1 files changed, 148 insertions, 0 deletions
diff --git a/src/rocksdb/util/filelock_test.cc b/src/rocksdb/util/filelock_test.cc
new file mode 100644
index 000000000..69947a732
--- /dev/null
+++ b/src/rocksdb/util/filelock_test.cc
@@ -0,0 +1,148 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+//
+#include <fcntl.h>
+
+#include "rocksdb/env.h"
+#include "rocksdb/status.h"
+#ifdef __FreeBSD__
+#include <sys/types.h>
+#include <sys/wait.h>
+#endif
+#include <vector>
+
+#include "test_util/testharness.h"
+#include "util/coding.h"
+#include "util/string_util.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+class LockTest : public testing::Test {
+ public:
+  static LockTest* current_;
+  std::string file_;
+  ROCKSDB_NAMESPACE::Env* env_;
+
+  LockTest()
+      : file_(test::PerThreadDBPath("db_testlock_file")),
+        env_(ROCKSDB_NAMESPACE::Env::Default()) {
+    current_ = this;
+  }
+
+  ~LockTest() override {}
+
+  Status LockFile(FileLock** db_lock) { return env_->LockFile(file_, db_lock); }
+
+  Status UnlockFile(FileLock* db_lock) { return env_->UnlockFile(db_lock); }
+
+  bool AssertFileIsLocked() {
+    return CheckFileLock(/* lock_expected = */ true);
+  }
+
+  bool AssertFileIsNotLocked() {
+    return CheckFileLock(/* lock_expected = */ false);
+  }
+
+  bool CheckFileLock(bool lock_expected) {
+    // We need to fork to check the fcntl lock as we need
+    // to open and close the file from a different process
+    // to avoid either releasing the lock on close, or not
+    // contending for it when requesting a lock.
+
+#ifdef OS_WIN
+
+    // WaitForSingleObject and GetExitCodeProcess can do what waitpid does.
+    // TODO - implement on Windows
+    return true;
+
+#else
+
+    pid_t pid = fork();
+    if (0 == pid) {
+      // child process
+      int exit_val = EXIT_FAILURE;
+      int fd = open(file_.c_str(), O_RDWR | O_CREAT, 0644);
+      if (fd < 0) {
+        // could not open file, could not check if it was locked
+        fprintf(stderr, "Open on on file %s failed.\n", file_.c_str());
+        exit(exit_val);
+      }
+
+      struct flock f;
+      memset(&f, 0, sizeof(f));
+      f.l_type = (F_WRLCK);
+      f.l_whence = SEEK_SET;
+      f.l_start = 0;
+      f.l_len = 0;  // Lock/unlock entire file
+      int value = fcntl(fd, F_SETLK, &f);
+      if (value == -1) {
+        if (lock_expected) {
+          exit_val = EXIT_SUCCESS;
+        }
+      } else {
+        if (!lock_expected) {
+          exit_val = EXIT_SUCCESS;
+        }
+      }
+      close(fd);  // lock is released for child process
+      exit(exit_val);
+    } else if (pid > 0) {
+      // parent process
+      int status;
+      while (-1 == waitpid(pid, &status, 0))
+        ;
+      if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
+        // child process exited with non success status
+        return false;
+      } else {
+        return true;
+      }
+    } else {
+      fprintf(stderr, "Fork failed\n");
+      return false;
+    }
+    return false;
+
+#endif
+  }
+};
+LockTest* LockTest::current_;
+
+TEST_F(LockTest, LockBySameThread) {
+  FileLock* lock1;
+  FileLock* lock2;
+
+  // acquire a lock on a file
+  ASSERT_OK(LockFile(&lock1));
+
+  // check the file is locked
+  ASSERT_TRUE(AssertFileIsLocked());
+
+  // re-acquire the lock on the same file. This should fail.
+  Status s = LockFile(&lock2);
+  ASSERT_TRUE(s.IsIOError());
+#ifndef OS_WIN
+  // Validate that error message contains current thread ID.
+  ASSERT_TRUE(s.ToString().find(std::to_string(
+                  Env::Default()->GetThreadID())) != std::string::npos);
+#endif
+
+  // check the file is locked
+  ASSERT_TRUE(AssertFileIsLocked());
+
+  // release the lock
+  ASSERT_OK(UnlockFile(lock1));
+
+  // check the file is not locked
+  ASSERT_TRUE(AssertFileIsNotLocked());
+}
+
+}  // namespace ROCKSDB_NAMESPACE
+
+int main(int argc, char** argv) {
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}