From 8dfc73af1d58f2a9e3fecf6e2a5623b78f6b68e2 Mon Sep 17 00:00:00 2001
From: nxshock <nxshock@gmail.com>
Date: Fri, 9 Dec 2022 20:05:30 +0500
Subject: [PATCH] Speedup small writes by using write buffer

---
 defaults.go |  3 ++-
 options.go  |  7 +++++--
 zkv.go      | 17 +++++++++++++----
 zkv_test.go |  4 ++--
 4 files changed, 22 insertions(+), 9 deletions(-)

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++ {