pikiwidb/include/pika_cmd_table_manager.h at 1f046c452eeb7b30827bcd7fbb09d8ac68c4b13b · OpenAtomFoundation/pikiwidb · 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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
// Copyright (c) 2018-present, Qihoo, Inc.  All rights reserved.
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree. An additional grant
// of patent rights can be found in the PATENTS file in the same directory.

#ifndef PIKA_CMD_TABLE_MANAGER_H_
#define PIKA_CMD_TABLE_MANAGER_H_

#include <shared_mutex>
#include <thread>
#include <prometheus/exposer.h>
#include <prometheus/registry.h>
#include <prometheus/counter.h>
#include <prometheus/histogram.h>

#include "include/acl.h"
#include "include/pika_command.h"
#include "include/pika_data_distribution.h"

using namespace prometheus;

struct CommandStatistics {
  CommandStatistics() = default;
  CommandStatistics(const CommandStatistics& other) {
    cmd_time_consuming.store(other.cmd_time_consuming.load());
    cmd_count.store(other.cmd_count.load());
  }
  std::atomic<uint64_t> cmd_count = 0;
  std::atomic<uint64_t> cmd_time_consuming = 0;
};

struct HistogramData {
  std::shared_ptr<prometheus::Registry> registry;
  prometheus::Family<prometheus::Histogram>* family;
  std::unordered_map<std::string, prometheus::Histogram*> histograms;

  HistogramData() {
    registry = std::make_shared<prometheus::Registry>();
    family = &prometheus::BuildHistogram()
        .Name("pika_command_duration_seconds")
        .Help("Execution time of Pika commands in seconds")
        .Register(*registry);
  }

  HistogramData(const HistogramData&) = delete;
  HistogramData& operator=(const HistogramData&) = delete;
};

class PikaCmdTableManager {
  friend AclSelector;

 public:
  PikaCmdTableManager();
  virtual ~PikaCmdTableManager() = default;
  void InitCmdTable(void);
  void RenameCommand(const std::string before, const std::string after);
  std::shared_ptr<Cmd> GetCmd(const std::string& opt);
  bool CmdExist(const std::string& cmd) const;
  CmdTable* GetCmdTable();
  uint32_t GetMaxCmdId();

  std::vector<std::string> GetAclCategoryCmdNames(uint32_t flag);

  /*
  * Info Commandstats used
  */
  std::unordered_map<std::string, CommandStatistics>* GetCommandStatMap();
  std::unordered_map<std::string, CommandStatistics> GetSlowCommandCount();
  void UpdateSlowCommandCount(const std::string& opt);
  void ResetCommandCount();
  prometheus::Histogram& GetHistogram(const std::string& opt);
  std::shared_ptr<HistogramData> GetHistogramsData();

 private:
  std::shared_ptr<Cmd> NewCommand(const std::string& opt);

  void InsertCurrentThreadDistributionMap();
  bool CheckCurrentThreadDistributionMapExist(const std::thread::id& tid);

  std::unique_ptr<CmdTable> cmds_;

  uint32_t cmdId_ = 0;

  std::shared_mutex map_protector_;
  std::unordered_map<std::thread::id, std::unique_ptr<PikaDataDistribution>> thread_distribution_map_;

  /*
  * Info Commandstats used
  */
  std::unordered_map<std::string, CommandStatistics> cmdstat_map_;
  std::unordered_map<std::string, CommandStatistics> slow_command_count_;
  std::shared_mutex slow_command_mutex_;
  std::shared_mutex histograms_mutex_;
  std::mutex data_mutex_;
  std::shared_ptr<HistogramData> data_;
};
#endif