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
28namespace proxima {
29namespace be {
30namespace agent {
31
32class CollectionCounter;
33class CollectionCounterMap;
34
35using CollectionCounterPtr = std::shared_ptr<CollectionCounter>;
36using CollectionCounterMapPtr = std::shared_ptr<CollectionCounterMap>;
37
38/*! CollectionCounter
39 */
40class 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 */
75class 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