components/sqlite_vfs/shared_locks_unittest.cc - codesearch/chromium/src - Git at Google

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

 #include "components/sqlite_vfs/shared_locks.h"

 #include "base/barrier_closure.h"
 #include "base/functional/bind.h"
 #include "base/functional/callback.h"
 #include "base/memory/unsafe_shared_memory_region.h"
 #include "base/numerics/clamped_math.h"
 #include "base/run_loop.h"
 #include "base/task/thread_pool.h"
 #include "base/test/bind.h"
 #include "base/test/gmock_expected_support.h"
 #include "base/test/task_environment.h"
 #include "base/threading/platform_thread.h"
 #include "components/sqlite_vfs/lock_state.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/sqlite/sqlite3.h"

 namespace sqlite_vfs {

 namespace {

 using ::testing::Ne;

 TEST(SharedLocksTest, CreateRegion) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());
   // Size should be at least sizeof(uint32_t) for the atomic lock.
   EXPECT_GE(shared_region.GetSize(), sizeof(uint32_t));
 }

 TEST(SharedLocksTest, Create) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());
   ASSERT_THAT(SharedLocks::Create(shared_region, /*wal_mode=*/false),
               Ne(std::nullopt));
 }

 TEST(SharedLocksTest, LockBasics) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());
   ASSERT_OK_AND_ASSIGN(auto locks,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int current_mode = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks.Lock(SQLITE_LOCK_SHARED, current_mode), SQLITE_OK);
   EXPECT_EQ(current_mode, SQLITE_LOCK_SHARED);

   EXPECT_EQ(locks.Lock(SQLITE_LOCK_RESERVED, current_mode), SQLITE_OK);
   EXPECT_EQ(current_mode, SQLITE_LOCK_RESERVED);

   EXPECT_EQ(locks.Lock(SQLITE_LOCK_EXCLUSIVE, current_mode), SQLITE_OK);
   EXPECT_EQ(current_mode, SQLITE_LOCK_EXCLUSIVE);

   EXPECT_EQ(locks.Unlock(SQLITE_LOCK_SHARED, current_mode), SQLITE_OK);
   EXPECT_EQ(current_mode, SQLITE_LOCK_SHARED);

   EXPECT_EQ(locks.Unlock(SQLITE_LOCK_NONE, current_mode), SQLITE_OK);
   EXPECT_EQ(current_mode, SQLITE_LOCK_NONE);
 }

 TEST(SharedLocksTest, MultipleLocks) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks3,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;
   int mode2 = SQLITE_LOCK_NONE;
   int mode3 = SQLITE_LOCK_NONE;

   // Reader 1 takes SHARED lock.
   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

   // Reader 2 takes SHARED lock.
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);
   EXPECT_EQ(mode2, SQLITE_LOCK_SHARED);

   // Writer 1 tries to take RESERVED lock, should succeed.
   EXPECT_EQ(locks3.Lock(SQLITE_LOCK_SHARED, mode3), SQLITE_OK);
   EXPECT_EQ(mode3, SQLITE_LOCK_SHARED);

   EXPECT_EQ(locks3.Lock(SQLITE_LOCK_RESERVED, mode3), SQLITE_OK);
   EXPECT_EQ(mode3, SQLITE_LOCK_RESERVED);

   // Writer 2 (we reuse locks2 as another connection) tries to take RESERVED,
   // should fail.
   int mode2_temp = mode2;  // Save current mode (SHARED)
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_RESERVED, mode2), SQLITE_BUSY);
   EXPECT_EQ(mode2, mode2_temp);  // Mode should not change

   // Writer 1 tries to upgrade to EXCLUSIVE, should fail because of other
   // readers.
   EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_BUSY);
   EXPECT_EQ(mode3, SQLITE_LOCK_PENDING);  // Should be upgraded to PENDING

   // Reader 1 releases SHARED lock.
   EXPECT_EQ(locks1.Unlock(SQLITE_LOCK_NONE, mode1), SQLITE_OK);
   EXPECT_EQ(mode1, SQLITE_LOCK_NONE);

   // Writer 1 tries to upgrade to EXCLUSIVE again, should still fail because
   // reader 2 is still there.
   EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_BUSY);
   EXPECT_EQ(mode3, SQLITE_LOCK_PENDING);

   // Reader 2 releases SHARED lock.
   EXPECT_EQ(locks2.Unlock(SQLITE_LOCK_NONE, mode2), SQLITE_OK);
   EXPECT_EQ(mode2, SQLITE_LOCK_NONE);

   // Writer 1 should now be able to upgrade to EXCLUSIVE.
   EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_OK);
   EXPECT_EQ(mode3, SQLITE_LOCK_EXCLUSIVE);

   // Unlock writer 1.
   EXPECT_EQ(locks3.Unlock(SQLITE_LOCK_NONE, mode3), SQLITE_OK);
   EXPECT_EQ(mode3, SQLITE_LOCK_NONE);
 }

 TEST(SharedLocksTest, Abandon) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;
   int mode2 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

   // Abandon locks1.
   LockState state = locks1.Abandon();
   EXPECT_EQ(state, LockState::kReading);

   // locks2 should fail to acquire lock because it's abandoned.
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_IOERR_LOCK);
 }

 TEST(SharedLocksTest, IsReserved) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;

   EXPECT_FALSE(locks1.IsReserved());

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);
   EXPECT_EQ(mode1, SQLITE_LOCK_RESERVED);

   EXPECT_TRUE(locks2.IsReserved());
 }

 TEST(SharedLocksTest, AbandonWithPending) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;
   int mode2 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

   // locks1 tries to upgrade to EXCLUSIVE, fails because locks2 has SHARED.
   // But it should become PENDING.
   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_EXCLUSIVE, mode1), SQLITE_BUSY);
   EXPECT_EQ(mode1, SQLITE_LOCK_PENDING);

   // Abandon locks1. Should return kWriting because it was PENDING.
   LockState state = locks1.Abandon();
   EXPECT_EQ(state, LockState::kWriting);
 }

 TEST(SharedLocksTest, AbandonWithReserved) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);

   // Abandon locks1. Should return kWriting because it was RESERVED.
   LockState state = locks1.Abandon();
   EXPECT_EQ(state, LockState::kWriting);
 }

 TEST(SharedLocksTest, AbandonWithExclusive) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);
   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_EXCLUSIVE, mode1), SQLITE_OK);

   // Abandon locks1. Should return kWriting because it was EXCLUSIVE.
   LockState state = locks1.Abandon();
   EXPECT_EQ(state, LockState::kWriting);
 }

 TEST(SharedLocksTest, LockReservedAfterAbandon) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;
   int mode2 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

   locks1.Abandon();

   // locks2 has SHARED lock, tries to acquire RESERVED.
   // Should fail with IOERR_LOCK because locks1 abandoned.
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_RESERVED, mode2), SQLITE_IOERR_LOCK);
 }

 TEST(SharedLocksTest, LockExclusiveAfterAbandon) {
   auto shared_region = SharedLocks::CreateRegion(false);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/false));

   int mode1 = SQLITE_LOCK_NONE;
   int mode2 = SQLITE_LOCK_NONE;

   EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

   locks1.Abandon();

   // locks2 has SHARED lock, tries to acquire EXCLUSIVE.
   // Should fail with IOERR_LOCK because locks1 abandoned.
   EXPECT_EQ(locks2.Lock(SQLITE_LOCK_EXCLUSIVE, mode2), SQLITE_IOERR_LOCK);
 }

 TEST(SharedLocksTest, ShmLockShared) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Locks 1 acquires shared lock on index 0.
   EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_OK);

   // Locks 2 acquires shared lock on index 0 (should succeed).
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_OK);

   // Locks 1 unlocks.
   EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
             SQLITE_OK);

   // Locks 2 unlocks.
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
             SQLITE_OK);
 }

 TEST(SharedLocksTest, ShmLockExclusive) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Locks 1 acquires exclusive lock on index 0.
   EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_OK);

   // Locks 2 tries to acquire shared lock on index 0 (should fail).
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_BUSY);

   // Locks 2 tries to acquire exclusive lock on index 0 (should fail).
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_BUSY);

   // Locks 1 unlocks.
   EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kExclusive),
             SQLITE_OK);

   // Locks 2 can now acquire exclusive lock.
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_OK);
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kExclusive),
             SQLITE_OK);
 }

 TEST(SharedLocksTest, ShmLockRange) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Locks 1 acquires exclusive lock on index 1 to 3.
   EXPECT_EQ(locks1.ShmLock(1, 3, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_OK);

   // Locks 2 tries to acquire shared lock on index 2 (should fail).
   EXPECT_EQ(locks2.ShmLock(2, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_BUSY);

   // Locks 2 acquires shared lock on index 0 (should succeed).
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_OK);

   // Locks 1 unlocks range.
   EXPECT_EQ(locks1.ShmLock(1, 3, LockOperation::kRelease, LockType::kExclusive),
             SQLITE_OK);

   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
             SQLITE_OK);
 }

 TEST(SharedLocksTest, ShmLockRollback) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Locks 1 acquires exclusive lock on index 2.
   EXPECT_EQ(locks1.ShmLock(2, 1, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_OK);

   // Locks 2 tries to acquire exclusive lock on index 1 to 3.
   // It will succeed on 1, but fail on 2!
   // It should rollback lock on 1!
   EXPECT_EQ(locks2.ShmLock(1, 3, LockOperation::kAcquire, LockType::kExclusive),
             SQLITE_BUSY);

   // Locks 2 tries to acquire shared lock on index 1 (should succeed if rolled
   // back).
   EXPECT_EQ(locks2.ShmLock(1, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_OK);

   EXPECT_EQ(locks2.ShmLock(1, 1, LockOperation::kRelease, LockType::kShared),
             SQLITE_OK);
   EXPECT_EQ(locks1.ShmLock(2, 1, LockOperation::kRelease, LockType::kExclusive),
             SQLITE_OK);
 }

 TEST(SharedLocksTest, ShmLockAbandon) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));
   ASSERT_OK_AND_ASSIGN(auto locks2,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   locks1.Abandon();

   // locks2 should fail to acquire WAL lock because it's abandoned.
   EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_IOERR_LOCK);
 }

 TEST(SharedLocksTest, ShmLockReleaseAfterAbandon) {
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks1,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Acquire a lock first.
   ASSERT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
             SQLITE_OK);

   locks1.Abandon();

   // Releasing should still succeed even though abandoned.
   EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
             SQLITE_OK);
 }

 TEST(SharedLocksTest, WalLocksThreadingAndDataRaces) {
   base::test::TaskEnvironment task_environment;
   using LockOperation = SharedLocks::LockOperation;
   using LockType = SharedLocks::LockType;

   auto shared_region = SharedLocks::CreateRegion(true);
   ASSERT_TRUE(shared_region.IsValid());

   ASSERT_OK_AND_ASSIGN(auto locks,
                        SharedLocks::Create(shared_region, /*wal_mode=*/true));

   // Protected 8-byte vector.
   std::vector<uint8_t> protected_vector(8, 0);

   const int kNumWorkers = 4;
   const int kIterations = 1000;

   base::RunLoop run_loop;
   base::RepeatingClosure barrier_closure =
       base::BarrierClosure(kNumWorkers, run_loop.QuitClosure());

   for (int i = 0; i < kNumWorkers; ++i) {
     base::ThreadPool::PostTask(
         FROM_HERE, {base::MayBlock(), base::TaskPriority::USER_BLOCKING},
         base::BindLambdaForTesting([&locks, &protected_vector, barrier_closure,
                                     i]() {
           uint32_t local_sum = 0;
           for (int iter = 0; iter < kIterations; ++iter) {
             int byte_index = (i + iter) % 8;
             bool is_write = (iter % 4) == 0;

             LockType type = is_write ? LockType::kExclusive : LockType::kShared;

             while (true) {
               int rv =
                   locks.ShmLock(byte_index, 1, LockOperation::kAcquire, type);
               if (rv == SQLITE_OK) {
                 if (is_write) {
                   protected_vector[byte_index]++;
                 } else {
                   local_sum =
                       base::ClampAdd(local_sum, protected_vector[byte_index]);
                 }
                 EXPECT_EQ(
                     locks.ShmLock(byte_index, 1, LockOperation::kRelease, type),
                     SQLITE_OK);
                 break;
               } else if (rv == SQLITE_BUSY) {
                 base::PlatformThread::YieldCurrentThread();
               } else {
                 ADD_FAILURE() << "Unexpected result: " << rv;
                 break;
               }
             }
           }
           EXPECT_GE(local_sum, 0u);
           barrier_closure.Run();
         }));
   }

   run_loop.Run();
 }

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

	#include "components/sqlite_vfs/shared_locks.h"

	#include "base/barrier_closure.h"
	#include "base/functional/bind.h"
	#include "base/functional/callback.h"
	#include "base/memory/unsafe_shared_memory_region.h"
	#include "base/numerics/clamped_math.h"
	#include "base/run_loop.h"
	#include "base/task/thread_pool.h"
	#include "base/test/bind.h"
	#include "base/test/gmock_expected_support.h"
	#include "base/test/task_environment.h"
	#include "base/threading/platform_thread.h"
	#include "components/sqlite_vfs/lock_state.h"
	#include "testing/gmock/include/gmock/gmock.h"
	#include "testing/gtest/include/gtest/gtest.h"
	#include "third_party/sqlite/sqlite3.h"

	namespace sqlite_vfs {

	namespace {

	using ::testing::Ne;

	TEST(SharedLocksTest, CreateRegion) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());
	// Size should be at least sizeof(uint32_t) for the atomic lock.
	EXPECT_GE(shared_region.GetSize(), sizeof(uint32_t));
	}

	TEST(SharedLocksTest, Create) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());
	ASSERT_THAT(SharedLocks::Create(shared_region, /wal_mode=/false),
	Ne(std::nullopt));
	}

	TEST(SharedLocksTest, LockBasics) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());
	ASSERT_OK_AND_ASSIGN(auto locks,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int current_mode = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks.Lock(SQLITE_LOCK_SHARED, current_mode), SQLITE_OK);
	EXPECT_EQ(current_mode, SQLITE_LOCK_SHARED);

	EXPECT_EQ(locks.Lock(SQLITE_LOCK_RESERVED, current_mode), SQLITE_OK);
	EXPECT_EQ(current_mode, SQLITE_LOCK_RESERVED);

	EXPECT_EQ(locks.Lock(SQLITE_LOCK_EXCLUSIVE, current_mode), SQLITE_OK);
	EXPECT_EQ(current_mode, SQLITE_LOCK_EXCLUSIVE);

	EXPECT_EQ(locks.Unlock(SQLITE_LOCK_SHARED, current_mode), SQLITE_OK);
	EXPECT_EQ(current_mode, SQLITE_LOCK_SHARED);

	EXPECT_EQ(locks.Unlock(SQLITE_LOCK_NONE, current_mode), SQLITE_OK);
	EXPECT_EQ(current_mode, SQLITE_LOCK_NONE);
	}

	TEST(SharedLocksTest, MultipleLocks) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks3,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;
	int mode2 = SQLITE_LOCK_NONE;
	int mode3 = SQLITE_LOCK_NONE;

	// Reader 1 takes SHARED lock.
	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

	// Reader 2 takes SHARED lock.
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);
	EXPECT_EQ(mode2, SQLITE_LOCK_SHARED);

	// Writer 1 tries to take RESERVED lock, should succeed.
	EXPECT_EQ(locks3.Lock(SQLITE_LOCK_SHARED, mode3), SQLITE_OK);
	EXPECT_EQ(mode3, SQLITE_LOCK_SHARED);

	EXPECT_EQ(locks3.Lock(SQLITE_LOCK_RESERVED, mode3), SQLITE_OK);
	EXPECT_EQ(mode3, SQLITE_LOCK_RESERVED);

	// Writer 2 (we reuse locks2 as another connection) tries to take RESERVED,
	// should fail.
	int mode2_temp = mode2; // Save current mode (SHARED)
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_RESERVED, mode2), SQLITE_BUSY);
	EXPECT_EQ(mode2, mode2_temp); // Mode should not change

	// Writer 1 tries to upgrade to EXCLUSIVE, should fail because of other
	// readers.
	EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_BUSY);
	EXPECT_EQ(mode3, SQLITE_LOCK_PENDING); // Should be upgraded to PENDING

	// Reader 1 releases SHARED lock.
	EXPECT_EQ(locks1.Unlock(SQLITE_LOCK_NONE, mode1), SQLITE_OK);
	EXPECT_EQ(mode1, SQLITE_LOCK_NONE);

	// Writer 1 tries to upgrade to EXCLUSIVE again, should still fail because
	// reader 2 is still there.
	EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_BUSY);
	EXPECT_EQ(mode3, SQLITE_LOCK_PENDING);

	// Reader 2 releases SHARED lock.
	EXPECT_EQ(locks2.Unlock(SQLITE_LOCK_NONE, mode2), SQLITE_OK);
	EXPECT_EQ(mode2, SQLITE_LOCK_NONE);

	// Writer 1 should now be able to upgrade to EXCLUSIVE.
	EXPECT_EQ(locks3.Lock(SQLITE_LOCK_EXCLUSIVE, mode3), SQLITE_OK);
	EXPECT_EQ(mode3, SQLITE_LOCK_EXCLUSIVE);

	// Unlock writer 1.
	EXPECT_EQ(locks3.Unlock(SQLITE_LOCK_NONE, mode3), SQLITE_OK);
	EXPECT_EQ(mode3, SQLITE_LOCK_NONE);
	}

	TEST(SharedLocksTest, Abandon) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;
	int mode2 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

	// Abandon locks1.
	LockState state = locks1.Abandon();
	EXPECT_EQ(state, LockState::kReading);

	// locks2 should fail to acquire lock because it's abandoned.
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_IOERR_LOCK);
	}

	TEST(SharedLocksTest, IsReserved) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;

	EXPECT_FALSE(locks1.IsReserved());

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(mode1, SQLITE_LOCK_SHARED);

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);
	EXPECT_EQ(mode1, SQLITE_LOCK_RESERVED);

	EXPECT_TRUE(locks2.IsReserved());
	}

	TEST(SharedLocksTest, AbandonWithPending) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;
	int mode2 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

	// locks1 tries to upgrade to EXCLUSIVE, fails because locks2 has SHARED.
	// But it should become PENDING.
	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_EXCLUSIVE, mode1), SQLITE_BUSY);
	EXPECT_EQ(mode1, SQLITE_LOCK_PENDING);

	// Abandon locks1. Should return kWriting because it was PENDING.
	LockState state = locks1.Abandon();
	EXPECT_EQ(state, LockState::kWriting);
	}

	TEST(SharedLocksTest, AbandonWithReserved) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);

	// Abandon locks1. Should return kWriting because it was RESERVED.
	LockState state = locks1.Abandon();
	EXPECT_EQ(state, LockState::kWriting);
	}

	TEST(SharedLocksTest, AbandonWithExclusive) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_RESERVED, mode1), SQLITE_OK);
	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_EXCLUSIVE, mode1), SQLITE_OK);

	// Abandon locks1. Should return kWriting because it was EXCLUSIVE.
	LockState state = locks1.Abandon();
	EXPECT_EQ(state, LockState::kWriting);
	}

	TEST(SharedLocksTest, LockReservedAfterAbandon) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;
	int mode2 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

	locks1.Abandon();

	// locks2 has SHARED lock, tries to acquire RESERVED.
	// Should fail with IOERR_LOCK because locks1 abandoned.
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_RESERVED, mode2), SQLITE_IOERR_LOCK);
	}

	TEST(SharedLocksTest, LockExclusiveAfterAbandon) {
	auto shared_region = SharedLocks::CreateRegion(false);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/false));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/false));

	int mode1 = SQLITE_LOCK_NONE;
	int mode2 = SQLITE_LOCK_NONE;

	EXPECT_EQ(locks1.Lock(SQLITE_LOCK_SHARED, mode1), SQLITE_OK);
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_SHARED, mode2), SQLITE_OK);

	locks1.Abandon();

	// locks2 has SHARED lock, tries to acquire EXCLUSIVE.
	// Should fail with IOERR_LOCK because locks1 abandoned.
	EXPECT_EQ(locks2.Lock(SQLITE_LOCK_EXCLUSIVE, mode2), SQLITE_IOERR_LOCK);
	}

	TEST(SharedLocksTest, ShmLockShared) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Locks 1 acquires shared lock on index 0.
	EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_OK);

	// Locks 2 acquires shared lock on index 0 (should succeed).
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_OK);

	// Locks 1 unlocks.
	EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
	SQLITE_OK);

	// Locks 2 unlocks.
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
	SQLITE_OK);
	}

	TEST(SharedLocksTest, ShmLockExclusive) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Locks 1 acquires exclusive lock on index 0.
	EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_OK);

	// Locks 2 tries to acquire shared lock on index 0 (should fail).
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_BUSY);

	// Locks 2 tries to acquire exclusive lock on index 0 (should fail).
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_BUSY);

	// Locks 1 unlocks.
	EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kExclusive),
	SQLITE_OK);

	// Locks 2 can now acquire exclusive lock.
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_OK);
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kExclusive),
	SQLITE_OK);
	}

	TEST(SharedLocksTest, ShmLockRange) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Locks 1 acquires exclusive lock on index 1 to 3.
	EXPECT_EQ(locks1.ShmLock(1, 3, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_OK);

	// Locks 2 tries to acquire shared lock on index 2 (should fail).
	EXPECT_EQ(locks2.ShmLock(2, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_BUSY);

	// Locks 2 acquires shared lock on index 0 (should succeed).
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_OK);

	// Locks 1 unlocks range.
	EXPECT_EQ(locks1.ShmLock(1, 3, LockOperation::kRelease, LockType::kExclusive),
	SQLITE_OK);

	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
	SQLITE_OK);
	}

	TEST(SharedLocksTest, ShmLockRollback) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Locks 1 acquires exclusive lock on index 2.
	EXPECT_EQ(locks1.ShmLock(2, 1, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_OK);

	// Locks 2 tries to acquire exclusive lock on index 1 to 3.
	// It will succeed on 1, but fail on 2!
	// It should rollback lock on 1!
	EXPECT_EQ(locks2.ShmLock(1, 3, LockOperation::kAcquire, LockType::kExclusive),
	SQLITE_BUSY);

	// Locks 2 tries to acquire shared lock on index 1 (should succeed if rolled
	// back).
	EXPECT_EQ(locks2.ShmLock(1, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_OK);

	EXPECT_EQ(locks2.ShmLock(1, 1, LockOperation::kRelease, LockType::kShared),
	SQLITE_OK);
	EXPECT_EQ(locks1.ShmLock(2, 1, LockOperation::kRelease, LockType::kExclusive),
	SQLITE_OK);
	}

	TEST(SharedLocksTest, ShmLockAbandon) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));
	ASSERT_OK_AND_ASSIGN(auto locks2,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	locks1.Abandon();

	// locks2 should fail to acquire WAL lock because it's abandoned.
	EXPECT_EQ(locks2.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_IOERR_LOCK);
	}

	TEST(SharedLocksTest, ShmLockReleaseAfterAbandon) {
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks1,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Acquire a lock first.
	ASSERT_EQ(locks1.ShmLock(0, 1, LockOperation::kAcquire, LockType::kShared),
	SQLITE_OK);

	locks1.Abandon();

	// Releasing should still succeed even though abandoned.
	EXPECT_EQ(locks1.ShmLock(0, 1, LockOperation::kRelease, LockType::kShared),
	SQLITE_OK);
	}

	TEST(SharedLocksTest, WalLocksThreadingAndDataRaces) {
	base::test::TaskEnvironment task_environment;
	using LockOperation = SharedLocks::LockOperation;
	using LockType = SharedLocks::LockType;

	auto shared_region = SharedLocks::CreateRegion(true);
	ASSERT_TRUE(shared_region.IsValid());

	ASSERT_OK_AND_ASSIGN(auto locks,
	SharedLocks::Create(shared_region, /wal_mode=/true));

	// Protected 8-byte vector.
	std::vector<uint8_t> protected_vector(8, 0);

	const int kNumWorkers = 4;
	const int kIterations = 1000;

	base::RunLoop run_loop;
	base::RepeatingClosure barrier_closure =
	base::BarrierClosure(kNumWorkers, run_loop.QuitClosure());

	for (int i = 0; i < kNumWorkers; ++i) {
	base::ThreadPool::PostTask(
	FROM_HERE, {base::MayBlock(), base::TaskPriority::USER_BLOCKING},
	base::BindLambdaForTesting([&locks, &protected_vector, barrier_closure,
	i]() {
	uint32_t local_sum = 0;
	for (int iter = 0; iter < kIterations; ++iter) {
	int byte_index = (i + iter) % 8;
	bool is_write = (iter % 4) == 0;

	LockType type = is_write ? LockType::kExclusive : LockType::kShared;

	while (true) {
	int rv =
	locks.ShmLock(byte_index, 1, LockOperation::kAcquire, type);
	if (rv == SQLITE_OK) {
	if (is_write) {
	protected_vector[byte_index]++;
	} else {
	local_sum =
	base::ClampAdd(local_sum, protected_vector[byte_index]);
	}
	EXPECT_EQ(
	locks.ShmLock(byte_index, 1, LockOperation::kRelease, type),
	SQLITE_OK);
	break;
	} else if (rv == SQLITE_BUSY) {
	base::PlatformThread::YieldCurrentThread();
	} else {
	ADD_FAILURE() << "Unexpected result: " << rv;
	break;
	}
	}
	}
	EXPECT_GE(local_sum, 0u);
	barrier_closure.Run();
	}));
	}

	run_loop.Run();
	}

	} // namespace
	} // namespace sqlite_vfs