chore: smarter loading scheme

apps/workspace-engine/pkg/db/workspaces.go

@@ -89,14 +89,16 @@ func GetAllWorkspaceIDs(ctx context.Context) ([]string, error) {
 }
 
 type WorkspaceSnapshot struct {
+	WorkspaceID   string
 	Path          string
 	Timestamp     time.Time
 	Partition     int32
 	NumPartitions int32
+	Offset        int64
 }
 
 const WORKSPACE_SNAPSHOT_SELECT_QUERY = `
-	SELECT path, timestamp, partition, num_partitions FROM workspace_snapshot WHERE workspace_id = $1 ORDER BY timestamp DESC LIMIT 1
+	SELECT workspace_id, path, timestamp, partition, num_partitions, offset FROM workspace_snapshot WHERE workspace_id = $1 ORDER BY timestamp DESC LIMIT 1
 `
 
 func GetWorkspaceSnapshot(ctx context.Context, workspaceID string) (*WorkspaceSnapshot, error) {
@@ -108,10 +110,12 @@ func GetWorkspaceSnapshot(ctx context.Context, workspaceID string) (*WorkspaceSn
 
 	workspaceSnapshot := &WorkspaceSnapshot{}
 	err = db.QueryRow(ctx, WORKSPACE_SNAPSHOT_SELECT_QUERY, workspaceID).Scan(
+		&workspaceSnapshot.WorkspaceID,
 		&workspaceSnapshot.Path,
 		&workspaceSnapshot.Timestamp,
 		&workspaceSnapshot.Partition,
 		&workspaceSnapshot.NumPartitions,
+		&workspaceSnapshot.Offset,
 	)
 	if err != nil {
 		if err == pgx.ErrNoRows {
@@ -122,19 +126,58 @@ func GetWorkspaceSnapshot(ctx context.Context, workspaceID string) (*WorkspaceSn
 	return workspaceSnapshot, nil
 }
 
+func GetLatestWorkspaceSnapshots(ctx context.Context, workspaceIDs []string) (map[string]*WorkspaceSnapshot, error) {
+	db, err := GetDB(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer db.Release()
+
+	const query = `
+		SELECT DISTINCT ON (workspace_id) workspace_id, path, timestamp, partition, num_partitions, offset 
+		FROM workspace_snapshot 
+		WHERE workspace_id = ANY($1)
+		ORDER BY workspace_id, timestamp DESC
+	`
+	rows, err := db.Query(ctx, query, workspaceIDs)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var snapshots []*WorkspaceSnapshot
+	for rows.Next() {
+		var snapshot WorkspaceSnapshot
+		err := rows.Scan(&snapshot.WorkspaceID, &snapshot.Path, &snapshot.Timestamp, &snapshot.Partition, &snapshot.NumPartitions, &snapshot.Offset)
+		if err != nil {
+			return nil, err
+		}
+		snapshots = append(snapshots, &snapshot)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	snapshotMap := make(map[string]*WorkspaceSnapshot)
+	for _, snapshot := range snapshots {
+		snapshotMap[snapshot.WorkspaceID] = snapshot
+	}
+	return snapshotMap, nil
+}
+
 const WORKSPACE_SNAPSHOT_INSERT_QUERY = `
-	INSERT INTO workspace_snapshot (workspace_id, path, timestamp, partition, num_partitions)
-	VALUES ($1, $2, $3, $4, $5)
+	INSERT INTO workspace_snapshot (workspace_id, path, timestamp, partition, num_partitions, offset)
+	VALUES ($1, $2, $3, $4, $5, $6)
 `
 
-func WriteWorkspaceSnapshot(ctx context.Context, workspaceID string, snapshot *WorkspaceSnapshot) error {
+func WriteWorkspaceSnapshot(ctx context.Context, snapshot *WorkspaceSnapshot) error {
 	db, err := GetDB(ctx)
 	if err != nil {
 		return err
 	}
 	defer db.Release()
 
-	_, err = db.Exec(ctx, WORKSPACE_SNAPSHOT_INSERT_QUERY, workspaceID, snapshot.Path, snapshot.Timestamp, snapshot.Partition, snapshot.NumPartitions)
+	_, err = db.Exec(ctx, WORKSPACE_SNAPSHOT_INSERT_QUERY, snapshot.WorkspaceID, snapshot.Path, snapshot.Timestamp, snapshot.Partition, snapshot.NumPartitions, snapshot.Offset)
 	if err != nil {
 		return err
 	}

apps/workspace-engine/pkg/events/handler/handler.go

@@ -102,8 +102,14 @@ func NewEventListener(handlers HandlerRegistry) *EventListener {
 	return &EventListener{handlers: handlers}
 }
 
+type OffsetTracker struct {
+	LastCommittedOffset int64
+	LastWorkspaceOffset int64
+	MessageOffset       int64
+}
+
 // ListenAndRoute processes incoming Kafka messages and routes them to the appropriate handler
-func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message) (*workspace.Workspace, error) {
+func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message, offsetTracker OffsetTracker) (*workspace.Workspace, error) {
 	ctx, span := tracer.Start(ctx, "ListenAndRoute",
 		trace.WithAttributes(
 			attribute.String("kafka.topic", *msg.TopicPartition.Topic),
@@ -149,6 +155,13 @@ func (el *EventListener) ListenAndRoute(ctx context.Context, msg *kafka.Message)
 		return nil, fmt.Errorf("workspace not found: %s", rawEvent.WorkspaceID)
 	}
 
+	isReplay := offsetTracker.MessageOffset <= offsetTracker.LastCommittedOffset
+	ws.Store().SetIsReplay(isReplay)
+
+	if offsetTracker.MessageOffset <= offsetTracker.LastWorkspaceOffset {
+		return ws, nil
+	}
+
 	ctx = changeset.WithChangeSet(ctx, changeSet)
 
 	if err := handler(ctx, ws, rawEvent); err != nil {

apps/workspace-engine/pkg/kafka/consumer.go

@@ -1,6 +1,10 @@
 package kafka
 
 import (
+	"context"
+	"math"
+	"workspace-engine/pkg/db"
+
 	"github.com/charmbracelet/log"
 	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
 )
@@ -27,3 +31,41 @@ func createConsumer() (*kafka.Consumer, error) {
 
 	return c, nil
 }
+
+func getEarliestOffset(snapshots map[string]*db.WorkspaceSnapshot) int64 {
+	beginning := int64(kafka.OffsetBeginning)
+	if len(snapshots) == 0 {
+		return beginning
+	}
+
+	earliestOffset := int64(math.MaxInt64)
+	for _, snapshot := range snapshots {
+		if snapshot.Offset < earliestOffset {
+			earliestOffset = snapshot.Offset
+		}
+	}
+	if earliestOffset == math.MaxInt64 {
+		return beginning
+	}
+	return earliestOffset
+}
+
+func setOffsets(ctx context.Context, consumer *kafka.Consumer, partitionWorkspaceMap map[int32][]string) {
+	for partition, workspaceIDs := range partitionWorkspaceMap {
+		snapshots, err := db.GetLatestWorkspaceSnapshots(ctx, workspaceIDs)
+		if err != nil {
+			log.Error("Failed to get latest workspace snapshots", "error", err)
+			continue
+		}
+
+		earliestOffset := getEarliestOffset(snapshots)
+		if err := consumer.Seek(kafka.TopicPartition{
+			Topic:     &Topic,
+			Partition: partition,
+			Offset:    kafka.Offset(earliestOffset),
+		}, 0); err != nil {
+			log.Error("Failed to seek to earliest offset", "error", err)
+			continue
+		}
+	}
+}

apps/workspace-engine/pkg/kafka/kafka.go

@@ -2,25 +2,30 @@ package kafka
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 	"time"
 
 	"workspace-engine/pkg/db"
 	"workspace-engine/pkg/events"
+	eventHanlder "workspace-engine/pkg/events/handler"
 	"workspace-engine/pkg/oapi"
 	"workspace-engine/pkg/workspace"
 	wskafka "workspace-engine/pkg/workspace/kafka"
 
 	"github.com/aws/smithy-go/ptr"
 	"github.com/charmbracelet/log"
+	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
 )
 
 // Configuration variables loaded from environment
 var (
-	Topic   = getEnv("KAFKA_TOPIC", "workspace-events")
-	GroupID = getEnv("KAFKA_GROUP_ID", "workspace-engine")
-	Brokers = getEnv("KAFKA_BROKERS", "localhost:9092")
+	Topic               = getEnv("KAFKA_TOPIC", "workspace-events")
+	GroupID             = getEnv("KAFKA_GROUP_ID", "workspace-engine")
+	Brokers             = getEnv("KAFKA_BROKERS", "localhost:9092")
+	MinSnapshotDistance = getEnvInt("SNAPSHOT_DISTANCE_MINUTES", 60)
 )
 
 // getEnv retrieves an environment variable or returns a default value
@@ -32,6 +37,40 @@ func getEnv(varName string, defaultValue string) string {
 	return v
 }
 
+// getEnvInt retrieves an integer environment variable or returns a default value
+func getEnvInt(varName string, defaultValue int64) int64 {
+	v := os.Getenv(varName)
+	if v == "" {
+		return defaultValue
+	}
+	i, err := strconv.ParseInt(v, 10, 64)
+	if err != nil {
+		log.Warn("Failed to parse environment variable as integer, using default", "var", varName, "value", v, "default", defaultValue)
+		return defaultValue
+	}
+	return i
+}
+
+func getLastSnapshot(ctx context.Context, msg *kafka.Message) (*db.WorkspaceSnapshot, error) {
+	var rawEvent eventHanlder.RawEvent
+	if err := json.Unmarshal(msg.Value, &rawEvent); err != nil {
+		log.Error("Failed to unmarshal event", "error", err, "message", string(msg.Value))
+		return nil, fmt.Errorf("failed to unmarshal event: %w", err)
+	}
+
+	return db.GetWorkspaceSnapshot(ctx, rawEvent.WorkspaceID)
+}
+
+func getLastWorkspaceOffset(snapshot *db.WorkspaceSnapshot) int64 {
+	beginning := int64(kafka.OffsetBeginning)
+
+	if snapshot == nil {
+		return beginning
+	}
+
+	return snapshot.Offset
+}
+
 // RunConsumerWithWorkspaceLoader starts the Kafka consumer with workspace-based offset resume
 //
 // Flow:
@@ -83,11 +122,17 @@ func RunConsumer(ctx context.Context) error {
 	}
 	log.Info("Partition assignment complete", "assigned", assignedPartitions)
 
-	allWorkspaceIDs, err := wskafka.GetAssignedWorkspaceIDs(ctx, assignedPartitions, numPartitions)
+	partitionWorkspaceMap, err := wskafka.GetAssignedWorkspaceIDs(ctx, assignedPartitions, numPartitions)
 	if err != nil {
 		return fmt.Errorf("failed to get assigned workspace IDs: %w", err)
 	}
 
+	// Flatten the map to get all workspace IDs
+	var allWorkspaceIDs []string
+	for _, workspaceIDs := range partitionWorkspaceMap {
+		allWorkspaceIDs = append(allWorkspaceIDs, workspaceIDs...)
+	}
+
 	storage := workspace.NewFileStorage("./state")
 	if workspace.IsGCSStorageEnabled() {
 		storage, err = workspace.NewGCSStorageClient(ctx)
@@ -118,6 +163,8 @@ func RunConsumer(ctx context.Context) error {
 	// Start consuming messages
 	handler := events.NewEventHandler()
 
+	setOffsets(ctx, consumer, partitionWorkspaceMap)
+
 	for {
 		// Check for cancellation
 		select {
@@ -134,21 +181,46 @@ func RunConsumer(ctx context.Context) error {
 			continue
 		}
 
-		ws, err := handler.ListenAndRoute(ctx, msg)
+		lastSnapshot, err := getLastSnapshot(ctx, msg)
 		if err != nil {
-			log.Error("Failed to route message", "error", err)
+			log.Error("Failed to get last snapshot", "error", err)
 			continue
 		}
 
-		snapshot := &db.WorkspaceSnapshot{
-			Path:          fmt.Sprintf("%s.gob", ws.ID),
-			Timestamp:     msg.Timestamp,
-			Partition:     int32(msg.TopicPartition.Partition),
-			NumPartitions: numPartitions,
+		messageOffset := int64(msg.TopicPartition.Offset)
+		lastCommittedOffset, err := getCommittedOffset(consumer, msg.TopicPartition.Partition)
+		if err != nil {
+			log.Error("Failed to get committed offset", "error", err)
+			continue
+		}
+		lastWorkspaceOffset := getLastWorkspaceOffset(lastSnapshot)
+
+		offsetTracker := eventHanlder.OffsetTracker{
+			LastCommittedOffset: lastCommittedOffset,
+			LastWorkspaceOffset: lastWorkspaceOffset,
+			MessageOffset:       messageOffset,
+		}
+
+		ws, err := handler.ListenAndRoute(ctx, msg, offsetTracker)
+		if err != nil {
+			log.Error("Failed to route message", "error", err)
+			continue
 		}
 
-		if err := workspace.Save(ctx, storage, ws, snapshot); err != nil {
-			log.Error("Failed to save workspace", "workspaceID", ws.ID, "snapshotPath", snapshot.Path, "error", err)
+		shouldSaveSnapshot := lastSnapshot == nil || lastSnapshot.Timestamp.Before(msg.Timestamp.Add(-time.Duration(MinSnapshotDistance)*time.Minute))
+		if shouldSaveSnapshot {
+			snapshot := &db.WorkspaceSnapshot{
+				WorkspaceID:   ws.ID,
+				Path:          fmt.Sprintf("%s.gob", ws.ID),
+				Timestamp:     msg.Timestamp,
+				Partition:     int32(msg.TopicPartition.Partition),
+				Offset:        int64(msg.TopicPartition.Offset),
+				NumPartitions: numPartitions,
+			}
+
+			if err := workspace.Save(ctx, storage, ws, snapshot); err != nil {
+				log.Error("Failed to save workspace", "workspaceID", ws.ID, "snapshotPath", snapshot.Path, "error", err)
+			}
 		}
 
 		// Commit offset to Kafka

apps/workspace-engine/pkg/kafka/offset.go

@@ -139,3 +139,26 @@ func getTopicPartitionCount(c *kafka.Consumer) (int32, error) {
 
 	return numPartitions, nil
 }
+
+// getCommittedOffset retrieves the last committed offset for a partition
+func getCommittedOffset(consumer *kafka.Consumer, partition int32) (int64, error) {
+	partitions := []kafka.TopicPartition{
+		{
+			Topic:     &Topic,
+			Partition: partition,
+			Offset:    kafka.OffsetStored, // This fetches the committed offset
+		},
+	}
+
+	committed, err := consumer.Committed(partitions, 5000)
+	if err != nil {
+		return int64(kafka.OffsetInvalid), err
+	}
+
+	if len(committed) == 0 || committed[0].Offset == kafka.OffsetInvalid {
+		// No committed offset yet, this is the beginning
+		return int64(kafka.OffsetBeginning), nil
+	}
+
+	return int64(committed[0].Offset), nil
+}

apps/workspace-engine/pkg/workspace/kafka/state.go

@@ -29,7 +29,7 @@ func FilterWorkspaceIDsForPartition(workspaceIDs []string, targetPartition int32
 
 type WorkspaceIDDiscoverer func(ctx context.Context, targetPartition int32, numPartitions int32) ([]string, error)
 
-func GetAssignedWorkspaceIDs(ctx context.Context, assignedPartitions []int32, numPartitions int32) ([]string, error) {
+func GetAssignedWorkspaceIDs(ctx context.Context, assignedPartitions []int32, numPartitions int32) (map[int32][]string, error) {
 	workspaceIDs, err := db.GetAllWorkspaceIDs(ctx)
 	if err != nil {
 		return nil, err
@@ -40,10 +40,11 @@ func GetAssignedWorkspaceIDs(ctx context.Context, assignedPartitions []int32, nu
 		assignedSet[p] = true
 	}
 
-	var result []string
+	result := make(map[int32][]string)
 	for _, workspaceID := range workspaceIDs {
-		if assignedSet[PartitionForWorkspace(workspaceID, numPartitions)] {
-			result = append(result, workspaceID)
+		partition := PartitionForWorkspace(workspaceID, numPartitions)
+		if assignedSet[partition] {
+			result[partition] = append(result[partition], workspaceID)
 		}
 	}
 

apps/workspace-engine/pkg/workspace/loader.go

@@ -36,7 +36,7 @@ func Save(ctx context.Context, storage StorageClient, workspace *Workspace, snap
 		return fmt.Errorf("failed to write workspace to disk: %w", err)
 	}
 
-	if err := db.WriteWorkspaceSnapshot(ctx, workspace.ID, snapshot); err != nil {
+	if err := db.WriteWorkspaceSnapshot(ctx, snapshot); err != nil {
 		span.RecordError(err)
 		span.SetStatus(codes.Error, "Failed to write workspace snapshot")
 		return fmt.Errorf("failed to write workspace snapshot: %w", err)