Merge pull request #7 from 0xsequence/timeflusher

time interval based file roll policy + refactor

Author: marino39

writer.go

@@ -19,6 +19,7 @@ const defaultBufferSize = 8 * datasize.MB
 type Writer[T any] interface {
 	Write(b Block[T]) error
 	BlockNum() uint64
+	RollFile() error
 	Close() error
 }
 
@@ -79,6 +80,7 @@ func NewWriter[T any](opt Options) (Writer[T], error) {
 		lastBlockNum = walFiles[len(walFiles)-1].LastBlockNum
 	}
 
+	// create new writer
 	return &writer[T]{
 		options:       opt,
 		path:          walPath,
@@ -109,13 +111,16 @@ func (w *writer[T]) Write(b Block[T]) error {
 	}
 
 	w.lastBlockNum = b.Number
-	if p, ok := w.options.FileRollPolicy.(LastBlockNumberRollPolicy); ok {
-		p.LastBlockNum(w.lastBlockNum)
-	}
-
+	w.options.FileRollPolicy.onBlockProcessed(w.lastBlockNum)
 	return nil
 }
 
+func (w *writer[T]) RollFile() error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.rollFile()
+}
+
 func (w *writer[T]) BlockNum() uint64 {
 	w.mu.Lock()
 	defer w.mu.Unlock()
@@ -209,9 +214,7 @@ func (w *writer[T]) newFile() error {
 
 	// create new buffer writer
 	bufferWriter := io.Writer(w.buffer)
-	if policy, ok := w.options.FileRollPolicy.(FileSizeRollPolicy); ok {
-		bufferWriter = policy.WrapWriter(bufferWriter)
-	}
+	bufferWriter = &writerWrapper{Writer: bufferWriter, fsrp: w.options.FileRollPolicy}
 
 	// create new buffer closer
 	w.bufferCloser = &funcCloser{

writer_file_roll_policy.go

@@ -1,56 +1,53 @@
 package ethwal
 
-import "io"
+import (
+	"context"
+	"io"
+	"sync"
+	"time"
+)
 
 type FileRollPolicy interface {
 	ShouldRoll() bool
 	Reset()
-}
-
-type FileSizeRollPolicy interface {
-	FileRollPolicy
-
-	WrapWriter(w io.Writer) io.Writer
-}
 
-type LastBlockNumberRollPolicy interface {
-	FileRollPolicy
-
-	LastBlockNum(blockNum uint64)
+	onWrite(data []byte)
+	onBlockProcessed(blockNum uint64)
 }
 
 type fileSizeRollPolicy struct {
-	maxSize uint64
-	stats   *fileStats
+	maxSize      uint64
+	bytesWritten uint64
 }
 
-func NewFileSizeRollPolicy(maxSize uint64) FileSizeRollPolicy {
+func NewFileSizeRollPolicy(maxSize uint64) FileRollPolicy {
 	return &fileSizeRollPolicy{maxSize: maxSize}
 }
 
-func (p *fileSizeRollPolicy) WrapWriter(w io.Writer) io.Writer {
-	p.stats = &fileStats{Writer: w}
-	return p.stats
-}
-
 func (p *fileSizeRollPolicy) ShouldRoll() bool {
-	return p.stats.BytesWritten >= p.maxSize
+	return p.bytesWritten >= p.maxSize
 }
 
 func (p *fileSizeRollPolicy) Reset() {
-	p.stats = &fileStats{}
+	p.bytesWritten = 0
+}
+
+func (p *fileSizeRollPolicy) onWrite(data []byte) {
+	p.bytesWritten += uint64(len(data))
 }
 
+func (p *fileSizeRollPolicy) onBlockProcessed(blockNum uint64) {}
+
 // fileStats is a writer that keeps track of the number of bytes written to it.
-type fileStats struct {
+type writerWrapper struct {
 	io.Writer
-	BytesWritten uint64
+
+	fsrp FileRollPolicy
 }
 
-func (w *fileStats) Write(p []byte) (n int, err error) {
-	n, err = w.Writer.Write(p)
-	w.BytesWritten += uint64(n)
-	return
+func (w *writerWrapper) Write(p []byte) (n int, err error) {
+	defer w.fsrp.onWrite(p)
+	return w.Writer.Write(p)
 }
 
 type lastBlockNumberRollPolicy struct {
@@ -59,7 +56,9 @@ type lastBlockNumberRollPolicy struct {
 	lastBlockNum uint64
 }
 
-func NewLastBlockNumberRollPolicy(rollInterval uint64) LastBlockNumberRollPolicy {
+func (l *lastBlockNumberRollPolicy) onWrite(data []byte) {}
+
+func NewLastBlockNumberRollPolicy(rollInterval uint64) FileRollPolicy {
 	return &lastBlockNumberRollPolicy{rollInterval: rollInterval}
 }
 
@@ -71,27 +70,70 @@ func (l *lastBlockNumberRollPolicy) Reset() {
 	// noop
 }
 
-func (l *lastBlockNumberRollPolicy) LastBlockNum(blockNum uint64) {
+func (l *lastBlockNumberRollPolicy) onBlockProcessed(blockNum uint64) {
 	l.lastBlockNum = blockNum
 }
 
-type fileSizeOrLastBlockNumberRollPolicy struct {
-	FileSizeRollPolicy
-	LastBlockNumberRollPolicy
+type timeBasedRollPolicy struct {
+	rollInterval time.Duration
+	onError      func(err error)
+
+	rollFunc func() error
+
+	bgCtx    context.Context
+	bgCancel context.CancelFunc
+
+	lastTimeRolled time.Time
+
+	mu sync.Mutex
+}
+
+func NewTimeBasedRollPolicy(rollInterval time.Duration, onError func(err error)) FileRollPolicy {
+	return &timeBasedRollPolicy{rollInterval: rollInterval, lastTimeRolled: time.Now(), onError: onError}
+}
+
+func (t *timeBasedRollPolicy) ShouldRoll() bool {
+	if time.Since(t.lastTimeRolled) >= t.rollInterval {
+		return true
+	}
+	return false
+}
+
+func (t *timeBasedRollPolicy) Reset() {
+	t.lastTimeRolled = time.Now()
+}
+
+func (t *timeBasedRollPolicy) onWrite(data []byte) {}
+
+func (t *timeBasedRollPolicy) onBlockProcessed(blockNum uint64) {}
+
+type FileRollPolicies []FileRollPolicy
+
+func (policies FileRollPolicies) ShouldRoll() bool {
+	for _, p := range policies {
+		if p.ShouldRoll() {
+			return true
+		}
+	}
+	return false
 }
 
-func NewFileSizeOrLastBlockNumberRollPolicy(maxSize, rollInterval uint64) FileRollPolicy {
-	return &fileSizeOrLastBlockNumberRollPolicy{
-		FileSizeRollPolicy:        NewFileSizeRollPolicy(maxSize),
-		LastBlockNumberRollPolicy: NewLastBlockNumberRollPolicy(rollInterval),
+func (policies FileRollPolicies) Reset() {
+	for _, p := range policies {
+		p.Reset()
 	}
 }
 
-func (f *fileSizeOrLastBlockNumberRollPolicy) ShouldRoll() bool {
-	return f.FileSizeRollPolicy.ShouldRoll() || f.LastBlockNumberRollPolicy.ShouldRoll()
+func (policies FileRollPolicies) onWrite(data []byte) {
+	for _, p := range policies {
+		p.onWrite(data)
+	}
 }
 
-func (f *fileSizeOrLastBlockNumberRollPolicy) Reset() {
-	f.FileSizeRollPolicy.Reset()
-	f.LastBlockNumberRollPolicy.Reset()
+func (policies FileRollPolicies) onBlockProcessed(blockNum uint64) {
+	for _, p := range policies {
+		p.onBlockProcessed(blockNum)
+	}
 }
+
+var _ FileRollPolicy = &fileSizeRollPolicy{}

writer_file_roll_policy_test.go

@@ -3,17 +3,17 @@ package ethwal
 import (
 	"bytes"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestFileSizeRollPolicy(t *testing.T) {
-	var p FileSizeRollPolicy
 	var buff = bytes.NewBuffer(nil)
 
-	p = NewFileSizeRollPolicy(10)
-	w := p.WrapWriter(buff)
+	p := NewFileSizeRollPolicy(10)
+	w := writerWrapper{buff, p}
 
 	_, err := w.Write([]byte("hello"))
 	require.NoError(t, err)
@@ -26,7 +26,7 @@ func TestFileSizeRollPolicy(t *testing.T) {
 	assert.True(t, p.ShouldRoll())
 
 	p.Reset()
-	w = p.WrapWriter(buff)
+	w = writerWrapper{buff, p}
 	assert.False(t, p.ShouldRoll())
 
 	_, err = w.Write([]byte("hello world"))
@@ -38,50 +38,59 @@ func TestFileSizeRollPolicy(t *testing.T) {
 }
 
 func TestLastBlockNumberRollPolicy(t *testing.T) {
-	var p LastBlockNumberRollPolicy
+	p := NewLastBlockNumberRollPolicy(10)
+	assert.False(t, p.ShouldRoll())
+
+	p.onBlockProcessed(5)
+	assert.False(t, p.ShouldRoll())
 
-	p = NewLastBlockNumberRollPolicy(10)
+	p.onBlockProcessed(10)
+	assert.True(t, p.ShouldRoll())
+
+	p.onBlockProcessed(11)
 	assert.False(t, p.ShouldRoll())
+}
 
-	p.LastBlockNum(5)
+func TestTimeBasedRollPolicy(t *testing.T) {
+	p := NewTimeBasedRollPolicy(1500*time.Millisecond, nil)
 	assert.False(t, p.ShouldRoll())
 
-	p.LastBlockNum(10)
+	time.Sleep(1500 * time.Millisecond)
 	assert.True(t, p.ShouldRoll())
 
-	p.LastBlockNum(11)
+	p.Reset()
 	assert.False(t, p.ShouldRoll())
+
+	time.Sleep(1500 * time.Millisecond)
+	assert.True(t, p.ShouldRoll())
 }
 
 func TestNewFileSizeOrLastBlockNumberRollPolicy(t *testing.T) {
 	var buff = bytes.NewBuffer(nil)
 
-	fol := NewFileSizeOrLastBlockNumberRollPolicy(10, 10)
-
-	fs := fol.(FileSizeRollPolicy)
-	lb := fol.(LastBlockNumberRollPolicy)
-
-	require.NotNil(t, fs)
-	require.NotNil(t, lb)
+	fol := FileRollPolicies{
+		NewFileSizeRollPolicy(10),
+		NewLastBlockNumberRollPolicy(10),
+	}
 
-	w := fs.WrapWriter(buff)
+	w := writerWrapper{buff, fol}
 
-	assert.False(t, fs.ShouldRoll())
+	assert.False(t, fol.ShouldRoll())
 
-	lb.LastBlockNum(10)
-	assert.True(t, lb.ShouldRoll())
+	fol.onBlockProcessed(10)
+	assert.True(t, fol.ShouldRoll())
 
-	lb.LastBlockNum(11)
-	assert.False(t, lb.ShouldRoll())
+	fol.onBlockProcessed(11)
+	assert.False(t, fol.ShouldRoll())
 
 	_, err := w.Write([]byte("hello world"))
 	require.NoError(t, err)
 
-	assert.True(t, fs.ShouldRoll())
+	assert.True(t, fol.ShouldRoll())
 
-	fs.Reset()
-	assert.False(t, fs.ShouldRoll())
+	fol.Reset()
+	assert.False(t, fol.ShouldRoll())
 
-	lb.LastBlockNum(20)
-	assert.True(t, lb.ShouldRoll())
+	fol.onBlockProcessed(20)
+	assert.True(t, fol.ShouldRoll())
 }

writer_no_gap.go

@@ -33,6 +33,10 @@ func (n *noGapWriter[T]) Write(b Block[T]) error {
 	return n.w.Write(b)
 }
 
+func (n *noGapWriter[T]) RollFile() error {
+	return n.w.RollFile()
+}
+
 func (n *noGapWriter[T]) BlockNum() uint64 {
 	return n.w.BlockNum()
 }