// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "content/browser/indexed_db/instance/bucket_context.h"

#include <memory>
#include <string>

#include "base/run_loop.h"
#include "base/strings/stringprintf.h"
#include "base/test/metrics/histogram_tester.h"
#include "base/test/test_future.h"
#include "components/services/storage/privileged/mojom/indexed_db_internals_types.mojom.h"
#include "content/browser/indexed_db/indexed_db_test_base.h"
#include "content/browser/indexed_db/instance/fake_transaction.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/blink/public/common/storage_key/storage_key.h"
#include "url/gurl.h"

namespace content::indexed_db {

using testing::SizeIs;
using ITS = storage::mojom::IdbTransactionState;

class BucketContextTest : public IndexedDBTestBase {
 public:
  BucketContextTest()
      : IndexedDBTestBase(/*use_default_buckets=*/true, /*use_sqlite=*/false) {}

  void SetUp() override {
    IndexedDBTestBase::SetUp();
    bucket_context_ = InitBucketContext(GetOrCreateBucket(GetTestStorageKey()),
                                        /*create_backing_store=*/false);
  }

  void SetQuotaLeft(int64_t quota_manager_response) {
    quota_manager_->SetQuota(bucket_context_->bucket_locator().storage_key,
                             quota_manager_response);
  }

 protected:
  base::WeakPtr<BucketContext> bucket_context_;
};

TEST_F(BucketContextTest, CanUseDiskSpaceQueuing) {
  base::HistogramTester tester;
  // Request space 3 times consecutively. The requests should coalesce.
  SetQuotaLeft(100);

  base::test::TestFuture<bool> success_future;
  base::test::TestFuture<bool> success_future2;
  base::test::TestFuture<bool> success_future3;
  bucket_context_->CheckCanUseDiskSpace(30, success_future.GetCallback());
  bucket_context_->CheckCanUseDiskSpace(30, success_future2.GetCallback());
  bucket_context_->CheckCanUseDiskSpace(50, success_future3.GetCallback());
  ASSERT_FALSE(success_future.IsReady());
  ASSERT_FALSE(success_future2.IsReady());
  ASSERT_FALSE(success_future3.IsReady());

  EXPECT_TRUE(success_future.Get());

  // We know these requests coalesced because only the first request waited
  // (via `Get()`), yet all 3 requests are now ready. The first two requests
  // succeed but the third fails.
  ASSERT_TRUE(success_future2.IsReady());
  ASSERT_TRUE(success_future3.IsReady());
  EXPECT_TRUE(success_future2.Get());
  EXPECT_FALSE(success_future3.Get());

  tester.ExpectTotalCount("IndexedDB.QuotaCheckTime2.Success", 1);
}

TEST_F(BucketContextTest, CanUseDiskSpaceCaching) {
  // Verify the limited authority that BucketContext has to approve
  // disk usage without checking the quota manager. First set the quota manager
  // to report a large amount of disk space, but request even more --- the usage
  // shouldn't be approved.
  constexpr const int64_t kLargeSpace = 120;
  SetQuotaLeft(kLargeSpace);
  {
    base::test::TestFuture<bool> success_future;
    bucket_context_->CheckCanUseDiskSpace(kLargeSpace + 1,
                                          success_future.GetCallback());
    EXPECT_FALSE(success_future.IsReady());
    EXPECT_FALSE(success_future.Get());
  }

  // Second, simulate something using up a lot of the quota.
  // `CheckCanUseDiskSpace` will fudge and not check with the QuotaManager, so
  // this usage is also approved.
  SetQuotaLeft(10);
  {
    base::test::TestFuture<bool> success_future;
    bucket_context_->CheckCanUseDiskSpace(kLargeSpace / 2 + 1,
                                          success_future.GetCallback());
    EXPECT_TRUE(success_future.IsReady());
    EXPECT_TRUE(success_future.Get());
  }
  // Next, request the same amount of space again. `CheckCanUseDiskSpace` does
  // need to double check with the QuotaManager this time as its limited
  // authority has been exhausted, and hence this usage will not be approved.
  {
    base::test::TestFuture<bool> success_future;
    bucket_context_->CheckCanUseDiskSpace(kLargeSpace / 2 + 1,
                                          success_future.GetCallback());
    EXPECT_FALSE(success_future.IsReady());
    EXPECT_FALSE(success_future.Get());
  }

  // Set a large amount of disk space again, request a little.
  SetQuotaLeft(kLargeSpace);
  {
    base::test::TestFuture<bool> success_future;
    bucket_context_->CheckCanUseDiskSpace(kLargeSpace / 4,
                                          success_future.GetCallback());
    EXPECT_FALSE(success_future.IsReady());
    EXPECT_TRUE(success_future.Get());
  }
  // Wait for the cached value to expire. The next request should be approved,
  // but only after going to the QuotaManager again.
  task_environment_.FastForwardBy(BucketContext::kBucketSpaceCacheTimeLimit *
                                  2);
  {
    base::test::TestFuture<bool> success_future;
    bucket_context_->CheckCanUseDiskSpace(kLargeSpace / 4,
                                          success_future.GetCallback());
    EXPECT_FALSE(success_future.IsReady());
    EXPECT_TRUE(success_future.Get());
  }
}

TEST_F(BucketContextTest, CanUseDiskSpaceWarmUp) {
  constexpr const int64_t kLargeSpace = 120;
  SetQuotaLeft(kLargeSpace);

  bucket_context_->CheckCanUseDiskSpace(120, {});
  base::RunLoop().RunUntilIdle();

  base::test::TestFuture<bool> success_future;
  bucket_context_->CheckCanUseDiskSpace(120, success_future.GetCallback());
  EXPECT_TRUE(success_future.IsReady());
  EXPECT_TRUE(success_future.Get());
}

TEST_F(BucketContextTest, BucketSpaceDecay) {
  constexpr const int64_t kLargeSpace = 120;
  SetQuotaLeft(kLargeSpace);

  base::test::TestFuture<bool> success_future;
  bucket_context_->CheckCanUseDiskSpace(1, success_future.GetCallback());
  EXPECT_TRUE(success_future.Get());

  const int64_t can_allot = bucket_context_->GetBucketSpaceToAllot();
  EXPECT_LE(can_allot, 120);

  task_environment_.FastForwardBy(BucketContext::kBucketSpaceCacheTimeLimit /
                                  2);
  const int64_t decayed_can_allot = bucket_context_->GetBucketSpaceToAllot();
  EXPECT_LT(decayed_can_allot, can_allot);
  EXPECT_GT(decayed_can_allot, 0);

  task_environment_.FastForwardBy(BucketContext::kBucketSpaceCacheTimeLimit /
                                  2);
  EXPECT_EQ(bucket_context_->GetBucketSpaceToAllot(), 0);
}

// Verifies state history is calculated correctly based on snapshots.
TEST_F(BucketContextTest, MetadataRecordingStateHistory) {
  bucket_context_->StartMetadataRecording();

  // Helper function to make a idb internals metadata snapshot with a single
  // transaction.
  auto make_snapshot = [](std::u16string db_name, int64_t tid, ITS state,
                          double age) {
    storage::mojom::IdbBucketMetadataPtr metadata =
        storage::mojom::IdbBucketMetadata::New();
    auto db = storage::mojom::IdbDatabaseMetadata::New();
    db->name = db_name;
    auto tx = storage::mojom::IdbTransactionMetadata::New();
    tx->tid = tid;
    tx->state = state;
    tx->age = age;
    tx->connection_id = 0;
    db->transactions.push_back(std::move(tx));
    metadata->databases.push_back(std::move(db));

    // Add to a second parallel connection_id to test that it doesn't interfere.
    db = storage::mojom::IdbDatabaseMetadata::New();
    db->name = db_name;
    tx = storage::mojom::IdbTransactionMetadata::New();
    tx->connection_id = 1;
    tx->tid = tid;
    tx->state = state;
    tx->age = age * 2;
    db->transactions.push_back(std::move(tx));
    metadata->databases.push_back(std::move(db));

    return metadata;
  };

  bucket_context_->metadata_recording_buffer_.push_back(
      make_snapshot(u"database0", /* tid = */ 1, ITS::kStarted,
                    /* age = */ 0));

  // Add another transaction with a different tid to ensure it does not
  // interfere.
  auto tx = storage::mojom::IdbTransactionMetadata::New();
  tx->tid = 2;
  tx->state = ITS::kRunning;
  tx->age = 4;
  bucket_context_->metadata_recording_buffer_[1]
      ->databases[0]
      ->transactions.push_back(std::move(tx));

  bucket_context_->metadata_recording_buffer_.push_back(
      make_snapshot(u"database0", /* tid = */ 1, ITS::kRunning,
                    /* age = */ 10));
  bucket_context_->metadata_recording_buffer_.push_back(
      make_snapshot(u"database0", /* tid = */ 1, ITS::kCommitting,
                    /* age = */ 20));
  bucket_context_->metadata_recording_buffer_.push_back(
      make_snapshot(u"database0", /* tid = */ 1, ITS::kRunning,
                    /* age = */ 30));
  bucket_context_->metadata_recording_buffer_.push_back(
      make_snapshot(u"database0", /* tid = */ 1, ITS::kFinished,
                    /* age = */ 50));

  std::vector<storage::mojom::IdbBucketMetadataPtr> result =
      bucket_context_->StopMetadataRecording();
  EXPECT_THAT(result, SizeIs(6));

  // Snapshot 0
  EXPECT_THAT(result[0]->databases, SizeIs(0));

  // Snapshot 1.
  tx = std::move(result[1]->databases[0]->transactions[0]);
  EXPECT_THAT(tx->state_history, SizeIs(1));
  EXPECT_EQ(tx->state_history[0]->state, ITS::kStarted);
  EXPECT_EQ(tx->state_history[0]->duration, 0);

  // Snapshot 2.
  tx = std::move(result[2]->databases[0]->transactions[0]);
  EXPECT_THAT(tx->state_history, SizeIs(2));
  EXPECT_EQ(tx->state_history[0]->state, ITS::kStarted);
  EXPECT_EQ(tx->state_history[0]->duration, 10);
  EXPECT_EQ(tx->state_history[1]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[1]->duration, 0);

  // Snapshot 3.
  tx = std::move(result[3]->databases[0]->transactions[0]);
  EXPECT_THAT(tx->state_history, SizeIs(3));
  EXPECT_EQ(tx->state_history[0]->state, ITS::kStarted);
  EXPECT_EQ(tx->state_history[0]->duration, 10);
  EXPECT_EQ(tx->state_history[1]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[1]->duration, 10);
  EXPECT_EQ(tx->state_history[2]->state, ITS::kCommitting);
  EXPECT_EQ(tx->state_history[2]->duration, 0);

  // Snapshot 4.
  tx = std::move(result[4]->databases[0]->transactions[0]);
  EXPECT_THAT(tx->state_history, SizeIs(4));
  EXPECT_EQ(tx->state_history[0]->state, ITS::kStarted);
  EXPECT_EQ(tx->state_history[0]->duration, 10);
  EXPECT_EQ(tx->state_history[1]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[1]->duration, 10);
  EXPECT_EQ(tx->state_history[2]->state, ITS::kCommitting);
  EXPECT_EQ(tx->state_history[2]->duration, 10);
  EXPECT_EQ(tx->state_history[3]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[3]->duration, 0);

  // Snapshot 5.
  tx = std::move(result[5]->databases[0]->transactions[0]);
  EXPECT_THAT(tx->state_history, SizeIs(5));
  EXPECT_EQ(tx->state_history[0]->state, ITS::kStarted);
  EXPECT_EQ(tx->state_history[0]->duration, 10);
  EXPECT_EQ(tx->state_history[1]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[1]->duration, 10);
  EXPECT_EQ(tx->state_history[2]->state, ITS::kCommitting);
  EXPECT_EQ(tx->state_history[2]->duration, 10);
  EXPECT_EQ(tx->state_history[3]->state, ITS::kRunning);
  EXPECT_EQ(tx->state_history[3]->duration, 20);
  EXPECT_EQ(tx->state_history[4]->state, ITS::kFinished);
  EXPECT_EQ(tx->state_history[4]->duration, 0);
}

TEST_F(BucketContextTest, OverrideShouldUseSqliteForTesting) {
  {
    // Toss out the override from the test fixture.
    auto release = std::move(sqlite_override_);
  }
  int bucket_counter = 0;
  auto create_new_bucket_context = [&]() {
    storage::BucketInfo info =
        GetOrCreateBucket(blink::StorageKey::CreateFromStringForTesting(
            base::StringPrintf("https://test%d.example/", bucket_counter++)));
    return InitBucketContext(info);
  };
  // `bucket_context_` was created before deleting `sqlite_override_`, so we
  // have to make a new one to verify what happens without any override.
  base::WeakPtr<BucketContext> bucket_context_no_override =
      create_new_bucket_context();
  {
    base::AutoReset<std::optional<bool>> scoped_override =
        BucketContext::OverrideShouldUseSqliteForTesting(false);
    EXPECT_FALSE(create_new_bucket_context()->IsUsingSqlite());
  }
  EXPECT_EQ(bucket_context_no_override->IsUsingSqlite(),
            create_new_bucket_context()->IsUsingSqlite());
  {
    base::AutoReset<std::optional<bool>> scoped_override =
        BucketContext::OverrideShouldUseSqliteForTesting(true);
    EXPECT_TRUE(create_new_bucket_context()->IsUsingSqlite());
  }
}

TEST_F(BucketContextTest, InMemoryContextUsesSqliteByDefault) {
  {
    // Toss out the override from the test fixture.
    auto release = std::move(sqlite_override_);
  }
  SetUpInMemoryContext();

  EXPECT_TRUE(InitBucketContext()->IsUsingSqlite());
}

}  // namespace content::indexed_db