aboutsummaryrefslogtreecommitdiffstats
path: root/searchcore/src/tests/proton/documentdb/documentbucketmover/bucketmover_common.h
blob: 9d5c4e0e4ec907921bd89ca8c05df65f15f088f0 (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
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include <vespa/searchcore/proton/bucketdb/bucketdbhandler.h>
#include <vespa/searchcore/proton/bucketdb/bucket_create_notifier.h>
#include <vespa/searchcore/proton/test/bucketfactory.h>
#include <vespa/searchcore/proton/feedoperation/moveoperation.h>
#include <vespa/searchcore/proton/server/i_move_operation_limiter.h>
#include <vespa/searchcore/proton/server/idocumentmovehandler.h>
#include <vespa/searchcore/proton/server/imaintenancejobrunner.h>
#include <vespa/searchcore/proton/server/maintenancedocumentsubdb.h>
#include <vespa/searchcore/proton/server/ibucketmodifiedhandler.h>
#include <vespa/searchcore/proton/server/i_maintenance_job.h>
#include <vespa/searchcore/proton/test/buckethandler.h>
#include <vespa/searchcore/proton/test/clusterstatehandler.h>
#include <vespa/searchcore/proton/test/disk_mem_usage_notifier.h>
#include <vespa/searchcore/proton/test/test.h>
#include <vespa/document/test/make_bucket_space.h>

namespace proton {
    class DocumentMetaStore;
}
namespace proton::move::test {

struct MyMoveOperationLimiter : public IMoveOperationLimiter {
    uint32_t beginOpCount;

    MyMoveOperationLimiter() : beginOpCount(0) {}

    vespalib::IDestructorCallback::SP beginOperation() override {
        ++beginOpCount;
        return {};
    }
    size_t numPending() const override { return beginOpCount; }
};

struct MyMoveHandler : public IDocumentMoveHandler {
    using MoveOperationVector = std::vector<MoveOperation>;
    bucketdb::BucketDBOwner &_bucketDb;
    MoveOperationVector _moves;
    size_t _numCachedBuckets;
    bool _storeMoveDoneContexts;
    std::vector<vespalib::IDestructorCallback::SP> _moveDoneContexts;

    MyMoveHandler(bucketdb::BucketDBOwner &bucketDb, bool storeMoveDoneContext = false);
    ~MyMoveHandler() override;
    void handleMove(MoveOperation &op, vespalib::IDestructorCallback::SP moveDoneCtx) override;

    void reset() {
        _moves.clear();
        _numCachedBuckets = 0;
    }

    void clearMoveDoneContexts() {
        _moveDoneContexts.clear();
    }
};

struct MyDocumentRetriever : public DocumentRetrieverBaseForTest {
    using DocumentTypeRepo = document::DocumentTypeRepo;
    using DocumentMetaData = search::DocumentMetaData;
    using Document = document::Document;
    using DocumentId = document::DocumentId;
    using DocumentIdT = search::DocumentIdT;
    using DocumentVector = std::vector<Document::SP>;
    std::shared_ptr<const DocumentTypeRepo> _repo;
    DocumentVector _docs;
    uint32_t _lid2Fail;

    MyDocumentRetriever(std::shared_ptr<const DocumentTypeRepo> repo)
        : _repo(std::move(repo)),
          _docs(),
          _lid2Fail(0)
    {
        _docs.push_back(Document::UP()); // lid 0 invalid
    }

    const DocumentTypeRepo &getDocumentTypeRepo() const override { return *_repo; }

    void getBucketMetaData(const storage::spi::Bucket &, DocumentMetaData::Vector &) const override {}

    DocumentMetaData getDocumentMetaData(const DocumentId &) const override { return DocumentMetaData(); }

    Document::UP getFullDocument(DocumentIdT lid) const override {
        return (lid != _lid2Fail) ? Document::UP(_docs[lid]->clone()) : Document::UP();
    }

    void failRetrieveForLid(uint32_t lid) { _lid2Fail = lid; }

    CachedSelect::SP parseSelect(const vespalib::string &) const override {
        return {};
    }
};

struct MyBucketModifiedHandler : public IBucketModifiedHandler {
    using BucketId = document::BucketId;
    std::vector<BucketId> _modified;

    ~MyBucketModifiedHandler() override;
    void notifyBucketModified(const BucketId &bucket) override;

    void reset() { _modified.clear(); }
};

struct MySubDb {
    using BucketId = document::BucketId;
    using Document = document::Document;
    using DocumentTypeRepo = document::DocumentTypeRepo;
    using DocumentVector = proton::test::DocumentVector;
    using UserDocuments = proton::test::UserDocuments;
    std::shared_ptr<DocumentMetaStore>    _metaStoreSP;
    DocumentMetaStore                    &_metaStore;
    std::shared_ptr<MyDocumentRetriever>  _realRetriever;
    std::shared_ptr<IDocumentRetriever>   _retriever;
    MaintenanceDocumentSubDB              _subDb;
    UserDocuments                         _docs;
    bucketdb::BucketDBHandler             _bucketDBHandler;

    MySubDb(const std::shared_ptr<const DocumentTypeRepo> &repo, std::shared_ptr<bucketdb::BucketDBOwner> bucketDB,
            uint32_t subDbId, SubDbType subDbType);

    ~MySubDb();

    void insertDocs(const UserDocuments &docs_);

    void failRetrieveForLid(uint32_t lid) {
        _realRetriever->failRetrieveForLid(lid);
    }

    BucketId bucket(uint32_t userId) const {
        return _docs.getBucket(userId);
    }

    DocumentVector docs(uint32_t userId) {
        return _docs.getGidOrderDocs(userId);
    }

    void setBucketState(const BucketId &bucketId, bool active);
};

struct MyCountJobRunner : public IMaintenanceJobRunner {
    uint32_t runCount;
    explicit MyCountJobRunner(IMaintenanceJob &job) : runCount(0) {
        job.registerRunner(this);
    }
    void run() override { ++runCount; }
};

bool
assertEqual(const document::BucketId &bucket, const proton::test::Document &doc,
            uint32_t sourceSubDbId, uint32_t targetSubDbId, const MoveOperation &op);

}