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
31namespace proxima {
32namespace be {
33namespace agent {
34
35class IndexAgent;
36using IndexAgentPtr = std::shared_ptr<IndexAgent>;
37
38/*! Index Agent
39 */
40class 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