DBImpl.cpp 25.7 KB
Newer Older
X
Xu Peng 已提交
1 2 3 4 5
/*******************************************************************************
 * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited.
 * Proprietary and confidential.
 ******************************************************************************/
6
#include "DBImpl.h"
S
starlord 已提交
7
#include "src/db/meta/SqliteMetaImpl.h"
G
groot 已提交
8
#include "Log.h"
9
#include "Utils.h"
S
starlord 已提交
10
#include "engine/EngineFactory.h"
Z
update  
zhiru 已提交
11
#include "Factories.h"
G
groot 已提交
12
#include "metrics/Metrics.h"
G
groot 已提交
13
#include "scheduler/TaskScheduler.h"
J
jinhai 已提交
14

G
groot 已提交
15
#include "scheduler/context/DeleteContext.h"
G
groot 已提交
16
#include "utils/TimeRecorder.h"
S
starlord 已提交
17
#include "meta/MetaConsts.h"
X
Xu Peng 已提交
18

X
Xu Peng 已提交
19
#include <assert.h>
X
Xu Peng 已提交
20
#include <chrono>
X
Xu Peng 已提交
21
#include <thread>
22
#include <iostream>
X
xj.lin 已提交
23
#include <cstring>
X
Xu Peng 已提交
24
#include <cache/CpuCacheMgr.h>
G
groot 已提交
25
#include <boost/filesystem.hpp>
W
wxyu 已提交
26
#include "scheduler/SchedInst.h"
Y
Yu Kun 已提交
27
#include <src/cache/GpuCacheMgr.h>
X
Xu Peng 已提交
28

