DBImpl.cpp 97.2 KB
Newer Older
1
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
J
jinhai 已提交
2
//
3 4
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
J
jinhai 已提交
5
//
6 7 8 9 10
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
J
jinhai 已提交
11

S
starlord 已提交
12
#include "db/DBImpl.h"
Z
Zhiru Zhu 已提交
13 14

#include <assert.h>
15
#include <fiu-local.h>
Z
Zhiru Zhu 已提交
16 17 18 19 20

#include <algorithm>
#include <boost/filesystem.hpp>
#include <chrono>
#include <cstring>
21
#include <functional>
Z
Zhiru Zhu 已提交
22
#include <iostream>
23
#include <limits>
24
#include <queue>
Z
Zhiru Zhu 已提交
25 26
#include <set>
#include <thread>
27
#include <unordered_map>
Z
Zhiru Zhu 已提交
28 29
#include <utility>

S
starlord 已提交
30
#include "Utils.h"
S
starlord 已提交
31 32
#include "cache/CpuCacheMgr.h"
#include "cache/GpuCacheMgr.h"
33
#include "db/IDGenerator.h"
G
groot 已提交
34
#include "db/merge/MergeManagerFactory.h"
S
starlord 已提交
35
#include "engine/EngineFactory.h"
36
#include "index/thirdparty/faiss/utils/distances.h"
37
#include "insert/MemManagerFactory.h"
S
starlord 已提交
38
#include "meta/MetaConsts.h"
S
starlord 已提交
39 40
#include "meta/MetaFactory.h"
#include "meta/SqliteMetaImpl.h"
G
groot 已提交
41
#include "metrics/Metrics.h"
G
groot 已提交
42
#include "scheduler/Definition.h"
S
starlord 已提交
43
#include "scheduler/SchedInst.h"
Y
Yu Kun 已提交
44
#include "scheduler/job/BuildIndexJob.h"
S
starlord 已提交
45 46
#include "scheduler/job/DeleteJob.h"
#include "scheduler/job/SearchJob.h"
47 48 49
#include "segment/SegmentReader.h"
#include "segment/SegmentWriter.h"
#include "utils/Exception.h"
S
starlord 已提交
50
#include "utils/Log.h"
G
groot 已提交
51
#include "utils/StringHelpFunctions.h"
S
starlord 已提交
52
#include "utils/TimeRecorder.h"
53 54
#include "utils/ValidationUtil.h"
#include "wal/WalDefinations.h"
X
Xu Peng 已提交
55

56 57
#include "search/TaskInst.h"

J
jinhai 已提交
58
namespace milvus {
X
Xu Peng 已提交
59
namespace engine {
X
Xu Peng 已提交
60

G
groot 已提交
61
namespace {
G
groot 已提交
62 63
constexpr uint64_t BACKGROUND_METRIC_INTERVAL = 1;
constexpr uint64_t BACKGROUND_INDEX_INTERVAL = 1;
G
groot 已提交
64
constexpr uint64_t WAIT_BUILD_INDEX_INTERVAL = 5;
G
groot 已提交
65

66 67 68 69 70 71 72 73
constexpr const char* JSON_ROW_COUNT = "row_count";
constexpr const char* JSON_PARTITIONS = "partitions";
constexpr const char* JSON_PARTITION_TAG = "tag";
constexpr const char* JSON_SEGMENTS = "segments";
constexpr const char* JSON_SEGMENT_NAME = "name";
constexpr const char* JSON_INDEX_NAME = "index_name";
constexpr const char* JSON_DATA_SIZE = "data_size";

G
groot 已提交
74
static const Status SHUTDOWN_ERROR = Status(DB_ERROR, "Milvus server is shutdown!");
G
groot 已提交
75

S
starlord 已提交
76
}  // namespace
G
groot 已提交
77

Y
Yu Kun 已提交
78
DBImpl::DBImpl(const DBOptions& options)
79
    : options_(options), initialized_(false), merge_thread_pool_(1, 1), index_thread_pool_(1, 1) {
S
starlord 已提交
80
    meta_ptr_ = MetaFactory::Build(options.meta_, options.mode_);
Z
zhiru 已提交
81
    mem_mgr_ = MemManagerFactory::Build(meta_ptr_, options_);
G
groot 已提交
82
    merge_mgr_ptr_ = MergeManagerFactory::Build(meta_ptr_, options_);
83 84 85 86 87 88 89 90 91 92

    if (options_.wal_enable_) {
        wal::MXLogConfiguration mxlog_config;
        mxlog_config.recovery_error_ignore = options_.recovery_error_ignore_;
        // 2 buffers in the WAL
        mxlog_config.buffer_size = options_.buffer_size_ / 2;
        mxlog_config.mxlog_path = options_.mxlog_path_;
        wal_mgr_ = std::make_shared<wal::WalManager>(mxlog_config);
    }

93 94
    SetIdentity("DBImpl");
    AddCacheInsertDataListener();
95
    AddUseBlasThresholdListener();
96

S
starlord 已提交
97 98 99 100 101 102 103
    Start();
}

DBImpl::~DBImpl() {
    Stop();
}

S
starlord 已提交
104
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
105
// external api
S
starlord 已提交
106
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
107 108
Status
DBImpl::Start() {
109
    if (initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
110 111 112
        return Status::OK();
    }

113
    // LOG_ENGINE_TRACE_ << "DB service start";
114
    initialized_.store(true, std::memory_order_release);
S
starlord 已提交
115

116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
    // wal
    if (options_.wal_enable_) {
        auto error_code = DB_ERROR;
        if (wal_mgr_ != nullptr) {
            error_code = wal_mgr_->Init(meta_ptr_);
        }
        if (error_code != WAL_SUCCESS) {
            throw Exception(error_code, "Wal init error!");
        }

        // recovery
        while (1) {
            wal::MXLogRecord record;
            auto error_code = wal_mgr_->GetNextRecovery(record);
            if (error_code != WAL_SUCCESS) {
                throw Exception(error_code, "Wal recovery error!");
            }
            if (record.type == wal::MXLogType::None) {
                break;
            }
            ExecWalRecord(record);
        }

        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
G
groot 已提交
141 142
            // background wal thread
            bg_wal_thread_ = std::thread(&DBImpl::BackgroundWalThread, this);
143 144 145 146
        }
    } else {
        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
G
groot 已提交
147 148
            // background flush thread
            bg_flush_thread_ = std::thread(&DBImpl::BackgroundFlushThread, this);
149
        }
Z
update  
zhiru 已提交
150
    }
S
starlord 已提交
151

G
groot 已提交
152 153 154 155 156 157 158 159 160
    // for distribute version, some nodes are read only
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        // background build index thread
        bg_index_thread_ = std::thread(&DBImpl::BackgroundIndexThread, this);
    }

    // background metric thread
    bg_metric_thread_ = std::thread(&DBImpl::BackgroundMetricThread, this);

S
starlord 已提交
161 162 163
    return Status::OK();
}

S
starlord 已提交
164 165
Status
DBImpl::Stop() {
166
    if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
167 168
        return Status::OK();
    }
169

170
    initialized_.store(false, std::memory_order_release);
S
starlord 已提交
171

172 173
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        if (options_.wal_enable_) {
G
groot 已提交
174 175
            // wait wal thread finish
            swn_wal_.Notify();
176 177
            bg_wal_thread_.join();
        } else {
G
groot 已提交
178
            // flush all without merge
179 180 181 182
            wal::MXLogRecord record;
            record.type = wal::MXLogType::Flush;
            ExecWalRecord(record);

G
groot 已提交
183 184 185
            // wait flush thread finish
            swn_flush_.Notify();
            bg_flush_thread_.join();
186
        }
S
starlord 已提交
187

188 189
        WaitMergeFileFinish();

G
groot 已提交
190 191 192
        swn_index_.Notify();
        bg_index_thread_.join();

193
        meta_ptr_->CleanUpShadowFiles();
S
starlord 已提交
194 195
    }

G
groot 已提交
196 197 198 199
    // wait metric thread exit
    swn_metric_.Notify();
    bg_metric_thread_.join();

200
    // LOG_ENGINE_TRACE_ << "DB service stop";
S
starlord 已提交
201
    return Status::OK();
X
Xu Peng 已提交
202 203
}

S
starlord 已提交
204 205
Status
DBImpl::DropAll() {
S
starlord 已提交
206 207 208
    return meta_ptr_->DropAll();
}

S
starlord 已提交
209
Status
210
DBImpl::CreateCollection(meta::CollectionSchema& collection_schema) {
211
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
212
        return SHUTDOWN_ERROR;
S
starlord 已提交
213 214
    }

215
    meta::CollectionSchema temp_schema = collection_schema;
B
bigbraver 已提交
216
    temp_schema.index_file_size_ *= MB;  // store as MB
217
    if (options_.wal_enable_) {
218
        temp_schema.flush_lsn_ = wal_mgr_->CreateCollection(collection_schema.collection_id_);
219 220
    }

221
    return meta_ptr_->CreateCollection(temp_schema);
222 223
}

224 225 226 227 228 229 230
Status
DBImpl::CreateHybridCollection(meta::CollectionSchema& collection_schema, meta::hybrid::FieldsSchema& fields_schema) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    meta::CollectionSchema temp_schema = collection_schema;
Y
yukun 已提交
231
    temp_schema.index_file_size_ *= MB;
232
    if (options_.wal_enable_) {
Y
yukun 已提交
233
        temp_schema.flush_lsn_ = wal_mgr_->CreateHybridCollection(collection_schema.collection_id_);
234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
    }

    return meta_ptr_->CreateHybridCollection(temp_schema, fields_schema);
}

Status
DBImpl::DescribeHybridCollection(meta::CollectionSchema& collection_schema,
                                 milvus::engine::meta::hybrid::FieldsSchema& fields_schema) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    auto stat = meta_ptr_->DescribeHybridCollection(collection_schema, fields_schema);
    return stat;
}

S
starlord 已提交
250
Status
251
DBImpl::DropCollection(const std::string& collection_id) {
252
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
253
        return SHUTDOWN_ERROR;
S
starlord 已提交
254 255
    }

256
    if (options_.wal_enable_) {
257
        wal_mgr_->DropCollection(collection_id);
258 259
    }

260
    return DropCollectionRecursively(collection_id);
G
groot 已提交
261 262
}

S
starlord 已提交
263
Status
264
DBImpl::DescribeCollection(meta::CollectionSchema& collection_schema) {
265
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
266
        return SHUTDOWN_ERROR;
S
starlord 已提交
267 268
    }

269
    auto stat = meta_ptr_->DescribeCollection(collection_schema);
B
bigbraver 已提交
270
    collection_schema.index_file_size_ /= MB;  // return as MB
S
starlord 已提交
271
    return stat;
272 273
}

S
starlord 已提交
274
Status
275
DBImpl::HasCollection(const std::string& collection_id, bool& has_or_not) {
276
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
277
        return SHUTDOWN_ERROR;
S
starlord 已提交
278 279
    }

G
groot 已提交
280
    return meta_ptr_->HasCollection(collection_id, has_or_not, false);
281 282
}

283
Status
G
groot 已提交
284
DBImpl::HasNativeCollection(const std::string& collection_id, bool& has_or_not) {
285 286 287 288
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
289
    return meta_ptr_->HasCollection(collection_id, has_or_not, true);
290 291
}

S
starlord 已提交
292
Status
293
DBImpl::AllCollections(std::vector<meta::CollectionSchema>& collection_schema_array) {
294
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
295
        return SHUTDOWN_ERROR;
S
starlord 已提交
296 297
    }

298 299
    std::vector<meta::CollectionSchema> all_collections;
    auto status = meta_ptr_->AllCollections(all_collections);
300

301 302 303 304 305
    // only return real collections, dont return partition collections
    collection_schema_array.clear();
    for (auto& schema : all_collections) {
        if (schema.owner_collection_.empty()) {
            collection_schema_array.push_back(schema);
306 307 308 309
        }
    }

    return status;
G
groot 已提交
310 311
}

