tair-kvcache/kv_cache_manager/data_storage/data_storage_manager.h at e6ee490350244407e4a4337a205d785cdb39dd90 · alibaba/tair-kvcache · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#pragma once

#include <map>
#include <memory>
#include <shared_mutex>
#include <string>
#include <thread>
#include <vector>

#include "kv_cache_manager/common/request_context.h"
#include "kv_cache_manager/data_storage/data_storage_backend.h"

namespace kv_cache_manager {

class MetricsRegistry;

class DataStorageManager {
public:
    DataStorageManager() = delete;
    explicit DataStorageManager(std::shared_ptr<MetricsRegistry> metrics_registry);
    ~DataStorageManager() = default;

    DataStorageManager(const DataStorageManager &) = delete;
    DataStorageManager &operator=(const DataStorageManager &) = delete;

public:
    std::vector<std::string> GetAllStorageNames() const;
    std::vector<std::shared_ptr<DataStorageBackend>> GetAvailableStorages();
    std::shared_ptr<DataStorageBackend> GetDataStorageBackend(const std::string &name);
    std::vector<StorageConfig> ListStorageConfig();

    ErrorCode EnableStorage(const std::string &name);
    ErrorCode DisableStorage(const std::string &name);
    ErrorCode
    RegisterStorage(RequestContext *request_context, const std::string &name, const StorageConfig &storage_config);
    ErrorCode UnRegisterStorage(const std::string &name);
    ErrorCode DoCleanup();

    std::vector<SpecCreateResult> Create(RequestContext *request_context,
                                        const std::string &unique_name,
                                        const CreateBlocksRequest &request,
                                        std::function<void()> cb);

    std::vector<ErrorCode> Delete(RequestContext *request_context,
                                  const std::string &unique_name,
                                  const std::vector<DataStorageUri> &storage_uris,
                                  std::function<void()> cb);
    std::vector<bool>
    Exist(const std::string &unique_name, const std::vector<DataStorageUri> &storage_uris, bool fastpath = false);
    std::vector<ErrorCode> Lock(const std::string &unique_name, const std::vector<DataStorageUri> &storage_uris);
    std::vector<ErrorCode> UnLock(const std::string &unique_name, const std::vector<DataStorageUri> &storage_uris);

private:
    std::shared_ptr<DataStorageBackend> CreateStorageBackend(const DataStorageType &type);
    std::string ToString(DataStorageType type);

private:
    mutable std::shared_mutex rw_lock_;
    std::thread heartbeat_thread_;
    // stroage unique name -> storage_backend
    std::map<std::string, std::shared_ptr<DataStorageBackend>> storage_map_;
    std::shared_ptr<MetricsRegistry> metrics_registry_;
};

} // namespace kv_cache_manager