X
Xu Peng 已提交
29
namespace zilliz {
J
jinhai 已提交
30
namespace milvus {
X
Xu Peng 已提交
31
namespace engine {
X
Xu Peng 已提交
32

G
groot 已提交
33 34
namespace {

J
jinhai 已提交
35 36 37
constexpr uint64_t METRIC_ACTION_INTERVAL = 1;
constexpr uint64_t COMPACT_ACTION_INTERVAL = 1;
constexpr uint64_t INDEX_ACTION_INTERVAL = 1;
G
groot 已提交
38

G
groot 已提交
39
}
Y
yu yunfeng 已提交
40

G
groot 已提交
41 42

DBImpl::DBImpl(const Options& options)
G
groot 已提交
43
    : options_(options),
X
Xu Peng 已提交
44
      shutting_down_(false),
G
groot 已提交
45 46
      compact_thread_pool_(1, 1),
      index_thread_pool_(1, 1) {
Z
update  
zhiru 已提交
47
    meta_ptr_ = DBMetaImplFactory::Build(options.meta, options.mode);
Z
zhiru 已提交
48
    mem_mgr_ = MemManagerFactory::Build(meta_ptr_, options_);
Z
update  
zhiru 已提交
49
    if (options.mode != Options::MODE::READ_ONLY) {
50
        ENGINE_LOG_TRACE << "StartTimerTasks";
Z
update  
zhiru 已提交
51 52
        StartTimerTasks();
    }
S
starlord 已提交
53 54


X
Xu Peng 已提交
55 56
}

G
groot 已提交
57
Status DBImpl::CreateTable(meta::TableSchema& table_schema) {
58 59 60
    meta::TableSchema temp_schema = table_schema;
    temp_schema.index_file_size_ *= ONE_MB;
    return meta_ptr_->CreateTable(temp_schema);
61 62
}

G
groot 已提交
63
Status DBImpl::DeleteTable(const std::string& table_id, const meta::DatesT& dates) {
G
groot 已提交
64
    //dates partly delete files of the table but currently we don't support
S
starlord 已提交
65
    ENGINE_LOG_DEBUG << "Prepare to delete table " << table_id;
G
groot 已提交
66

67 68 69 70 71 72
    if (dates.empty()) {
        mem_mgr_->EraseMemVector(table_id); //not allow insert
        meta_ptr_->DeleteTable(table_id); //soft delete table

        //scheduler will determine when to delete table files
        TaskScheduler& scheduler = TaskScheduler::GetInstance();
W
wxyu 已提交
73 74 75
        DeleteContextPtr context = std::make_shared<DeleteContext>(table_id,
                                                               meta_ptr_,
                                                               ResMgrInst::GetInstance()->GetNumOfComputeResource());
76
        scheduler.Schedule(context);
W
wxyu 已提交
77
        context->WaitAndDelete();
78 79 80
    } else {
        meta_ptr_->DropPartitionsByDates(table_id, dates);
    }
G
groot 已提交
81

G
groot 已提交
82 83 84
    return Status::OK();
}

G
groot 已提交
85
Status DBImpl::DescribeTable(meta::TableSchema& table_schema) {
G
groot 已提交
86
    return meta_ptr_->DescribeTable(table_schema);
87 88
}

G
groot 已提交
89
Status DBImpl::HasTable(const std::string& table_id, bool& has_or_not) {
G
groot 已提交
90
    return meta_ptr_->HasTable(table_id, has_or_not);
91 92
}

G
groot 已提交
93
Status DBImpl::AllTables(std::vector<meta::TableSchema>& table_schema_array) {
G
groot 已提交
94
    return meta_ptr_->AllTables(table_schema_array);
G
groot 已提交
95 96
}

Y
Yu Kun 已提交
97
Status DBImpl::PreloadTable(const std::string &table_id) {
Y
Yu Kun 已提交
98
    meta::DatePartionedTableFilesSchema files;
Y
Yu Kun 已提交
99

Y
Yu Kun 已提交
100 101 102 103 104
    meta::DatesT dates;
    auto status = meta_ptr_->FilesToSearch(table_id, dates, files);
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
105

Y
Yu Kun 已提交
106 107
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
108 109
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
110 111 112

    for(auto &day_files : files) {
        for (auto &file : day_files.second) {
S
starlord 已提交
113
            ExecutionEnginePtr engine = EngineFactory::Build(file.dimension_, file.location_, (EngineType)file.engine_type_, (MetricType)file.metric_type_, file.nlist_);
Y
Yu Kun 已提交
114 115 116 117
            if(engine == nullptr) {
                ENGINE_LOG_ERROR << "Invalid engine type";
                return Status::Error("Invalid engine type");
            }
Y
Yu Kun 已提交
118

Y
Yu Kun 已提交
119
            size += engine->PhysicalSize();
Y
Yu Kun 已提交
120
            if (size > available_size) {
Y
Yu Kun 已提交
121 122 123 124
                break;
            } else {
                try {
                    //step 1: load index
Y
Yu Kun 已提交
125
                    engine->Load(true);
Y
Yu Kun 已提交
126 127 128 129 130
                } catch (std::exception &ex) {
                    std::string msg = "load to cache exception" + std::string(ex.what());
                    ENGINE_LOG_ERROR << msg;
                    return Status::Error(msg);
                }
Y
Yu Kun 已提交
131 132 133
            }
        }
    }
Y
Yu Kun 已提交
134
    return Status::OK();
Y
Yu Kun 已提交
135 136
}

S
starlord 已提交
137 138 139 140
Status DBImpl::UpdateTableFlag(const std::string &table_id, int64_t flag) {
    return meta_ptr_->UpdateTableFlag(table_id, flag);
}

G
groot 已提交
141
Status DBImpl::GetTableRowCount(const std::string& table_id, uint64_t& row_count) {
G
groot 已提交
142
    return meta_ptr_->Count(table_id, row_count);
G
groot 已提交
143 144
}

G
groot 已提交
145
Status DBImpl::InsertVectors(const std::string& table_id_,
G
groot 已提交
146
        uint64_t n, const float* vectors, IDNumbers& vector_ids_) {
S
starlord 已提交
147
    ENGINE_LOG_DEBUG << "Insert " << n << " vectors to cache";
Y
yu yunfeng 已提交
148

149 150 151
    Status status;
    zilliz::milvus::server::CollectInsertMetrics metrics(n, status);
    status = mem_mgr_->InsertVectors(table_id_, n, vectors, vector_ids_);
Y
yu yunfeng 已提交
152 153 154
//    std::chrono::microseconds time_span = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
//    double average_time = double(time_span.count()) / n;

S
starlord 已提交
155 156
    ENGINE_LOG_DEBUG << "Insert vectors to cache finished";

G
groot 已提交
157
    return status;
X
Xu Peng 已提交
158 159
}

Y
Yu Kun 已提交
160
Status DBImpl::Query(const std::string &table_id, uint64_t k, uint64_t nq, uint64_t nprobe,
X
xj.lin 已提交
161
                      const float *vectors, QueryResults &results) {
Y
Yu Kun 已提交
162 163
    server::CollectQueryMetrics metrics(nq);

X
Xu Peng 已提交
164
    meta::DatesT dates = {meta::Meta::GetDate()};
Y
Yu Kun 已提交
165
    Status result = Query(table_id, k, nq, nprobe, vectors, dates, results);
Y
yu yunfeng 已提交
166

Y
yu yunfeng 已提交
167
    return result;
X
Xu Peng 已提交
168 169
}

Y
Yu Kun 已提交
170
Status DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe,
X
Xu Peng 已提交
171
        const float* vectors, const meta::DatesT& dates, QueryResults& results) {
S
starlord 已提交
172
    ENGINE_LOG_DEBUG << "Query by vectors " << table_id;
S
starlord 已提交
173

174 175
    //get all table files from table
    meta::DatePartionedTableFilesSchema files;
G
groot 已提交
176
    auto status = meta_ptr_->FilesToSearch(table_id, dates, files);
177 178 179 180 181 182 183 184 185
    if (!status.ok()) { return status; }

    meta::TableFilesSchema file_id_array;
    for (auto &day_files : files) {
        for (auto &file : day_files.second) {
            file_id_array.push_back(file);
        }
    }

S
starlord 已提交
186
    cache::CpuCacheMgr::GetInstance()->PrintInfo(); //print cache info before query
Y
Yu Kun 已提交
187
    status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, dates, results);
S
starlord 已提交
188 189
    cache::CpuCacheMgr::GetInstance()->PrintInfo(); //print cache info after query
    return status;
G
groot 已提交
190
}
X
Xu Peng 已提交
191

192
Status DBImpl::Query(const std::string& table_id, const std::vector<std::string>& file_ids,
Y
Yu Kun 已提交
193
        uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
194
        const meta::DatesT& dates, QueryResults& results) {
S
starlord 已提交
195 196
    ENGINE_LOG_DEBUG << "Query by file ids";

197
    //get specified files
198
    std::vector<size_t> ids;
199 200
    for (auto &id : file_ids) {
        meta::TableFileSchema table_file;
201 202
        table_file.table_id_ = table_id;
        std::string::size_type sz;
J
jinhai 已提交
203
        ids.push_back(std::stoul(id, &sz));
204 205
    }

X
xj.lin 已提交
206 207
    meta::DatePartionedTableFilesSchema files_array;
    auto status = meta_ptr_->FilesToSearch(table_id, ids, dates, files_array);
208 209
    if (!status.ok()) {
        return status;
210 211
    }

X
xj.lin 已提交
212 213 214 215 216 217 218 219
    meta::TableFilesSchema file_id_array;
    for (auto &day_files : files_array) {
        for (auto &file : day_files.second) {
            file_id_array.push_back(file);
        }
    }

    if(file_id_array.empty()) {
G
groot 已提交
220 221 222
        return Status::Error("Invalid file id");
    }

S
starlord 已提交
223
    cache::CpuCacheMgr::GetInstance()->PrintInfo(); //print cache info before query
Y
Yu Kun 已提交
224
    status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, dates, results);
S
starlord 已提交
225 226
    cache::CpuCacheMgr::GetInstance()->PrintInfo(); //print cache info after query
    return status;
227 228 229
}

