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
24namespace proxima {
25namespace be {
26namespace agent {
27
28IndexAgentPtr IndexAgent::Create(meta::MetaServicePtr meta_service) {
29 return std::make_shared<IndexAgent>(meta_service);
30}
31
32IndexAgent::IndexAgent(meta::MetaServicePtr meta_service)
33 : meta_service_(std::move(meta_service)) {}
34
35IndexAgent::~IndexAgent() {}
36
37int 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
61int 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
100int 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
114int 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
126bool 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
138int 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
186int 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
198int 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
235int IndexAgent::cleanup() {
236 if (index_service_) {
237 index_service_->cleanup();
238 }
239
240 LOG_INFO("IndexAgent cleanup complete.");
241 return 0;
242}
243
244int 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
265int 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
278int 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
292int 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
309void 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
322int 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