312
Status
313
DBImpl::GetCollectionInfo(const std::string& collection_id, std::string& collection_info) {
314 315 316 317 318
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // step1: get all partition ids
J
Jin Hai 已提交
319 320
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
321

J
Jin Hai 已提交
322 323
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::INDEX};
324

325 326 327
    milvus::json json_info;
    milvus::json json_partitions;
    size_t total_row_count = 0;
328

329
    auto get_info = [&](const std::string& col_id, const std::string& tag) {
G
groot 已提交
330 331
        meta::FilesHolder files_holder;
        status = meta_ptr_->FilesByType(col_id, file_types, files_holder);
332
        if (!status.ok()) {
J
Jin Hai 已提交
333
            std::string err_msg = "Failed to get collection info: " + status.ToString();
334
            LOG_ENGINE_ERROR_ << err_msg;
335 336 337
            return Status(DB_ERROR, err_msg);
        }

338 339 340 341 342
        milvus::json json_partition;
        json_partition[JSON_PARTITION_TAG] = tag;

        milvus::json json_segments;
        size_t row_count = 0;
G
groot 已提交
343
        milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
344
        for (auto& file : collection_files) {
345 346 347 348 349 350 351 352 353
            milvus::json json_segment;
            json_segment[JSON_SEGMENT_NAME] = file.segment_id_;
            json_segment[JSON_ROW_COUNT] = file.row_count_;
            json_segment[JSON_INDEX_NAME] = utils::GetIndexName(file.engine_type_);
            json_segment[JSON_DATA_SIZE] = (int64_t)file.file_size_;
            json_segments.push_back(json_segment);

            row_count += file.row_count_;
            total_row_count += file.row_count_;
354 355
        }

356 357 358 359 360 361 362
        json_partition[JSON_ROW_COUNT] = row_count;
        json_partition[JSON_SEGMENTS] = json_segments;

        json_partitions.push_back(json_partition);

        return Status::OK();
    };
363

364 365 366 367
    // step2: get default partition info
    status = get_info(collection_id, milvus::engine::DEFAULT_PARTITON_TAG);
    if (!status.ok()) {
        return status;
368 369
    }

370 371 372 373 374 375 376 377 378 379 380 381 382
    // step3: get partitions info
    for (auto& schema : partition_array) {
        status = get_info(schema.collection_id_, schema.partition_tag_);
        if (!status.ok()) {
            return status;
        }
    }

    json_info[JSON_ROW_COUNT] = total_row_count;
    json_info[JSON_PARTITIONS] = json_partitions;

    collection_info = json_info.dump();

383 384 385
    return Status::OK();
}

S
starlord 已提交
386
Status
387
DBImpl::PreloadCollection(const std::string& collection_id) {
388
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
389
        return SHUTDOWN_ERROR;
S
starlord 已提交
390 391
    }

J
Jin Hai 已提交
392
    // step 1: get all collection files from parent collection
G
groot 已提交
393 394
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesToSearch(collection_id, files_holder);
Y
Yu Kun 已提交
395 396 397
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
398

399
    // step 2: get files from partition collections
J
Jin Hai 已提交
400 401
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
402
    for (auto& schema : partition_array) {
G
groot 已提交
403
        status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
G
groot 已提交
404 405
    }

Y
Yu Kun 已提交
406 407
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
408 409
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
410

411
    // step 3: load file one by one
G
groot 已提交
412
    milvus::engine::meta::SegmentsSchema& files_array = files_holder.HoldFiles();
413 414
    LOG_ENGINE_DEBUG_ << "Begin pre-load collection:" + collection_id + ", totally " << files_array.size()
                      << " files need to be pre-loaded";
J
Jin Hai 已提交
415
    TimeRecorderAuto rc("Pre-load collection:" + collection_id);
G
groot 已提交
416
    for (auto& file : files_array) {
417
        EngineType engine_type;
J
Jin Hai 已提交
418 419 420
        if (file.file_type_ == meta::SegmentSchema::FILE_TYPE::RAW ||
            file.file_type_ == meta::SegmentSchema::FILE_TYPE::TO_INDEX ||
            file.file_type_ == meta::SegmentSchema::FILE_TYPE::BACKUP) {
421 422
            engine_type =
                utils::IsBinaryMetricType(file.metric_type_) ? EngineType::FAISS_BIN_IDMAP : EngineType::FAISS_IDMAP;
423 424 425
        } else {
            engine_type = (EngineType)file.engine_type_;
        }
426 427 428 429

        auto json = milvus::json::parse(file.index_params_);
        ExecutionEnginePtr engine =
            EngineFactory::Build(file.dimension_, file.location_, engine_type, (MetricType)file.metric_type_, json);
430
        fiu_do_on("DBImpl.PreloadCollection.null_engine", engine = nullptr);
G
groot 已提交
431
        if (engine == nullptr) {
432
            LOG_ENGINE_ERROR_ << "Invalid engine type";
G
groot 已提交
433 434
            return Status(DB_ERROR, "Invalid engine type");
        }
Y
Yu Kun 已提交
435

436
        fiu_do_on("DBImpl.PreloadCollection.exceed_cache", size = available_size + 1);
437 438

        try {
439
            fiu_do_on("DBImpl.PreloadCollection.engine_throw_exception", throw std::exception());
440 441
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
442 443 444 445
            status = engine->Load(true);
            if (!status.ok()) {
                return status;
            }
446 447 448

            size += engine->Size();
            if (size > available_size) {
449
                LOG_ENGINE_DEBUG_ << "Pre-load cancelled since cache is almost full";
450
                return Status(SERVER_CACHE_FULL, "Cache is full");
Y
Yu Kun 已提交
451
            }
452
        } catch (std::exception& ex) {
J
Jin Hai 已提交
453
            std::string msg = "Pre-load collection encounter exception: " + std::string(ex.what());
454
            LOG_ENGINE_ERROR_ << msg;
455
            return Status(DB_ERROR, msg);
Y
Yu Kun 已提交
456 457
        }
    }
G
groot 已提交
458

Y
Yu Kun 已提交
459
    return Status::OK();
Y
Yu Kun 已提交
460 461
}

S
starlord 已提交
462
Status
463
DBImpl::UpdateCollectionFlag(const std::string& collection_id, int64_t flag) {
464
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
465
        return SHUTDOWN_ERROR;
S
starlord 已提交
466 467
    }

468
    return meta_ptr_->UpdateCollectionFlag(collection_id, flag);
S
starlord 已提交
469 470
}

S
starlord 已提交
471
Status
472
DBImpl::GetCollectionRowCount(const std::string& collection_id, uint64_t& row_count) {
473
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
474 475 476
        return SHUTDOWN_ERROR;
    }

477
    return GetCollectionRowCountRecursively(collection_id, row_count);
G
groot 已提交
478 479 480
}

Status
J
Jin Hai 已提交
481
DBImpl::CreatePartition(const std::string& collection_id, const std::string& partition_name,
G
groot 已提交
482
                        const std::string& partition_tag) {
483
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
484 485 486
        return SHUTDOWN_ERROR;
    }

487
    uint64_t lsn = 0;
488
    meta_ptr_->GetCollectionFlushLSN(collection_id, lsn);
J
Jin Hai 已提交
489
    return meta_ptr_->CreatePartition(collection_id, partition_name, partition_tag, lsn);
G
groot 已提交
490 491 492 493
}

Status
DBImpl::DropPartition(const std::string& partition_name) {
494
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
495
        return SHUTDOWN_ERROR;
S
starlord 已提交
496 497
    }

498
    mem_mgr_->EraseMemVector(partition_name);                // not allow insert
J
Jin Hai 已提交
499
    auto status = meta_ptr_->DropPartition(partition_name);  // soft delete collection
500
    if (!status.ok()) {
501
        LOG_ENGINE_ERROR_ << status.message();
502 503
        return status;
    }
G
groot 已提交
504

J
Jin Hai 已提交
505
    // scheduler will determine when to delete collection files
G
groot 已提交
506 507 508 509 510 511
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(partition_name, meta_ptr_, nres);
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

    return Status::OK();
G
groot 已提交
512 513
}

S
starlord 已提交
514
Status
J
Jin Hai 已提交
515
DBImpl::DropPartitionByTag(const std::string& collection_id, const std::string& partition_tag) {
516
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
517 518 519 520
        return SHUTDOWN_ERROR;
    }

    std::string partition_name;
J
Jin Hai 已提交
521
    auto status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
522
    if (!status.ok()) {
523
        LOG_ENGINE_ERROR_ << status.message();
524 525 526
        return status;
    }

G
groot 已提交
527 528 529 530
    return DropPartition(partition_name);
}

Status
J
Jin Hai 已提交
531
DBImpl::ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) {
532
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
533 534 535
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
536
    return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
G
groot 已提交
537 538 539
}

Status
J
Jin Hai 已提交
540
DBImpl::InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) {
541
    //    LOG_ENGINE_DEBUG_ << "Insert " << n << " vectors to cache";
542
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
543
        return SHUTDOWN_ERROR;
S
starlord 已提交
544
    }
Y
yu yunfeng 已提交
545

J
Jin Hai 已提交
546
    // insert vectors into target collection
547 548
    // (zhiru): generate ids
    if (vectors.id_array_.empty()) {
J
Jin Hai 已提交
549 550 551
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(vectors.vector_count_, vectors.id_array_);
        if (!status.ok()) {
552
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get next id number fail: %s", "insert", 0, status.message().c_str());
J
Jin Hai 已提交
553 554
            return status;
        }
555 556
    }

557
    Status status;
558
    if (options_.wal_enable_) {
559 560
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
G
groot 已提交
561
        if (!status.ok()) {
562
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
G
groot 已提交
563 564
            return status;
        }
565 566

        if (!vectors.float_data_.empty()) {
J
Jin Hai 已提交
567
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.float_data_);
568
        } else if (!vectors.binary_data_.empty()) {
J
Jin Hai 已提交
569
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.binary_data_);
570
        }
G
groot 已提交
571
        swn_wal_.Notify();
572 573 574
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
J
Jin Hai 已提交
575
        record.collection_id = collection_id;
576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591
        record.partition_tag = partition_tag;
        record.ids = vectors.id_array_.data();
        record.length = vectors.vector_count_;
        if (vectors.binary_data_.empty()) {
            record.type = wal::MXLogType::InsertVector;
            record.data = vectors.float_data_.data();
            record.data_size = vectors.float_data_.size() * sizeof(float);
        } else {
            record.type = wal::MXLogType::InsertBinary;
            record.ids = vectors.id_array_.data();
            record.length = vectors.vector_count_;
            record.data = vectors.binary_data_.data();
            record.data_size = vectors.binary_data_.size() * sizeof(uint8_t);
        }

        status = ExecWalRecord(record);
G
groot 已提交
592 593
    }

594 595 596
    return status;
}

597
Status
Y
yukun 已提交
598 599 600 601 602
CopyToAttr(std::vector<uint8_t>& record, uint64_t row_num, const std::vector<std::string>& field_names,
           std::unordered_map<std::string, meta::hybrid::DataType>& attr_types,
           std::unordered_map<std::string, std::vector<uint8_t>>& attr_datas,
           std::unordered_map<std::string, uint64_t>& attr_nbytes,
           std::unordered_map<std::string, uint64_t>& attr_data_size) {
603
    uint64_t offset = 0;
Y
yukun 已提交
604 605
    for (auto name : field_names) {
        switch (attr_types.at(name)) {
606 607
            case meta::hybrid::DataType::INT8: {
                std::vector<uint8_t> data;
Y
yukun 已提交
608
                data.resize(row_num * sizeof(int8_t));
609

Y
yukun 已提交
610 611
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
612

Y
yukun 已提交
613 614
                std::vector<int8_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
615 616 617
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
618 619
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int8_t));
                attr_datas.insert(std::make_pair(name, data));
620

Y
yukun 已提交
621 622 623
                attr_nbytes.insert(std::make_pair(name, sizeof(int8_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int8_t)));
                offset += row_num * sizeof(int64_t);
624 625 626 627
                break;
            }
            case meta::hybrid::DataType::INT16: {
                std::vector<uint8_t> data;
Y
yukun 已提交
628
                data.resize(row_num * sizeof(int16_t));
629

Y
yukun 已提交
630 631
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
632

Y
yukun 已提交
633 634
                std::vector<int16_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
635 636 637
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
638 639
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int16_t));
                attr_datas.insert(std::make_pair(name, data));