Status DBImpl::QueryAsync(const std::string& table_id, const meta::TableFilesSchema& files,
Y
Yu Kun 已提交
230
                          uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
231
                          const meta::DatesT& dates, QueryResults& results) {
Y
Yu Kun 已提交
232 233
    server::CollectQueryMetrics metrics(nq);

S
starlord 已提交
234
    server::TimeRecorder rc("");
G
groot 已提交
235 236

    //step 1: get files to search
S
starlord 已提交
237
    ENGINE_LOG_DEBUG << "Engine query begin, index file count:" << files.size() << " date range count:" << dates.size();
Y
Yu Kun 已提交
238
    SearchContextPtr context = std::make_shared<SearchContext>(k, nq, nprobe, vectors);
239 240 241
    for (auto &file : files) {
        TableFileSchemaPtr file_ptr = std::make_shared<meta::TableFileSchema>(file);
        context->AddIndexFile(file_ptr);
G
groot 已提交
242 243
    }

G
groot 已提交
244
    //step 2: put search task to scheduler
G
groot 已提交
245 246
    TaskScheduler& scheduler = TaskScheduler::GetInstance();
    scheduler.Schedule(context);
G
groot 已提交
247 248

    context->WaitResult();
G
groot 已提交
249

S
starlord 已提交
250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272
    //step 3: print time cost information
    double load_cost = context->LoadCost();
    double search_cost = context->SearchCost();
    double reduce_cost = context->ReduceCost();
    std::string load_info = server::TimeRecorder::GetTimeSpanStr(load_cost);
    std::string search_info = server::TimeRecorder::GetTimeSpanStr(search_cost);
    std::string reduce_info = server::TimeRecorder::GetTimeSpanStr(reduce_cost);
    if(search_cost > 0.0 || reduce_cost > 0.0) {
        double total_cost = load_cost + search_cost + reduce_cost;
        double load_percent = load_cost/total_cost;
        double search_percent = search_cost/total_cost;
        double reduce_percent = reduce_cost/total_cost;

        ENGINE_LOG_DEBUG << "Engine load index totally cost:" << load_info << " percent: " << load_percent*100 << "%";
        ENGINE_LOG_DEBUG << "Engine search index totally cost:" << search_info << " percent: " << search_percent*100 << "%";
        ENGINE_LOG_DEBUG << "Engine reduce topk totally cost:" << reduce_info << " percent: " << reduce_percent*100 << "%";
    } else {
        ENGINE_LOG_DEBUG << "Engine load cost:" << load_info
            << " search cost: " << search_info
            << " reduce cost: " << reduce_info;
    }

    //step 4: construct results
J
jinhai 已提交
273
    results = context->GetResult();
S
starlord 已提交
274
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
275 276 277 278

    return Status::OK();
}

