diff --git a/defaults.go b/defaults.go index ad476f8..f22add0 100644 --- a/defaults.go +++ b/defaults.go @@ -9,5 +9,6 @@ import ( var defaultOptions = Options{ MaxParallelReads: runtime.NumCPU(), CompressionLevel: zstd.SpeedDefault, - BufferSize: 4 * 1024 * 1024, + MemoryBufferSize: 4 * 1024 * 1024, + DiskBufferSize: 1 * 1024 * 1024, } diff --git a/options.go b/options.go index 69b3c28..1471014 100644 --- a/options.go +++ b/options.go @@ -9,8 +9,11 @@ type Options struct { // Compression level CompressionLevel zstd.EncoderLevel - // Write buffer size in bytes - BufferSize int + // Memory write buffer size in bytes + MemoryBufferSize int + + // Diwk write buffer size in bytes + DiskBufferSize int } func (o *Options) setDefaults() { diff --git a/zkv.go b/zkv.go index 1bbd4b4..11d2cee 100644 --- a/zkv.go +++ b/zkv.go @@ -1,6 +1,7 @@ package zkv import ( + "bufio" "bytes" "crypto/sha256" "encoding/base64" @@ -124,7 +125,7 @@ func (s *Store) Delete(key interface{}) error { return err } - if s.buffer.Len() > s.options.BufferSize { + if s.buffer.Len() > s.options.MemoryBufferSize { err = s.flush() if err != nil { @@ -209,7 +210,7 @@ func (s *Store) setBytes(keyHash [sha256.Size224]byte, valueBytes []byte) error return err } - if s.buffer.Len() > s.options.BufferSize { + if s.buffer.Len() > s.options.MemoryBufferSize { err = s.flush() if err != nil { @@ -238,7 +239,7 @@ func (s *Store) set(key, value interface{}) error { return err } - if s.buffer.Len() > s.options.BufferSize { + if s.buffer.Len() > s.options.MemoryBufferSize { err = s.flush() if err != nil { @@ -377,7 +378,9 @@ func (s *Store) flush() error { return fmt.Errorf("open store file: %v", err) } - encoder, err := zstd.NewWriter(f, zstd.WithEncoderLevel(s.options.CompressionLevel)) + diskWriteBuffer := bufio.NewWriterSize(f, s.options.DiskBufferSize) + + encoder, err := zstd.NewWriter(diskWriteBuffer, zstd.WithEncoderLevel(s.options.CompressionLevel)) if err != nil { f.Close() return fmt.Errorf("open store file: %v", err) @@ -402,6 +405,12 @@ func (s *Store) flush() error { return err } + err = diskWriteBuffer.Flush() + if err != nil { + // TODO: truncate file to previous state + return err + } + err = f.Close() if err != nil { return err diff --git a/zkv_test.go b/zkv_test.go index a252f9d..2b8005f 100644 --- a/zkv_test.go +++ b/zkv_test.go @@ -165,7 +165,7 @@ func TestBufferBasic(t *testing.T) { const filePath = "TestBuffer.zkv" defer os.Remove(filePath) - db, err := OpenWithOptions(filePath, Options{BufferSize: 100}) + db, err := OpenWithOptions(filePath, Options{MemoryBufferSize: 100}) assert.NoError(t, err) err = db.Set(1, make([]byte, 100)) @@ -190,7 +190,7 @@ func TestBufferRead(t *testing.T) { const recordCount = 100 defer os.Remove(filePath) - db, err := OpenWithOptions(filePath, Options{BufferSize: 100}) + db, err := OpenWithOptions(filePath, Options{MemoryBufferSize: 100}) assert.NoError(t, err) for i := 1; i <= recordCount; i++ {