diff --git a/compact_log_test.go b/compact_log_test.go index f7be3f299..695a68f4c 100644 --- a/compact_log_test.go +++ b/compact_log_test.go @@ -76,7 +76,8 @@ func TestCompactLogBasic(t *testing.T) { opt := getTestOptions(dir) { - kv, _ := NewKV(opt) + kv, err := NewKV(opt) + require.NoError(t, err) n := 5000 for i := 0; i < n; i++ { if (i % 10000) == 0 { @@ -87,17 +88,18 @@ func TestCompactLogBasic(t *testing.T) { } kv.Set([]byte("testkey"), []byte("testval")) kv.validate() - kv.Close() + require.NoError(t, kv.Close()) } - kv, _ := NewKV(opt) + kv, err := NewKV(opt) + require.NoError(t, err) var item KVItem if err := kv.Get([]byte("testkey"), &item); err != nil { t.Error(err) } require.EqualValues(t, "testval", string(item.Value())) - kv.Close() + require.NoError(t, kv.Close()) } func key(prefix string, i int) string { diff --git a/kv.go b/kv.go index 0eebbd5f9..1d686d116 100644 --- a/kv.go +++ b/kv.go @@ -17,9 +17,11 @@ package badger import ( + "fmt" "log" "math" "os" + "path/filepath" "sync" "time" @@ -145,6 +147,9 @@ func NewKV(opt *Options) (out *KV, err error) { return nil, ErrInvalidDir } } + if err := createLockFile(filepath.Join(opt.Dir, lockFile)); err != nil { + return nil, err + } if !(opt.ValueLogFileSize <= 2<<30 && opt.ValueLogFileSize >= 1<<20) { return nil, ErrValueLogSize } @@ -259,7 +264,7 @@ func (s *KV) Close() error { // Now close the value log. if err := s.vlog.Close(); err != nil { - return errors.Wrapf(err, "Close()") + return errors.Wrapf(err, "KV.Close") } // Make sure that block writer is done pushing stuff into memtable! @@ -310,9 +315,52 @@ func (s *KV) Close() error { s.closer.SignalAll() s.closer.WaitForAll() s.elog.Finish() + if err := os.Remove(filepath.Join(s.opt.Dir, lockFile)); err != nil { + return errors.Wrap(err, "KV.Close") + } + // Sync Dir so that pid file is guaranteed removed from directory entries. + if err := syncDir(s.opt.Dir); err != nil { + return errors.Wrap(err, "KV.Close cannot sync Dir") + } + return nil +} + +const ( + lockFile = "LOCK" +) + +// Opens a file, errors if it exists, and writes the process id to the file +func createLockFile(path string) error { + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0666) + if err != nil { + return errors.Wrap(err, "cannot create pid lock file") + } + _, err = fmt.Fprintf(f, "%d\n", os.Getpid()) + closeErr := f.Close() + if err != nil { + return errors.Wrap(err, "cannot write to pid lock file") + } + if closeErr != nil { + return errors.Wrap(closeErr, "cannot close pid lock file") + } return nil } +// When you create or delete a file, you have to ensure the directory entry for the file is synced +// in order to guarantee the file is visible (if the system crashes). +func syncDir(dir string) error { + f, err := os.Open(dir) + if err != nil { + return err + } + err = f.Sync() + closeErr := f.Close() + if err != nil { + return err + } + return closeErr +} + // getMemtables returns the current memtables and get references. func (s *KV) getMemTables() ([]*skl.Skiplist, func()) { s.RLock() diff --git a/kv_test.go b/kv_test.go index 97536a237..f818f30a7 100644 --- a/kv_test.go +++ b/kv_test.go @@ -794,3 +794,16 @@ func BenchmarkExists(b *testing.B) { fmt.Println("Done and closing") } + +func TestPidFile(t *testing.T) { + dir, err := ioutil.TempDir("", "badger") + require.NoError(t, err) + defer os.RemoveAll(dir) + options := getTestOptions(dir) + kv1, err := NewKV(options) + require.NoError(t, err) + defer kv1.Close() + _, err = NewKV(options) + require.Error(t, err) + require.Contains(t, err.Error(), "cannot create pid lock file") +}