G
groot 已提交
279 280
void DBImpl::StartTimerTasks() {
    bg_timer_thread_ = std::thread(&DBImpl::BackgroundTimerTask, this);
X
Xu Peng 已提交
281 282
}

G
groot 已提交
283
void DBImpl::BackgroundTimerTask() {
X
Xu Peng 已提交
284
    Status status;
Y
yu yunfeng 已提交
285
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
286
    while (true) {
G
groot 已提交
287 288 289 290 291 292 293
        if (shutting_down_.load(std::memory_order_acquire)){
            for(auto& iter : compact_thread_results_) {
                iter.wait();
            }
            for(auto& iter : index_thread_results_) {
                iter.wait();
            }
S
starlord 已提交
294 295

            ENGINE_LOG_DEBUG << "DB background thread exit";
G
groot 已提交
296 297
            break;
        }
X
Xu Peng 已提交
298

G
groot 已提交
299
        std::this_thread::sleep_for(std::chrono::seconds(1));
X
Xu Peng 已提交
300

G
groot 已提交
301
        StartMetricTask();
G
groot 已提交
302 303 304
        StartCompactionTask();
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
305 306
}

G
groot 已提交
307 308 309 310 311 312 313
void DBImpl::StartMetricTask() {
    static uint64_t metric_clock_tick = 0;
    metric_clock_tick++;
    if(metric_clock_tick%METRIC_ACTION_INTERVAL != 0) {
        return;
    }

314
    ENGINE_LOG_TRACE << "Start metric task";
S
starlord 已提交
315

G
groot 已提交
316 317 318
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(METRIC_ACTION_INTERVAL);
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
319
    server::Metrics::GetInstance().CpuCacheUsageGaugeSet(cache_usage*100/cache_total);
Y
Yu Kun 已提交
320
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
321 322 323 324 325 326 327 328
    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 已提交
329

K
kun yu 已提交
330
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
331 332
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
K
kun yu 已提交
333

334
    ENGINE_LOG_TRACE << "Metric task finished";
G
groot 已提交
335 336
}

G
groot 已提交
337
void DBImpl::StartCompactionTask() {
G
groot 已提交
338 339 340 341 342 343
    static uint64_t compact_clock_tick = 0;
    compact_clock_tick++;
    if(compact_clock_tick%COMPACT_ACTION_INTERVAL != 0) {
        return;
    }

G
groot 已提交
344
    //serialize memory data
G
groot 已提交
345
    std::set<std::string> temp_table_ids;
G
groot 已提交
346
    mem_mgr_->Serialize(temp_table_ids);
G
groot 已提交
347 348 349
    for(auto& id : temp_table_ids) {
        compact_table_ids_.insert(id);
    }
X
Xu Peng 已提交
350

351 352 353
    if(!temp_table_ids.empty()) {
        SERVER_LOG_DEBUG << "Insert cache serialized";
    }
S
starlord 已提交
354

G
groot 已提交
355 356 357 358 359 360 361
    //compactiong has been finished?
    if(!compact_thread_results_.empty()) {
        std::chrono::milliseconds span(10);
        if (compact_thread_results_.back().wait_for(span) == std::future_status::ready) {
            compact_thread_results_.pop_back();
        }
    }
X
Xu Peng 已提交
362

G
groot 已提交
363 364 365 366 367 368
    //add new compaction task
    if(compact_thread_results_.empty()) {
        compact_thread_results_.push_back(
                compact_thread_pool_.enqueue(&DBImpl::BackgroundCompaction, this, compact_table_ids_));
        compact_table_ids_.clear();
    }
X
Xu Peng 已提交
369 370
}

G
groot 已提交
371
Status DBImpl::MergeFiles(const std::string& table_id, const meta::DateT& date,
372
        const meta::TableFilesSchema& files) {
S
starlord 已提交
373
    ENGINE_LOG_DEBUG << "Merge files for table " << table_id;
S
starlord 已提交
374

S
starlord 已提交
375
    //step 1: create table file
X
Xu Peng 已提交
376
    meta::TableFileSchema table_file;
G
groot 已提交
377 378
    table_file.table_id_ = table_id;
    table_file.date_ = date;
379
    table_file.file_type_ = meta::TableFileSchema::NEW_MERGE;
G
groot 已提交
380
    Status status = meta_ptr_->CreateTableFile(table_file);
X
Xu Peng 已提交
381

382
    if (!status.ok()) {
S
starlord 已提交
383
        ENGINE_LOG_ERROR << "Failed to create table: " << status.ToString();
384 385 386
        return status;
    }

S
starlord 已提交
387
    //step 2: merge files
G
groot 已提交
388
    ExecutionEnginePtr index =
S
starlord 已提交
389 390
            EngineFactory::Build(table_file.dimension_, table_file.location_, (EngineType)table_file.engine_type_,
                    (MetricType)table_file.metric_type_, table_file.nlist_);
391

392
    meta::TableFilesSchema updated;
X
Xu Peng 已提交
393
    long  index_size = 0;
394 395

    for (auto& file : files) {
Y
Yu Kun 已提交
396
        server::CollectMergeFilesMetrics metrics;
Y
yu yunfeng 已提交
397

G
groot 已提交
398
        index->Merge(file.location_);
399
        auto file_schema = file;
G
groot 已提交
400
        file_schema.file_type_ = meta::TableFileSchema::TO_DELETE;
401
        updated.push_back(file_schema);
G
groot 已提交
402
        ENGINE_LOG_DEBUG << "Merging file " << file_schema.file_id_;
G
groot 已提交
403
        index_size = index->Size();
X
Xu Peng 已提交
404

S
starlord 已提交
405
        if (index_size >= file_schema.index_file_size_) break;
406 407
    }

S
starlord 已提交
408 409 410 411 412 413 414
    //step 3: serialize to disk
    try {
        index->Serialize();
    } catch (std::exception& ex) {
        //typical error: out of disk space or permition denied
        std::string msg = "Serialize merged index encounter exception" + std::string(ex.what());
        ENGINE_LOG_ERROR << msg;
Y
yu yunfeng 已提交
415

S
starlord 已提交
416 417 418
        table_file.file_type_ = meta::TableFileSchema::TO_DELETE;
        status = meta_ptr_->UpdateTableFile(table_file);
        ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";
X
Xu Peng 已提交
419

S
starlord 已提交
420 421 422 423 424 425 426
        std::cout << "ERROR: failed to persist merged index file: " << table_file.location_
                  << ", possible out of disk space" << std::endl;

        return Status::Error(msg);
    }

    //step 4: update table files state
427 428 429 430 431 432 433 434
    //if index type isn't IDMAP, set file type to TO_INDEX if file size execeed index_file_size
    //else set file type to RAW, no need to build index
    if (table_file.engine_type_ != (int)EngineType::FAISS_IDMAP) {
        table_file.file_type_ = (index->PhysicalSize() >= table_file.index_file_size_) ?
                                meta::TableFileSchema::TO_INDEX : meta::TableFileSchema::RAW;
    } else {
        table_file.file_type_ = meta::TableFileSchema::RAW;
    }
435 436
    table_file.file_size_ = index->PhysicalSize();
    table_file.row_count_ = index->Count();
X
Xu Peng 已提交
437
    updated.push_back(table_file);
G
groot 已提交
438 439
    status = meta_ptr_->UpdateTableFiles(updated);
    ENGINE_LOG_DEBUG << "New merged file " << table_file.file_id_ <<
S
starlord 已提交
440
        " of size " << index->PhysicalSize() << " bytes";
441

S
starlord 已提交
442 443 444
    if(options_.insert_cache_immediately_) {
        index->Cache();
    }
X
Xu Peng 已提交
445

446 447 448
    return status;
}

G
groot 已提交
449
Status DBImpl::BackgroundMergeFiles(const std::string& table_id) {
450
    meta::DatePartionedTableFilesSchema raw_files;
G
groot 已提交
451
    auto status = meta_ptr_->FilesToMerge(table_id, raw_files);
X
Xu Peng 已提交
452
    if (!status.ok()) {
S
starlord 已提交
453
        ENGINE_LOG_ERROR << "Failed to get merge files for table: " << table_id;
X
Xu Peng 已提交
454 455
        return status;
    }
456

X
Xu Peng 已提交
457
    bool has_merge = false;
458
    for (auto& kv : raw_files) {
X
Xu Peng 已提交
459
        auto files = kv.second;
S
starlord 已提交
460 461
        if (files.size() < options_.merge_trigger_number) {
            ENGINE_LOG_DEBUG << "Files number not greater equal than merge trigger number, skip merge action";
X
Xu Peng 已提交
462 463
            continue;
        }
X
Xu Peng 已提交
464
        has_merge = true;
X
Xu Peng 已提交
465
        MergeFiles(table_id, kv.first, kv.second);
G
groot 已提交
466 467

        if (shutting_down_.load(std::memory_order_acquire)){
S
starlord 已提交
468
            ENGINE_LOG_DEBUG << "Server will shutdown, skip merge action for table " << table_id;
G
groot 已提交
469 470
            break;
        }
471
    }
X
Xu Peng 已提交
472

G
groot 已提交
473 474
    return Status::OK();
}
475

G
groot 已提交
476
void DBImpl::BackgroundCompaction(std::set<std::string> table_ids) {
477
    ENGINE_LOG_TRACE << " Background compaction thread start";
S
starlord 已提交
478

G
groot 已提交
479
    Status status;
J
jinhai 已提交
480
    for (auto& table_id : table_ids) {
G
groot 已提交
481 482
        status = BackgroundMergeFiles(table_id);
        if (!status.ok()) {
S
starlord 已提交
483
            ENGINE_LOG_ERROR << "Merge files for table " << table_id << " failed: " << status.ToString();
S
starlord 已提交
484
            continue;//let other table get chance to merge
G
groot 已提交
485
        }
S
starlord 已提交
486 487 488 489 490

        if (shutting_down_.load(std::memory_order_acquire)){
            ENGINE_LOG_DEBUG << "Server will shutdown, skip merge action";
            break;
        }
G
groot 已提交
491
    }
X
Xu Peng 已提交
492

G
groot 已提交
493
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
494

495
    int ttl = 5*meta::M_SEC;//default: file will be deleted after 5 minutes
Z
update  
zhiru 已提交
496
    if (options_.mode == Options::MODE::CLUSTER) {
Z
update  
zhiru 已提交
497 498 499
        ttl = meta::D_SEC;
    }
    meta_ptr_->CleanUpFilesWithTTL(ttl);
S
starlord 已提交
500

501
    ENGINE_LOG_TRACE << " Background compaction thread exit";
G
groot 已提交
502
}
X
Xu Peng 已提交
503

P
peng.xu 已提交
504
void DBImpl::StartBuildIndexTask(bool force) {
G
groot 已提交
505 506
    static uint64_t index_clock_tick = 0;
    index_clock_tick++;
P
peng.xu 已提交
507
    if(!force && (index_clock_tick%INDEX_ACTION_INTERVAL != 0)) {
G
groot 已提交
508 509 510
        return;
    }

G
groot 已提交
511 512 513 514 515 516 517 518 519 520 521 522 523
    //build index has been finished?
    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();
        }
    }

    //add new build index task
    if(index_thread_results_.empty()) {
        index_thread_results_.push_back(
                index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
    }
X
Xu Peng 已提交
524 525
}