640

Y
yukun 已提交
641 642 643
                attr_nbytes.insert(std::make_pair(name, sizeof(int16_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int16_t)));
                offset += row_num * sizeof(int64_t);
644 645 646 647
                break;
            }
            case meta::hybrid::DataType::INT32: {
                std::vector<uint8_t> data;
Y
yukun 已提交
648
                data.resize(row_num * sizeof(int32_t));
649

Y
yukun 已提交
650 651
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
652

Y
yukun 已提交
653 654
                std::vector<int32_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
655 656 657
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
658 659
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int32_t));
                attr_datas.insert(std::make_pair(name, data));
660

Y
yukun 已提交
661 662 663
                attr_nbytes.insert(std::make_pair(name, sizeof(int32_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int32_t)));
                offset += row_num * sizeof(int64_t);
664 665 666 667
                break;
            }
            case meta::hybrid::DataType::INT64: {
                std::vector<uint8_t> data;
Y
yukun 已提交
668 669 670 671 672 673
                data.resize(row_num * sizeof(int64_t));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(int64_t));
                attr_datas.insert(std::make_pair(name, data));

                std::vector<int64_t> test_data(row_num);
                memcpy(test_data.data(), record.data(), row_num * sizeof(int64_t));
674

Y
yukun 已提交
675 676 677
                attr_nbytes.insert(std::make_pair(name, sizeof(int64_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int64_t)));
                offset += row_num * sizeof(int64_t);
678 679 680 681
                break;
            }
            case meta::hybrid::DataType::FLOAT: {
                std::vector<uint8_t> data;
Y
yukun 已提交
682
                data.resize(row_num * sizeof(float));
683

Y
yukun 已提交
684 685
                std::vector<double> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(double));
686

Y
yukun 已提交
687 688
                std::vector<float> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
689 690 691
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
692 693
                memcpy(data.data(), raw_value.data(), row_num * sizeof(float));
                attr_datas.insert(std::make_pair(name, data));
694

Y
yukun 已提交
695 696 697
                attr_nbytes.insert(std::make_pair(name, sizeof(float)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(float)));
                offset += row_num * sizeof(double);
698 699 700 701
                break;
            }
            case meta::hybrid::DataType::DOUBLE: {
                std::vector<uint8_t> data;
Y
yukun 已提交
702 703 704
                data.resize(row_num * sizeof(double));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(double));
                attr_datas.insert(std::make_pair(name, data));
705

Y
yukun 已提交
706 707 708
                attr_nbytes.insert(std::make_pair(name, sizeof(double)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(double)));
                offset += row_num * sizeof(double);
709 710
                break;
            }
711 712
            default:
                break;
713 714
        }
    }
Y
yukun 已提交
715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764
    return Status::OK();
}

Status
DBImpl::InsertEntities(const std::string& collection_id, const std::string& partition_tag,
                       const std::vector<std::string>& field_names, Entity& entity,
                       std::unordered_map<std::string, meta::hybrid::DataType>& attr_types) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // Generate id
    if (entity.id_array_.empty()) {
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(entity.entity_count_, entity.id_array_);
        if (!status.ok()) {
            return status;
        }
    }

    Status status;
    std::unordered_map<std::string, std::vector<uint8_t>> attr_data;
    std::unordered_map<std::string, uint64_t> attr_nbytes;
    std::unordered_map<std::string, uint64_t> attr_data_size;
    status = CopyToAttr(entity.attr_value_, entity.entity_count_, field_names, attr_types, attr_data, attr_nbytes,
                        attr_data_size);
    if (!status.ok()) {
        return status;
    }

    wal::MXLogRecord record;
    record.lsn = 0;
    record.collection_id = collection_id;
    record.partition_tag = partition_tag;
    record.ids = entity.id_array_.data();
    record.length = entity.entity_count_;

    auto vector_it = entity.vector_data_.begin();
    if (vector_it->second.binary_data_.empty()) {
        record.type = wal::MXLogType::Entity;
        record.data = vector_it->second.float_data_.data();
        record.data_size = vector_it->second.float_data_.size() * sizeof(float);
        record.attr_data = attr_data;
        record.attr_nbytes = attr_nbytes;
        record.attr_data_size = attr_data_size;
    } else {
        //        record.type = wal::MXLogType::InsertBinary;
        //        record.data = entities.vector_data_[0].binary_data_.data();
        //        record.length = entities.vector_data_[0].binary_data_.size() * sizeof(uint8_t);
    }
765 766

    status = ExecWalRecord(record);
Y
yukun 已提交
767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812

#if 0
    if (options_.wal_enable_) {
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
        if (!status.ok()) {
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
            return status;
        }

        auto vector_it = entity.vector_data_.begin();
        if (!vector_it->second.binary_data_.empty()) {
            wal_mgr_->InsertEntities(collection_id, partition_tag, entity.id_array_, vector_it->second.binary_data_,
                                     attr_nbytes, attr_data);
        } else if (!vector_it->second.float_data_.empty()) {
            wal_mgr_->InsertEntities(collection_id, partition_tag, entity.id_array_, vector_it->second.float_data_,
                                     attr_nbytes, attr_data);
        }
        swn_wal_.Notify();
    } else {
        // insert entities: collection_name is field id
        wal::MXLogRecord record;
        record.lsn = 0;
        record.collection_id = collection_id;
        record.partition_tag = partition_tag;
        record.ids = entity.id_array_.data();
        record.length = entity.entity_count_;

        auto vector_it = entity.vector_data_.begin();
        if (vector_it->second.binary_data_.empty()) {
            record.type = wal::MXLogType::Entity;
            record.data = vector_it->second.float_data_.data();
            record.data_size = vector_it->second.float_data_.size() * sizeof(float);
            record.attr_data = attr_data;
            record.attr_nbytes = attr_nbytes;
            record.attr_data_size = attr_data_size;
        } else {
            //        record.type = wal::MXLogType::InsertBinary;
            //        record.data = entities.vector_data_[0].binary_data_.data();
            //        record.length = entities.vector_data_[0].binary_data_.size() * sizeof(uint8_t);
        }

        status = ExecWalRecord(record);
    }
#endif

813 814 815
    return status;
}

816
Status
J
Jin Hai 已提交
817
DBImpl::DeleteVector(const std::string& collection_id, IDNumber vector_id) {
818 819
    IDNumbers ids;
    ids.push_back(vector_id);
J
Jin Hai 已提交
820
    return DeleteVectors(collection_id, ids);
821 822 823
}

Status
J
Jin Hai 已提交
824
DBImpl::DeleteVectors(const std::string& collection_id, IDNumbers vector_ids) {
825 826 827 828 829 830
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
J
Jin Hai 已提交
831
        wal_mgr_->DeleteById(collection_id, vector_ids);
G
groot 已提交
832
        swn_wal_.Notify();
833 834 835 836
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
J
Jin Hai 已提交
837
        record.collection_id = collection_id;
838 839 840 841 842 843 844 845 846 847
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
J
Jin Hai 已提交
848
DBImpl::Flush(const std::string& collection_id) {
849 850 851 852 853
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
854 855
    bool has_collection;
    status = HasCollection(collection_id, has_collection);
856 857 858
    if (!status.ok()) {
        return status;
    }
859
    if (!has_collection) {
860
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_id;
J
Jin Hai 已提交
861
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
862 863
    }

864
    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_id;
865 866

    if (options_.wal_enable_) {
867
        LOG_ENGINE_DEBUG_ << "WAL flush";
J
Jin Hai 已提交
868
        auto lsn = wal_mgr_->Flush(collection_id);
869
        if (lsn != 0) {
G
groot 已提交
870 871
            swn_wal_.Notify();
            flush_req_swn_.Wait();
872 873 874
        }

    } else {
875
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
876
        InternalFlush(collection_id);
877 878
    }

879
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_id;
880 881 882 883 884 885 886 887 888 889

    return status;
}

Status
DBImpl::Flush() {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

890
    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
891 892 893

    Status status;
    if (options_.wal_enable_) {
894
        LOG_ENGINE_DEBUG_ << "WAL flush";
895 896
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
G
groot 已提交
897 898
            swn_wal_.Notify();
            flush_req_swn_.Wait();
899 900
        }
    } else {
901
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
902
        InternalFlush();
903 904
    }

905
    LOG_ENGINE_DEBUG_ << "End flush all collections";
906 907 908 909 910

    return status;
}

Status
G
groot 已提交
911
DBImpl::Compact(const std::string& collection_id, double threshold) {
912 913 914 915
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

916 917 918
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
919 920
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
921
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
922
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
923 924 925 926
        } else {
            return status;
        }
    } else {
927
        if (!collection_schema.owner_collection_.empty()) {
928
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
929
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
930 931 932
        }
    }

933
    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
934

Z
update  
Zhiru Zhu 已提交
935
    // WaitBuildIndexFinish();
936

Z
update  
Zhiru Zhu 已提交
937
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
938
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
939

940
    LOG_ENGINE_DEBUG_ << "Compacting collection: " << collection_id;
Z
Zhiru Zhu 已提交
941

942
    // Get files to compact from meta.
J
Jin Hai 已提交
943 944
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
G
groot 已提交
945 946
    meta::FilesHolder files_holder;
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
947 948
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
949
        LOG_ENGINE_ERROR_ << err_msg;
950 951 952
        return Status(DB_ERROR, err_msg);
    }

G
groot 已提交
953
    LOG_ENGINE_DEBUG_ << "Found " << files_holder.HoldFiles().size() << " segment to compact";
Z
Zhiru Zhu 已提交
954 955

    Status compact_status;
G
groot 已提交
956 957
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files_to_compact = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
958
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
J
Jin Hai 已提交
959
        meta::SegmentSchema file = *iter;
G
groot 已提交
960 961
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
962 963 964
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
965

Z
Zhiru Zhu 已提交
966
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
967 968
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
969
        if (!status.ok()) {
G
groot 已提交
970
            files_holder.UnmarkFile(file);
G
groot 已提交
971
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
972 973
        }

J
Jin Hai 已提交
974
        meta::SegmentsSchema files_to_update;
Z
Zhiru Zhu 已提交
975
        if (deleted_docs_size != 0) {
G
groot 已提交
976
            compact_status = CompactFile(collection_id, threshold, file, files_to_update);
Z
Zhiru Zhu 已提交
977 978

            if (!compact_status.ok()) {
979 980
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << file.segment_id_ << ": "
                                  << compact_status.message();
G
groot 已提交
981
                files_holder.UnmarkFile(file);
G
groot 已提交
982
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
983 984
            }
        } else {
G
groot 已提交
985
            files_holder.UnmarkFile(file);
986
            LOG_ENGINE_DEBUG_ << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
987
            continue;  // skip this file and try compact next one
988
        }
Z
Zhiru Zhu 已提交
989

990
        LOG_ENGINE_DEBUG_ << "Updating meta after compaction...";
991
        status = meta_ptr_->UpdateCollectionFiles(files_to_update);
G
groot 已提交
992
        files_holder.UnmarkFile(file);
G
groot 已提交
993 994 995 996
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
997 998
    }

G
groot 已提交
999
    if (compact_status.ok()) {
1000
        LOG_ENGINE_DEBUG_ << "Finished compacting collection: " << collection_id;
G
groot 已提交
1001
    }
1002

G
groot 已提交
1003
    return compact_status;
1004 1005 1006
}

