Fix flushing logic and add functional test case. (#113)

persist/fs/manager.go

@@ -48,7 +48,7 @@ func NewPersistManager(opts *Options) persist.Manager {
 		Persist: persist.Fns{
 			WriteFields: pm.writeFields,
 		},
-		Close: pm.close,
+		Close: pm.writer.Finish,
 	}
 	return pm
 }
@@ -82,13 +82,13 @@ func (pm *persistManager) Prepare(opts persist.PrepareOptions) (persist.Prepared
 		return prepared, errPersistManagerCannotPrepareDataNotPersisting
 	}
 
-	writerOpts := writerOpenOptions{
+	writerOpts := writerStartOptions{
 		Namespace:    opts.Namespace,
 		Shard:        opts.Shard,
 		NumDocuments: opts.NumDocuments,
 		SegmentMeta:  opts.SegmentMeta,
 	}
-	if err := pm.writer.Open(writerOpts); err != nil {
+	if err := pm.writer.Start(writerOpts); err != nil {
 		return prepared, err
 	}
 
@@ -99,12 +99,8 @@ func (pm *persistManager) writeFields(fields []field.DocsField) error {
 	return pm.writer.WriteFields(fields...)
 }
 
-func (pm *persistManager) close() error {
-	return pm.writer.Close()
-}
-
-// Done is called to finish the data persistence process.
-func (pm *persistManager) Done() error {
+// Finish is called to finish the data persistence process.
+func (pm *persistManager) Finish() error {
 	pm.Lock()
 	defer pm.Unlock()
 
@@ -117,3 +113,8 @@ func (pm *persistManager) Done() error {
 
 	return nil
 }
+
+// Close is a no-op for now.
+func (pm *persistManager) Close() error {
+	return nil
+}

persist/fs/reader_bench_test.go

@@ -41,7 +41,7 @@ func BenchmarkReadField(b *testing.B) {
 		}
 	}
 	docsField := builder.Seal(int32(totDocs))
-	err := writer.Open(writerOpenOptions{
+	err := writer.Start(writerStartOptions{
 		Shard:        0,
 		NumDocuments: numDocs,
 		Namespace:    namespace,

persist/fs/writer.go

@@ -3,7 +3,6 @@ package fs
 import (
 	"bytes"
 	"encoding/binary"
-	"errors"
 	"fmt"
 	"os"
 	"time"
@@ -21,24 +20,20 @@ import (
 	xbytes "github.com/xichen2020/eventdb/x/bytes"
 )
 
-var (
-	errSegmentWriterClosed = errors.New("segment writer is closed")
-)
-
 // segmentWriter is responsible for writing segments to filesystem.
 type segmentWriter interface {
-	// Open opens the writer.
-	Open(opts writerOpenOptions) error
+	// Start starts persisting a segment.
+	Start(opts writerStartOptions) error
+
+	// Finish finishes persisting a segment and performs cleanups as necessary.
+	Finish() error
 
 	// WriteFields writes a set of document fields.
 	WriteFields(fields ...indexfield.DocsField) error
-
-	// Close closes the writer.
-	Close() error
 }
 
-// writerOpenOptions provide a set of options for opening a writer.
-type writerOpenOptions struct {
+// writerStartOptions provide a set of options for opening a writer.
+type writerStartOptions struct {
 	Namespace    []byte
 	Shard        uint32
 	NumDocuments int32
@@ -66,8 +61,7 @@ type writer struct {
 	tw     encoding.TimeEncoder
 	values valuesUnion
 
-	err    error
-	closed bool
+	err error
 }
 
 // newSegmentWriter creates a new segment writer.
@@ -93,11 +87,7 @@ func newSegmentWriter(opts *Options) segmentWriter {
 	return w
 }
 
-func (w *writer) Open(opts writerOpenOptions) error {
-	if w.closed {
-		return errSegmentWriterClosed
-	}
-
+func (w *writer) Start(opts writerStartOptions) error {
 	var (
 		namespace   = opts.Namespace
 		shard       = opts.Shard
@@ -122,10 +112,6 @@ func (w *writer) Open(opts writerOpenOptions) error {
 }
 
 func (w *writer) WriteFields(fields ...indexfield.DocsField) error {
-	if w.closed {
-		return errSegmentWriterClosed
-	}
-
 	for _, field := range fields {
 		if err := w.writeField(field); err != nil {
 			return err
@@ -134,29 +120,12 @@ func (w *writer) WriteFields(fields ...indexfield.DocsField) error {
 	return nil
 }
 
-func (w *writer) Close() error {
-	if w.closed {
-		return errSegmentWriterClosed
-	}
-
+func (w *writer) Finish() error {
 	if w.err != nil {
 		return w.err
 	}
-	// NB(xichen): only write out the checkpoint file if there are no errors
-	// encountered between calling writer.Open() and writer.Close().
-	if err := w.writeCheckpointFile(w.segmentDir); err != nil {
-		w.err = err
-		return err
-	}
-
-	w.closed = true
-	w.info = nil
-	w.bw = nil
-	w.iw = nil
-	w.dw = nil
-	w.sw = nil
-	w.tw = nil
-	return nil
+	w.err = w.writeCheckpointFile(w.segmentDir)
+	return w.err
 }
 
 func (w *writer) writeInfoFile(

persist/persist.go

@@ -16,8 +16,11 @@ type Persister interface {
 	// Prepare prepares for data persistence.
 	Prepare(opts PrepareOptions) (PreparedPersister, error)
 
-	// Done marks the persistence as complete.
-	Done() error
+	// Finish marks the persistence as complete.
+	Finish() error
+
+	// Close all resources owned by the persist manager.
+	Close() error
 }
 
 // SegmentMetadata contains the metadata for a segment.

storage/database.go

@@ -250,6 +250,9 @@ func (d *db) Close() error {
 	d.state = databaseClosed
 
 	// Close database-level resources.
+	if err := d.mediator.Close(); err != nil {
+		return err
+	}
 
 	// Close namespaces.
 	var multiErr xerrors.MultiError

storage/database_test.go

@@ -0,0 +1,107 @@
+package storage
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"github.com/xichen2020/eventdb/document"
+	"github.com/xichen2020/eventdb/parser/json"
+	"github.com/xichen2020/eventdb/parser/json/value"
+	"github.com/xichen2020/eventdb/sharding"
+
+	"github.com/m3db/m3cluster/shard"
+	"github.com/pborman/uuid"
+)
+
+const (
+	testJSONDataFilePath  = "./testdata/testdata.json"
+	testDataPath          = "./testdata/eventdb"
+	testMaxDocsPerSegment = 1 // Keep this # low for flushing purposes.
+	testDataSeparator     = '\n'
+)
+
+var (
+	testNamespace = []byte("testNamespace")
+	testNumShards = 8
+)
+
+type testDatabase struct {
+	Database
+}
+
+func newTestDatabase() (*testDatabase, error) {
+	namespaces := []NamespaceMetadata{
+		NamespaceMetadata{
+			id:   testNamespace,
+			opts: NewNamespaceOptions(),
+		},
+	}
+
+	shardIDs := make([]uint32, 0, testNumShards)
+	for i := 0; i < testNumShards; i++ {
+		shardIDs = append(shardIDs, uint32(i))
+	}
+	shards := sharding.NewShards(shardIDs, shard.Available)
+	hashFn := sharding.DefaultHashFn(testNumShards)
+	shardSet, err := sharding.NewShardSet(shards, hashFn)
+	if err != nil {
+		return nil, err
+	}
+
+	dbOpts := NewOptions().
+		SetNamespaceFieldName(string(testNamespace)).
+		SetMaxNumDocsPerSegment(testMaxDocsPerSegment)
+
+	db := NewDatabase(namespaces, shardSet, dbOpts)
+	if err := db.Open(); err != nil {
+		return nil, err
+	}
+
+	return &testDatabase{db}, nil
+}
+
+// Close removes all filesystem resources owned by the test database.
+func (tdb *testDatabase) Close() error {
+	tdb.Database.Close()
+	return os.RemoveAll(testDataPath)
+}
+
+func createTestDocuments() ([]byte, []document.Document, error) {
+	p := json.NewParser(json.NewOptions())
+
+	data, err := ioutil.ReadFile(testJSONDataFilePath)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	var (
+		end  = bytes.IndexByte(data, testDataSeparator)
+		docs = make([]document.Document, 0)
+	)
+	if end < 0 {
+		end = len(data)
+	}
+	for end > 0 {
+		v, err := p.ParseBytes(data[:end])
+		if err != nil {
+			err = fmt.Errorf("cannot parse document %s: %v", data, err)
+			return nil, nil, err
+		}
+
+		fieldIter := value.NewFieldIterator(v)
+		doc := document.Document{
+			ID:        uuid.NewUUID(),
+			TimeNanos: time.Now().UnixNano(), // This doesn't need to line up to what's in the raw data.
+			FieldIter: fieldIter,
+			RawData:   data,
+		}
+		docs = append(docs, doc)
+
+		data = data[end+1:]
+		end = bytes.IndexByte(data, testDataSeparator)
+	}
+	return []byte(testNamespace), docs, nil
+}

storage/flush.go

@@ -66,7 +66,7 @@ func (m *flushManager) Flush() error {
 	if err != nil {
 		return err
 	}
-	defer persister.Done()
+	defer persister.Finish()
 
 	namespaces, err := m.database.GetOwnedNamespaces()
 	if err != nil {

storage/flush_test.go

@@ -0,0 +1,26 @@
+package storage
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFlush(t *testing.T) {
+	tdb, err := newTestDatabase()
+	require.Nil(t, err)
+	defer tdb.Close()
+
+	ns, docs, err := createTestDocuments()
+	require.Nil(t, err)
+
+	err = tdb.WriteBatch(ns, docs)
+	require.Nil(t, err)
+
+	// Get a reference to the underlying fs manager and force a flush.
+	db := tdb.Database.(*db)
+	mediator := db.mediator.(*mediator)
+	fsMgr := mediator.databaseFileSystemManager.(*fileSystemManager)
+	err = fsMgr.databaseFlushManager.Flush()
+	require.Nil(t, err)
+}