526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551
Status DBImpl::CreateIndex(const std::string& table_id, const TableIndex& index) {
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

        //step 1: check index difference
        TableIndex old_index;
        auto status = DescribeIndex(table_id, old_index);
        if(!status.ok()) {
            ENGINE_LOG_ERROR << "Failed to get table index info";
            return status;
        }

        if(utils::IsSameIndex(old_index, index)) {
            ENGINE_LOG_DEBUG << "Same index setting, no need to create index again";
            return Status::OK();
        }

        //step 2: drop old index files
        DropIndex(table_id);

        //step 3: update index info
        status = meta_ptr_->UpdateTableIndexParam(table_id, index);
        if (!status.ok()) {
            ENGINE_LOG_ERROR << "Failed to update table index info";
            return status;
        }
552 553 554 555 556

        if(index.engine_type_ == (int)EngineType::FAISS_IDMAP) {
            ENGINE_LOG_DEBUG << "index type = IDMAP, no need to build index";
            return Status::OK();
        }
557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579
    }

    bool has = false;
    auto status = meta_ptr_->HasNonIndexFiles(table_id, has);
    int times = 1;

    while (has) {
        ENGINE_LOG_DEBUG << "Non index files detected! Will build index " << times;
        status = meta_ptr_->UpdateTableFilesToIndex(table_id);
        /* StartBuildIndexTask(true); */
        std::this_thread::sleep_for(std::chrono::milliseconds(std::min(10*1000, times*100)));
        status = meta_ptr_->HasNonIndexFiles(table_id, has);
        times++;
    }
    return Status::OK();
}

