text: // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
  #include <vespa/document/base/documentid.h>
  #include <vespa/document/test/make_document_bucket.h>
  #include <vespa/storage/distributor/distributor.h>
  #include <vespa/storage/distributor/idealstatemanager.h>
  #include <vespa/storage/distributor/operation_sequencer.h>
  #include <vespa/storage/distributor/operations/idealstate/splitoperation.h>
  #include <vespa/storageapi/message/bucketsplitting.h>
  #include <vespa/storageapi/message/persistence.h>
  #include <tests/common/dummystoragelink.h>
  #include <tests/distributor/distributortestutil.h>
  #include <vespa/vespalib/gtest/gtest.h>
  
  using document::test::makeDocumentBucket;
  using namespace document;
  using namespace ::testing;
  
  namespace storage::distributor {
  
  struct SplitOperationTest : Test, DistributorTestUtil {
      uint32_t splitByteSize;
      uint32_t tooLargeBucketSize;
      uint32_t splitCount;
      uint32_t maxSplitBits;
  
      SplitOperationTest();
  
      void SetUp() override {
          createLinks();
          getConfig().setSplitCount(splitCount);
          getConfig().setSplitSize(splitByteSize);
  
      }
  
      void TearDown() override {
          close();
      }
  };
  
  SplitOperationTest::SplitOperationTest()
      : splitByteSize(10*1024*1024),
        tooLargeBucketSize(splitByteSize * 1.1),
        splitCount(UINT32_MAX),
        maxSplitBits(58)
  {
  }
  
  namespace {
      vespalib::string _Storage("storage");
      api::StorageMessageAddress _Storage0Address(&_Storage, lib::NodeType::STORAGE, 0);
  }
  
  TEST_F(SplitOperationTest, simple) {
      enableDistributorClusterState("distributor:1 storage:1");
  
      insertBucketInfo(document::BucketId(16, 1), 0, 0xabc, 1000,
                       tooLargeBucketSize, 250);
  
      SplitOperation op(&_Storage,
                        BucketAndNodes(makeDocumentBucket(document::BucketId(16, 1)),
                                       toVector<uint16_t>(0)),
                        maxSplitBits,
                        splitCount,
                        splitByteSize);
  
      op.setIdealStateManager(&getIdealStateManager());
      op.start(_sender, framework::MilliSecTime(0));
  
      {
          ASSERT_EQ(1, _sender.commands().size());
  
          std::shared_ptr<api::StorageCommand> msg  = _sender.command(0);
          ASSERT_EQ(msg->getType(), api::MessageType::SPLITBUCKET);
          EXPECT_EQ(_Storage0Address.toString(),
                    msg->getAddress()->toString());
  
          std::shared_ptr<api::StorageReply> reply(msg->makeReply().release());
          auto* sreply = static_cast<api::SplitBucketReply*>(reply.get());
  
          sreply->getSplitInfo().emplace_back(document::BucketId(17, 1),
                                              api::BucketInfo(100, 600, 5000000));
  
          sreply->getSplitInfo().emplace_back(document::BucketId(17, 0x10001),
                                              api::BucketInfo(110, 400, 6000000));
  
          op.receive(_sender, reply);
      }
  
      ASSERT_FALSE(getBucket(document::BucketId(16, 1)).valid());
  
      {
          BucketDatabase::Entry entry = getBucket(document::BucketId(17, 1));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(1, entry->getNodeCount());
          EXPECT_EQ(0, entry->getNodeRef(0).getNode());
          EXPECT_EQ(100, entry->getNodeRef(0).getChecksum());
          EXPECT_EQ(5000000, entry->getNodeRef(0).getTotalDocumentSize());
          EXPECT_EQ(600, entry->getNodeRef(0).getDocumentCount());
      }
  
      {
          BucketDatabase::Entry entry(getBucket(document::BucketId(17, 0x10001)));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(1, entry->getNodeCount());
          EXPECT_EQ(0, entry->getNodeRef(0).getNode());
          EXPECT_EQ(110, entry->getNodeRef(0).getChecksum());
          EXPECT_EQ(6000000, entry->getNodeRef(0).getTotalDocumentSize());
          EXPECT_EQ(400, entry->getNodeRef(0).getDocumentCount());
      }
  }
  
  TEST_F(SplitOperationTest, multi_node_failure) {
      {
          BucketDatabase::Entry entry(document::BucketId(16, 1));
  
          BucketCopy copy(0, 0, api::BucketInfo(250, 1000, tooLargeBucketSize));
          entry->addNode(copy, toVector<uint16_t>(0));
  
          entry->addNode(BucketCopy(0, 1, copy.getBucketInfo()),
                         toVector<uint16_t>(0));
          getBucketDatabase().update(entry);
      }
  
      enableDistributorClusterState("distributor:1 storage:2");
  
      SplitOperation op(&_Storage,
                        BucketAndNodes(makeDocumentBucket(document::BucketId(16, 1)),
                                       toVector<uint16_t>(0,1)),
                        maxSplitBits,
                        splitCount,
                        splitByteSize);
  
      op.setIdealStateManager(&getIdealStateManager());
      op.start(_sender, framework::MilliSecTime(0));
  
      {
          ASSERT_EQ(2, _sender.commands().size());
  
          {
              std::shared_ptr<api::StorageCommand> msg  = _sender.command(0);
              ASSERT_EQ(msg->getType(), api::MessageType::SPLITBUCKET);
              EXPECT_EQ(_Storage0Address.toString(),
                        msg->getAddress()->toString());
  
              auto* sreply = static_cast<api::SplitBucketReply*>(msg->makeReply().release());
              sreply->setResult(api::ReturnCode::OK);
  
              sreply->getSplitInfo().emplace_back(document::BucketId(17, 1),
                                                  api::BucketInfo(100, 600, 5000000));
  
              sreply->getSplitInfo().emplace_back(document::BucketId(17, 0x10001),
                                                  api::BucketInfo(110, 400, 6000000));
  
              op.receive(_sender, std::shared_ptr<api::StorageReply>(sreply));
          }
  
          sendReply(op, 1, api::ReturnCode::NOT_CONNECTED);
      }
  
      {
          BucketDatabase::Entry entry = getBucket(document::BucketId(16, 1));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(1, entry->getNodeCount());
  
          EXPECT_EQ(1, entry->getNodeRef(0).getNode());
          EXPECT_EQ(250, entry->getNodeRef(0).getChecksum());
          EXPECT_EQ(tooLargeBucketSize, entry->getNodeRef(0).getTotalDocumentSize());
          EXPECT_EQ(1000, entry->getNodeRef(0).getDocumentCount());
      }
  
      {
          BucketDatabase::Entry entry = getBucket(document::BucketId(17, 1));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(1, entry->getNodeCount());
  
          EXPECT_EQ(0, entry->getNodeRef(0).getNode());
          EXPECT_EQ(100, entry->getNodeRef(0).getChecksum());
          EXPECT_EQ(5000000, entry->getNodeRef(0).getTotalDocumentSize());
          EXPECT_EQ(600, entry->getNodeRef(0).getDocumentCount());
      }
  
      {
          BucketDatabase::Entry entry(getBucket(document::BucketId(17, 0x10001)));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(1, entry->getNodeCount());
  
          EXPECT_EQ(0, entry->getNodeRef(0).getNode());
          EXPECT_EQ(110, entry->getNodeRef(0).getChecksum());
          EXPECT_EQ(6000000, entry->getNodeRef(0).getTotalDocumentSize());
          EXPECT_EQ(400, entry->getNodeRef(0).getDocumentCount());
      }
  }
  
  TEST_F(SplitOperationTest, copy_trusted_status_not_carried_over_after_split) {
      enableDistributorClusterState("distributor:1 storage:2");
  
      document::BucketId sourceBucket(16, 1);
      /*
       * Need 3 nodes to reproduce bug 6418516. Otherwise, the source bucket is
       * left with only 1 copy which implicitly becomes trusted. When this copy
       * is then split, the distributor db will automatically un-trust all buckets
       * since it sees that multiple copies are trusted that are not consistent
       * with each other. This prevents the bug from being visible.
       */
      addNodesToBucketDB(sourceBucket, "0=150/20/30000000/t,1=450/50/60000/u,"
                                       "2=550/60/70000");
  
      SplitOperation op(&_Storage,
                        BucketAndNodes(makeDocumentBucket(sourceBucket), toVector<uint16_t>(0, 1)),
                        maxSplitBits,
                        splitCount,
                        splitByteSize);
  
      op.setIdealStateManager(&getIdealStateManager());
      op.start(_sender, framework::MilliSecTime(0));
  
      ASSERT_EQ(3, _sender.commands().size());
  
      std::vector<document::BucketId> childBuckets;
      childBuckets.emplace_back(17, 1);
      childBuckets.emplace_back(17, 0x10001);
  
      // Note: only 2 out of 3 requests replied to!
      for (int i = 0; i < 2; ++i) {
          std::shared_ptr<api::StorageCommand> msg  = _sender.command(i);
          ASSERT_EQ(msg->getType(), api::MessageType::SPLITBUCKET);
          std::shared_ptr<api::StorageReply> reply(msg->makeReply().release());
          auto* sreply = static_cast<api::SplitBucketReply*>(reply.get());
  
          // Make sure copies differ so they cannot become implicitly trusted.
          sreply->getSplitInfo().emplace_back(childBuckets[0],
                                              api::BucketInfo(100 + i, 600, 5000000));
          sreply->getSplitInfo().emplace_back(childBuckets[1],
                                              api::BucketInfo(110 + i, 400, 6000000));
  
          op.receive(_sender, reply);
      }
  
      ASSERT_TRUE(getBucket(sourceBucket).valid()); // Still alive
  
      for (uint32_t i = 0; i < 2; ++i) {
          BucketDatabase::Entry entry(getBucket(childBuckets[i]));
  
          ASSERT_TRUE(entry.valid());
          ASSERT_EQ(2, entry->getNodes().size());
  
          for (uint16_t j = 0; j < 2; ++j) {
              EXPECT_FALSE(entry->getNodeRef(i).trusted());
          }
      }
  }
  
  TEST_F(SplitOperationTest, operation_blocked_by_pending_join) {
      StorageComponentRegisterImpl compReg;
      framework::defaultimplementation::FakeClock clock;
      compReg.setClock(clock);
      clock.setAbsoluteTimeInSeconds(1);
      PendingMessageTracker tracker(compReg);
      OperationSequencer op_seq;
  
      enableDistributorClusterState("distributor:1 storage:2");
  
      document::BucketId joinTarget(2, 1);
      std::vector<document::BucketId> joinSources = {
          document::BucketId(3, 1), document::BucketId(3, 5)
      };
      auto joinCmd = std::make_shared<api::JoinBucketsCommand>(makeDocumentBucket(joinTarget));
      joinCmd->getSourceBuckets() = joinSources;
      joinCmd->setAddress(_Storage0Address);
  
      tracker.insert(joinCmd);
  
      insertBucketInfo(joinTarget, 0, 0xabc, 1000, 1234, true);
  
      SplitOperation op(&_Storage,
                        BucketAndNodes(makeDocumentBucket(joinTarget), toVector<uint16_t>(0)),
                        maxSplitBits,
                        splitCount,
                        splitByteSize);
  
      EXPECT_TRUE(op.isBlocked(tracker, op_seq));
  
      // Now, pretend there's a join for another node in the same bucket. This
      // will happen when a join is partially completed.
      tracker.clearMessagesForNode(0);
      EXPECT_FALSE(op.isBlocked(tracker, op_seq));
  
      joinCmd->setAddress(api::StorageMessageAddress::create(&_Storage, lib::NodeType::STORAGE, 1));
      tracker.insert(joinCmd);
  
      EXPECT_TRUE(op.isBlocked(tracker, op_seq));
  }
  
  TEST_F(SplitOperationTest, split_is_blocked_by_locked_bucket) {
      StorageComponentRegisterImpl compReg;
      framework::defaultimplementation::FakeClock clock;
      compReg.setClock(clock);
      clock.setAbsoluteTimeInSeconds(1);
      PendingMessageTracker tracker(compReg);
      OperationSequencer op_seq;
  
      enableDistributorClusterState("distributor:1 storage:2");
  
      document::BucketId source_bucket(16, 1);
      insertBucketInfo(source_bucket, 0, 0xabc, 1000, tooLargeBucketSize, 250);
  
      SplitOperation op(&_Storage, BucketAndNodes(makeDocumentBucket(source_bucket), toVector<uint16_t>(0)),
                        maxSplitBits, splitCount, splitByteSize);
  
      EXPECT_FALSE(op.isBlocked(tracker, op_seq));
      auto token = op_seq.try_acquire(makeDocumentBucket(source_bucket), "foo");
      EXPECT_TRUE(token.valid());
      EXPECT_TRUE(op.isBlocked(tracker, op_seq));
  }
  
  } // storage::distributor
  