Status
G
groot 已提交
1007
DBImpl::CompactFile(const std::string& collection_id, double threshold, const meta::SegmentSchema& file,
J
Jin Hai 已提交
1008
                    meta::SegmentsSchema& files_to_update) {
1009
    LOG_ENGINE_DEBUG_ << "Compacting segment " << file.segment_id_ << " for collection: " << collection_id;
1010

G
groot 已提交
1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029
    std::string segment_dir_to_merge;
    utils::GetParentPath(file.location_, segment_dir_to_merge);

    // no need to compact if deleted vectors are too few(less than threashold)
    if (file.row_count_ > 0 && threshold > 0.0) {
        segment::SegmentReader segment_reader_to_merge(segment_dir_to_merge);
        segment::DeletedDocsPtr deleted_docs_ptr;
        auto status = segment_reader_to_merge.LoadDeletedDocs(deleted_docs_ptr);
        if (status.ok()) {
            auto delete_items = deleted_docs_ptr->GetDeletedDocs();
            double delete_rate = (double)delete_items.size() / (double)file.row_count_;
            if (delete_rate < threshold) {
                LOG_ENGINE_DEBUG_ << "Delete rate less than " << threshold << ", no need to compact for"
                                  << segment_dir_to_merge;
                return Status::OK();
            }
        }
    }

J
Jin Hai 已提交
1030 1031 1032
    // Create new collection file
    meta::SegmentSchema compacted_file;
    compacted_file.collection_id_ = collection_id;
1033
    // compacted_file.date_ = date;
J
Jin Hai 已提交
1034
    compacted_file.file_type_ = meta::SegmentSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
G
groot 已提交
1035
    auto status = meta_ptr_->CreateCollectionFile(compacted_file);
1036 1037

    if (!status.ok()) {
1038
        LOG_ENGINE_ERROR_ << "Failed to create collection file: " << status.message();
1039 1040 1041
        return status;
    }

J
Jin Hai 已提交
1042
    // Compact (merge) file to the newly created collection file
1043 1044 1045 1046 1047

    std::string new_segment_dir;
    utils::GetParentPath(compacted_file.location_, new_segment_dir);
    auto segment_writer_ptr = std::make_shared<segment::SegmentWriter>(new_segment_dir);

1048
    LOG_ENGINE_DEBUG_ << "Compacting begin...";
1049 1050 1051
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
1052
    LOG_ENGINE_DEBUG_ << "Serializing compacted segment...";
1053 1054
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
1055
        LOG_ENGINE_ERROR_ << "Failed to serialize compacted segment: " << status.message();
J
Jin Hai 已提交
1056
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1057
        auto mark_status = meta_ptr_->UpdateCollectionFile(compacted_file);
1058
        if (mark_status.ok()) {
1059
            LOG_ENGINE_DEBUG_ << "Mark file: " << compacted_file.file_id_ << " to to_delete";
1060
        }
G
groot 已提交
1061

1062 1063 1064
        return status;
    }

1065 1066 1067 1068 1069 1070 1071
    // Update compacted file state, if origin file is backup or to_index, set compected file to to_index
    compacted_file.file_size_ = segment_writer_ptr->Size();
    compacted_file.row_count_ = segment_writer_ptr->VectorCount();
    if ((file.file_type_ == (int32_t)meta::SegmentSchema::BACKUP ||
         file.file_type_ == (int32_t)meta::SegmentSchema::TO_INDEX) &&
        (compacted_file.row_count_ > meta::BUILD_INDEX_THRESHOLD)) {
        compacted_file.file_type_ = meta::SegmentSchema::TO_INDEX;
1072
    } else {
J
Jin Hai 已提交
1073
        compacted_file.file_type_ = meta::SegmentSchema::RAW;
1074 1075 1076
    }

    if (compacted_file.row_count_ == 0) {
1077
        LOG_ENGINE_DEBUG_ << "Compacted segment is empty. Mark it as TO_DELETE";
J
Jin Hai 已提交
1078
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1079 1080
    }

Z
Zhiru Zhu 已提交
1081
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
1082

Z
Zhiru Zhu 已提交
1083 1084
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
G
groot 已提交
1085 1086
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
Z
Zhiru Zhu 已提交
1087 1088 1089
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
1090 1091

    milvus::engine::meta::SegmentsSchema& segment_files = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1092
    for (auto& f : segment_files) {
J
Jin Hai 已提交
1093
        f.file_type_ = meta::SegmentSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
1094 1095
        files_to_update.emplace_back(f);
    }
G
groot 已提交
1096
    files_holder.ReleaseFiles();
1097

1098 1099 1100
    LOG_ENGINE_DEBUG_ << "Compacted segment " << compacted_file.segment_id_ << " from "
                      << std::to_string(file.file_size_) << " bytes to " << std::to_string(compacted_file.file_size_)
                      << " bytes";
1101 1102 1103 1104 1105 1106 1107 1108 1109

    if (options_.insert_cache_immediately_) {
        segment_writer_ptr->Cache();
    }

    return status;
}

Status
1110 1111
DBImpl::GetVectorsByID(const std::string& collection_id, const IDNumbers& id_array,
                       std::vector<engine::VectorsData>& vectors) {
1112 1113 1114 1115
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

1116 1117 1118
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1119
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1120
        return Status(DB_NOT_FOUND, "Collection does not exist");
1121 1122 1123 1124 1125
    }
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1126
    meta::FilesHolder files_holder;
J
Jin Hai 已提交
1127 1128
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
1129

G
groot 已提交
1130
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
1131
    if (!status.ok()) {
1132
        std::string err_msg = "Failed to get files for GetVectorsByID: " + status.message();
1133
        LOG_ENGINE_ERROR_ << err_msg;
1134 1135 1136
        return status;
    }

J
Jin Hai 已提交
1137 1138
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
1139
    for (auto& schema : partition_array) {
G
groot 已提交
1140
        status = meta_ptr_->FilesByType(schema.collection_id_, file_types, files_holder);
1141 1142
        if (!status.ok()) {
            std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
1143
            LOG_ENGINE_ERROR_ << err_msg;
1144 1145 1146 1147
            return status;
        }
    }

G
groot 已提交
1148
    if (files_holder.HoldFiles().empty()) {
1149
        LOG_ENGINE_DEBUG_ << "No files to get vector by id from";
J
Jin Hai 已提交
1150
        return Status(DB_NOT_FOUND, "Collection is empty");
1151 1152 1153
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
G
groot 已提交
1154
    status = GetVectorsByIdHelper(collection_id, id_array, vectors, files_holder);
1155 1156 1157 1158 1159 1160
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

    return status;
}

Status
J
Jin Hai 已提交
1161
DBImpl::GetVectorIDs(const std::string& collection_id, const std::string& segment_id, IDNumbers& vector_ids) {
1162 1163 1164 1165
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
1166
    // step 1: check collection existence
1167 1168 1169
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1170
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1171
        return Status(DB_NOT_FOUND, "Collection does not exist");
1172 1173 1174 1175 1176 1177
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
G
groot 已提交
1178 1179
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
1180 1181 1182 1183
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1184
    milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
1185
    if (collection_files.empty()) {
1186 1187 1188
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

J
Jin Hai 已提交
1189
    // check the segment is belong to this collection
1190
    if (collection_files[0].collection_id_ != collection_id) {
J
Jin Hai 已提交
1191
        // the segment could be in a partition under this collection
1192 1193 1194 1195
        meta::CollectionSchema collection_schema;
        collection_schema.collection_id_ = collection_files[0].collection_id_;
        status = DescribeCollection(collection_schema);
        if (collection_schema.owner_collection_ != collection_id) {
J
Jin Hai 已提交
1196
            return Status(DB_NOT_FOUND, "Segment does not belong to this collection");
1197 1198 1199 1200 1201
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
1202
    engine::utils::GetParentPath(collection_files[0].location_, segment_dir);
1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227
    segment::SegmentReader segment_reader(segment_dir);

    std::vector<segment::doc_id_t> uids;
    status = segment_reader.LoadUids(uids);
    if (!status.ok()) {
        return status;
    }

    segment::DeletedDocsPtr deleted_docs_ptr;
    status = segment_reader.LoadDeletedDocs(deleted_docs_ptr);
    if (!status.ok()) {
        return status;
    }

    // step 4: construct id array
    // avoid duplicate offset and erase from max offset to min offset
    auto& deleted_offset = deleted_docs_ptr->GetDeletedDocs();
    std::set<segment::offset_t, std::greater<segment::offset_t>> ordered_offset;
    for (segment::offset_t offset : deleted_offset) {
        ordered_offset.insert(offset);
    }
    for (segment::offset_t offset : ordered_offset) {
        uids.erase(uids.begin() + offset);
    }
    vector_ids.swap(uids);
S
starlord 已提交
1228

G
groot 已提交
1229
    return status;
X
Xu Peng 已提交
1230 1231
}

1232
Status
1233
DBImpl::GetVectorsByIdHelper(const std::string& collection_id, const IDNumbers& id_array,
G
groot 已提交
1234 1235 1236
                             std::vector<engine::VectorsData>& vectors, meta::FilesHolder& files_holder) {
    // attention: this is a copy, not a reference, since the files_holder.UnMarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
1237
    LOG_ENGINE_DEBUG_ << "Getting vector by id in " << files.size() << " files, id count = " << id_array.size();
J
Jin Hai 已提交
1238

1239 1240 1241 1242 1243 1244 1245
    // sometimes not all of id_array can be found, we need to return empty vector for id not found
    // for example:
    // id_array = [1, -1, 2, -1, 3]
    // vectors should return [valid_vector, empty_vector, valid_vector, empty_vector, valid_vector]
    // the ID2RAW is to ensure returned vector sequence is consist with id_array
    using ID2VECTOR = std::map<int64_t, VectorsData>;
    ID2VECTOR map_id2vector;
1246

1247 1248 1249
    vectors.clear();

    IDNumbers temp_ids = id_array;
1250 1251 1252 1253 1254 1255 1256 1257
    for (auto& file : files) {
        // Load bloom filter
        std::string segment_dir;
        engine::utils::GetParentPath(file.location_, segment_dir);
        segment::SegmentReader segment_reader(segment_dir);
        segment::IdBloomFilterPtr id_bloom_filter_ptr;
        segment_reader.LoadBloomFilter(id_bloom_filter_ptr);

1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268
        for (IDNumbers::iterator it = temp_ids.begin(); it != temp_ids.end();) {
            int64_t vector_id = *it;
            // each id must has a VectorsData
            // if vector not found for an id, its VectorsData's vector_count = 0, else 1
            VectorsData& vector_ref = map_id2vector[vector_id];

            // Check if the id is present in bloom filter.
            if (id_bloom_filter_ptr->Check(vector_id)) {
                // Load uids and check if the id is indeed present. If yes, find its offset.
                std::vector<segment::doc_id_t> uids;
                auto status = segment_reader.LoadUids(uids);
1269 1270 1271
                if (!status.ok()) {
                    return status;
                }
1272 1273 1274 1275 1276 1277 1278 1279

                auto found = std::find(uids.begin(), uids.end(), vector_id);
                if (found != uids.end()) {
                    auto offset = std::distance(uids.begin(), found);

                    // Check whether the id has been deleted
                    segment::DeletedDocsPtr deleted_docs_ptr;
                    status = segment_reader.LoadDeletedDocs(deleted_docs_ptr);
1280
                    if (!status.ok()) {
J
Jin Hai 已提交
1281
                        LOG_ENGINE_ERROR_ << status.message();
1282 1283
                        return status;
                    }
1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309
                    auto& deleted_docs = deleted_docs_ptr->GetDeletedDocs();

                    auto deleted = std::find(deleted_docs.begin(), deleted_docs.end(), offset);
                    if (deleted == deleted_docs.end()) {
                        // Load raw vector
                        bool is_binary = utils::IsBinaryMetricType(file.metric_type_);
                        size_t single_vector_bytes = is_binary ? file.dimension_ / 8 : file.dimension_ * sizeof(float);
                        std::vector<uint8_t> raw_vector;
                        status =
                            segment_reader.LoadVectors(offset * single_vector_bytes, single_vector_bytes, raw_vector);
                        if (!status.ok()) {
                            LOG_ENGINE_ERROR_ << status.message();
                            return status;
                        }

                        vector_ref.vector_count_ = 1;
                        if (is_binary) {
                            vector_ref.binary_data_.swap(raw_vector);
                        } else {
                            std::vector<float> float_vector;
                            float_vector.resize(file.dimension_);
                            memcpy(float_vector.data(), raw_vector.data(), single_vector_bytes);
                            vector_ref.float_data_.swap(float_vector);
                        }
                        temp_ids.erase(it);
                        continue;
1310 1311 1312
                    }
                }
            }
1313 1314 1315

            it++;
        }
G
groot 已提交
1316 1317 1318

        // unmark file, allow the file to be deleted
        files_holder.UnmarkFile(file);
1319 1320 1321 1322 1323 1324 1325 1326
    }

    for (auto id : id_array) {
        VectorsData& vector_ref = map_id2vector[id];

        VectorsData data;
        data.vector_count_ = vector_ref.vector_count_;
        if (data.vector_count_ > 0) {
G
groot 已提交
1327 1328
            data.float_data_ = vector_ref.float_data_;    // copy data since there could be duplicated id
            data.binary_data_ = vector_ref.binary_data_;  // copy data since there could be duplicated id
1329
        }
1330
        vectors.emplace_back(data);
1331 1332
    }

1333 1334
    if (vectors.empty()) {
        std::string msg = "Vectors not found in collection " + collection_id;
J
Jin Hai 已提交
1335 1336 1337
        LOG_ENGINE_DEBUG_ << msg;
    }

1338 1339 1340
    return Status::OK();
}

S
starlord 已提交
1341
Status
G
groot 已提交
1342 1343
DBImpl::CreateIndex(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const CollectionIndex& index) {
1344
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1345 1346 1347
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1348
    // serialize memory data
1349 1350
    //    std::set<std::string> sync_collection_ids;
    //    auto status = SyncMemData(sync_collection_ids);
1351
    auto status = Flush();
G
groot 已提交
1352

S
starlord 已提交
1353 1354 1355
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

S
starlord 已提交
1356
        // step 1: check index difference
1357
        CollectionIndex old_index;
J
Jin Hai 已提交
1358
        status = DescribeIndex(collection_id, old_index);
S
starlord 已提交
1359
        if (!status.ok()) {
1360
            LOG_ENGINE_ERROR_ << "Failed to get collection index info for collection: " << collection_id;
S
starlord 已提交
1361 1362 1363
            return status;
        }

S
starlord 已提交
1364
        // step 2: update index info
1365 1366
        CollectionIndex new_index = index;
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateCollection
S
starlord 已提交
1367
        if (!utils::IsSameIndex(old_index, new_index)) {
1368
            status = UpdateCollectionIndexRecursively(collection_id, new_index);
S
starlord 已提交
1369 1370 1371 1372 1373 1374
            if (!status.ok()) {
                return status;
            }
        }
    }

S
starlord 已提交
1375 1376
    // step 3: let merge file thread finish
    // to avoid duplicate data bug
1377 1378
    WaitMergeFileFinish();

S
starlord 已提交
1379
    // step 4: wait and build index
1380
    status = index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
G
groot 已提交
1381
    status = WaitCollectionIndexRecursively(context, collection_id, index);
S
starlord 已提交
1382

G
groot 已提交
1383
    return status;
S
starlord 已提交
1384 1385
}

S
starlord 已提交
1386
Status
1387
DBImpl::DescribeIndex(const std::string& collection_id, CollectionIndex& index) {
1388
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1389 1390 1391
        return SHUTDOWN_ERROR;
    }

1392
    return meta_ptr_->DescribeCollectionIndex(collection_id, index);
S
starlord 已提交
1393 1394
}

S
starlord 已提交
1395
Status
J
Jin Hai 已提交
1396
DBImpl::DropIndex(const std::string& collection_id) {
1397
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1398 1399 1400
        return SHUTDOWN_ERROR;
    }

1401
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
1402
    return DropCollectionIndexRecursively(collection_id);
S
starlord 已提交
1403 1404
}

