Carried stream subscriptions through reset and failover.

chasm/lib/stream/cursor.go

@@ -139,3 +139,23 @@ func (c *Cursor) AdvanceKnownHead(_ chasm.MutableContext, head int64) {
 func (c *Cursor) StartOffset() int64 {
 	return c.State.StartOffset
 }
+
+// Restore puts the cursor where a completed task's event says it stood, for a
+// run being rebuilt from its history. It only moves forward: the events are
+// applied in order, and a range folded in earlier is never taken back.
+func (c *Cursor) Restore(_ chasm.MutableContext, offset int64) {
+	if offset > c.State.Offset {
+		c.State.Offset = offset
+	}
+	if offset > c.State.KnownHead {
+		c.State.KnownHead = offset
+	}
+}
+
+// MarkExternal says the stream lives in another execution, which the subscribe
+// event alone cannot tell a rebuilt run, and carries over the frontier that
+// stream last pushed at the run this one was rebuilt from.
+func (c *Cursor) MarkExternal(mctx chasm.MutableContext, knownHead int64) {
+	c.State.External = true
+	c.AdvanceKnownHead(mctx, knownHead)
+}

chasm/lib/stream/doc.go

@@ -86,4 +86,13 @@
 // re-supply the budget cuts short is refused as a whole rather than marked:
 // paging it needs a short flag on the poll response, or on the last
 // StreamSlice, that this package does not have on the wire yet.
+//
+// # Reset
+//
+// A reset run keeps its subscriptions. A range recorded before the reset
+// point is re-supplied from the run reset from; the range the reset-point task
+// had been given is delivered again to the reset run's first task, from a
+// stream of its own that begins with a copy of that input (see [Stream.Seed]).
+// Outside consumers are not told that a reset happened: reporting it as a
+// record needs a record shape on the wire, which this package does not have.
 package stream

chasm/lib/stream/seed.go

@@ -0,0 +1,78 @@
+package stream
+
+import (
+	"time"
+
+	commonpb "go.temporal.io/api/common/v1"
+	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/server/chasm"
+	streamlib "go.temporal.io/server/chasm/lib/stream/gen/streampb/v1"
+	"google.golang.org/protobuf/proto"
+)
+
+// Seed copies the records of another stream's window [from, to), read with
+// ReadBatches, into this one, which has to begin at from and hold nothing yet.
+// The records keep their offsets, so a history that names them stays true. No
+// producer entry comes along: the copy is this stream's own past, not an
+// append anyone could retry.
+//
+// Made for a reset. The run reset from keeps its stream, and what it held past
+// the inherited cursor is the input the reset-point task had been given; the
+// run that re-runs that task needs the same input in a stream of its own.
+func (s *Stream) Seed(
+	mctx chasm.MutableContext,
+	blobs []*commonpb.DataBlob,
+	starts []int64,
+	from int64,
+	to int64,
+) error {
+	if s.State.BaseOffset != from || s.State.HeadOffset != from || len(s.Batches) != 0 {
+		return serviceerror.NewInternalf(
+			"a stream holding [%d,%d) cannot be seeded from offset %d",
+			s.State.BaseOffset, s.State.HeadOffset, from)
+	}
+	if s.Batches == nil {
+		s.Batches = make(chasm.Map[int64, *commonpb.DataBlob])
+	}
+	for i, blob := range blobs {
+		var batch streamlib.StreamRecordBatch
+		if err := proto.Unmarshal(blob.GetData(), &batch); err != nil {
+			return err
+		}
+		// A batch straddling the start is trimmed, so the copy begins exactly
+		// at the inherited offset and is keyed there.
+		var kept []*streamlib.StreamRecord
+		for j, record := range batch.GetRecords() {
+			if offset := starts[i] + int64(j); offset >= from && offset < to {
+				kept = append(kept, record)
+			}
+		}
+		if len(kept) == 0 {
+			continue
+		}
+		first := max(starts[i], from)
+		if first != s.State.HeadOffset {
+			return serviceerror.NewInternalf(
+				"seeding a stream at %d with a batch starting at %d", s.State.HeadOffset, first)
+		}
+		// The copy keeps the source batch's age: the records are as old as
+		// they were, whichever stream holds them.
+		var appendedAt time.Time
+		if batch.GetAppendedAt() != nil {
+			appendedAt = batch.GetAppendedAt().AsTime()
+		}
+		encoded, err := marshalBatch(kept, appendedAt)
+		if err != nil {
+			return err
+		}
+		s.Batches[first] = chasm.NewDataField(mctx, encoded)
+		s.State.HeadOffset = first + int64(len(kept))
+		s.State.AppendedBytes += int64(len(encoded.Data))
+		s.State.HeldBytes += int64(len(encoded.Data))
+	}
+	if s.State.HeadOffset != to {
+		return serviceerror.NewInternalf(
+			"seeding stopped at offset %d, short of %d", s.State.HeadOffset, to)
+	}
+	return nil
+}

chasm/lib/stream/seed_test.go

@@ -0,0 +1,81 @@
+package stream
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	commonpb "go.temporal.io/api/common/v1"
+	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/server/chasm"
+	streamlib "go.temporal.io/server/chasm/lib/stream/gen/streampb/v1"
+)
+
+func bodiesOf(records []*streamlib.StreamRecord) []string {
+	out := make([]string, len(records))
+	for i, r := range records {
+		out[i] = string(r.GetBody().GetData())
+	}
+	return out
+}
+
+// A copy that begins inside a batch trims it and keys what it kept at the
+// copy's start, so the seeded stream begins exactly at the inherited offset.
+func TestSeedTrimsTheBatchStraddlingTheStart(t *testing.T) {
+	// A clock on the source, so its batches carry an age the copy must keep.
+	written := time.Unix(1_700_000_000, 0).UTC()
+	mctx := &chasm.MockMutableContext{MockContext: chasm.MockContext{
+		HandleNow: func(chasm.Component) time.Time { return written },
+	}}
+	source := newTestStream(t)
+	source.Batches = make(chasm.Map[int64, *commonpb.DataBlob])
+	for _, batch := range [][]string{{"a", "b", "c"}, {"d", "e"}} {
+		_, err := source.AddMessages(mctx, AddMessagesRequest{Records: msgs(batch...)})
+		require.NoError(t, err)
+	}
+	blobs, starts, err := source.ReadBatches(nil, 1, 4, 0)
+	require.NoError(t, err)
+
+	target := &Stream{State: &streamlib.StreamState{BaseOffset: 1, HeadOffset: 1}}
+	require.NoError(t, target.Seed(nil, blobs, starts, 1, 4))
+	require.Equal(t, int64(4), target.State.HeadOffset)
+	require.Equal(t, []int64{1, 3}, target.batchStarts())
+	require.Positive(t, target.State.AppendedBytes, "the copy counts against the budget")
+	require.Equal(t, target.State.AppendedBytes, target.State.HeldBytes, "and is held")
+	require.Empty(t, target.State.Producers)
+	appendedAt, ok, err := target.batchAppendedAt(nil, 1)
+	require.NoError(t, err)
+	require.True(t, ok, "the copy keeps the source batch's age")
+	require.Equal(t, written, appendedAt)
+
+	copied, copiedStarts, err := target.ReadBatches(nil, 1, 4, 0)
+	require.NoError(t, err)
+	records, next, err := CollectRecords(copied, copiedStarts, 1, 4, 10, nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(4), next)
+	require.Equal(t, []string{"b", "c", "d"}, bodiesOf(records))
+	require.Equal(t, []int64{1, 2, 3}, []int64{
+		records[0].GetOffset(), records[1].GetOffset(), records[2].GetOffset()})
+}
+
+func TestSeedRefusesAStreamThatIsNotEmptyAtTheStart(t *testing.T) {
+	source := newTestStream(t)
+	_, err := source.AddMessages(nil, AddMessagesRequest{Records: msgs("a", "b")})
+	require.NoError(t, err)
+	blobs, starts, err := source.ReadBatches(nil, 0, 2, 0)
+	require.NoError(t, err)
+
+	var internal *serviceerror.Internal
+	occupied := newTestStream(t)
+	_, err = occupied.AddMessages(nil, AddMessagesRequest{Records: msgs("x")})
+	require.NoError(t, err)
+	require.ErrorAs(t, occupied.Seed(nil, blobs, starts, 0, 2), &internal)
+
+	elsewhere := &Stream{State: &streamlib.StreamState{BaseOffset: 5, HeadOffset: 5}}
+	require.ErrorAs(t, elsewhere.Seed(nil, blobs, starts, 0, 2), &internal)
+
+	// A window that does not reach the end it was asked for is refused rather
+	// than left as a stream shorter than its history claims.
+	short := newTestStream(t)
+	require.ErrorAs(t, short.Seed(nil, blobs, starts, 0, 3), &internal)
+}

