Restores segments on start up

Author: ashishkujoy

internal/queue/queue.go

@@ -17,6 +17,14 @@ func NewQueue(config *config.Config) (*Queue, error) {
 	return &Queue{segments: segments}, nil
 }
 
+func RestoreQueue(config *config.Config) (*Queue, error) {
+	segments, err := storage.RestoreSegments(config, storage.NewIndex())
+	if err != nil {
+		return nil, err
+	}
+	return &Queue{segments: segments}, nil
+}
+
 func (q *Queue) Enqueue(data []byte) (int, error) {
 	return q.segments.Append(data)
 }

internal/storage/segment.go

@@ -29,6 +29,17 @@ func NewSegment(id int, config *config.Config) (*Segment, error) {
 	return &Segment{store: store, id: id, mu: &sync.RWMutex{}}, nil
 }
 
+// RestoreSegment restores a segment from the given ID and configuration.
+func RestoreSegment(id int, config *config.Config) (*Segment, error) {
+	filePath := fmt.Sprintf("%s/segment-%d", config.SegmentsRoot(), id)
+	store, err := RestoreStore(filePath)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Segment{store: store, id: id, mu: &sync.RWMutex{}}, nil
+}
+
 // Append appends data to the segment.
 // It locks the segment for writing to ensure thread safety.
 // It returns the offset of the appended data or an error if the operation fails.

internal/storage/segments.go

@@ -3,6 +3,10 @@ package storage
 import (
 	"ashishkujoy/queue/internal/config"
 	"fmt"
+	"os"
+	"slices"
+	"strconv"
+	"strings"
 	"sync"
 )
 
@@ -35,6 +39,77 @@ func NewSegments(config *config.Config, index *Index) (*Segments, error) {
 	}, nil
 }
 
