Repository: petermattis/pebble
Pull request #157 (2019-06-01 15:53:11)

Provide guard rails for various invalid operations
Return errors from operations on snapshots that have been closed.

petermattis: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.png" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/petermattis/pebble/157)
<!-- Reviewable:end -->


Patch:
diff --git a/compaction.go b/compaction.go
index 98447aa..b09bebd 100644
--- a/compaction.go
+++ b/compaction.go
@@ -11,6 +11,7 @@ import (
 	"math"
 	"os"
 	"sort"
+	"sync/atomic"
 	"unsafe"
 
 	"github.com/petermattis/pebble/internal/base"
@@ -576,7 +577,7 @@ type manualCompaction struct {
 //
 // d.mu must be held when calling this.
 func (d *DB) maybeScheduleFlush() {
-	if d.mu.compact.flushing || d.mu.closed {
+	if d.mu.compact.flushing || atomic.LoadInt32(&d.closed) != 0 {
 		return
 	}
 	if len(d.mu.mem.queue) <= 1 {
@@ -701,7 +702,7 @@ func (d *DB) flush1() error {
 //
 // d.mu must be held when calling this.
 func (d *DB) maybeScheduleCompaction() {
-	if d.mu.compact.compacting || d.mu.closed {
+	if d.mu.compact.compacting || atomic.LoadInt32(&d.closed) != 0 {
 		return
 	}
 
diff --git a/db.go b/db.go
index deafb03..002ca8e 100644
--- a/db.go
+++ b/db.go
@@ -32,6 +32,9 @@ var (
 	// ErrNotFound is returned when a get operation does not find the requested
 	// key.
 	ErrNotFound = base.ErrNotFound
+	// ErrClosed is returned when an operation is performed on a closed snapshot
+	// or DB.
+	ErrClosed = errors.New("pebble: closed")
 )
 
 type flushable interface {
@@ -165,12 +168,13 @@ type DB struct {
 
 	logRecycler logRecycler
 
+	closed int32 // updated atomically
+
 	// TODO(peter): describe exactly what this mutex protects. So far: every
 	// field in the struct.
 	mu struct {
 		sync.Mutex
 
-		closed    bool
 		nextJobID int
 
 		versions versionSet
@@ -227,6 +231,10 @@ func (d *DB) Get(key []byte) ([]byte, error) {
 }
 
 func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, error) {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	// Grab and reference the current readState. This prevents the underlying
 	// files in the associated version from being deleted if there is a current
 	// compaction. The readState is unref'd by Iterator.Close().
@@ -343,6 +351,10 @@ func (d *DB) LogData(data []byte, opts *WriteOptions) error {
 //
 // It is safe to modify the contents of the arguments after Apply returns.
 func (d *DB) Apply(batch *Batch, opts *WriteOptions) error {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	sync := opts.GetSync()
 	if sync && d.opts.DisableWAL {
 		return errors.New("pebble: WAL disabled")
@@ -437,6 +449,10 @@ func (d *DB) newIterInternal(
 	s *Snapshot,
 	o *IterOptions,
 ) *Iterator {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	// Grab and reference the current readState. This prevents the underlying
 	// files in the associated version from being deleted if there is a current
 	// compaction. The readState is unref'd by Iterator.Close().
@@ -578,6 +594,10 @@ func (d *DB) NewIter(o *IterOptions) *Iterator {
 // deleted. Instead, a snapshot prevents deletion of sequence numbers
 // referenced by the snapshot.
 func (d *DB) NewSnapshot() *Snapshot {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	s := &Snapshot{
 		db:     d,
 		seqNum: atomic.LoadUint64(&d.mu.versions.visibleSeqNum),
@@ -596,9 +616,10 @@ func (d *DB) NewSnapshot() *Snapshot {
 func (d *DB) Close() error {
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	if d.mu.closed {
-		return nil
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
 	}
+	atomic.StoreInt32(&d.closed, 1)
 	for d.mu.compact.compacting || d.mu.compact.flushing {
 		d.mu.compact.cond.Wait()
 	}
@@ -606,7 +627,6 @@ func (d *DB) Close() error {
 	err = firstError(err, d.mu.log.Close())
 	err = firstError(err, d.fileLock.Close())
 	d.commit.Close()
-	d.mu.closed = true
 
 	err = firstError(err, d.dataDir.Close())
 
@@ -632,6 +652,10 @@ func (d *DB) Close() error {
 
 // Compact the specified range of keys in the database.
 func (d *DB) Compact(start, end []byte /* CompactionOptions */) error {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	iStart := base.MakeInternalKey(start, InternalKeySeqNumMax, InternalKeyKindMax)
 	iEnd := base.MakeInternalKey(end, 0, 0)
 	meta := []*fileMetadata{&fileMetadata{smallest: iStart, largest: iEnd}}
@@ -703,6 +727,10 @@ func (d *DB) manualCompact(manual *manualCompaction) error {
 
 // Flush the memtable to stable storage.
 func (d *DB) Flush() error {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	d.mu.Lock()
 	mem := d.mu.mem.mutable
 	err := d.makeRoomForWrite(nil)
@@ -718,6 +746,10 @@ func (d *DB) Flush() error {
 //
 // TODO(peter): untested
 func (d *DB) AsyncFlush() error {
+	if atomic.LoadInt32(&d.closed) != 0 {
+		panic(ErrClosed)
+	}
+
 	d.mu.Lock()
 	err := d.makeRoomForWrite(nil)
 	d.mu.Unlock()
diff --git a/iterator.go b/iterator.go
index 22ea9e8..1697105 100644
--- a/iterator.go
+++ b/iterator.go
@@ -279,14 +279,14 @@ func (i *Iterator) SeekGE(key []byte) bool {
 // the Comparer. Also note that the iterator will not observe keys not matching
 // the prefix.
 func (i *Iterator) SeekPrefixGE(key []byte) bool {
-	if i.split == nil {
-		panic("pebble: split must be provided for SeekPrefixGE")
-	}
-
 	if i.err != nil {
 		return false
 	}
 
+	if i.split == nil {
+		panic("pebble: split must be provided for SeekPrefixGE")
+	}
+
 	// Make a copy of the prefix so that modifications to the key after
 	// SeekPrefixGE returns does not affect the stored prefix.
 	prefixLen := i.split(key)
diff --git a/snapshot.go b/snapshot.go
index 6d889e3..7834640 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -25,6 +25,9 @@ var _ Reader = (*Snapshot)(nil)
 // The caller should not modify the contents of the returned slice, but it is
 // safe to modify the contents of the argument after Get returns.
 func (s *Snapshot) Get(key []byte) ([]byte, error) {
+	if s.db == nil {
+		panic(ErrClosed)
+	}
 	return s.db.getInternal(key, nil /* batch */, s)
 }
 
@@ -32,6 +35,9 @@ func (s *Snapshot) Get(key []byte) ([]byte, error) {
 // return false). The iterator can be positioned via a call to SeekGE,
 // SeekLT, First or Last.
 func (s *Snapshot) NewIter(o *IterOptions) *Iterator {
+	if s.db == nil {
+		panic(ErrClosed)
+	}
 	return s.db.newIterInternal(nil /* batchIter */, nil /* batchRangeDelIter */, s, o)
 }
 
@@ -40,9 +46,13 @@ func (s *Snapshot) NewIter(o *IterOptions) *Iterator {
 // leak of resources on disk due to the entries the snapshot is preventing from
 // being deleted.
 func (s *Snapshot) Close() error {
+	if s.db == nil {
+		panic(ErrClosed)
+	}
 	s.db.mu.Lock()
 	s.db.mu.snapshots.remove(s)
 	s.db.mu.Unlock()
+	s.db = nil
 	return nil
 }
 


Test patch:
diff --git a/db_test.go b/db_test.go
index 569b448..2ba3a4b 100644
--- a/db_test.go
+++ b/db_test.go
@@ -690,3 +690,45 @@ func TestRollManifest(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestDBClosed(t *testing.T) {
+	d, err := Open("", &Options{
+		FS: vfs.NewMem(),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := d.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	catch := func(f func()) (err error) {
+		defer func() {
+			if r := recover(); r != nil {
+				err = r.(error)
+			}
+		}()
+		f()
+		return nil
+	}
+
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Close() }))
+
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Compact(nil, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Flush() }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.AsyncFlush() }))
+
+	require.EqualValues(t, ErrClosed, catch(func() { _, _ = d.Get(nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Delete(nil, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.DeleteRange(nil, nil, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.LogData(nil, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Merge(nil, nil, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Set(nil, nil, nil) }))
+
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.NewSnapshot() }))
+
+	b := d.NewIndexedBatch()
+	require.EqualValues(t, ErrClosed, catch(func() { _ = b.Commit(nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = d.Apply(b, nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { _ = b.NewIter(nil) }))
+}
diff --git a/snapshot_test.go b/snapshot_test.go
index f4ca494..b3b2f37 100644
--- a/snapshot_test.go
+++ b/snapshot_test.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/petermattis/pebble/internal/datadriven"
 	"github.com/petermattis/pebble/vfs"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSnapshotListToSlice(t *testing.T) {
@@ -163,3 +164,30 @@ func TestSnapshot(t *testing.T) {
 		}
 	})
 }
+
+func TestSnapshotClosed(t *testing.T) {
+	d, err := Open("", &Options{
+		FS: vfs.NewMem(),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	catch := func(f func()) (err error) {
+		defer func() {
+			if r := recover(); r != nil {
+				err = r.(error)
+			}
+		}()
+		f()
+		return nil
+	}
+
+	snap := d.NewSnapshot()
+	if err := snap.Close(); err != nil {
+		t.Fatal(err)
+	}
+	require.EqualValues(t, ErrClosed, catch(func() { _ = snap.Close() }))
+	require.EqualValues(t, ErrClosed, catch(func() { _, _ = snap.Get(nil) }))
+	require.EqualValues(t, ErrClosed, catch(func() { snap.NewIter(nil) }))
+}