chasm/lib/stream/stream.go

@@ -56,6 +56,12 @@ type NewStreamRequest struct {
 	// child of the root, so an attached stream carries none and is found
 	// through its owner instead of through ListStreams.
 	Attached bool
+
+	// StartOffset is where the stream's offsets begin. Zero for a new stream.
+	// A run created by a reset inherits a cursor whose position is written in
+	// its History, so the stream it goes on reading has to continue that
+	// offset space rather than start over at zero.
+	StartOffset int64
 }
 
 type AddMessagesRequest struct {
@@ -97,6 +103,9 @@ type AddMessagesResult struct {
 }
 
 func NewStream(ctx chasm.MutableContext, req NewStreamRequest) (*Stream, error) {
+	if req.StartOffset < 0 {
+		return nil, serviceerror.NewInvalidArgument("start offset cannot be negative")
+	}
 	visibility := chasm.NewEmptyField[*chasm.Visibility]()
 	if !req.Attached {
 		visibility = chasm.NewComponentField(ctx, chasm.NewVisibility(ctx))
@@ -105,10 +114,12 @@ func NewStream(ctx chasm.MutableContext, req NewStreamRequest) (*Stream, error)
 		Visibility: visibility,
 		Batches:    make(chasm.Map[int64, *commonpb.DataBlob]),
 		State: &streamlib.StreamState{
-			Lifecycle: req.Lifecycle,
-			Budget:    req.Budget,
-			Producers: make(map[string]*streamlib.ProducerCursor),
-			Consumers: make(map[string]*streamlib.ConsumerCursor),
+			HeadOffset: req.StartOffset,
+			BaseOffset: req.StartOffset,
+			Lifecycle:  req.Lifecycle,
+			Budget:     req.Budget,
+			Producers:  make(map[string]*streamlib.ProducerCursor),
+			Consumers:  make(map[string]*streamlib.ConsumerCursor),
 		},
 	}, nil
 }

chasm/lib/workflow/stream_admission_test.go

@@ -164,6 +164,50 @@ func TestSubscriptionsPerWorkflowAreBounded(t *testing.T) {
 	require.NoError(t, subscribe("a"))
 }
 
+// A reset run's stream continues the base run's offset space, so it is born
+// with a head offset well above zero while holding nothing. Measuring the
+// budget against that head would put it over the moment it exists, and every
+// publish on the reset run would be refused.
+func TestAnInheritedStreamIsNotBornOverItsItemBudget(t *testing.T) {
+	ctx := newStreamBudgetTestContext()
+	w := &Workflow{}
+	limits := stream.Limits{
+		MaxOwnedStreamsPerWorkflow: 10,
+		MaxConsumersPerStream:      10,
+		OwnedStreamMaxItems:        3,
+		OwnedStreamMaxBytes:        1 << 20,
+	}
+
+	// The inherited cursor stands far past the item budget.
+	created, createErr := w.ownStreamFrom(ctx, DefaultStreamName, 100, limits)
+	require.NoError(t, createErr)
+	require.True(t, created)
+
+	_, err := w.AppendToOwnedStream(ctx, DefaultStreamName, stream.AddMessagesRequest{
+		Records: budgetTestRecords(8),
+		Limits:  limits,
+	})
+	require.NoError(t, err, "the stream holds nothing, so its whole budget is free")
+
+	owned := w.Streams[DefaultStreamName].Get(ctx)
+	require.Equal(t, int64(101), owned.State.GetHeadOffset())
+
+	// Still bounded against what it holds.
+	for range 2 {
+		_, err = w.AppendToOwnedStream(ctx, DefaultStreamName, stream.AddMessagesRequest{
+			Records: budgetTestRecords(8),
+			Limits:  limits,
+		})
+		require.NoError(t, err)
+	}
+	_, err = w.AppendToOwnedStream(ctx, DefaultStreamName, stream.AddMessagesRequest{
+		Records: budgetTestRecords(8),
+		Limits:  limits,
+	})
+	var exhausted *serviceerror.ResourceExhausted
+	require.ErrorAs(t, err, &exhausted)
+}
+
 // An activity id may contain a slash, so the reserved key escapes it. Without
 // that two different (activity, name) pairs would land on one stream.
 func TestActivityStreamKeysDoNotCollide(t *testing.T) {

chasm/lib/workflow/stream_commands.go

@@ -400,6 +400,13 @@ func (w *Workflow) RecordStreamRecordsAppended(
 // streamNamed returns the workflow's stream of that name, creating it on first
 // use. Implicit creation is deliberate: a workflow publishing to its own output
 // should not have to coordinate with anyone about who creates it.
+//
+// The stream belongs to this run. After a reset the new run publishes to and
+// subscribes on streams of its own, created here on first use, and the run it
+// was reset from keeps the records its own history refers to. A stream the
+// reset run inherited a subscription to is created by the reset itself, at the
+// offset that subscription stood at, so it is already here by the time a
+// command names it.
 func (w *Workflow) streamNamed(
 	ctx chasm.MutableContext,
 	name string,