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 guonix
17 * \date Dec 2020
18 * \brief
19 */
20
21#include "admin_agent.h"
22#include <utility>
23#include "admin/admin_proto_converter.h"
24#include "common/logger.h"
25
26namespace proxima {
27namespace be {
28namespace admin {
29
30
31/**
32 * Proxima BE Admin module
33 */
34class AdminAgentImpl : public AdminAgent {
35 public:
36 //! Constructor
37 explicit AdminAgentImpl(meta::MetaAgentPtr meta, agent::IndexAgentPtr index,
38 query::QueryAgentPtr query)
39 : meta_agent_(std::move(meta)),
40 index_agent_(std::move(index)),
41 query_agent_(std::move(query)) {}
42 //! Destructor
43 ~AdminAgentImpl() override = default;
44
45 public:
46 //! Init Meta Agent
47 int init() override {
48 LOG_INFO("AdminAgent initialize complete.");
49 return 0;
50 }
51
52 //! Clean up object
53 int cleanup() override {
54 LOG_INFO("AdminAgent cleanup complete.");
55 return 0;
56 }
57
58 //! Start background service
59 int start() override {
60 LOG_INFO("AdminAgent start complete.");
61 return 0;
62 }
63
64 //! Stop background service
65 int stop() override {
66 LOG_INFO("AdminAgent stopped.");
67 return 0;
68 }
69
70 //! Create collection
71 int create_collection(const proto::CollectionConfig &request) override {
72 meta::CollectionBase param;
73 int code = AdminProtoConverter::PBToCollectionBase(request, &param);
74 if (code != 0) {
75 LOG_ERROR("Deserialize collection meta from pb failed.");
76 return code;
77 }
78
79 meta::CollectionMetaPtr collection;
80 code = meta_agent_->create_collection(param, &collection);
81 if (code != 0) {
82 LOG_ERROR("MetaAgent create collection failed. code[%d] what[%s]", code,
83 ErrorCode::What(code));
84 return code;
85 }
86
87 code = index_agent_->create_collection(collection->name());
88 if (code == 0) {
89 LOG_INFO("Create collection success. collection_config[%s]",
90 request.ShortDebugString().c_str());
91 meta_agent_->enable_collection(collection->name(),
92 collection->revision());
93 } else {
94 meta_agent_->delete_collection(collection->name());
95 LOG_ERROR("IndexAgent create collection failed. code[%d] what[%s]", code,
96 ErrorCode::What(code));
97 }
98 return code;
99 }
100
101 //! Describe collection
102 int describe_collection(
103 const std::string &collection_name,
104 proto::DescribeCollectionResponse *collection_info) override {
105 auto collection = meta_agent_->get_collection(collection_name);
106 if (!collection) {
107 LOG_ERROR("Failed to describe collection. collection[%s]",
108 collection_name.c_str());
109 return PROXIMA_BE_ERROR_CODE(InexistentCollection);
110 }
111 return fill_collection_info(collection,
112 collection_info->mutable_collection());
113 }
114
115 //! Delete collection
116 int drop_collection(const std::string &collection_name) override {
117 int code = index_agent_->drop_collection(collection_name);
118 if (code != 0) {
119 LOG_ERROR("IndexAgent delete collection failed. code[%d] what[%s]", code,
120 ErrorCode::What(code));
121 return code;
122 }
123
124 // Continue to remove collection in meta
125 code = meta_agent_->delete_collection(collection_name);
126 if (code != 0) {
127 LOG_ERROR("MetaAgent delete collection failed. code[%d] what[%s]", code,
128 ErrorCode::What(code));
129 return code;
130 }
131 return code;
132 }
133
134 //! Retrieve collections
135 int list_collections(const proto::ListCondition &condition,
136 proto::ListCollectionsResponse *response) override {
137 meta::CollectionMetaPtrList collections;
138 int code = meta_agent_->list_collections(&collections);
139 if (code != 0) {
140 LOG_ERROR("Failed to list collections. code[%d] what[%s].", code,
141 ErrorCode::What(code));
142 return code;
143 }
144 for (auto &collection : collections) {
145 if (condition.repository_name().empty() ||
146 (condition.repository_name() == collection->repository_name())) {
147 auto *pb_collection = response->add_collections();
148 int ret = fill_collection_info(collection, pb_collection);
149 if (ret != 0) {
150 return ret;
151 }
152 }
153 }
154
155 return 0;
156 }
157
158 int stats_collection(const std::string &collection_name,
159 proto::StatsCollectionResponse *stats) override {
160 index::CollectionStats collection_stats;
161 int code =
162 index_agent_->get_collection_stats(collection_name, &collection_stats);
163 if (code != 0) {
164 LOG_ERROR(
165 "Failed to get collection stats, collection[%s] code[%d], "
166 "what[%s].",
167 collection_name.c_str(), code, ErrorCode::What(code));
168 return code;
169 }
170 AdminProtoConverter::CollectionStatsToPB(collection_stats,
171 stats->mutable_collection_stats());
172 return 0;
173 }
174
175 int reload_meta() override {
176 return meta_agent_->reload();
177 }
178
179 int start_query_service() override {
180 return query_agent_->start();
181 }
182
183 int stop_query_service() override {
184 return query_agent_->stop();
185 }
186
187 int get_query_service_status() override {
188 return query_agent_->is_running();
189 }
190
191 private:
192 //! only called in proxy write scenario
193 int fill_lsn_context_and_magic_number(
194 const std::string &collection_name,
195 proto::CollectionInfo *collection_info) {
196 collection_info->set_magic_number(index_agent_->get_magic_number());
197 std::string context;
198 uint64_t lsn;
199 int ret = index_agent_->get_latest_lsn(collection_name, &lsn, &context);
200 if (ret != 0) {
201 LOG_WARN("Get latest lsn failed. collection_name[%s] ret[%d]",
202 collection_name.c_str(), ret);
203 return ret;
204 }
205 auto *lsn_context = collection_info->mutable_latest_lsn_context();
206 lsn_context->set_lsn(lsn);
207 lsn_context->set_context(std::move(context));
208 return 0;
209 }
210
211 //! fill protobuf collection info
212 int fill_collection_info(const meta::CollectionMetaPtr &meta,
213 proto::CollectionInfo *collection) {
214 int ret = 0;
215 AdminProtoConverter::CollectionMetaToPB(*meta, collection);
216 if (meta->repository()) {
217 ret = fill_lsn_context_and_magic_number(meta->name(), collection);
218 }
219 return ret;
220 }
221
222 private:
223 //! Meta agent
224 meta::MetaAgentPtr meta_agent_{nullptr};
225
226 //! Index Agent
227 agent::IndexAgentPtr index_agent_{nullptr};
228
229 //! Query agent
230 query::QueryAgentPtr query_agent_{nullptr};
231};
232
233AdminAgentPtr AdminAgent::Create(const meta::MetaAgentPtr &meta,
234 const agent::IndexAgentPtr &agent,
235 const query::QueryAgentPtr &query) {
236 return std::make_shared<AdminAgentImpl>(meta, agent, query);
237}
238
239} // namespace admin
240} // namespace be
241} // namespace proxima
242