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 IndexAgent service implementation for proxima search engine |
19 | */ |
20 | |
21 | #include "index_agent.h" |
22 | #include "common/config.h" |
23 | |
24 | namespace proxima { |
25 | namespace be { |
26 | namespace agent { |
27 | |
28 | IndexAgentPtr IndexAgent::Create(meta::MetaServicePtr meta_service) { |
29 | return std::make_shared<IndexAgent>(meta_service); |
30 | } |
31 | |
32 | IndexAgent::IndexAgent(meta::MetaServicePtr meta_service) |
33 | : meta_service_(std::move(meta_service)) {} |
34 | |
35 | IndexAgent::~IndexAgent() {} |
36 | |
37 | int IndexAgent::create_collection(const std::string &collection_name) { |
38 | // Add counter |
39 | counter_map_->add_counter(collection_name); |
40 | |
41 | // Add column order |
42 | auto schema = meta_service_->get_current_collection(collection_name); |
43 | if (!schema) { |
44 | LOG_ERROR("Get latest collection for meta service failed. collection[%s]" , |
45 | collection_name.c_str()); |
46 | return ErrorCode_InexistentCollection; |
47 | } |
48 | column_order_map_->add_column_order(*schema); |
49 | |
50 | // Create collection |
51 | int ret = index_service_->create_collection(collection_name, schema); |
52 | if (ret != 0) { |
53 | LOG_ERROR("Index service create collection failed. collection[%s]" , |
54 | collection_name.c_str()); |
55 | return ret; |
56 | } |
57 | |
58 | return 0; |
59 | } |
60 | |
61 | int IndexAgent::update_collection(const std::string &collection_name, |
62 | uint32_t revision) { |
63 | // get collection's counter |
64 | CollectionCounterPtr counter = counter_map_->get_counter(collection_name); |
65 | if (!counter) { |
66 | LOG_ERROR("Get collection counter failed. collection[%s]" , |
67 | collection_name.c_str()); |
68 | return ErrorCode_RuntimeError; |
69 | } |
70 | |
71 | // wait collection's all records to be processed |
72 | while (true) { |
73 | if (counter->active_count() == 0) { |
74 | break; |
75 | } |
76 | std::this_thread::sleep_for(std::chrono::seconds(1)); |
77 | } |
78 | |
79 | // get specified revision's collection meta |
80 | auto collection = meta_service_->get_collection(collection_name, revision); |
81 | if (!collection) { |
82 | LOG_ERROR("Meta service update collection failed. collection[%s]" , |
83 | collection_name.c_str()); |
84 | return ErrorCode_InexistentCollection; |
85 | } |
86 | |
87 | column_order_map_->update_column_order(*collection); |
88 | |
89 | // update collection |
90 | int ret = index_service_->update_collection(collection_name, collection); |
91 | if (ret != 0) { |
92 | LOG_ERROR("Index service update collection failed. collection[%s]" , |
93 | collection_name.c_str()); |
94 | return ret; |
95 | } |
96 | |
97 | return 0; |
98 | } |
99 | |
100 | int IndexAgent::drop_collection(const std::string &collection_name) { |
101 | int ret = index_service_->drop_collection(collection_name); |
102 | if (ret != 0) { |
103 | LOG_ERROR("Index service drop collection failed. collection[%s]" , |
104 | collection_name.c_str()); |
105 | return ret; |
106 | } |
107 | |
108 | counter_map_->remove_counter(collection_name); |
109 | column_order_map_->remove_column_order(collection_name); |
110 | |
111 | return 0; |
112 | } |
113 | |
114 | int IndexAgent::get_collection_stats( |
115 | const std::string &name, index::CollectionStats *collection_stats) const { |
116 | int ret = index_service_->get_collection_stats(name, collection_stats); |
117 | if (ret != 0) { |
118 | LOG_ERROR("Index service get collection stats failed. collection[%s]" , |
119 | name.c_str()); |
120 | return ret; |
121 | } |
122 | |
123 | return 0; |
124 | } |
125 | |
126 | bool IndexAgent::is_collection_suspend(const std::string &collection) { |
127 | meta::CollectionMetaPtr meta = |
128 | meta_service_->get_current_collection(collection); |
129 | if (!meta) { |
130 | LOG_ERROR("Meta service get latest collection failed. collection[%s]" , |
131 | collection.c_str()); |
132 | return false; |
133 | } |
134 | |
135 | return !meta->writable(); |
136 | } |
137 | |
138 | int IndexAgent::write(const WriteRequest &request) { |
139 | // check request empty |
140 | int row_count = request.row_count(); |
141 | if (!row_count) { |
142 | return 0; |
143 | } |
144 | |
145 | // check if suspend |
146 | const std::string &collection = request.collection_name(); |
147 | if (is_collection_suspend(collection)) { |
148 | return ErrorCode_SuspendedCollection; |
149 | } |
150 | |
151 | // check magic number |
152 | bool proxy_request = request.is_proxy_request(); |
153 | if (proxy_request && agent_timestamp_ != request.magic_number()) { |
154 | LOG_ERROR("Write request magic number mismatched." ); |
155 | return ErrorCode_MismatchedMagicNumber; |
156 | } |
157 | |
158 | // acquire permits |
159 | if (!rate_limiter_->try_acquire(row_count, acquire_timeout_)) { |
160 | LOG_WARN("Acquire permits failed. count[%d] collection[%s]" , row_count, |
161 | collection.c_str()); |
162 | return ErrorCode_ExceedRateLimit; |
163 | } |
164 | |
165 | CollectionCounterPtr counter = counter_map_->get_counter(collection); |
166 | if (!counter) { |
167 | LOG_ERROR("Get collection counter failed. collection[%s]" , |
168 | collection.c_str()); |
169 | return ErrorCode_InexistentCollection; |
170 | } |
171 | counter->add_active_count(row_count); |
172 | |
173 | // double check collection is suspended |
174 | if (is_collection_suspend(collection)) { |
175 | counter->sub_active_count(row_count); |
176 | return ErrorCode_SuspendedCollection; |
177 | } |
178 | |
179 | if (proxy_request) { |
180 | return proxy_write(request, counter.get()); |
181 | } else { |
182 | return direct_write(request, counter.get()); |
183 | } |
184 | } |
185 | |
186 | int IndexAgent::get_latest_lsn(const std::string &collection_name, |
187 | uint64_t *lsn, std::string *lsn_context) { |
188 | int ret = index_service_->get_latest_lsn(collection_name, lsn, lsn_context); |
189 | if (ret != 0) { |
190 | LOG_ERROR("Index service get collection latest lsn failed. collection[%s]" , |
191 | collection_name.c_str()); |
192 | return ret; |
193 | } |
194 | |
195 | return 0; |
196 | } |
197 | |
198 | int IndexAgent::init() { |
199 | if (!meta_service_) { |
200 | LOG_ERROR("Meta service is nullptr." ); |
201 | return ErrorCode_RuntimeError; |
202 | } |
203 | |
204 | // init index service |
205 | index_service_ = std::make_shared<index::IndexService>(); |
206 | int ret = index_service_->init(); |
207 | if (ret != 0) { |
208 | LOG_ERROR("Init index service failed." ); |
209 | return ret; |
210 | } |
211 | |
212 | // init rate limiter |
213 | const Config &config = Config::Instance(); |
214 | rate_limiter_ = ailego::RateLimiter::Create(config.get_index_max_build_qps()); |
215 | if (!rate_limiter_) { |
216 | LOG_ERROR("Create rate limiter failed." ); |
217 | return ErrorCode_RuntimeError; |
218 | } |
219 | // Set to default value |
220 | acquire_timeout_ = 500; |
221 | |
222 | // init agent stat timestamp |
223 | agent_timestamp_ = ailego::Monotime::MicroSeconds(); |
224 | |
225 | // create counter map |
226 | counter_map_ = std::make_shared<CollectionCounterMap>(); |
227 | |
228 | // create columns order map |
229 | column_order_map_ = std::make_shared<ColumnOrderMap>(); |
230 | |
231 | LOG_INFO("IndexAgent initialzie complete." ); |
232 | return 0; |
233 | } |
234 | |
235 | int IndexAgent::cleanup() { |
236 | if (index_service_) { |
237 | index_service_->cleanup(); |
238 | } |
239 | |
240 | LOG_INFO("IndexAgent cleanup complete." ); |
241 | return 0; |
242 | } |
243 | |
244 | int IndexAgent::start() { |
245 | // Start index service |
246 | int ret = index_service_->start(); |
247 | if (ret != 0) { |
248 | LOG_ERROR("Start index service failed." ); |
249 | return ret; |
250 | } |
251 | |
252 | ret = load_index_service(); |
253 | if (ret != 0) { |
254 | LOG_ERROR("Load index service failed." ); |
255 | return ret; |
256 | } |
257 | |
258 | thread_pool_ = std::make_shared<ailego::ThreadQueue>( |
259 | Config::Instance().get_index_build_thread_count()); |
260 | |
261 | LOG_INFO("IndexAgent start complete." ); |
262 | return 0; |
263 | } |
264 | |
265 | int IndexAgent::stop() { |
266 | if (thread_pool_) { |
267 | thread_pool_->stop(); |
268 | } |
269 | |
270 | if (index_service_) { |
271 | index_service_->stop(); |
272 | } |
273 | |
274 | LOG_INFO("IndexAgent stopped." ); |
275 | return 0; |
276 | } |
277 | |
278 | int IndexAgent::proxy_write(const WriteRequest &request, |
279 | CollectionCounter *counter) { |
280 | auto &collection = request.collection_name(); |
281 | int dataset_count = (int)request.collection_dataset_count(); |
282 | for (int i = 0; i < dataset_count; ++i) { |
283 | auto &dataset = request.get_collection_dataset(i); |
284 | thread_pool_->execute(dataset->get(0).primary_key, this, |
285 | &IndexAgent::write_dataset, collection, dataset, |
286 | counter); |
287 | } |
288 | |
289 | return 0; |
290 | } |
291 | |
292 | int IndexAgent::direct_write(const WriteRequest &request, |
293 | CollectionCounter *counter) { |
294 | auto &collection = request.collection_name(); |
295 | int row_count = request.row_count(); |
296 | auto &dataset = request.get_collection_dataset(0); |
297 | int ret = index_service_->write_records(collection, dataset); |
298 | if (ret != 0) { |
299 | counter->sub_active_count(row_count); |
300 | LOG_ERROR("Index service write records failed. collection[%s]" , |
301 | collection.c_str()); |
302 | return ret; |
303 | } |
304 | counter->sub_active_count(row_count); |
305 | |
306 | return 0; |
307 | } |
308 | |
309 | void IndexAgent::write_dataset(const std::string &collection_name, |
310 | const index::CollectionDatasetPtr &record, |
311 | CollectionCounter *counter) { |
312 | int ret = index_service_->write_records(collection_name, record); |
313 | if (ret != 0) { |
314 | LOG_ERROR( |
315 | "Index service write record failed. " |
316 | "code[%d] reason[%s] collection[%s]" , |
317 | ret, ErrorCode::What(ret), collection_name.c_str()); |
318 | } |
319 | counter->dec_active_count(); |
320 | } |
321 | |
322 | int IndexAgent::load_index_service() { |
323 | // 1.Get all valid collection schemas |
324 | meta::CollectionMetaPtrList schemas; |
325 | int ret = meta_service_->get_latest_collections(&schemas); |
326 | if (ret != 0) { |
327 | LOG_ERROR("Meta service get latest collections failed." ); |
328 | return ret; |
329 | } |
330 | |
331 | // 2.Get all valid collection names, and create all collections columns order |
332 | std::vector<std::string> collection_names; |
333 | for (auto &schema : schemas) { |
334 | collection_names.emplace_back(schema->name()); |
335 | column_order_map_->add_column_order(*schema); |
336 | } |
337 | |
338 | // 3.Create all collections counter |
339 | for (auto &collection_name : collection_names) { |
340 | counter_map_->add_counter(collection_name); |
341 | } |
342 | |
343 | // 4.Load index service |
344 | ret = index_service_->load_collections(collection_names, schemas); |
345 | if (ret != 0) { |
346 | LOG_ERROR("Index service load collections failed." ); |
347 | return ret; |
348 | } |
349 | |
350 | return 0; |
351 | } |
352 | |
353 | } // end namespace agent |
354 | } // namespace be |
355 | } // end namespace proxima |
356 | |