YugabyteDB (2.13.0.0-b42, bfc6a6643e7399ac8a0e81d06a3ee6d6571b33ab)

Coverage Report

Created: 2022-03-09 17:30

/Users/deen/code/yugabyte-db/src/yb/master/cluster_balance_mocked.h
Line
Count
Source (jump to first uncovered line)
1
// Copyright (c) YugaByte, Inc.
2
3
#ifndef YB_MASTER_CLUSTER_BALANCE_MOCKED_H
4
#define YB_MASTER_CLUSTER_BALANCE_MOCKED_H
5
6
#include "yb/master/cluster_balance.h"
7
#include "yb/master/cluster_balance_util.h"
8
#include "yb/master/ts_manager.h"
9
10
namespace yb {
11
namespace master {
12
13
class ClusterLoadBalancerMocked : public ClusterLoadBalancer {
14
 public:
15
2
  explicit ClusterLoadBalancerMocked(Options* options) : ClusterLoadBalancer(nullptr)  {
16
2
    const int kHighNumber = 100;
17
2
    options->kMaxConcurrentAdds = kHighNumber;
18
2
    options->kMaxConcurrentRemovals = kHighNumber;
19
2
    options->kAllowLimitStartingTablets = false;
20
2
    options->kAllowLimitOverReplicatedTablets = false;
21
22
2
    auto table_state = std::make_unique<PerTableLoadState>(global_state_.get());
23
2
    table_state->options_ = options;
24
2
    state_ = table_state.get();
25
2
    per_table_states_[""] = std::move(table_state);
26
27
2
    SetOptions(LIVE, "");
28
29
2
    InitTablespaceManager();
30
2
  }
31
32
  // Overrides for base class functionality to bypass calling CatalogManager.
33
73
  void GetAllReportedDescriptors(TSDescriptorVector* ts_descs) const override {
34
73
    *ts_descs = ts_descs_;
35
73
  }
36
37
73
  void GetAllAffinitizedZones(AffinitizedZonesSet* affinitized_zones) const override {
38
73
    *affinitized_zones = affinitized_zones_;
39
73
  }
40
41
2.46k
  const TabletInfoMap& GetTabletMap() const override { return tablet_map_; }
42
43
0
  const TableInfoMap& GetTableMap() const override { return table_map_; }
44
45
73
  const scoped_refptr<TableInfo> GetTableInfo(const TableId& table_uuid) const override {
46
73
    return FindPtrOrNull(table_map_, table_uuid);
47
73
  }
48
49
62
  const ReplicationInfoPB& GetClusterReplicationInfo() const override {
50
62
    return replication_info_;
51
62
  }
52
53
3
  const PlacementInfoPB& GetClusterPlacementInfo() const override {
54
3
    return state_->options_->type == LIVE ?
55
3
        replication_info_.live_replicas() : replication_info_.read_replicas(0);
56
3
  }
57
58
73
  const BlacklistPB& GetServerBlacklist() const override { return blacklist_; }
59
73
  const BlacklistPB& GetLeaderBlacklist() const override { return leader_blacklist_; }
60
61
  Status SendReplicaChanges(scoped_refptr<TabletInfo> tablet, const TabletServerId& ts_uuid,
62
                          const bool is_add, const bool should_remove,
63
132
                          const TabletServerId& new_leader_uuid) override {
64
    // Do nothing.
65
132
    return Status::OK();
66
132
  }
67
68
  void GetPendingTasks(const TableId& table_uuid,
69
                       TabletToTabletServerMap* pending_add_replica_tasks,
70
                       TabletToTabletServerMap* pending_remove_replica_tasks,
71
4
                       TabletToTabletServerMap* pending_stepdown_leader_tasks) override {
72
8
    for (const auto& tablet_id : pending_add_replica_tasks_) {
73
8
      (*pending_add_replica_tasks)[tablet_id] = "";
74
8
    }
75
8
    for (const auto& tablet_id : pending_remove_replica_tasks_) {
76
8
      (*pending_remove_replica_tasks)[tablet_id] = "";
77
8
    }
78
0
    for (const auto& tablet_id : pending_stepdown_leader_tasks_) {
79
0
      (*pending_stepdown_leader_tasks)[tablet_id] = "";
80
0
    }
81
4
  }
82
83
68
  void ResetTableStatePtr(const TableId& table_id, Options* options) override {
84
68
    if (state_) {
85
68
      options = state_->options_;
86
68
    }
87
68
    auto table_state = std::make_unique<PerTableLoadState>(global_state_.get());
88
68
    table_state->options_ = options;
89
68
    table_state->check_ts_liveness_ = false;
90
68
    state_ = table_state.get();
91
92
68
    per_table_states_[table_id] = std::move(table_state);
93
68
  }
94
95
2
  void InitTablespaceManager() override {
96
2
    tablespace_manager_ = std::make_shared<YsqlTablespaceManager>(nullptr, nullptr);
97
2
  }
98
99
4
  void SetOptions(ReplicaType type, const string& placement_uuid) {
100
4
    state_->options_->type = type;
101
4
    state_->options_->placement_uuid = placement_uuid;
102
4
  }
103
104
  TSDescriptorVector ts_descs_;
105
  AffinitizedZonesSet affinitized_zones_;
106
  TabletInfoMap tablet_map_;
107
  TableInfoMap table_map_;
108
  ReplicationInfoPB replication_info_;
109
  BlacklistPB blacklist_;
110
  BlacklistPB leader_blacklist_;
111
  vector<TabletId> pending_add_replica_tasks_;
112
  vector<TabletId> pending_remove_replica_tasks_;
113
  vector<TabletId> pending_stepdown_leader_tasks_;
114
115
  friend class TestLoadBalancerEnterprise;
116
};
117
118
} // namespace master
119
} // namespace yb
120
121
#endif // YB_MASTER_CLUSTER_BALANCE_MOCKED_H