S
starlord 已提交
1405
Status
1406 1407 1408
DBImpl::QueryByIDs(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                   const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
                   const IDNumbers& id_array, ResultIds& result_ids, ResultDistances& result_distances) {
1409
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1410
        return SHUTDOWN_ERROR;
S
starlord 已提交
1411 1412
    }

1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529
    if (id_array.empty()) {
        return Status(DB_ERROR, "Empty id array during query by id");
    }

    TimeRecorder rc("Query by id in collection:" + collection_id);

    // get collection schema
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
            std::string msg = "Collection to search does not exist: " + collection_id;
            LOG_ENGINE_ERROR_ << msg;
            return Status(DB_NOT_FOUND, msg);
        } else {
            return status;
        }
    } else {
        if (!collection_schema.owner_collection_.empty()) {
            std::string msg = "Collection to search does not exist: " + collection_id;
            LOG_ENGINE_ERROR_ << msg;
            return Status(DB_NOT_FOUND, msg);
        }
    }

    rc.RecordSection("get collection schema");

    // get target vectors data
    std::vector<milvus::engine::VectorsData> vectors;
    status = GetVectorsByID(collection_id, id_array, vectors);
    if (!status.ok()) {
        std::string msg = "Failed to get vector data for collection: " + collection_id;
        LOG_ENGINE_ERROR_ << msg;
        return status;
    }

    // some vectors could not be found, no need to search them
    uint64_t valid_count = 0;
    bool is_binary = utils::IsBinaryMetricType(collection_schema.metric_type_);
    for (auto& vector : vectors) {
        if (vector.vector_count_ > 0) {
            valid_count++;
        }
    }

    // copy valid vectors data for search input
    uint64_t dimension = collection_schema.dimension_;
    VectorsData valid_vectors;
    valid_vectors.vector_count_ = valid_count;
    if (is_binary) {
        valid_vectors.binary_data_.resize(valid_count * dimension / 8);
    } else {
        valid_vectors.float_data_.resize(valid_count * dimension * sizeof(float));
    }

    int64_t valid_index = 0;
    for (size_t i = 0; i < vectors.size(); i++) {
        if (vectors[i].vector_count_ == 0) {
            continue;
        }
        if (is_binary) {
            memcpy(valid_vectors.binary_data_.data() + valid_index * dimension / 8, vectors[i].binary_data_.data(),
                   vectors[i].binary_data_.size());
        } else {
            memcpy(valid_vectors.float_data_.data() + valid_index * dimension, vectors[i].float_data_.data(),
                   vectors[i].float_data_.size() * sizeof(float));
        }
        valid_index++;
    }

    rc.RecordSection("construct query input");

    // search valid vectors
    ResultIds valid_result_ids;
    ResultDistances valid_result_distances;
    status = Query(context, collection_id, partition_tags, k, extra_params, valid_vectors, valid_result_ids,
                   valid_result_distances);
    if (!status.ok()) {
        std::string msg = "Failed to query by id in collection " + collection_id + ", error: " + status.message();
        LOG_ENGINE_ERROR_ << msg;
        return status;
    }

    if (valid_result_ids.size() != valid_count * k || valid_result_distances.size() != valid_count * k) {
        std::string msg = "Failed to query by id in collection " + collection_id + ", result doesn't match id count";
        return Status(DB_ERROR, msg);
    }

    rc.RecordSection("query vealid vectors");

    // construct result
    if (valid_count == id_array.size()) {
        result_ids.swap(valid_result_ids);
        result_distances.swap(valid_result_distances);
    } else {
        result_ids.resize(vectors.size() * k);
        result_distances.resize(vectors.size() * k);
        int64_t valid_index = 0;
        for (uint64_t i = 0; i < vectors.size(); i++) {
            if (vectors[i].vector_count_ > 0) {
                memcpy(result_ids.data() + i * k, valid_result_ids.data() + valid_index * k, k * sizeof(int64_t));
                memcpy(result_distances.data() + i * k, valid_result_distances.data() + valid_index * k,
                       k * sizeof(float));
                valid_index++;
            } else {
                memset(result_ids.data() + i * k, -1, k * sizeof(int64_t));
                for (uint64_t j = i * k; j < i * k + k; j++) {
                    result_distances[j] = std::numeric_limits<float>::max();
                }
            }
        }
    }

    rc.RecordSection("construct result");

    return status;
X
Xu Peng 已提交
1530 1531
}

1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544
Status
DBImpl::HybridQuery(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const std::vector<std::string>& partition_tags,
                    context::HybridSearchContextPtr hybrid_search_context, query::GeneralQueryPtr general_query,
                    std::unordered_map<std::string, engine::meta::hybrid::DataType>& attr_type, uint64_t& nq,
                    ResultIds& result_ids, ResultDistances& result_distances) {
    auto query_ctx = context->Child("Query");

    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
G
groot 已提交
1545
    meta::FilesHolder files_holder;
1546 1547 1548
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
G
groot 已提交
1549
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
1550 1551 1552 1553 1554 1555 1556 1557 1558 1559
        if (!status.ok()) {
            return status;
        }

        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
        if (!status.ok()) {
            return status;
        }
        for (auto& schema : partition_array) {
G
groot 已提交
1560
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1561
            if (!status.ok()) {
G
groot 已提交
1562
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1563 1564 1565
            }
        }

G
groot 已提交
1566 1567
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1568 1569 1570 1571 1572 1573 1574
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
        GetPartitionsByTags(collection_id, partition_tags, partition_name_array);

        for (auto& partition_name : partition_name_array) {
G
groot 已提交
1575
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1576
            if (!status.ok()) {
G
groot 已提交
1577
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1578 1579 1580
            }
        }

G
groot 已提交
1581
        if (files_holder.HoldFiles().empty()) {
1582 1583 1584 1585 1586
            return Status::OK();
        }
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1587
    status = HybridQueryAsync(query_ctx, collection_id, files_holder, hybrid_search_context, general_query, attr_type,
1588 1589 1590 1591 1592 1593 1594 1595 1596 1597 1598
                              nq, result_ids, result_distances);
    if (!status.ok()) {
        return status;
    }
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query

    query_ctx->GetTraceContext()->GetSpan()->Finish();

    return status;
}

S
starlord 已提交
1599
Status
J
Jin Hai 已提交
1600
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
1601 1602
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
1603
    milvus::server::ContextChild tracer(context, "Query");
Z
Zhiru Zhu 已提交
1604

1605
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1606
        return SHUTDOWN_ERROR;
S
starlord 已提交
1607 1608
    }

G
groot 已提交
1609
    Status status;
G
groot 已提交
1610
    meta::FilesHolder files_holder;
G
groot 已提交
1611
    if (partition_tags.empty()) {
J
Jin Hai 已提交
1612 1613
        // no partition tag specified, means search in whole collection
        // get all collection files from parent collection
G
groot 已提交
1614
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
G
groot 已提交
1615 1616 1617 1618
        if (!status.ok()) {
            return status;
        }

J
Jin Hai 已提交
1619 1620
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
1621
        for (auto& schema : partition_array) {
G
groot 已提交
1622
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1623 1624
        }

G
groot 已提交
1625 1626
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
G
groot 已提交
1627 1628 1629 1630
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
J
Jin Hai 已提交
1631
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
T
Tinkerrr 已提交
1632 1633 1634
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
G
groot 已提交
1635 1636

        for (auto& partition_name : partition_name_array) {
G
groot 已提交
1637
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1638 1639
        }

G
groot 已提交
1640 1641
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1642 1643 1644
        }
    }

S
starlord 已提交
1645
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1646
    status = QueryAsync(tracer.Context(), files_holder, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1647
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1648

S
starlord 已提交
1649
    return status;
G
groot 已提交
1650
}
X
Xu Peng 已提交
1651

S
starlord 已提交
1652
Status
1653 1654 1655
DBImpl::QueryByFileID(const std::shared_ptr<server::Context>& context, const std::vector<std::string>& file_ids,
                      uint64_t k, const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                      ResultDistances& result_distances) {
1656
    milvus::server::ContextChild tracer(context, "Query by file id");
Z
Zhiru Zhu 已提交
1657

1658
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1659
        return SHUTDOWN_ERROR;
S
starlord 已提交
1660 1661
    }

