update the code

pkg/actor/disk_storage.go

@@ -27,6 +27,7 @@ type DiskStorage[V any] struct {
 
 type LogStorage[V any] interface {
 	LoadEvents(ctx context.Context, id uint64, grain Grain[V]) error
+	LoadEventsFunc(ctx context.Context, id uint64, grain Grain[V], condition func(msg proto.Message, index int, timeStamp time.Time) bool) error
 	AppendMutations(id uint64, msg ...proto.Message) error
 }
 
@@ -87,6 +88,27 @@ func (s *DiskStorage[V]) logPath(id uint64) string {
 	return filepath.Join(s.path, fmt.Sprintf("%d.events.log", id))
 }
 
+func (s *DiskStorage[V]) LoadEventsFunc(ctx context.Context, id uint64, grain Grain[V], condition func(msg proto.Message, index int, timeStamp time.Time) bool) error {
+	path := s.logPath(id)
+	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
+		// No log -> nothing to replay
+		return nil
+	}
+
+	fh, err := os.Open(path)
+	if err != nil {
+		return fmt.Errorf("open replay file: %w", err)
+	}
+	defer fh.Close()
+	index := 0
+	return s.Load(fh, func(msg proto.Message, when time.Time) {
+		if condition(msg, index, when) {
+			s.registry.Apply(ctx, grain, msg)
+		}
+		index++
+	})
+}
+
 func (s *DiskStorage[V]) LoadEvents(ctx context.Context, id uint64, grain Grain[V]) error {
 	path := s.logPath(id)
 	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
@@ -99,7 +121,7 @@ func (s *DiskStorage[V]) LoadEvents(ctx context.Context, id uint64, grain Grain[
 		return fmt.Errorf("open replay file: %w", err)
 	}
 	defer fh.Close()
-	return s.Load(fh, func(msg proto.Message) {
+	return s.Load(fh, func(msg proto.Message, _ time.Time) {
 		s.registry.Apply(ctx, grain, msg)
 	})
 }

pkg/actor/state.go

@@ -34,14 +34,14 @@ func NewState(registry MutationRegistry) *StateStorage {
 
 var ErrUnknownType = errors.New("unknown type")
 
-func (s *StateStorage) Load(r io.Reader, onMessage func(msg proto.Message)) error {
+func (s *StateStorage) Load(r io.Reader, onMessage func(msg proto.Message, timeStamp time.Time)) error {
 	var err error
 	var evt *StorageEvent
 	scanner := bufio.NewScanner(r)
 	for err == nil {
 		evt, err = s.Read(scanner)
 		if err == nil {
-			onMessage(evt.Mutation)
+			onMessage(evt.Mutation, evt.TimeStamp)
 		}
 	}
 	if err == io.EOF {