Periodically persists consumer records

ashishkujoy · ashishkujoy · commit 97d9151a5534 · 2025-05-22T17:14:09.000+05:30
diff --git a/cmd/server/main.go b/cmd/server/main.go
@@ -4,10 +4,11 @@ import (
 	"ashishkujoy/queue/internal/config"
 	netinternal "ashishkujoy/queue/internal/net"
 	"log"
+	"time"
 )
 
 func main() {
-	conf := config.NewConfigWithMetadataPath("data/segments", "data/metadata", 1024*1024)
+	conf := config.NewConfig("data/segments", "data/metadata", 1024*1024, time.Second*100)
 	server, err := netinternal.NewQueueServer(conf, ":50051")
 	if err != nil {
 		log.Fatalf("Failed to create server: %v", err)
diff --git a/internal/arrayUtils.go b/internal/arrayUtils.go
@@ -20,3 +20,11 @@ func Filter[T any](elements []T, f func(T) bool) []T {
 
 	return filtered
 }
+
+func Map[I any, O any](elements []I, f func(I) O) []O {
+	var mapped []O
+	for _, element := range elements {
+		mapped = append(mapped, f(element))
+	}
+	return mapped
+}
diff --git a/internal/config/base_config.go b/internal/config/base_config.go
@@ -1,9 +1,12 @@
 package config
 
+import "time"
+
 type Config struct {
-	segmentsRoot          string
-	MetadataPath          string
-	maxSegmentSizeInBytes int
+	segmentsRoot              string
+	MetadataPath              string
+	consumerIndexSyncInterval time.Duration
+	maxSegmentSizeInBytes     int
 }
 
 func (c *Config) MaxSegmentSizeInBytes() int {
@@ -14,17 +17,20 @@ func (c *Config) SegmentsRoot() string {
 	return c.segmentsRoot
 }
 
-func NewConfig(segmentsRoot string, maxSegmentSizeInBytes int) *Config {
-	return &Config{
-		segmentsRoot:          segmentsRoot,
-		maxSegmentSizeInBytes: maxSegmentSizeInBytes,
-	}
+func (c *Config) ConsumerIndexSyncInterval() time.Duration {
+	return c.consumerIndexSyncInterval
 }
 
-func NewConfigWithMetadataPath(segmentsRoot, metadataPath string, maxSegmentSizeInBytes int) *Config {
+func NewConfig(
+	segmentsRoot string,
+	metadataPath string,
+	maxSegmentSizeInBytes int,
+	consumerIndexSyncInterval time.Duration,
+) *Config {
 	return &Config{
-		segmentsRoot:          segmentsRoot,
-		MetadataPath:          metadataPath,
-		maxSegmentSizeInBytes: maxSegmentSizeInBytes,
+		segmentsRoot:              segmentsRoot,
+		maxSegmentSizeInBytes:     maxSegmentSizeInBytes,
+		MetadataPath:              metadataPath,
+		consumerIndexSyncInterval: consumerIndexSyncInterval,
 	}
 }
diff --git a/internal/consumer/consumer_index.go b/internal/consumer/consumer_index.go
@@ -1,6 +1,7 @@
 package consumer
 
 import (
+	"ashishkujoy/queue/internal"
 	"ashishkujoy/queue/internal/config"
 	"encoding/binary"
 	"fmt"
@@ -30,12 +31,30 @@ func NewConsumerIndex(config *config.Config) (*ConsumerIndex, error) {
 		return nil, err
 	}
 
-	return &ConsumerIndex{
+	consumerIndex := &ConsumerIndex{
 		writer:  writer,
 		mu:      &sync.RWMutex{},
 		indexes: make(map[int]int),
 		config:  config,
-	}, nil
+	}
+	return consumerIndex, nil
+}
+
+func (ci *ConsumerIndex) schedulePersist() {
+	ticker := time.NewTicker(ci.config.ConsumerIndexSyncInterval())
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				ci.mu.Lock()
+				err := ci.Persist()
+				ci.mu.Unlock()
+				if err != nil {
+					fmt.Printf("Error persisting consumer index: %v\n", err)
+				}
+			}
+		}
+	}()
 }
 
 func extractTimestamp(filename string) int64 {
@@ -46,12 +65,20 @@ func extractTimestamp(filename string) int64 {
 }
 
 func getLastIndexFile(config *config.Config) (*os.File, error) {
-	enteries, err := os.ReadDir(config.MetadataPath)
+	files, err := getSortedIndexFiles(config)
+	if err != nil {
+		return nil, err
+	}
+	return files[0], nil
+}
+
+func getSortedIndexFiles(config *config.Config) ([]*os.File, error) {
+	entries, err := os.ReadDir(config.MetadataPath)
 	if err != nil {
 		return nil, err
 	}
 	var metadataTimestamp []int64
-	for _, entry := range enteries {
+	for _, entry := range entries {
 		if strings.Contains(entry.Name(), "consumer_index_") {
 			metadataTimestamp = append(metadataTimestamp, extractTimestamp(entry.Name()))
 		}
@@ -62,10 +89,14 @@ func getLastIndexFile(config *config.Config) (*os.File, error) {
 	sort.Slice(metadataTimestamp, func(i, j int) bool {
 		return metadataTimestamp[i] > metadataTimestamp[j]
 	})
-	lastTimestamp := metadataTimestamp[0]
-	lastIndexFile := fmt.Sprintf("%s/consumer_index_%d", config.MetadataPath, lastTimestamp)
-
-	return os.OpenFile(lastIndexFile, os.O_RDWR, 0666)
+	files := internal.Map(metadataTimestamp, func(i int64) *os.File {
+		file, err := os.OpenFile(fmt.Sprintf("%s/consumer_index_%d", config.MetadataPath, i), os.O_RDWR, 0666)
+		if err != nil {
+			return nil
+		}
+		return file
+	})
+	return files, nil
 }
 
 func restoreIndexesFromFile(file *os.File) (map[int]int, error) {
@@ -78,11 +109,17 @@ func restoreIndexesFromFile(file *os.File) (map[int]int, error) {
 	indexes := make(map[int]int, stat.Size()/8)
 	for i := int64(0); i < indexSize; i++ {
 		buf := make([]byte, 4)
-		file.ReadAt(buf, int64(offset))
+		_, err := file.ReadAt(buf, int64(offset))
+		if err != nil {
+			return nil, err
+		}
 		offset += 4
 		consumerId := binary.BigEndian.Uint32(buf)
 		buf = make([]byte, 4)
-		file.ReadAt(buf, int64(offset))
+		_, err = file.ReadAt(buf, int64(offset))
+		if err != nil {
+			return nil, err
+		}
 		offset += 4
 		consumerIndex := binary.BigEndian.Uint32(buf)
 		indexes[int(consumerId)] = int(consumerIndex)
@@ -192,12 +229,27 @@ func (ci *ConsumerIndex) CreateSnapshot() []byte {
 }
 
 func (ci *ConsumerIndex) Close() error {
+	return ci.Persist()
+}
+
+func (ci *ConsumerIndex) Persist() error {
 	snapshot := ci.CreateSnapshot()
 	indexFile, err := createIndexFile(ci.config)
 	if err != nil {
 		return err
 	}
 	_, err = indexFile.Write(snapshot)
-	indexFile.Sync()
+	err = indexFile.Sync()
+	if err != nil {
+		return err
+	}
+	ci.writer = indexFile
+	files, err := getSortedIndexFiles(ci.config)
+	for _, file := range files[1:] {
+		_ = file.Close()
+	}
+	if err != nil {
+		return err
+	}
 	return err
 }
diff --git a/internal/consumer/consumer_index_test.go b/internal/consumer/consumer_index_test.go
@@ -4,6 +4,7 @@ import (
 	"ashishkujoy/queue/internal/config"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -22,7 +23,7 @@ func TestReadAndWriteIndex(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(metadataDir)
 
-	config := config.NewConfigWithMetadataPath("/tmp", metadataDir, 1234)
+	config := config.NewConfig("/tmp", metadataDir, 1234, time.Second*100)
 	index, err := NewConsumerIndex(config)
 	assert.NoError(t, err)
 
@@ -41,17 +42,17 @@ func TestReadFromARestoredIndex(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(metadataDir)
 
-	config := config.NewConfigWithMetadataPath("/tmp", metadataDir, 1234)
-	index, err := NewConsumerIndex(config)
+	cfg := config.NewConfig("/tmp", metadataDir, 1234, time.Second*100)
+	index, err := NewConsumerIndex(cfg)
 	assert.NoError(t, err)
 
 	index.WriteIndex(11, 10)
 	index.WriteIndex(12, 20)
 	index.WriteIndex(13, 30)
 
-	index.Close()
+	assert.NoError(t, index.Close())
 
-	restoredIndex, err := RestoreConsumerIndex(config)
+	restoredIndex, err := RestoreConsumerIndex(cfg)
 	assert.NoError(t, err)
 
 	assert.Equal(t, 10, restoredIndex.ReadIndex(11))
@@ -65,7 +66,7 @@ func TestRestoreIndexUsesTheLatestSnapshot(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.RemoveAll(metadataDir)
 
-	config := config.NewConfigWithMetadataPath("/tmp", metadataDir, 1234)
+	config := config.NewConfig("/tmp", metadataDir, 1234, time.Second*100)
 
 	index1, err := NewConsumerIndex(config)
 	assert.NoError(t, err)
diff --git a/internal/queue/queue_service.go b/internal/queue/queue_service.go
@@ -15,7 +15,7 @@ func NewQueueService(config *config.Config) (*QueueService, error) {
 	if err != nil {
 		return nil, err
 	}
-	consumerIndex, err := consumer.NewConsumerIndex(config)
+	consumerIndex, err := consumer.RestoreConsumerIndex(config)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/queue/queue_service_test.go b/internal/queue/queue_service_test.go
@@ -4,6 +4,7 @@ import (
 	"ashishkujoy/queue/internal/config"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -13,7 +14,7 @@ func TestEnqueue(t *testing.T) {
 	defer os.RemoveAll(segmentPath)
 	metaDataPath := createTempDir("testEnqueue/metadata")
 	defer os.RemoveAll(metaDataPath)
-	config := config.NewConfigWithMetadataPath(segmentPath, metaDataPath, 1024)
+	config := config.NewConfig(segmentPath, metaDataPath, 1024, time.Second)
 
 	queueService, err := NewQueueService(config)
 	assert.NoError(t, err)
diff --git a/internal/queue/queue_test.go b/internal/queue/queue_test.go
@@ -4,6 +4,7 @@ import (
 	"ashishkujoy/queue/internal/config"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -20,9 +21,13 @@ func removeTempDir(suffix string) {
 }
 
 func TestEnqueueSingleElement(t *testing.T) {
-	config := config.NewConfig(createTempDir("TestEnqueueSingleElement"), 1000)
+	cfg := config.NewConfig(
+		createTempDir("TestEnqueueSingleElement"),
+		"/tmp/metadata",
+		1000,
+		time.Second)
 	defer removeTempDir("TestEnqueueSingleElement")
-	queue, err := NewQueue(config)
+	queue, err := NewQueue(cfg)
 	assert.NoError(t, err)
 	defer queue.Close()
 
diff --git a/internal/storage/segment_test.go b/internal/storage/segment_test.go
@@ -4,17 +4,18 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"time"
 
 	"ashishkujoy/queue/internal/config"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestAppendToTheSegment(t *testing.T) {
-	config := config.NewConfig(os.TempDir(), 1024)
+	cfg := config.NewConfig(os.TempDir(), "/tmp", 1024, time.Second*5)
 	defer os.Remove(fmt.Sprintf("%s/segment_10", os.TempDir()))
 
-	segment, err := NewSegment(10, config)
+	segment, err := NewSegment(10, cfg)
 	assert.NoError(t, err)
 	defer segment.Close()
 
@@ -23,10 +24,10 @@ func TestAppendToTheSegment(t *testing.T) {
 }
 
 func TestReadFromTheSegment(t *testing.T) {
-	config := config.NewConfig(os.TempDir(), 1024)
+	cfg := config.NewConfig(os.TempDir(), "/tmp", 1024, time.Second)
 	defer os.Remove(fmt.Sprintf("%s/segment_10", os.TempDir()))
 
-	segment, err := NewSegment(10, config)
+	segment, err := NewSegment(10, cfg)
 	assert.NoError(t, err)
 	defer segment.Close()
 
@@ -46,10 +47,10 @@ func TestReadFromTheSegment(t *testing.T) {
 }
 
 func TestReadNonExistingMessage(t *testing.T) {
-	config := config.NewConfig(os.TempDir(), 1024)
+	cfg := config.NewConfig(os.TempDir(), "/tmp", 1024, time.Second)
 	defer os.Remove(fmt.Sprintf("%s/segment_12", os.TempDir()))
 
-	segment, err := NewSegment(12, config)
+	segment, err := NewSegment(12, cfg)
 	assert.NoError(t, err)
 	defer segment.Close()
 
@@ -58,18 +59,18 @@ func TestReadNonExistingMessage(t *testing.T) {
 }
 
 func TestReadFromReloadedSegment(t *testing.T) {
-	config := config.NewConfig(os.TempDir(), 1024)
+	cfg := config.NewConfig(os.TempDir(), "/tmp", 1024, time.Second)
 	defer os.Remove(fmt.Sprintf("%s/segment_13", os.TempDir()))
 
-	segment, err := NewSegment(12, config)
+	segment, err := NewSegment(12, cfg)
 	assert.NoError(t, err)
 
 	offset1, _ := segment.Append([]byte("Hello world"))
 	offset2, _ := segment.Append([]byte("Hello earth"))
 	offset3, _ := segment.Append([]byte("Hello India"))
 	segment.Close()
 
-	restoredSegment, err := NewSegment(12, config)
+	restoredSegment, err := NewSegment(12, cfg)
 	assert.NoError(t, err)
 
 	data1, _ := restoredSegment.Read(offset1)
diff --git a/internal/storage/segments_test.go b/internal/storage/segments_test.go

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ func NewQueueService(config config.Config) (QueueService, error) {`
`15`	`15`	`if err != nil {`
`16`	`16`	`return nil, err`
`17`	`17`	`}`
`18`		`- consumerIndex, err := consumer.NewConsumerIndex(config)`
	`18`	`+ consumerIndex, err := consumer.RestoreConsumerIndex(config)`
`19`	`19`	`if err != nil {`
`20`	`20`	`return nil, err`
`21`	`21`	`}`