S
starlord 已提交
1662
    // get specified files
1663
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1664
    for (auto& id : file_ids) {
1665
        std::string::size_type sz;
J
jinhai 已提交
1666
        ids.push_back(std::stoul(id, &sz));
1667 1668
    }

G
groot 已提交
1669 1670
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesByID(ids, files_holder);
1671 1672
    if (!status.ok()) {
        return status;
1673 1674
    }

G
groot 已提交
1675
    milvus::engine::meta::SegmentsSchema& search_files = files_holder.HoldFiles();
1676
    if (search_files.empty()) {
S
starlord 已提交
1677
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1678 1679
    }

S
starlord 已提交
1680
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1681
    status = QueryAsync(tracer.Context(), files_holder, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1682
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1683

S
starlord 已提交
1684
    return status;
1685 1686
}

S
starlord 已提交
1687
Status
Y
Yu Kun 已提交
1688
DBImpl::Size(uint64_t& result) {
1689
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1690
        return SHUTDOWN_ERROR;
S
starlord 已提交
1691 1692
    }

S
starlord 已提交
1693
    return meta_ptr_->Size(result);
S
starlord 已提交
1694 1695 1696
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1697
// internal methods
S
starlord 已提交
1698
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1699
Status
G
groot 已提交
1700
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, meta::FilesHolder& files_holder, uint64_t k,
1701 1702
                   const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                   ResultDistances& result_distances) {
1703
    milvus::server::ContextChild tracer(context, "Query Async");
G
groot 已提交
1704
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1705

G
groot 已提交
1706
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
G
groot 已提交
1707 1708 1709
    if (files.size() > milvus::scheduler::TASK_TABLE_MAX_COUNT) {
        std::string msg =
            "Search files count exceed scheduler limit: " + std::to_string(milvus::scheduler::TASK_TABLE_MAX_COUNT);
1710
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1711 1712 1713
        return Status(DB_ERROR, msg);
    }

S
starlord 已提交
1714
    TimeRecorder rc("");
G
groot 已提交
1715

1716
    // step 1: construct search job
1717
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files.size());
1718
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(tracer.Context(), k, extra_params, vectors);
Y
Yu Kun 已提交
1719
    for (auto& file : files) {
J
Jin Hai 已提交
1720
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
W
wxyu 已提交
1721
        job->AddIndexFile(file_ptr);
G
groot 已提交
1722 1723
    }

1724
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1725
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1726
    job->WaitResult();
1727

G
groot 已提交
1728
    files_holder.ReleaseFiles();
W
wxyu 已提交
1729 1730
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1731
    }
G
groot 已提交
1732

1733
    // step 3: construct results
G
groot 已提交
1734 1735
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1736
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1737 1738 1739 1740

    return Status::OK();
}

1741 1742
Status
DBImpl::HybridQueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
G
groot 已提交
1743
                         meta::FilesHolder& files_holder, context::HybridSearchContextPtr hybrid_search_context,
1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772
                         query::GeneralQueryPtr general_query,
                         std::unordered_map<std::string, engine::meta::hybrid::DataType>& attr_type, uint64_t& nq,
                         ResultIds& result_ids, ResultDistances& result_distances) {
    auto query_async_ctx = context->Child("Query Async");

#if 0
    // Construct tasks
    for (auto file : files) {
        std::unordered_map<std::string, engine::DataType> types;
        auto it = attr_type.begin();
        for (; it != attr_type.end(); it++) {
            types.insert(std::make_pair(it->first, (engine::DataType)it->second));
        }

        auto file_ptr = std::make_shared<meta::TableFileSchema>(file);
        search::TaskPtr
            task = std::make_shared<search::Task>(context, file_ptr, general_query, types, hybrid_search_context);
        search::TaskInst::GetInstance().load_queue().push(task);
        search::TaskInst::GetInstance().load_cv().notify_one();
        hybrid_search_context->tasks_.emplace_back(task);
    }

#endif

    //#if 0
    TimeRecorder rc("");

    // step 1: construct search job
    VectorsData vectors;
G
groot 已提交
1773 1774
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files_holder.HoldFiles().size());
1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785
    scheduler::SearchJobPtr job =
        std::make_shared<scheduler::SearchJob>(query_async_ctx, general_query, attr_type, vectors);
    for (auto& file : files) {
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
        job->AddIndexFile(file_ptr);
    }

    // step 2: put search job to scheduler and wait result
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitResult();

G
groot 已提交
1786
    files_holder.ReleaseFiles();
1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
    }

    // step 3: construct results
    nq = job->vector_count();
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
    rc.ElapseFromBegin("Engine query totally cost");

    query_async_ctx->GetTraceContext()->GetSpan()->Finish();
    //#endif

    return Status::OK();
}

S
starlord 已提交
1803
void
G
groot 已提交
1804
DBImpl::BackgroundIndexThread() {
Y
yu yunfeng 已提交
1805
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
1806
    while (true) {
1807
        if (!initialized_.load(std::memory_order_acquire)) {
1808 1809
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
1810

1811
            LOG_ENGINE_DEBUG_ << "DB background thread exit";
G
groot 已提交
1812 1813
            break;
        }
X
Xu Peng 已提交
1814

G
groot 已提交
1815
        swn_index_.Wait_For(std::chrono::seconds(BACKGROUND_INDEX_INTERVAL));
X
Xu Peng 已提交
1816

G
groot 已提交
1817
        WaitMergeFileFinish();
G
groot 已提交
1818 1819
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
1820 1821
}

S
starlord 已提交
1822 1823
void
DBImpl::WaitMergeFileFinish() {
1824
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
1825 1826
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
1827 1828
        iter.wait();
    }
1829
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
1830 1831
}

S
starlord 已提交
1832 1833
void
DBImpl::WaitBuildIndexFinish() {
1834
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
1835
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
1836
    for (auto& iter : index_thread_results_) {
1837 1838
        iter.wait();
    }
1839
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
1840 1841
}

S
starlord 已提交
1842 1843
void
DBImpl::StartMetricTask() {
G
groot 已提交
1844
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
G
groot 已提交
1845 1846
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
1847 1848
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
1849 1850 1851 1852 1853 1854 1855
    if (cache_total > 0) {
        double cache_usage_double = cache_usage;
        server::Metrics::GetInstance().CpuCacheUsageGaugeSet(cache_usage_double * 100 / cache_total);
    } else {
        server::Metrics::GetInstance().CpuCacheUsageGaugeSet(0);
    }

Y
Yu Kun 已提交
1856
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
1857 1858 1859 1860 1861 1862 1863 1864
    uint64_t size;
    Size(size);
    server::Metrics::GetInstance().DataFileSizeGaugeSet(size);
    server::Metrics::GetInstance().CPUUsagePercentSet();
    server::Metrics::GetInstance().RAMUsagePercentSet();
    server::Metrics::GetInstance().GPUPercentGaugeSet();
    server::Metrics::GetInstance().GPUMemoryUsageGaugeSet();
    server::Metrics::GetInstance().OctetsSet();
S
starlord 已提交
1865

K
kun yu 已提交
1866
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
1867 1868
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
1869
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
1870 1871
}

S
starlord 已提交
1872
void
1873
DBImpl::StartMergeTask() {
1874
    // LOG_ENGINE_DEBUG_ << "Begin StartMergeTask";
1875
    // merge task has been finished?
1876
    {
1877 1878
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
1879
            std::chrono::milliseconds span(10);
1880 1881
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
1882
            }
G
groot 已提交
1883 1884
        }
    }
X
Xu Peng 已提交
1885

1886
    // add new merge task
1887
    {
1888 1889
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
1890 1891
            // collect merge files for all collections(if merge_collection_ids_ is empty) for two reasons:
            // 1. other collections may still has un-merged files
1892
            // 2. server may be closed unexpected, these un-merge files need to be merged when server restart
1893 1894 1895 1896 1897
            if (merge_collection_ids_.empty()) {
                std::vector<meta::CollectionSchema> collection_schema_array;
                meta_ptr_->AllCollections(collection_schema_array);
                for (auto& schema : collection_schema_array) {
                    merge_collection_ids_.insert(schema.collection_id_);
1898 1899 1900 1901
                }
            }

            // start merge file thread
1902
            merge_thread_results_.push_back(
1903 1904
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_collection_ids_));
            merge_collection_ids_.clear();
1905
        }
G
groot 已提交
1906
    }
1907

1908
    // LOG_ENGINE_DEBUG_ << "End StartMergeTask";
X
Xu Peng 已提交
1909 1910
}

1911
Status
G
groot 已提交
1912
DBImpl::MergeHybridFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927 1928 1929 1930 1931 1932 1933 1934 1935 1936 1937 1938 1939
    // const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

    LOG_ENGINE_DEBUG_ << "Merge files for collection: " << collection_id;

    // step 1: create table file
    meta::SegmentSchema table_file;
    table_file.collection_id_ = collection_id;
    table_file.file_type_ = meta::SegmentSchema::NEW_MERGE;
    Status status = meta_ptr_->CreateHybridCollectionFile(table_file);

    if (!status.ok()) {
        LOG_ENGINE_ERROR_ << "Failed to create collection: " << status.ToString();
        return status;
    }

    // step 2: merge files
    /*
    ExecutionEnginePtr index =
        EngineFactory::Build(table_file.dimension_, table_file.location_, (EngineType)table_file.engine_type_,
                             (MetricType)table_file.metric_type_, table_file.nlist_);
*/
    meta::SegmentsSchema updated;

    std::string new_segment_dir;
    utils::GetParentPath(table_file.location_, new_segment_dir);
    auto segment_writer_ptr = std::make_shared<segment::SegmentWriter>(new_segment_dir);

G
groot 已提交
1940 1941
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
1942 1943 1944 1945 1946
    for (auto& file : files) {
        server::CollectMergeFilesMetrics metrics;
        std::string segment_dir_to_merge;
        utils::GetParentPath(file.location_, segment_dir_to_merge);
        segment_writer_ptr->Merge(segment_dir_to_merge, table_file.file_id_);
G
groot 已提交
1947 1948 1949

        files_holder.UnmarkFile(file);

1950 1951 1952 1953 1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998 1999 2000 2001 2002 2003 2004 2005
        auto file_schema = file;
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
        updated.push_back(file_schema);
        auto size = segment_writer_ptr->Size();
        if (size >= file_schema.index_file_size_) {
            break;
        }
    }

    // step 3: serialize to disk
    try {
        status = segment_writer_ptr->Serialize();
        fiu_do_on("DBImpl.MergeFiles.Serialize_ThrowException", throw std::exception());
        fiu_do_on("DBImpl.MergeFiles.Serialize_ErrorStatus", status = Status(DB_ERROR, ""));
    } catch (std::exception& ex) {
        std::string msg = "Serialize merged index encounter exception: " + std::string(ex.what());
        LOG_ENGINE_ERROR_ << msg;
        status = Status(DB_ERROR, msg);
    }

    if (!status.ok()) {
        LOG_ENGINE_ERROR_ << "Failed to persist merged segment: " << new_segment_dir << ". Error: " << status.message();

        // if failed to serialize merge file to disk
        // typical error: out of disk space, out of memory or permission denied
        table_file.file_type_ = meta::SegmentSchema::TO_DELETE;
        status = meta_ptr_->UpdateCollectionFile(table_file);
        LOG_ENGINE_DEBUG_ << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";

        return status;
    }

    // step 4: update table files state
    // if index type isn't IDMAP, set file type to TO_INDEX if file size exceed index_file_size
    // else set file type to RAW, no need to build index
    if (!utils::IsRawIndexType(table_file.engine_type_)) {
        table_file.file_type_ = (segment_writer_ptr->Size() >= table_file.index_file_size_)
                                    ? meta::SegmentSchema::TO_INDEX
                                    : meta::SegmentSchema::RAW;
    } else {
        table_file.file_type_ = meta::SegmentSchema::RAW;
    }
    table_file.file_size_ = segment_writer_ptr->Size();
    table_file.row_count_ = segment_writer_ptr->VectorCount();
    updated.push_back(table_file);
    status = meta_ptr_->UpdateCollectionFiles(updated);
    LOG_ENGINE_DEBUG_ << "New merged segment " << table_file.segment_id_ << " of size " << segment_writer_ptr->Size()
                      << " bytes";

    if (options_.insert_cache_immediately_) {
        segment_writer_ptr->Cache();
    }

    return status;
}

