aboutsummaryrefslogtreecommitdiffstats
path: root/storage/src/tests/distributor/distributortestutil.h
blob: 19da0483165b1c19d61cc8ec9c1dd6b836eb5c40 (plain) (blame)
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#pragma once

#include "messagesenderstub.h"
#include <tests/common/teststorageapp.h>
#include <tests/common/testhelper.h>
#include <tests/common/dummystoragelink.h>
#include <vespa/storage/common/hostreporter/hostinfo.h>
#include <vespa/storage/frameworkimpl/component/distributorcomponentregisterimpl.h>
#include <vespa/storage/storageutil/utils.h>
#include <vespa/storageframework/defaultimplementation/clock/fakeclock.h>
#include <vespa/storageapi/message/state.h>

namespace storage {

namespace framework { class TickingThreadPool; }

namespace distributor {

class BucketDBUpdater;
class Distributor;
class DistributorBucketSpace;
class DistributorBucketSpaceRepo;
class IdealStateManager;
class ExternalOperationHandler;
class Operation;

class DistributorTestUtil : private DoneInitializeHandler
{
public:
    DistributorTestUtil();
    ~DistributorTestUtil();

    /**
     * Sets up the storage link chain.
     */
    void createLinks();
    void setTypeRepo(const document::DocumentTypeRepo::SP &repo);

    void close();

    /**
     * Returns a string with the nodes currently stored in the bucket
     * database for the given bucket.
     */
    std::string getNodes(document::BucketId id);

    /**
     * Returns a string with the ideal state nodes for the given bucket.
     */
    std::string getIdealStr(document::BucketId id, const lib::ClusterState& state);

    /**
     * Adds the ideal nodes for the given bucket and the given cluster state
     * to the bucket database.
     */
    void addIdealNodes(const lib::ClusterState& state, const document::BucketId& id);

    /**
     * Adds all the ideal nodes for the given bucket to the bucket database.
     */
    void addIdealNodes(const document::BucketId& id);

    /**
     * Parses the given string to a set of node => bucket info data,
     * and inserts them as nodes in the given bucket.
     * Format:
     *   "node1=checksum/docs/size,node2=checksum/docs/size"
     */
    void addNodesToBucketDB(const document::BucketId& id, const std::string& nodeStr);

   /**
     * Removes the given bucket from the bucket database.
     */
    void removeFromBucketDB(const document::BucketId& id);

    /**
     * Inserts the given bucket information for the given bucket and node in
     * the bucket database.
     */
    void insertBucketInfo(document::BucketId id,
                          uint16_t node,
                          uint32_t checksum,
                          uint32_t count,
                          uint32_t size,
                          bool trusted = false,
                          bool active = false);

    /**
     * Inserts the given bucket information for the given bucket and node in
     * the bucket database.
     */
    void insertBucketInfo(document::BucketId id,
                          uint16_t node,
                          const api::BucketInfo& info,
                          bool trusted = false,
                          bool active = false);

    std::string dumpBucket(const document::BucketId& bucket);

    /**
     * Replies to message idx sent upwards with the given result code.
     * If idx = -1, replies to the last command received upwards.
     */
    void sendReply(Operation& op,
                   int idx = -1,
                   api::ReturnCode::Result result = api::ReturnCode::OK);

    BucketDBUpdater& getBucketDBUpdater();
    IdealStateManager& getIdealStateManager();
    ExternalOperationHandler& getExternalOperationHandler();

    Distributor& getDistributor() {
        return *_distributor;
    }

    bool tick();

    DistributorConfiguration& getConfig();

    vdstestlib::DirConfig& getDirConfig() {
        return _config;
    }

    // TODO explicit notion of bucket spaces for tests
    DistributorBucketSpace &getDistributorBucketSpace();
    BucketDatabase& getBucketDatabase();
    const BucketDatabase& getBucketDatabase() const;
    DistributorBucketSpaceRepo &getBucketSpaceRepo();
    const DistributorBucketSpaceRepo &getBucketSpaceRepo() const;

    const lib::Distribution& getDistribution() const;
    // "End to end" distribution change trigger, which will invoke the bucket
    // DB updater as expected based on the previous and new cluster state
    // and config.
    void triggerDistributionChange(lib::Distribution::SP distr);
    
    framework::defaultimplementation::FakeClock& getClock() { return _node->getClock(); }
    DistributorComponentRegister& getComponentRegister() { return _node->getComponentRegister(); }
    DistributorComponentRegisterImpl& getComponentRegisterImpl() { return _node->getComponentRegister(); }

    StorageComponent& getComponent() {
        if (_component.get() == 0) {
            _component.reset(new storage::DistributorComponent(
                    _node->getComponentRegister(), "distributor_test_utils"));
        }
        return *_component;
    }

    void setupDistributor(int redundancy,
                          int nodeCount,
                          const std::string& systemState,
                          uint32_t earlyReturn = false,
                          bool requirePrimaryToBeWritten = true);

    void setRedundancy(uint32_t redundancy);

    void notifyDoneInitializing() override {}

        // Must implement this for storage server interface for now
    virtual api::Timestamp getUniqueTimestamp() {
        return _component->getUniqueTimestamp();
    }

    void disableBucketActivationInConfig(bool disable);

    BucketDatabase::Entry getBucket(const document::BucketId& bId) const;

protected:
    vdstestlib::DirConfig _config;
    std::unique_ptr<TestDistributorApp> _node;
    std::unique_ptr<framework::TickingThreadPool> _threadPool;
    std::unique_ptr<Distributor> _distributor;
    std::unique_ptr<storage::DistributorComponent> _component;
    MessageSenderStub _sender;
    MessageSenderStub _senderDown;
    HostInfo _hostInfo;

    struct MessageSenderImpl : public ChainedMessageSender {
        MessageSenderStub& _sender;
        MessageSenderStub& _senderDown;
        MessageSenderImpl(MessageSenderStub& up, MessageSenderStub& down)
            : _sender(up), _senderDown(down) {}

        void sendUp(const std::shared_ptr<api::StorageMessage>& msg) override {
            _sender.send(msg);
        }
        void sendDown(const std::shared_ptr<api::StorageMessage>& msg) override {
            _senderDown.send(msg);
        }
    };
    MessageSenderImpl _messageSender;
};

}

}