1 | /** |
---|---|
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | |
16 | * \author Hongqing.hu |
17 | * \date Dec 2020 |
18 | * \brief Collection counter interface definition for proxima search engine |
19 | */ |
20 | |
21 | #pragma once |
22 | |
23 | #include <atomic> |
24 | #include <memory> |
25 | #include <mutex> |
26 | #include <unordered_map> |
27 | |
28 | namespace proxima { |
29 | namespace be { |
30 | namespace agent { |
31 | |
32 | class CollectionCounter; |
33 | class CollectionCounterMap; |
34 | |
35 | using CollectionCounterPtr = std::shared_ptr<CollectionCounter>; |
36 | using CollectionCounterMapPtr = std::shared_ptr<CollectionCounterMap>; |
37 | |
38 | /*! CollectionCounter |
39 | */ |
40 | class CollectionCounter { |
41 | public: |
42 | //! Constructor |
43 | CollectionCounter() = default; |
44 | |
45 | //! Destructor |
46 | ~CollectionCounter() = default; |
47 | |
48 | //! Add active count |
49 | void add_active_count(uint32_t count) { |
50 | active_count_.fetch_add(count); |
51 | } |
52 | |
53 | //! Sub active count |
54 | void sub_active_count(uint32_t count) { |
55 | active_count_.fetch_sub(count); |
56 | } |
57 | |
58 | //! Sub one active count |
59 | void dec_active_count() { |
60 | --active_count_; |
61 | } |
62 | |
63 | //! Get active count |
64 | uint32_t active_count() const { |
65 | return active_count_.load(); |
66 | } |
67 | |
68 | private: |
69 | //! Collection active record count |
70 | std::atomic<uint32_t> active_count_{0U}; |
71 | }; |
72 | |
73 | /*! CollectionCounterMap |
74 | */ |
75 | class CollectionCounterMap { |
76 | public: |
77 | //! Constructor |
78 | CollectionCounterMap() = default; |
79 | |
80 | //! Destructor |
81 | ~CollectionCounterMap() = default; |
82 | |
83 | //! Add collection counter |
84 | void add_counter(const std::string &name) { |
85 | CollectionCounterPtr counter = std::make_shared<CollectionCounter>(); |
86 | std::lock_guard<std::mutex> lock(mutex_); |
87 | counter_map_[name] = counter; |
88 | } |
89 | |
90 | //! Remove collection counter |
91 | void remove_counter(const std::string &name) { |
92 | std::lock_guard<std::mutex> lock(mutex_); |
93 | counter_map_.erase(name); |
94 | } |
95 | |
96 | //! Get collection counter |
97 | CollectionCounterPtr get_counter(const std::string &name) { |
98 | std::lock_guard<std::mutex> lock(mutex_); |
99 | auto it = counter_map_.find(name); |
100 | if (it != counter_map_.end()) { |
101 | return it->second; |
102 | } |
103 | return nullptr; |
104 | } |
105 | |
106 | private: |
107 | //! Mutex |
108 | std::mutex mutex_{}; |
109 | //! Collection counter map |
110 | std::unordered_map<std::string, CollectionCounterPtr> counter_map_{}; |
111 | }; |
112 | |
113 | } // end namespace agent |
114 | } // namespace be |
115 | } // end namespace proxima |
116 |