S
starlord 已提交
2006
void
2007
DBImpl::BackgroundMerge(std::set<std::string> collection_ids) {
2008
    // LOG_ENGINE_TRACE_ << " Background merge thread start";
S
starlord 已提交
2009

G
groot 已提交
2010
    Status status;
2011
    for (auto& collection_id : collection_ids) {
G
groot 已提交
2012 2013 2014
        const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

        auto status = merge_mgr_ptr_->MergeFiles(collection_id);
G
groot 已提交
2015
        if (!status.ok()) {
G
groot 已提交
2016 2017
            LOG_ENGINE_ERROR_ << "Failed to get merge files for collection: " << collection_id
                              << " reason:" << status.message();
G
groot 已提交
2018
        }
S
starlord 已提交
2019

2020
        if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2021
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action for collection: " << collection_id;
S
starlord 已提交
2022 2023
            break;
        }
G
groot 已提交
2024
    }
X
Xu Peng 已提交
2025

G
groot 已提交
2026
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
2027

2028
    {
G
groot 已提交
2029 2030
        uint64_t timeout = (options_.file_cleanup_timeout_ > 0) ? options_.file_cleanup_timeout_ : 10;
        uint64_t ttl = timeout * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
2031
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
2032
    }
S
starlord 已提交
2033

2034
    // LOG_ENGINE_TRACE_ << " Background merge thread exit";
G
groot 已提交
2035
}
X
Xu Peng 已提交
2036

S
starlord 已提交
2037
void
G
groot 已提交
2038
DBImpl::StartBuildIndexTask() {
S
starlord 已提交
2039
    // build index has been finished?
2040 2041 2042 2043 2044 2045 2046
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (!index_thread_results_.empty()) {
            std::chrono::milliseconds span(10);
            if (index_thread_results_.back().wait_for(span) == std::future_status::ready) {
                index_thread_results_.pop_back();
            }
G
groot 已提交
2047 2048 2049
        }
    }

S
starlord 已提交
2050
    // add new build index task
2051 2052 2053
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
S
starlord 已提交
2054
            index_thread_results_.push_back(index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
2055
        }
G
groot 已提交
2056
    }
X
Xu Peng 已提交
2057 2058
}

S
starlord 已提交
2059 2060
void
DBImpl::BackgroundBuildIndex() {
P
peng.xu 已提交
2061
    std::unique_lock<std::mutex> lock(build_index_mutex_);
G
groot 已提交
2062 2063 2064 2065
    meta::FilesHolder files_holder;
    meta_ptr_->FilesToIndex(files_holder);

    milvus::engine::meta::SegmentsSchema& to_index_files = files_holder.HoldFiles();
2066
    Status status = index_failed_checker_.IgnoreFailedIndexFiles(to_index_files);
2067

2068
    if (!to_index_files.empty()) {
G
groot 已提交
2069
        LOG_ENGINE_DEBUG_ << "Background build index thread begin " << to_index_files.size() << " files";
2070

2071
        // step 2: put build index task to scheduler
J
Jin Hai 已提交
2072
        std::vector<std::pair<scheduler::BuildIndexJobPtr, scheduler::SegmentSchemaPtr>> job2file_map;
2073
        for (auto& file : to_index_files) {
G
groot 已提交
2074
            scheduler::BuildIndexJobPtr job = std::make_shared<scheduler::BuildIndexJob>(meta_ptr_, options_);
J
Jin Hai 已提交
2075
            scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
2076
            job->AddToIndexFiles(file_ptr);
G
groot 已提交
2077
            scheduler::JobMgrInst::GetInstance()->Put(job);
G
groot 已提交
2078
            job2file_map.push_back(std::make_pair(job, file_ptr));
2079
        }
G
groot 已提交
2080

G
groot 已提交
2081
        // step 3: wait build index finished and mark failed files
G
groot 已提交
2082 2083
        for (auto iter = job2file_map.begin(); iter != job2file_map.end(); ++iter) {
            scheduler::BuildIndexJobPtr job = iter->first;
J
Jin Hai 已提交
2084
            meta::SegmentSchema& file_schema = *(iter->second.get());
G
groot 已提交
2085 2086 2087
            job->WaitBuildIndexFinish();
            if (!job->GetStatus().ok()) {
                Status status = job->GetStatus();
2088
                LOG_ENGINE_ERROR_ << "Building index job " << job->id() << " failed: " << status.ToString();
G
groot 已提交
2089

2090
                index_failed_checker_.MarkFailedIndexFile(file_schema, status.message());
G
groot 已提交
2091
            } else {
2092
                LOG_ENGINE_DEBUG_ << "Building index job " << job->id() << " succeed.";
G
groot 已提交
2093 2094

                index_failed_checker_.MarkSucceedIndexFile(file_schema);
G
groot 已提交
2095
            }
G
groot 已提交
2096 2097
            status = files_holder.UnmarkFile(file_schema);
            LOG_ENGINE_DEBUG_ << "Finish build index file " << file_schema.file_id_;
2098
        }
G
groot 已提交
2099

2100
        LOG_ENGINE_DEBUG_ << "Background build index thread finished";
G
groot 已提交
2101
        index_req_swn_.Notify();  // notify CreateIndex check circle
Y
Yu Kun 已提交
2102
    }
X
Xu Peng 已提交
2103 2104
}

G
groot 已提交
2105
Status
J
Jin Hai 已提交
2106
DBImpl::GetFilesToBuildIndex(const std::string& collection_id, const std::vector<int>& file_types,
G
groot 已提交
2107 2108 2109
                             meta::FilesHolder& files_holder) {
    files_holder.ReleaseFiles();
    auto status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
G
groot 已提交
2110

G
groot 已提交
2111 2112 2113 2114 2115 2116 2117 2118 2119 2120 2121
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
    for (const milvus::engine::meta::SegmentSchema& file : files) {
        if (file.file_type_ == static_cast<int>(meta::SegmentSchema::RAW) &&
            file.row_count_ < meta::BUILD_INDEX_THRESHOLD) {
            // skip build index for files that row count less than certain threshold
            files_holder.UnmarkFile(file);
        } else if (index_failed_checker_.IsFailedIndexFile(file)) {
            // skip build index for files that failed before
            files_holder.UnmarkFile(file);
        }
G
groot 已提交
2122 2123 2124 2125 2126
    }

    return Status::OK();
}

2127
Status
J
Jin Hai 已提交
2128 2129
DBImpl::GetPartitionByTag(const std::string& collection_id, const std::string& partition_tag,
                          std::string& partition_name) {
2130 2131 2132
    Status status;

    if (partition_tag.empty()) {
J
Jin Hai 已提交
2133
        partition_name = collection_id;
2134 2135 2136 2137 2138 2139 2140 2141

    } else {
        // trim side-blank of tag, only compare valid characters
        // for example: " ab cd " is treated as "ab cd"
        std::string valid_tag = partition_tag;
        server::StringHelpFunctions::TrimStringBlank(valid_tag);

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
J
Jin Hai 已提交
2142
            partition_name = collection_id;
2143 2144 2145
            return status;
        }

J
Jin Hai 已提交
2146
        status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
2147
        if (!status.ok()) {
2148
            LOG_ENGINE_ERROR_ << status.message();
2149 2150 2151 2152 2153 2154
        }
    }

    return status;
}

G
groot 已提交
2155
Status
J
Jin Hai 已提交
2156
DBImpl::GetPartitionsByTags(const std::string& collection_id, const std::vector<std::string>& partition_tags,
G
groot 已提交
2157
                            std::set<std::string>& partition_name_array) {
J
Jin Hai 已提交
2158 2159
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2160 2161

    for (auto& tag : partition_tags) {
2162 2163 2164 2165
        // trim side-blank of tag, only compare valid characters
        // for example: " ab cd " is treated as "ab cd"
        std::string valid_tag = tag;
        server::StringHelpFunctions::TrimStringBlank(valid_tag);
2166 2167

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
J
Jin Hai 已提交
2168
            partition_name_array.insert(collection_id);
2169 2170 2171
            return status;
        }

G
groot 已提交
2172
        for (auto& schema : partition_array) {
2173
            if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) {
J
Jin Hai 已提交
2174
                partition_name_array.insert(schema.collection_id_);
G
groot 已提交
2175 2176 2177 2178
            }
        }
    }

T
Tinkerrr 已提交
2179
    if (partition_name_array.empty()) {
G
groot 已提交
2180
        return Status(DB_PARTITION_NOT_FOUND, "The specified partiton does not exist");
T
Tinkerrr 已提交
2181 2182
    }

G
groot 已提交
2183 2184 2185 2186
    return Status::OK();
}

