1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 //
6 // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
7 // Use of this source code is governed by a BSD-style license that can be
8 // found in the LICENSE file. See the AUTHORS file for names of contributors.
9 
10 #pragma once
11 #include <atomic>
12 #include <memory>
13 
14 #include "rocksdb/env.h"
15 #include "rocksdb/concurrent_task_limiter.h"
16 
17 namespace ROCKSDB_NAMESPACE {
18 
19 class TaskLimiterToken;
20 
21 class ConcurrentTaskLimiterImpl : public ConcurrentTaskLimiter {
22  public:
23   explicit ConcurrentTaskLimiterImpl(const std::string& name,
24                                      int32_t max_outstanding_task);
25   // No copying allowed
26   ConcurrentTaskLimiterImpl(const ConcurrentTaskLimiterImpl&) = delete;
27   ConcurrentTaskLimiterImpl& operator=(const ConcurrentTaskLimiterImpl&) =
28       delete;
29 
30   virtual ~ConcurrentTaskLimiterImpl();
31 
32   virtual const std::string& GetName() const override;
33 
34   virtual void SetMaxOutstandingTask(int32_t limit) override;
35 
36   virtual void ResetMaxOutstandingTask() override;
37 
38   virtual int32_t GetOutstandingTask() const override;
39 
40   // Request token for adding a new task.
41   // If force == true, it requests a token bypassing throttle.
42   // Returns nullptr if it got throttled.
43   virtual std::unique_ptr<TaskLimiterToken> GetToken(bool force);
44 
45  private:
46   friend class TaskLimiterToken;
47 
48   std::string name_;
49   std::atomic<int32_t> max_outstanding_tasks_;
50   std::atomic<int32_t> outstanding_tasks_;
51 };
52 
53 class TaskLimiterToken {
54  public:
TaskLimiterToken(ConcurrentTaskLimiterImpl * limiter)55   explicit TaskLimiterToken(ConcurrentTaskLimiterImpl* limiter)
56       : limiter_(limiter) {}
57   ~TaskLimiterToken();
58 
59  private:
60   ConcurrentTaskLimiterImpl* limiter_;
61 
62   // no copying allowed
63   TaskLimiterToken(const TaskLimiterToken&) = delete;
64   void operator=(const TaskLimiterToken&) = delete;
65 };
66 
67 }  // namespace ROCKSDB_NAMESPACE
68