Status DBImpl::DescribeIndex(const std::string& table_id, TableIndex& index) {
    return meta_ptr_->DescribeTableIndex(table_id, index);
}

Status DBImpl::DropIndex(const std::string& table_id) {
    return meta_ptr_->DropTableIndex(table_id);
P
peng.xu 已提交
580 581
}

G
groot 已提交
582
Status DBImpl::BuildIndex(const meta::TableFileSchema& file) {
S
starlord 已提交
583 584 585
    ExecutionEnginePtr to_index =
            EngineFactory::Build(file.dimension_, file.location_, (EngineType)file.engine_type_,
                    (MetricType)file.metric_type_, file.nlist_);
G
groot 已提交
586
    if(to_index == nullptr) {
S
starlord 已提交
587
        ENGINE_LOG_ERROR << "Invalid engine type";
G
groot 已提交
588 589
        return Status::Error("Invalid engine type");
    }
590

G
groot 已提交
591
    try {
G
groot 已提交
592
        //step 1: load index
S
starlord 已提交
593
        to_index->Load(options_.insert_cache_immediately_);
G
groot 已提交
594 595 596 597 598

        //step 2: create table file
        meta::TableFileSchema table_file;
        table_file.table_id_ = file.table_id_;
        table_file.date_ = file.date_;
599
        table_file.file_type_ = meta::TableFileSchema::NEW_INDEX; //for multi-db-path, distribute index file averagely to each path
G
groot 已提交
600 601
        Status status = meta_ptr_->CreateTableFile(table_file);
        if (!status.ok()) {
S
starlord 已提交
602
            ENGINE_LOG_ERROR << "Failed to create table: " << status.ToString();
G
groot 已提交
603 604 605 606
            return status;
        }

        //step 3: build index
607 608 609
        std::shared_ptr<ExecutionEngine> index;

        try {
Y
Yu Kun 已提交
610
            server::CollectBuildIndexMetrics metrics;
S
starlord 已提交
611
            index = to_index->BuildIndex(table_file.location_, (EngineType)table_file.engine_type_);
612 613 614 615 616 617 618 619 620 621 622 623 624
        } catch (std::exception& ex) {
            //typical error: out of gpu memory
            std::string msg = "BuildIndex encounter exception" + std::string(ex.what());
            ENGINE_LOG_ERROR << msg;

            table_file.file_type_ = meta::TableFileSchema::TO_DELETE;
            status = meta_ptr_->UpdateTableFile(table_file);
            ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";

            std::cout << "ERROR: failed to build index, index file is too large or gpu memory is not enough" << std::endl;

            return Status::Error(msg);
        }
625

G
groot 已提交
626 627 628 629 630 631 632 633 634
        //step 4: if table has been deleted, dont save index file
        bool has_table = false;
        meta_ptr_->HasTable(file.table_id_, has_table);
        if(!has_table) {
            meta_ptr_->DeleteTableFiles(file.table_id_);
            return Status::OK();
        }

        //step 5: save index file
635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650
        try {
            index->Serialize();
        } catch (std::exception& ex) {
            //typical error: out of disk space or permition denied
            std::string msg = "Serialize index encounter exception" + std::string(ex.what());
            ENGINE_LOG_ERROR << msg;

            table_file.file_type_ = meta::TableFileSchema::TO_DELETE;
            status = meta_ptr_->UpdateTableFile(table_file);
            ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";

            std::cout << "ERROR: failed to persist index file: " << table_file.location_
                << ", possible out of disk space" << std::endl;

            return Status::Error(msg);
        }
G
groot 已提交
651 652

        //step 6: update meta
G
groot 已提交
653
        table_file.file_type_ = meta::TableFileSchema::INDEX;
654 655
        table_file.file_size_ = index->PhysicalSize();
        table_file.row_count_ = index->Count();
X
Xu Peng 已提交
656

657 658
        auto origin_file = file;
        origin_file.file_type_ = meta::TableFileSchema::BACKUP;
X
Xu Peng 已提交
659

660
        meta::TableFilesSchema update_files = {table_file, origin_file};
661 662 663 664
        status = meta_ptr_->UpdateTableFiles(update_files);
        if(status.ok()) {
            ENGINE_LOG_DEBUG << "New index file " << table_file.file_id_ << " of size "
                             << index->PhysicalSize() << " bytes"
665
                             << " from file " << origin_file.file_id_;
X
Xu Peng 已提交
666

667 668 669 670 671
            if(options_.insert_cache_immediately_) {
                index->Cache();
            }
        } else {
            //failed to update meta, mark the new file as to_delete, don't delete old file
672 673 674
            origin_file.file_type_ = meta::TableFileSchema::TO_INDEX;
            status = meta_ptr_->UpdateTableFile(origin_file);
            ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << origin_file.file_id_ << " to to_index";
675 676 677 678

            table_file.file_type_ = meta::TableFileSchema::TO_DELETE;
            status = meta_ptr_->UpdateTableFile(table_file);
            ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";
S
starlord 已提交
679
        }
G
groot 已提交
680 681

    } catch (std::exception& ex) {
S
starlord 已提交
682 683 684
        std::string msg = "Build index encounter exception" + std::string(ex.what());
        ENGINE_LOG_ERROR << msg;
        return Status::Error(msg);
G
groot 已提交
685
    }
X
Xu Peng 已提交
686

X
Xu Peng 已提交
687 688 689
    return Status::OK();
}