Status
2187
DBImpl::DropCollectionRecursively(const std::string& collection_id) {
J
Jin Hai 已提交
2188
    // dates partly delete files of the collection but currently we don't support
2189
    LOG_ENGINE_DEBUG_ << "Prepare to delete collection " << collection_id;
G
groot 已提交
2190 2191

    Status status;
2192
    if (options_.wal_enable_) {
2193
        wal_mgr_->DropCollection(collection_id);
G
groot 已提交
2194 2195
    }

2196 2197 2198
    status = mem_mgr_->EraseMemVector(collection_id);   // not allow insert
    status = meta_ptr_->DropCollection(collection_id);  // soft delete collection
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
2199

J
Jin Hai 已提交
2200
    // scheduler will determine when to delete collection files
2201
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
J
Jin Hai 已提交
2202
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(collection_id, meta_ptr_, nres);
2203 2204 2205
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

J
Jin Hai 已提交
2206 2207
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2208
    for (auto& schema : partition_array) {
2209 2210
        status = DropCollectionRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionRecursively.failed", status = Status(DB_ERROR, ""));
G
groot 已提交
2211 2212 2213 2214 2215 2216 2217 2218 2219
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2220
DBImpl::UpdateCollectionIndexRecursively(const std::string& collection_id, const CollectionIndex& index) {
J
Jin Hai 已提交
2221
    DropIndex(collection_id);
G
groot 已提交
2222

2223 2224
    auto status = meta_ptr_->UpdateCollectionIndex(collection_id, index);
    fiu_do_on("DBImpl.UpdateCollectionIndexRecursively.fail_update_collection_index",
S
shengjh 已提交
2225
              status = Status(DB_META_TRANSACTION_FAILED, ""));
G
groot 已提交
2226
    if (!status.ok()) {
2227
        LOG_ENGINE_ERROR_ << "Failed to update collection index info for collection: " << collection_id;
G
groot 已提交
2228 2229 2230
        return status;
    }

J
Jin Hai 已提交
2231 2232
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
A
AzAz 已提交
2233 2234 2235
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
2236
    for (auto& schema : partition_array) {
2237
        status = UpdateCollectionIndexRecursively(schema.collection_id_, index);
G
groot 已提交
2238 2239 2240 2241 2242 2243 2244 2245 2246
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
G
groot 已提交
2247 2248
DBImpl::WaitCollectionIndexRecursively(const std::shared_ptr<server::Context>& context,
                                       const std::string& collection_id, const CollectionIndex& index) {
G
groot 已提交
2249 2250 2251
    // for IDMAP type, only wait all NEW file converted to RAW file
    // for other type, wait NEW/RAW/NEW_MERGE/NEW_INDEX/TO_INDEX files converted to INDEX files
    std::vector<int> file_types;
2252
    if (utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
2253
        file_types = {
J
Jin Hai 已提交
2254 2255
            static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE),
G
groot 已提交
2256 2257 2258
        };
    } else {
        file_types = {
J
Jin Hai 已提交
2259 2260 2261
            static_cast<int32_t>(meta::SegmentSchema::RAW),       static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE), static_cast<int32_t>(meta::SegmentSchema::NEW_INDEX),
            static_cast<int32_t>(meta::SegmentSchema::TO_INDEX),
G
groot 已提交
2262 2263 2264 2265
        };
    }

    // get files to build index
G
groot 已提交
2266 2267 2268 2269
    {
        meta::FilesHolder files_holder;
        auto status = GetFilesToBuildIndex(collection_id, file_types, files_holder);
        int times = 1;
G
groot 已提交
2270
        uint64_t repeat = 0;
G
groot 已提交
2271
        while (!files_holder.HoldFiles().empty()) {
G
groot 已提交
2272 2273 2274 2275 2276 2277
            if (repeat % WAIT_BUILD_INDEX_INTERVAL == 0) {
                LOG_ENGINE_DEBUG_ << files_holder.HoldFiles().size() << " non-index files detected! Will build index "
                                  << times;
                if (!utils::IsRawIndexType(index.engine_type_)) {
                    status = meta_ptr_->UpdateCollectionFilesToIndex(collection_id);
                }
G
groot 已提交
2278
            }
G
groot 已提交
2279

G
groot 已提交
2280 2281 2282 2283 2284 2285 2286 2287 2288 2289 2290 2291 2292 2293
            index_req_swn_.Wait_For(std::chrono::seconds(1));

            // client break the connection, no need to block, check every 1 second
            if (context->IsConnectionBroken()) {
                LOG_ENGINE_DEBUG_ << "Client connection broken, build index in background";
                break;  // just break, not return, continue to update partitions files to to_index
            }

            // check to_index files every 5 seconds
            repeat++;
            if (repeat % WAIT_BUILD_INDEX_INTERVAL == 0) {
                GetFilesToBuildIndex(collection_id, file_types, files_holder);
                ++times;
            }
G
groot 已提交
2294 2295 2296 2297
        }
    }

    // build index for partition
J
Jin Hai 已提交
2298
    std::vector<meta::CollectionSchema> partition_array;
G
groot 已提交
2299
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2300
    for (auto& schema : partition_array) {
G
groot 已提交
2301
        status = WaitCollectionIndexRecursively(context, schema.collection_id_, index);
2302
        fiu_do_on("DBImpl.WaitCollectionIndexRecursively.fail_build_collection_Index_for_partition",
S
shengjh 已提交
2303
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2304 2305 2306 2307 2308
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
2309
    // failed to build index for some files, return error
2310
    std::string err_msg;
2311 2312
    index_failed_checker_.GetErrMsgForCollection(collection_id, err_msg);
    fiu_do_on("DBImpl.WaitCollectionIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
2313 2314
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
2315 2316
    }

G
groot 已提交
2317 2318
    LOG_ENGINE_DEBUG_ << "WaitCollectionIndexRecursively finished";

G
groot 已提交
2319 2320 2321 2322
    return Status::OK();
}

Status
2323
DBImpl::DropCollectionIndexRecursively(const std::string& collection_id) {
2324
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
2325 2326
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    auto status = meta_ptr_->DropCollectionIndex(collection_id);
G
groot 已提交
2327 2328 2329 2330 2331
    if (!status.ok()) {
        return status;
    }

    // drop partition index
J
Jin Hai 已提交
2332 2333
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2334
    for (auto& schema : partition_array) {
2335 2336
        status = DropCollectionIndexRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionIndexRecursively.fail_drop_collection_Index_for_partition",
S
shengjh 已提交
2337
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2338 2339 2340 2341 2342 2343 2344 2345 2346
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2347
DBImpl::GetCollectionRowCountRecursively(const std::string& collection_id, uint64_t& row_count) {
G
groot 已提交
2348
    row_count = 0;
J
Jin Hai 已提交
2349
    auto status = meta_ptr_->Count(collection_id, row_count);
G
groot 已提交
2350 2351 2352 2353 2354
    if (!status.ok()) {
        return status;
    }

    // get partition row count
J
Jin Hai 已提交
2355 2356
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2357
    for (auto& schema : partition_array) {
G
groot 已提交
2358
        uint64_t partition_row_count = 0;
2359 2360
        status = GetCollectionRowCountRecursively(schema.collection_id_, partition_row_count);
        fiu_do_on("DBImpl.GetCollectionRowCountRecursively.fail_get_collection_rowcount_for_partition",
S
shengjh 已提交
2361
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2362 2363 2364 2365 2366 2367 2368 2369 2370 2371
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

2372 2373 2374 2375
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

2376 2377
    auto collections_flushed = [&](const std::set<std::string>& collection_ids) -> uint64_t {
        if (collection_ids.empty()) {
2378 2379 2380 2381 2382
            return 0;
        }

        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
2383
            for (auto& collection : collection_ids) {
2384
                uint64_t lsn = 0;
2385 2386
                meta_ptr_->GetCollectionFlushLSN(collection, lsn);
                wal_mgr_->CollectionFlushed(collection, lsn);
2387 2388 2389 2390 2391 2392 2393
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
2394 2395
        for (auto& collection : collection_ids) {
            merge_collection_ids_.insert(collection);
2396 2397 2398 2399 2400 2401 2402
        }
        return max_lsn;
    };

    Status status;

    switch (record.type) {
2403 2404 2405 2406
        case wal::MXLogType::Entity: {
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
            if (!status.ok()) {
2407
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2408 2409 2410
                return status;
            }

2411
            std::set<std::string> flushed_collections;
2412 2413 2414
            status = mem_mgr_->InsertEntities(target_collection_name, record.length, record.ids,
                                              (record.data_size / record.length / sizeof(float)),
                                              (const float*)record.data, record.attr_nbytes, record.attr_data_size,
2415 2416
                                              record.attr_data, record.lsn, flushed_collections);
            collections_flushed(flushed_collections);
2417 2418 2419 2420

            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }
2421
        case wal::MXLogType::InsertBinary: {
2422 2423
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2424
            if (!status.ok()) {
2425
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2426 2427 2428
                return status;
            }

2429 2430
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2431
                                             (record.data_size / record.length / sizeof(uint8_t)),
2432
                                             (const u_int8_t*)record.data, record.lsn, flushed_collections);
2433
            // even though !status.ok, run
2434
            collections_flushed(flushed_collections);
2435 2436 2437 2438 2439 2440 2441

            // metrics
            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }

        case wal::MXLogType::InsertVector: {
2442 2443
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2444
            if (!status.ok()) {
2445
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2446 2447 2448
                return status;
            }

2449 2450
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2451
                                             (record.data_size / record.length / sizeof(float)),
2452
                                             (const float*)record.data, record.lsn, flushed_collections);
2453
            // even though !status.ok, run
2454
            collections_flushed(flushed_collections);
2455 2456 2457 2458 2459 2460 2461

            // metrics
            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }

        case wal::MXLogType::Delete: {
J
Jin Hai 已提交
2462 2463
            std::vector<meta::CollectionSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2464 2465 2466 2467
            if (!status.ok()) {
                return status;
            }

2468
            std::vector<std::string> collection_ids{record.collection_id};
2469
            for (auto& partition : partition_array) {
2470 2471
                auto& partition_collection_id = partition.collection_id_;
                collection_ids.emplace_back(partition_collection_id);
2472 2473 2474
            }

            if (record.length == 1) {
2475
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2476
                    status = mem_mgr_->DeleteVector(collection_id, *record.ids, record.lsn);
2477 2478 2479 2480 2481
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
2482
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2483
                    status = mem_mgr_->DeleteVectors(collection_id, record.length, record.ids, record.lsn);
2484 2485 2486 2487 2488 2489 2490 2491 2492
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
J
Jin Hai 已提交
2493 2494 2495 2496
            if (!record.collection_id.empty()) {
                // flush one collection
                std::vector<meta::CollectionSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2497 2498 2499 2500
                if (!status.ok()) {
                    return status;
                }

2501
                std::vector<std::string> collection_ids{record.collection_id};
2502
                for (auto& partition : partition_array) {
2503 2504
                    auto& partition_collection_id = partition.collection_id_;
                    collection_ids.emplace_back(partition_collection_id);
2505 2506
                }

2507 2508
                std::set<std::string> flushed_collections;
                for (auto& collection_id : collection_ids) {
2509
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
J
Jin Hai 已提交
2510
                    status = mem_mgr_->Flush(collection_id);
2511 2512 2513
                    if (!status.ok()) {
                        break;
                    }
2514
                    flushed_collections.insert(collection_id);
2515 2516
                }

2517
                collections_flushed(flushed_collections);
2518 2519

            } else {
2520 2521
                // flush all collections
                std::set<std::string> collection_ids;
2522 2523
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2524
                    status = mem_mgr_->Flush(collection_ids);
2525 2526
                }

2527
                uint64_t lsn = collections_flushed(collection_ids);
2528 2529 2530 2531 2532 2533 2534 2535 2536 2537 2538 2539
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
    }

    return status;
}

void
G
groot 已提交
2540 2541 2542 2543 2544 2545 2546 2547 2548 2549 2550
DBImpl::InternalFlush(const std::string& collection_id) {
    wal::MXLogRecord record;
    record.type = wal::MXLogType::Flush;
    record.collection_id = collection_id;
    ExecWalRecord(record);

    StartMergeTask();
}

void
DBImpl::BackgroundWalThread() {
2551
    SetThreadName("wal_thread");
2552 2553
    server::SystemInfo::GetInstance().Init();

2554
    std::chrono::system_clock::time_point next_auto_flush_time;
2555
    auto get_next_auto_flush_time = [&]() {
2556
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2557
    };
2558 2559 2560
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2561 2562

    while (true) {
2563 2564
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
G
groot 已提交
2565
                InternalFlush();
2566 2567
                next_auto_flush_time = get_next_auto_flush_time();
            }
2568 2569
        }

G
groot 已提交
2570
        wal::MXLogRecord record;
2571 2572
        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
2573
            LOG_ENGINE_ERROR_ << "WAL background GetNextRecord error";
2574 2575 2576 2577 2578 2579
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
G
groot 已提交
2580 2581
                // notify flush request to return
                flush_req_swn_.Notify();
2582 2583

                // if user flush all manually, update auto flush also
J
Jin Hai 已提交
2584
                if (record.collection_id.empty() && options_.auto_flush_interval_ > 0) {
2585 2586 2587 2588 2589 2590
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2591 2592
                InternalFlush();
                flush_req_swn_.Notify();
2593 2594
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
2595
                LOG_ENGINE_DEBUG_ << "WAL background thread exit";
2596 2597 2598
                break;
            }

2599
            if (options_.auto_flush_interval_ > 0) {
G
groot 已提交
2600
                swn_wal_.Wait_Until(next_auto_flush_time);
2601
            } else {
G
groot 已提交
2602
                swn_wal_.Wait();
2603
            }
2604 2605 2606 2607
        }
    }
}

G
groot 已提交
2608 2609
void
DBImpl::BackgroundFlushThread() {
2610
    SetThreadName("flush_thread");
G
groot 已提交
2611 2612 2613
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2614
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
G
groot 已提交
2615 2616 2617 2618 2619 2620 2621 2622 2623 2624 2625 2626 2627 2628 2629 2630 2631
            break;
        }

        InternalFlush();
        if (options_.auto_flush_interval_ > 0) {
            swn_flush_.Wait_For(std::chrono::seconds(options_.auto_flush_interval_));
        } else {
            swn_flush_.Wait();
        }
    }
}

void
DBImpl::BackgroundMetricThread() {
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2632
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
G
groot 已提交
2633 2634 2635 2636 2637
            break;
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
        StartMetricTask();
G
groot 已提交
2638
        meta::FilesHolder::PrintInfo();
G
groot 已提交
2639 2640 2641
    }
}

2642 2643 2644 2645 2646
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

2647 2648 2649 2650 2651
void
DBImpl::OnUseBlasThresholdChanged(int64_t threshold) {
    faiss::distance_compute_blas_threshold = threshold;
}

S
starlord 已提交
2652 2653
}  // namespace engine
}  // namespace milvus