+// RestoreSegments restores segments from the given configuration and index.
+func RestoreSegments(c *config.Config, index *Index) (*Segments, error) {
+	segmentIds, err := getSegmentIds(c.SegmentsRoot())
+	if err != nil {
+		return nil, err
+	}
+
+	closedSegments, err2 := restoreSegmentsById(c, segmentIds)
+	if err2 != nil {
+		return nil, err2
+	}
+
+	activeSegmentId, activeSegment, err := createActiveSegment(c, segmentIds, err)
+	if err != nil {
+		return nil, err
+	}
+	return &Segments{
+		config:         c,
+		active:         activeSegment,
+		id:             activeSegmentId,
+		index:          index,
+		closedSegments: closedSegments,
+		mu:             &sync.Mutex{},
+	}, nil
+}
+
+func createActiveSegment(c *config.Config, segmentIds []int, err error) (int, *Segment, error) {
+	activeSegmentId := 0
+	if len(segmentIds) != 0 {
+		activeSegmentId = segmentIds[len(segmentIds)-1]
+	}
+	activeSegment, err := NewSegment(activeSegmentId+1, c)
+	if err != nil {
+		return 0, nil, err
+	}
+	return activeSegmentId + 1, activeSegment, nil
+}
+
+func restoreSegmentsById(c *config.Config, segmentIds []int) ([]*Segment, error) {
+	var closedSegments []*Segment
+	for _, segmentId := range segmentIds {
+		segment, err := RestoreSegment(segmentId, c)
+		if err != nil {
+			return nil, err
+		}
+		closedSegments = append(closedSegments, segment)
+	}
+	return closedSegments, nil
+}
+
+func getSegmentIds(segmentsRoot string) ([]int, error) {
+	entries, err := os.ReadDir(segmentsRoot)
+	if err != nil {
+		return nil, err
+	}
+	var segmentIds []int
+	for _, entry := range entries {
+		if !strings.HasPrefix(entry.Name(), "segment-") {
+			continue
+		}
+		segmentIdStr := strings.TrimPrefix(entry.Name(), "segment-")
+		id, err := strconv.Atoi(segmentIdStr)
+		if err != nil {
+			return nil, err
+		}
+		segmentIds = append(segmentIds, id)
+	}
+	slices.Sort(segmentIds)
+	return segmentIds, nil
+}
+
 // Append appends data to the active segment.
 // If the active segment is full, it rolls over to a new segment.
 func (s *Segments) Append(data []byte) (int, error) {

internal/storage/segments_test.go

@@ -74,10 +74,10 @@ func TestSegmentRollOver(t *testing.T) {
 }
 
 func TestReadFromARolledOverSegment(t *testing.T) {
-	config := config.NewConfig(createTempDir("TestReadFromARolledOverSegment"), "tmp", 10, time.Second)
+	cfg := config.NewConfig(createTempDir("TestReadFromARolledOverSegment"), "tmp", 10, time.Second)
 	defer removeTempDir("TestReadFromARolledOverSegment")
 	index := NewIndex()
-	segments, err := NewSegments(config, index)
+	segments, err := NewSegments(cfg, index)
 	assert.NoError(t, err)
 
 	messageId1, _ := segments.Append([]byte("Hello Segments"))
@@ -90,3 +90,26 @@ func TestReadFromARolledOverSegment(t *testing.T) {
 	assert.Equal(t, []byte("Another Hello Segments"), data2)
 	assert.Equal(t, []byte("Hello Segments"), data1)
 }
+
+func TestRestoreASegment(t *testing.T) {
+	cfg := config.NewConfig(createTempDir("TestRestoreASegment"), "tmp", 10, time.Second)
+	defer removeTempDir("TestRestoreASegment")
+	index := NewIndex()
+	segments, err := NewSegments(cfg, index)
+	assert.NoError(t, err)
+
+	messageId1, _ := segments.Append([]byte("Hello Segments"))
+	messageId2, _ := segments.Append([]byte("Another Hello Segments"))
+	messageId3, _ := segments.Append([]byte("Yet Another Hello Segments"))
+	_ = segments.Close()
+	restoreSegments, err := RestoreSegments(cfg, index)
+	assert.NoError(t, err)
+	assert.Equal(t, 3, len(restoreSegments.closedSegments))
+
+	data2, _ := restoreSegments.Read(messageId2)
+	data1, _ := restoreSegments.Read(messageId1)
+	data3, _ := restoreSegments.Read(messageId3)
+	assert.Equal(t, []byte("Another Hello Segments"), data2)
+	assert.Equal(t, []byte("Hello Segments"), data1)
+	assert.Equal(t, []byte("Yet Another Hello Segments"), data3)
+}

internal/storage/store.go

@@ -38,6 +38,19 @@ func NewStore(filePath string) (*Store, error) {
 	}, nil
 }
 
+// RestoreStore restores a store from a file at the given filePath.
+// It opens the file for reading only
+// TODO: Add error handling of writing to a restored store.(closed segment)
+func RestoreStore(filePath string) (*Store, error) {
+	reader, err := os.OpenFile(filePath, os.O_RDONLY, 0644)
+	if err != nil {
+		return nil, err
+	}
+	return &Store{
+		reader: reader,
+	}, nil
+}
+
 func (s *Store) Append(data []byte) (int, error) {
 	currentOffset := s.offset
 	sizeBuff := make([]byte, 4)

internal/storage/store_test.go

@@ -65,3 +65,29 @@ func TestReadFromANewStoreAfterClose(t *testing.T) {
 	assert.Equal(t, data1, []byte("Hello World"))
 	assert.Equal(t, data2, []byte("Another Hello World"))
 }
+
+func TestRestoreStore(t *testing.T) {
+	filePath := fmt.Sprintf("%s/%s", os.TempDir(), "TestReadFromANewStoreAfterClose")
+	defer os.Remove(filePath)
+	store, err := NewStore(filePath)
+	defer store.Close()
+
+	assert.NoError(t, err)
+
+	offset1, _ := store.Append([]byte("Hello World"))
+	offset2, _ := store.Append([]byte("Another Hello World"))
+	err = store.Close()
+	assert.NoError(t, err)
+
+	restoreStore, err := RestoreStore(filePath)
+	assert.Nil(t, restoreStore.writer)
+	assert.NoError(t, err)
+
+	data1, err := restoreStore.Read(offset1)
+	assert.NoError(t, err)
+	assert.Equal(t, data1, []byte("Hello World"))
+	data2, err := restoreStore.Read(offset2)
+	assert.NoError(t, err)
+	assert.Equal(t, data2, []byte("Another Hello World"))
+	assert.Equal(t, data1, []byte("Hello World"))
+}