G
groot 已提交
690
void DBImpl::BackgroundBuildIndex() {
691
    ENGINE_LOG_TRACE << " Background build index thread start";
S
starlord 已提交
692

P
peng.xu 已提交
693
    std::unique_lock<std::mutex> lock(build_index_mutex_);
694
    meta::TableFilesSchema to_index_files;
G
groot 已提交
695
    meta_ptr_->FilesToIndex(to_index_files);
X
Xu Peng 已提交
696 697
    Status status;
    for (auto& file : to_index_files) {
X
Xu Peng 已提交
698
        status = BuildIndex(file);
X
Xu Peng 已提交
699
        if (!status.ok()) {
S
starlord 已提交
700
            ENGINE_LOG_ERROR << "Building index for " << file.id_ << " failed: " << status.ToString();
X
Xu Peng 已提交
701
            return;
X
Xu Peng 已提交
702
        }
703

G
groot 已提交
704
        if (shutting_down_.load(std::memory_order_acquire)){
S
starlord 已提交
705
            ENGINE_LOG_DEBUG << "Server will shutdown, skip build index action";
G
groot 已提交
706
            break;
X
Xu Peng 已提交
707
        }
708
    }
S
starlord 已提交
709

710
    ENGINE_LOG_TRACE << " Background build index thread exit";
X
Xu Peng 已提交
711 712
}

G
groot 已提交
713
Status DBImpl::DropAll() {
G
groot 已提交
714
    return meta_ptr_->DropAll();
X
Xu Peng 已提交
715 716
}

G
groot 已提交
717
Status DBImpl::Size(uint64_t& result) {
G
groot 已提交
718
    return  meta_ptr_->Size(result);
X
Xu Peng 已提交
719 720
}

G
groot 已提交
721
DBImpl::~DBImpl() {
G
groot 已提交
722
    shutting_down_.store(true, std::memory_order_release);
X
Xu Peng 已提交
723
    bg_timer_thread_.join();
G
groot 已提交
724
    std::set<std::string> ids;
G
groot 已提交
725
    mem_mgr_->Serialize(ids);
X
Xu Peng 已提交
726 727
}

X
Xu Peng 已提交
728
} // namespace engine
J
jinhai 已提交
729
} // namespace milvus
X
Xu Peng 已提交
730
} // namespace zilliz