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 Oct 2020 |
18 | * \brief Index agent interface definition for proxima search engine |
19 | */ |
20 | |
21 | #pragma once |
22 | |
23 | #include <ailego/algorithm/rate_limiter.h> |
24 | #include <ailego/parallel/thread_queue.h> |
25 | #include "index/index_service.h" |
26 | #include "meta/meta_service.h" |
27 | #include "collection_counter.h" |
28 | #include "column_order.h" |
29 | #include "write_request.h" |
30 | |
31 | namespace proxima { |
32 | namespace be { |
33 | namespace agent { |
34 | |
35 | class IndexAgent; |
36 | using IndexAgentPtr = std::shared_ptr<IndexAgent>; |
37 | |
38 | /*! Index Agent |
39 | */ |
40 | class IndexAgent { |
41 | public: |
42 | //! Create index agent |
43 | static IndexAgentPtr Create(meta::MetaServicePtr meta_service); |
44 | |
45 | //! Constructor |
46 | IndexAgent(meta::MetaServicePtr meta_service); |
47 | |
48 | //! Destructor |
49 | ~IndexAgent(); |
50 | |
51 | //! Create collection with schema |
52 | int create_collection(const std::string &collection_name); |
53 | |
54 | //! Update collection schema |
55 | int update_collection(const std::string &collection_name, uint32_t revision); |
56 | |
57 | //! Drop collection by name |
58 | int drop_collection(const std::string &collection_name); |
59 | |
60 | //! Get collection statstics |
61 | int get_collection_stats(const std::string &collection_name, |
62 | index::CollectionStats *collection_stats) const; |
63 | |
64 | //! Write records |
65 | int write(const WriteRequest &request); |
66 | |
67 | // Get latest lsn |
68 | int get_latest_lsn(const std::string &collection_name, uint64_t *lsn, |
69 | std::string *lsn_context); |
70 | |
71 | //! Get Index Service |
72 | const index::IndexServicePtr &get_service() const { |
73 | return index_service_; |
74 | } |
75 | |
76 | //! Get magic number |
77 | uint64_t get_magic_number() const { |
78 | return agent_timestamp_; |
79 | } |
80 | |
81 | //! Get collection column order |
82 | ColumnOrderPtr get_column_order(const std::string &collection_name) const { |
83 | return column_order_map_->get_column_order(collection_name); |
84 | } |
85 | |
86 | //! Get collection meta |
87 | meta::CollectionMetaPtr get_collection_meta( |
88 | const std::string &collection_name) const { |
89 | return meta_service_->get_current_collection(collection_name); |
90 | } |
91 | |
92 | public: |
93 | //! Init index agent |
94 | int init(); |
95 | |
96 | //! Cleanup index agent |
97 | int cleanup(); |
98 | |
99 | //! Start index agent |
100 | int start(); |
101 | |
102 | //! Stop index agent |
103 | int stop(); |
104 | |
105 | private: |
106 | //! Load index service |
107 | int load_index_service(); |
108 | |
109 | //! Is Index Agent Suspended |
110 | bool is_collection_suspend(const std::string &collection); |
111 | |
112 | //! Proxy write |
113 | int proxy_write(const WriteRequest &request, CollectionCounter *counter); |
114 | |
115 | //! Direct write |
116 | int direct_write(const WriteRequest &request, CollectionCounter *counter); |
117 | |
118 | //! Write the CollectionDataset to index service |
119 | void write_dataset(const std::string &collection_name, |
120 | const index::CollectionDatasetPtr &records, |
121 | CollectionCounter *counter); |
122 | |
123 | private: |
124 | IndexAgent(const IndexAgent &) = delete; |
125 | IndexAgent &operator=(const IndexAgent &) = delete; |
126 | |
127 | private: |
128 | //! Agent start timestamp |
129 | uint64_t agent_timestamp_{0}; |
130 | //! Request acquire timeout ms |
131 | int32_t acquire_timeout_{0}; |
132 | //! Meta service pointer |
133 | meta::MetaServicePtr meta_service_{}; |
134 | //! Index service pointer |
135 | index::IndexServicePtr index_service_{}; |
136 | //! Rate limiter ptr |
137 | ailego::RateLimiter::Pointer rate_limiter_{}; |
138 | //! Work thread queue |
139 | std::shared_ptr<ailego::ThreadQueue> thread_pool_{}; |
140 | //! Collection counter map |
141 | CollectionCounterMapPtr counter_map_{}; |
142 | //! Collection columns order map |
143 | ColumnOrderMapPtr column_order_map_{}; |
144 | }; |
145 | |
146 | } // end namespace agent |
147 | } // namespace be |
148 | } // end namespace proxima |
149 | |