From e6904d8dde3960f6f4ba06a2859220746cccf839 Mon Sep 17 00:00:00 2001
From: wdvxdr <wdvxdr1123@gmail.com>
Date: Sat, 19 Feb 2022 18:55:28 +0800
Subject: [PATCH] internal/btree: implement file lock

For #1366
---
 internal/btree/btree.go             | 18 ++++++++++++
 internal/btree/file_lock_unix.go    | 45 +++++++++++++++++++++++++++++
 internal/btree/file_lock_windows.go | 28 ++++++++++++++++++
 3 files changed, 91 insertions(+)
 create mode 100644 internal/btree/file_lock_unix.go
 create mode 100644 internal/btree/file_lock_windows.go

diff --git a/internal/btree/btree.go b/internal/btree/btree.go
index 7b3b49c..81adcb9 100644
--- a/internal/btree/btree.go
+++ b/internal/btree/btree.go
@@ -18,6 +18,10 @@ const (
 	tableStructSize = int(unsafe.Sizeof(table{}))
 )
 
+type fileLock interface {
+	release() error
+}
+
 type item struct {
 	hash   [hashSize]byte
 	offset int64
@@ -48,6 +52,7 @@ type DB struct {
 	alloc   int64
 	cache   [cacheSlots]cache
 
+	flock        fileLock
 	inAllocator  bool
 	deleteLarger bool
 	fqueue       [freeQueueLen]chunk
@@ -108,6 +113,10 @@ func (d *DB) flushSuper() {
 
 // Open opens an existed btree file
 func Open(name string) (*DB, error) {
+	lock, err := newFileLock(name + ".lock")
+	if err != nil {
+		return nil, errors.New("文件被其他进程占用")
+	}
 	btree := new(DB)
 	fd, err := os.OpenFile(name, os.O_RDWR, 0o644)
 	if err != nil {
@@ -120,17 +129,23 @@ func Open(name string) (*DB, error) {
 	btree.top = super.top
 	btree.freeTop = super.freeTop
 	btree.alloc = super.alloc
+	btree.flock = lock
 	return btree, errors.Wrap(err, "btree read meta info failed")
 }
 
 // Create creates a database
 func Create(name string) (*DB, error) {
+	lock, err := newFileLock(name + ".lock")
+	if err != nil {
+		return nil, errors.New("文件被其他进程占用")
+	}
 	btree := new(DB)
 	fd, err := os.OpenFile(name, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0o644)
 	if err != nil {
 		return nil, errors.Wrap(err, "btree open file failed")
 	}
 
+	btree.flock = lock
 	btree.fd = fd
 	btree.alloc = int64(superSize)
 	btree.flushSuper()
@@ -140,6 +155,9 @@ func Create(name string) (*DB, error) {
 // Close closes the database
 func (d *DB) Close() error {
 	_ = d.fd.Sync()
+	if err := d.flock.release(); err != nil {
+		return err
+	}
 	err := d.fd.Close()
 	for i := 0; i < cacheSlots; i++ {
 		d.cache[i] = cache{}
diff --git a/internal/btree/file_lock_unix.go b/internal/btree/file_lock_unix.go
new file mode 100644
index 0000000..8974284
--- /dev/null
+++ b/internal/btree/file_lock_unix.go
@@ -0,0 +1,45 @@
+//go:build darwin || dragonfly || freebsd || linux || netbsd || openbsd
+
+package btree
+
+import (
+	"os"
+	"syscall"
+)
+
+type unixFileLock struct {
+	f *os.File
+}
+
+func (fl *unixFileLock) release() error {
+	if err := setFileLock(fl.f, false); err != nil {
+		return err
+	}
+	return fl.f.Close()
+}
+
+func newFileLock(path string) (fl fileLock, err error) {
+	flag := os.O_RDWR
+	f, err := os.OpenFile(path, flag, 0)
+	if os.IsNotExist(err) {
+		f, err = os.OpenFile(path, flag|os.O_CREATE, 0644)
+	}
+	if err != nil {
+		return
+	}
+	err = setFileLock(f, true)
+	if err != nil {
+		f.Close()
+		return
+	}
+	fl = &unixFileLock{f: f}
+	return
+}
+
+func setFileLock(f *os.File, lock bool) error {
+	how := syscall.LOCK_UN
+	if lock {
+		how = syscall.LOCK_EX
+	}
+	return syscall.Flock(int(f.Fd()), how|syscall.LOCK_NB)
+}
diff --git a/internal/btree/file_lock_windows.go b/internal/btree/file_lock_windows.go
new file mode 100644
index 0000000..becdfad
--- /dev/null
+++ b/internal/btree/file_lock_windows.go
@@ -0,0 +1,28 @@
+package btree
+
+import "syscall"
+
+type windowsFileLock struct {
+	fd syscall.Handle
+}
+
+func (fl *windowsFileLock) release() error {
+	return syscall.Close(fl.fd)
+}
+
+func newFileLock(path string) (fileLock, error) {
+	pathp, err := syscall.UTF16PtrFromString(path)
+	if err != nil {
+		return nil, err
+	}
+
+	const access uint32 = syscall.GENERIC_READ | syscall.GENERIC_WRITE
+	fd, err := syscall.CreateFile(pathp, access, 0, nil, syscall.OPEN_EXISTING, syscall.FILE_ATTRIBUTE_NORMAL, 0)
+	if err == syscall.ERROR_FILE_NOT_FOUND {
+		fd, err = syscall.CreateFile(pathp, access, 0, nil, syscall.OPEN_ALWAYS, syscall.FILE_ATTRIBUTE_NORMAL, 0)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return &windowsFileLock{fd: fd}, nil
+}