diff --git a/purego/internal/stream/ackmodel.go b/purego/internal/stream/ackmodel.go index e4b9e8aa..3b17bcd8 100644 --- a/purego/internal/stream/ackmodel.go +++ b/purego/internal/stream/ackmodel.go @@ -1,28 +1,197 @@ package stream import ( + "fmt" + "time" + "github.com/databricks/zerobus-sdk/purego/internal/zerobuspb" ) -// ackModel extracts connection-local ack offsets and classifies other responses. +// SubmittedRange describes one logical item submitted on the active connection. +// UnitStart is inclusive and UnitEnd is the exclusive end actually submitted +// on this connection. ItemUnitEnd is the exclusive end of the complete logical +// item; it differs from UnitEnd only when a multi-frame Send fails after +// submitting a prefix. The core always sets it, and a zero value normalizes to +// UnitEnd so a hand-built range cannot report a silently truncated item. // -// TODO(arrow): the Arrow wire path will supply its own ackModel over a Flight -// response, mapping a cumulative record count back to an offset. +// The range is protocol-neutral: proto/JSON assign one unit to every atomic +// request, while a record-count protocol can assign one unit per record. +type SubmittedRange struct { + WireOffset int64 + LogicalOffset int64 + UnitStart uint64 + UnitEnd uint64 + ItemUnitEnd uint64 +} + +// AckState is the connection-local state supplied to an acknowledgment model. +// Ranges contains submitted logical items that have not been fully +// acknowledged, in submission order. Active is the item whose Send is still +// outstanding; it is held separately so that acknowledging mid-send costs +// nothing rather than copying the whole outstanding window into a combined +// slice, which under continuous ingestion would be most acknowledgments. +// Iterate both in submission order with NumRanges and RangeAt. A model must +// not retain or mutate Ranges. +type AckState struct { + Ranges []SubmittedRange + Active SubmittedRange + HasActive bool + AcknowledgedUnits uint64 + SubmittedUnits uint64 +} + +// NumRanges returns the number of submitted ranges, counting Active. +func (s AckState) NumRanges() int { + if s.HasActive { + return len(s.Ranges) + 1 + } + return len(s.Ranges) +} + +// RangeAt returns the i-th submitted range in submission order. Active sorts +// after everything in Ranges, since its Send started last. +func (s AckState) RangeAt(i int) SubmittedRange { + if s.HasActive && i == len(s.Ranges) { + return s.Active + } + return s.Ranges[i] +} + +// AckResolution translates one cumulative protocol acknowledgment into logical +// durability progress. FullyAcknowledgedOffset and PartialOffset are -1 when +// absent. PartialUnits is the acknowledged prefix of PartialOffset, measured +// from that item's UnitStart. +type AckResolution struct { + AcknowledgedUnits uint64 + FullyAcknowledgedOffset int64 + PartialOffset int64 + PartialUnits uint64 +} + +type invalidAcknowledgment struct { + cause error +} + +func (e *invalidAcknowledgment) Error() string { + return "stream: invalid acknowledgment: " + e.cause.Error() +} + +func (e *invalidAcknowledgment) Unwrap() error { return e.cause } + +func (*invalidAcknowledgment) IsRetryable() bool { return false } + +// ResolveAcknowledgedUnits maps a cumulative connection-local unit count onto +// logical offsets. It is exported within the internal package boundary so a +// record-count protocol can share the core's range validation and resolution. +func ResolveAcknowledgedUnits(ackedUnits uint64, state AckState) (AckResolution, error) { + resolution := AckResolution{ + AcknowledgedUnits: ackedUnits, + FullyAcknowledgedOffset: -1, + PartialOffset: -1, + } + if ackedUnits > state.SubmittedUnits { + return resolution, fmt.Errorf( + "stream: server ack claims %d units, but only %d units were submitted", + ackedUnits, state.SubmittedUnits, + ) + } + if state.AcknowledgedUnits > state.SubmittedUnits { + return resolution, fmt.Errorf( + "stream: acknowledged-unit watermark %d exceeds submitted units %d", + state.AcknowledgedUnits, state.SubmittedUnits, + ) + } + if ackedUnits <= state.AcknowledgedUnits { + return resolution, nil + } + +resolveLoop: + for i := range state.NumRanges() { + submitted := state.RangeAt(i) + if submitted.UnitStart >= submitted.UnitEnd { + return resolution, fmt.Errorf( + "stream: invalid submitted range [%d,%d) for logical offset %d", + submitted.UnitStart, submitted.UnitEnd, submitted.LogicalOffset, + ) + } + itemUnitEnd := submitted.ItemUnitEnd + if itemUnitEnd == 0 { + itemUnitEnd = submitted.UnitEnd + } + if itemUnitEnd < submitted.UnitEnd { + return resolution, fmt.Errorf( + "stream: logical item end %d precedes submitted range end %d for logical offset %d", + itemUnitEnd, submitted.UnitEnd, submitted.LogicalOffset, + ) + } + if i > 0 && submitted.UnitStart != state.RangeAt(i-1).UnitEnd { + return resolution, fmt.Errorf( + "stream: submitted unit range starts at %d after %d", + submitted.UnitStart, state.RangeAt(i-1).UnitEnd, + ) + } + switch { + case ackedUnits >= submitted.UnitEnd: + if submitted.UnitEnd < itemUnitEnd { + resolution.PartialOffset = submitted.LogicalOffset + resolution.PartialUnits = submitted.UnitEnd - submitted.UnitStart + return resolution, nil + } + resolution.FullyAcknowledgedOffset = submitted.LogicalOffset + case ackedUnits > submitted.UnitStart: + resolution.PartialOffset = submitted.LogicalOffset + resolution.PartialUnits = ackedUnits - submitted.UnitStart + return resolution, nil + default: + // The ack stops at or below this range's start, so neither this + // range nor any contiguous later one can resolve it. Whatever + // earlier ranges resolved still stands; if nothing did, the ack + // lands in a gap and the check below rejects it. + break resolveLoop + } + } + if resolution.FullyAcknowledgedOffset < 0 && resolution.PartialOffset < 0 { + return resolution, fmt.Errorf( + "stream: ack of %d units does not intersect an unacknowledged submitted range", + ackedUnits, + ) + } + return resolution, nil +} + +// responseClassification keeps durability progress and connection rotation +// orthogonal. A single server response may carry either signal or both. +type responseClassification struct { + hasAck bool + legacyOffset int64 + pause *pauseSignal + failure responseFailure +} + +// ackModel classifies responses and extracts the legacy connection-local wire +// offset used by proto/JSON. A record-count model additionally implements +// resolvingAckModel to translate its response against submitted ranges. type ackModel[Resp any] interface { - // classify reports what a server response means to the core: an ack offset, - // a pause request, or an unknown/malformed response the receiver must fail - // on. Detecting these here keeps the receiver blind to concrete proto types. - classify(resp Resp) (kind respKind, off int64, pause pauseSignal) + // classify reports independent ack and pause signals, or an + // unknown/malformed response the receiver must fail on. + classify(resp Resp) responseClassification } -// respKind is the category the ackModel assigns to a server response. -type respKind int +// resolvingAckModel is the record-count acknowledgment extension point. The +// returned AcknowledgedUnits must use the same connection-local unit domain as +// AckState.Ranges. +type resolvingAckModel[Resp any] interface { + resolve(resp Resp, state AckState) (AckResolution, error) +} + +// responseFailure is the unusable category assigned to a server response. +// Its zero value means the response is structurally usable. +type responseFailure int const ( - ackResponse respKind = iota // carries a durability ack offset - pauseResponse // server-requested pause (close-stream signal) - unknownResponse // unrecognized response type — receiver fails - malformedResponse // ack missing its offset field — receiver fails + usableResponse responseFailure = iota + unknownResponse + malformedResponse ) // ephemeralResp is the proto/JSON server response type. Aliased so the core's @@ -32,25 +201,101 @@ type ephemeralResp = *zerobuspb.EphemeralStreamResponse // offsetAckModel extracts proto/JSON physical offsets and pause signals. type offsetAckModel struct{} -func (offsetAckModel) classify(resp ephemeralResp) (respKind, int64, pauseSignal) { +func (offsetAckModel) classify(resp ephemeralResp) responseClassification { if resp == nil { - return unknownResponse, 0, pauseSignal{} + return responseClassification{failure: unknownResponse} } if sig := resp.GetCloseStreamSignal(); sig != nil { - return pauseResponse, 0, pauseSignal{duration: sig.GetDuration().AsDuration()} + pause := pauseSignal{duration: sig.GetDuration().AsDuration()} + return responseClassification{pause: &pause} } if ack := resp.GetIngestRecordResponse(); ack != nil { // Absent offset must be malformed, not a fabricated ack for offset 0. if ack.DurabilityAckUpToOffset == nil { - return malformedResponse, 0, pauseSignal{} + return responseClassification{failure: malformedResponse} } off := *ack.DurabilityAckUpToOffset if off < 0 { - return malformedResponse, 0, pauseSignal{} + return responseClassification{failure: malformedResponse} + } + return responseClassification{hasAck: true, legacyOffset: off} + } + return responseClassification{failure: unknownResponse} +} + +func resolveOffsetAck(offset int64, state AckState) (AckResolution, error) { + if offset < 0 { + return AckResolution{}, fmt.Errorf("stream: negative ack offset %d", offset) + } + return ResolveAcknowledgedUnits(uint64(offset)+1, state) +} + +// ResponseStatus is the exported/internal structural status used by +// ResponseClassification. +type ResponseStatus int + +const ( + // ResponseOK means the response contains at least one usable signal. + ResponseOK ResponseStatus = iota + // ResponseUnknown is an unrecognized response. + ResponseUnknown + // ResponseMalformed is a recognized response with invalid fields. + ResponseMalformed +) + +// ResponseClassification describes independent durability and rotation signals. +// A response with both HasAck and HasPause set applies the acknowledgment first. +type ResponseClassification struct { + Status ResponseStatus + HasAck bool + LegacyOffset int64 + HasPause bool + PauseDuration time.Duration +} + +// AckModelHooks adapts protocol functions into the stream's acknowledgment +// seam. Resolve is optional for atomic offset protocols and required for +// record-count protocols. +type AckModelHooks[Resp any] struct { + Classify func(resp Resp) ResponseClassification + Resolve func(resp Resp, state AckState) (AckResolution, error) +} + +type hookAckModel[Resp any] struct { + hooks AckModelHooks[Resp] +} + +func (m hookAckModel[Resp]) classify(resp Resp) responseClassification { + classified := m.hooks.Classify(resp) + switch classified.Status { + case ResponseUnknown: + return responseClassification{failure: unknownResponse} + case ResponseMalformed: + return responseClassification{failure: malformedResponse} + case ResponseOK: + if !classified.HasAck && !classified.HasPause { + return responseClassification{failure: unknownResponse} } - return ackResponse, off, pauseSignal{} + result := responseClassification{ + hasAck: classified.HasAck, + legacyOffset: classified.LegacyOffset, + } + if classified.HasPause { + pause := pauseSignal{duration: classified.PauseDuration} + result.pause = &pause + } + return result + default: + return responseClassification{failure: unknownResponse} } - return unknownResponse, 0, pauseSignal{} +} + +type resolvingHookAckModel[Resp any] struct { + hookAckModel[Resp] +} + +func (m resolvingHookAckModel[Resp]) resolve(resp Resp, state AckState) (AckResolution, error) { + return m.hooks.Resolve(resp, state) } // newAckModel returns the proto/JSON ack model for the given record type. diff --git a/purego/internal/stream/ackmodel_test.go b/purego/internal/stream/ackmodel_test.go index ca76434f..f954e4c5 100644 --- a/purego/internal/stream/ackmodel_test.go +++ b/purego/internal/stream/ackmodel_test.go @@ -18,12 +18,12 @@ func TestClassifyAck(t *testing.T) { }, }, } - kind, off, _ := offsetAckModel{}.classify(resp) - if kind != ackResponse { - t.Fatalf("want ackResponse, got %v", kind) + classified := offsetAckModel{}.classify(resp) + if !classified.hasAck || classified.failure != usableResponse { + t.Fatalf("want usable ack response, got %+v", classified) } - if off != 42 { - t.Fatalf("want offset 42, got %d", off) + if classified.legacyOffset != 42 { + t.Fatalf("want offset 42, got %d", classified.legacyOffset) } } @@ -37,9 +37,9 @@ func TestClassifyAckOffsetZero(t *testing.T) { }, }, } - kind, off, _ := offsetAckModel{}.classify(resp) - if kind != ackResponse || off != 0 { - t.Fatalf("want ackResponse offset 0, got kind=%v off=%d", kind, off) + classified := offsetAckModel{}.classify(resp) + if !classified.hasAck || classified.legacyOffset != 0 { + t.Fatalf("want ack response offset 0, got %+v", classified) } } @@ -51,12 +51,12 @@ func TestClassifyPause(t *testing.T) { }, }, } - kind, _, pause := offsetAckModel{}.classify(resp) - if kind != pauseResponse { - t.Fatalf("want pauseResponse, got %v", kind) + classified := offsetAckModel{}.classify(resp) + if classified.pause == nil || classified.failure != usableResponse { + t.Fatalf("want usable pause response, got %+v", classified) } - if pause.duration != 3*time.Second { - t.Fatalf("want 3s pause, got %v", pause.duration) + if classified.pause.duration != 3*time.Second { + t.Fatalf("want 3s pause, got %v", classified.pause.duration) } } @@ -68,12 +68,12 @@ func TestClassifyMalformedAckMissingOffset(t *testing.T) { IngestRecordResponse: &zerobuspb.IngestRecordResponse{}, // offset absent }, } - kind, off, _ := offsetAckModel{}.classify(resp) - if kind != malformedResponse { - t.Fatalf("want malformedResponse for absent offset, got %v", kind) + classified := offsetAckModel{}.classify(resp) + if classified.failure != malformedResponse { + t.Fatalf("want malformedResponse for absent offset, got %+v", classified) } - if off != 0 { - t.Fatalf("want offset 0 for malformed, got %d", off) + if classified.legacyOffset != 0 { + t.Fatalf("want offset 0 for malformed, got %d", classified.legacyOffset) } } @@ -85,12 +85,12 @@ func TestClassifyMalformedAckNegativeOffset(t *testing.T) { }, }, } - kind, off, _ := offsetAckModel{}.classify(resp) - if kind != malformedResponse { - t.Fatalf("want malformedResponse for negative offset, got %v", kind) + classified := offsetAckModel{}.classify(resp) + if classified.failure != malformedResponse { + t.Fatalf("want malformedResponse for negative offset, got %+v", classified) } - if off != 0 { - t.Fatalf("want offset 0 for malformed, got %d", off) + if classified.legacyOffset != 0 { + t.Fatalf("want offset 0 for malformed, got %d", classified.legacyOffset) } } @@ -98,9 +98,9 @@ func TestClassifyMalformedAckNegativeOffset(t *testing.T) { // can fail the stream on a wire-contract mismatch. func TestClassifyUnknown(t *testing.T) { for _, resp := range []ephemeralResp{nil, {}} { - kind, _, _ := offsetAckModel{}.classify(resp) - if kind != unknownResponse { - t.Fatalf("want unknownResponse, got %v", kind) + classified := offsetAckModel{}.classify(resp) + if classified.failure != unknownResponse { + t.Fatalf("want unknownResponse, got %+v", classified) } } } diff --git a/purego/internal/stream/buffer.go b/purego/internal/stream/buffer.go index 36dd4645..78e4a612 100644 --- a/purego/internal/stream/buffer.go +++ b/purego/internal/stream/buffer.go @@ -1,12 +1,8 @@ // Package stream is the generic ingestion core: offset assignment, send/recv // goroutines, ack watermark, Flush/WaitForOffset, and the recovery supervisor. -// Protocol-specific behaviour (encoding, ack parsing, wire transport) is -// injected through the encoder, ackModel, and wireStream interfaces, so -// proto and JSON share one implementation. -// -// Arrow Flight will reuse these seams but not unchanged: buffer entries carry no -// record count, and recovery replays whole entries, so a partially acknowledged -// batch cannot be sliced. Both are core changes, not encoder changes. +// Protocol-specific behaviour (encoding, ack parsing, payload slicing, and wire +// transport) is injected through narrow hooks, so atomic offset protocols and +// record-count protocols share one implementation. package stream import ( @@ -15,20 +11,24 @@ import ( "fmt" "math" "sync" + "time" ) // defaultMaxInflight is the fallback backpressure cap used when a non-positive // value reaches newBuffer. const defaultMaxInflight = 1_000_000 -// item is one unit of work in the buffer: an already-encoded wire message -// paired with the logical offset that identifies it for acknowledgment. Req is -// the wire request type the core is instantiated with (encodedMsg for -// proto/JSON; a Flight frame for Arrow). +// item is one logical unit of work in the buffer: an already-encoded payload +// paired with its SDK offset and protocol durability-unit count. Proto/JSON +// payloads are atomic and always carry one unit; a record-count protocol can +// carry multiple units and acknowledge a prefix. type item[Req any] struct { - offset int64 - payload Req - weight int64 + offset int64 + payload Req + units uint64 + ackedUnits uint64 + weight int64 + pendingAt time.Time } type discardResult struct { @@ -54,9 +54,9 @@ type capacityWaiter struct { // Concurrency model: // - Many goroutines may call enqueue concurrently. // - Exactly one sender goroutine calls next. -// - Exactly one receiver goroutine calls discardThrough as acks arrive. -// - The supervisor calls requeue and drain, but only while the sender is -// stopped — so next never runs concurrently with requeue or drain. +// - Exactly one receiver goroutine calls acknowledge as acks arrive. +// - The supervisor calls requeueWithSlicer and drain, but only while the +// sender is stopped — so next never runs concurrently with either. // // All state (queue, flight, capacity, cond) is private; the sender and receiver // interact only through these methods, never by touching the fields directly. @@ -74,6 +74,7 @@ type buffer[Req any] struct { usedBytes int64 accountingReset bool waiters *list.List + flightRevision uint64 } func newBuffer[Req any](maxInflight int, byteLimit int64) *buffer[Req] { @@ -184,8 +185,40 @@ func (b *buffer[Req]) release(weight int64) { b.mu.Unlock() } -// append adds an item after reserve succeeds. -func (b *buffer[Req]) append(offset int64, msg Req, weight int64) error { +// reconcileReservation replaces an admission estimate with the actual retained +// payload size. Shrinks are always applied immediately. An underestimated +// reservation may grow only when the unreserved byte budget is already +// available; it never waits while holding an item slot, which would deadlock if +// several concurrent builders all underestimated at the count limit. +func (b *buffer[Req]) reconcileReservation(estimated, actual int64) error { + if actual < 0 || actual > b.maxBufferedBytes { + return fmt.Errorf("%w: buffered payload weight %d exceeds limit %d", + ErrPayloadTooLarge, actual, b.maxBufferedBytes) + } + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return errClosed + } + delta := actual - estimated + if delta > b.maxBufferedBytes-b.usedBytes { + return fmt.Errorf( + "%w: actual retained payload weight %d exceeds reserved estimate %d and remaining byte capacity", + ErrPayloadTooLarge, + actual, + estimated, + ) + } + b.usedBytes += delta + if delta < 0 { + b.grantWaitersLocked() + } + return nil +} + +// appendUnits adds an item with its protocol durability-unit count after reserve +// succeeds. +func (b *buffer[Req]) appendUnits(offset int64, msg Req, units uint64, weight int64) error { b.mu.Lock() if b.closed { if !b.accountingReset { @@ -196,15 +229,17 @@ func (b *buffer[Req]) append(offset int64, msg Req, weight int64) error { b.cond.Broadcast() return errClosed } - b.queue = append(b.queue, item[Req]{offset: offset, payload: msg, weight: weight}) + b.queue = append(b.queue, item[Req]{ + offset: offset, payload: msg, units: units, weight: weight, + }) b.mu.Unlock() b.cond.Signal() return nil } // next blocks until a pending item is available and moves it to the in-flight -// list, returning the item. The sender must later call discard (on ack) or -// requeue (on reconnect) for every item returned by next. +// list, returning the item. Every item it returns is later either acknowledged +// or requeued for replay; see the concurrency model on buffer. // // Returns errClosed when the buffer has been closed and drained. // Returns ctx.Err() if ctx is cancelled while waiting. @@ -242,21 +277,60 @@ func (b *buffer[Req]) next(ctx context.Context) (item[Req], error) { it := b.queue[0] b.queue[0] = item[Req]{} // release the departed slot's payload for GC b.queue = b.queue[1:] + it.pendingAt = time.Now() b.flight = append(b.flight, it) + b.flightRevision++ b.mu.Unlock() return it, nil } -// discardThrough removes every in-flight item whose offset is <= offset (all -// now acknowledged by the server), releases its count and byte capacity, and -// grants queued admission waiters in FIFO order. It is the receiver's only hook -// for ack-driven eviction. Returns the contiguous discarded offset range -// without allocating per-item callback metadata. -func (b *buffer[Req]) discardThrough(offset int64) discardResult { +// acknowledge applies logical durability progress to the in-flight prefix, +// releasing count and byte capacity for every fully acknowledged item and +// granting queued admission waiters in FIFO order. A partial acknowledgment is +// retained on the first unacknowledged item so recovery or GetUnacked can slice +// its payload through the protocol hook. It returns the contiguous discarded +// offset range without allocating per-item callback metadata. +func (b *buffer[Req]) acknowledge( + resolution AckResolution, +) (discardResult, bool, error) { b.mu.Lock() + defer b.mu.Unlock() + + discardCount := 0 + if resolution.FullyAcknowledgedOffset >= 0 { + for discardCount < len(b.flight) && + b.flight[discardCount].offset <= resolution.FullyAcknowledgedOffset { + discardCount++ + } + if discardCount == 0 || + b.flight[discardCount-1].offset != resolution.FullyAcknowledgedOffset { + return discardResult{}, false, fmt.Errorf( + "stream: fully acknowledged logical offset %d is not in flight", + resolution.FullyAcknowledgedOffset, + ) + } + } + if resolution.PartialOffset >= 0 { + if discardCount >= len(b.flight) || + b.flight[discardCount].offset != resolution.PartialOffset { + return discardResult{}, false, fmt.Errorf( + "stream: partially acknowledged logical offset %d is not first in flight", + resolution.PartialOffset, + ) + } + partial := b.flight[discardCount] + if resolution.PartialUnits == 0 || + resolution.PartialUnits >= partial.units { + return discardResult{}, false, fmt.Errorf( + "stream: partial prefix %d is invalid for logical offset %d with %d units", + resolution.PartialUnits, resolution.PartialOffset, partial.units, + ) + } + } + var result discardResult var releasedBytes int64 - for len(b.flight) > 0 && b.flight[0].offset <= offset { + for range discardCount { if result.count == 0 { result.first = b.flight[0].offset } @@ -266,45 +340,97 @@ func (b *buffer[Req]) discardThrough(offset int64) discardResult { b.flight[0] = item[Req]{} // release the acked payload for GC b.flight = b.flight[1:] } + + progressed := result.count > 0 + if resolution.PartialOffset >= 0 && + resolution.PartialUnits > b.flight[0].ackedUnits { + b.flight[0].ackedUnits = resolution.PartialUnits + progressed = true + } + if progressed && len(b.flight) > 0 { + // Durable progress refreshes the head's lack-of-ack budget, whether it + // landed inside the head or promoted a new one: ordered acks mean a + // promoted item could not have been made durable sooner. A stale ack + // makes no progress and never reaches here, so a stalled server cannot + // postpone recovery. + b.flight[0].pendingAt = time.Now() + } + if progressed { + b.flightRevision++ + } b.usedItems -= result.count b.usedBytes -= releasedBytes b.grantWaitersLocked() - b.mu.Unlock() - return result + return result, progressed, nil } -// requeue moves all in-flight items back to the front of the pending queue so -// they are re-sent after a reconnect. Called by the supervisor on stream failure. -func (b *buffer[Req]) requeue() { - b.mu.Lock() - defer b.mu.Unlock() - if len(b.flight) == 0 { - return - } - // Prepend in-flight items (in order) before any still-pending ones. - requeued := make([]item[Req], 0, len(b.flight)+len(b.queue)) - requeued = append(requeued, b.flight...) - requeued = append(requeued, b.queue...) - b.queue = requeued - // Zero departed slots so payload references in the old backing array become - // GC-collectible after flight is reset. - for i := range b.flight { - b.flight[i] = item[Req]{} - } - b.flight = b.flight[:0] - b.cond.Broadcast() -} +// requeueWithSlicer moves all in-flight items back to the front of the pending +// queue so they are re-sent after a reconnect. Called by the supervisor on +// stream failure. +// +// Partially acknowledged items are sliced outside b.mu, then validated against +// a revision counter before the snapshot is installed, so concurrent queue +// admission stays available while an expensive re-encode runs. A nil slice +// function rejects partially acknowledged items, which is correct for protocols +// whose payloads are atomic. The retained byte charge stays conservative until +// the item is fully discarded. +func (b *buffer[Req]) requeueWithSlicer( + slice func(payload Req, acknowledgedPrefix uint64) (Req, error), +) error { + for { + b.mu.Lock() + if len(b.flight) == 0 { + b.mu.Unlock() + return nil + } + revision := b.flightRevision + snapshot := append([]item[Req](nil), b.flight...) + b.mu.Unlock() -// highestInFlight returns the greatest offset the sender has observed on the -// current connection. Pending records are deliberately excluded: the server -// cannot legitimately acknowledge work that has not entered the send path. -func (b *buffer[Req]) highestInFlight() (int64, bool) { - b.mu.Lock() - defer b.mu.Unlock() - if len(b.flight) == 0 { - return 0, false + for i := range snapshot { + if snapshot[i].ackedUnits == 0 { + snapshot[i].pendingAt = time.Time{} + continue + } + if slice == nil { + return fmt.Errorf( + "stream: no payload slicer for partially acknowledged logical offset %d", + snapshot[i].offset, + ) + } + payload, err := slice(snapshot[i].payload, snapshot[i].ackedUnits) + if err != nil { + return fmt.Errorf( + "stream: slice logical offset %d after %d acknowledged units: %w", + snapshot[i].offset, snapshot[i].ackedUnits, err, + ) + } + snapshot[i].payload = payload + snapshot[i].units -= snapshot[i].ackedUnits + snapshot[i].ackedUnits = 0 + snapshot[i].pendingAt = time.Time{} + } + + b.mu.Lock() + if b.flightRevision != revision { + b.mu.Unlock() + continue + } + // Prepend the validated in-flight snapshot before every item admitted + // while the transform ran. + requeued := make([]item[Req], 0, len(snapshot)+len(b.queue)) + requeued = append(requeued, snapshot...) + requeued = append(requeued, b.queue...) + b.queue = requeued + for i := range b.flight { + b.flight[i] = item[Req]{} + } + b.flight = b.flight[:0] + b.flightRevision++ + b.mu.Unlock() + b.cond.Broadcast() + return nil } - return b.flight[len(b.flight)-1].offset, true } // drain returns all items currently in the buffer (pending + in-flight) and @@ -321,6 +447,7 @@ func (b *buffer[Req]) drain() []item[Req] { all = append(all, b.queue...) b.queue = nil b.flight = nil + b.flightRevision++ b.closed = true b.accountingReset = true b.usedItems = 0 @@ -350,3 +477,16 @@ func (b *buffer[Req]) inFlight() int { defer b.mu.Unlock() return len(b.flight) } + +// oldestInFlightDeadline returns the absolute deadline of the oldest pending +// item. Durable progress refreshes that item's budget, while an ack repeating +// known progress does not. Replay assigns a fresh connection-local pendingAt +// when next observes the item again. +func (b *buffer[Req]) oldestInFlightDeadline(timeout time.Duration) (time.Time, bool) { + b.mu.Lock() + defer b.mu.Unlock() + if len(b.flight) == 0 || b.flight[0].pendingAt.IsZero() { + return time.Time{}, false + } + return b.flight[0].pendingAt.Add(timeout), true +} diff --git a/purego/internal/stream/buffer_test.go b/purego/internal/stream/buffer_test.go index e993f7f9..7ae2e599 100644 --- a/purego/internal/stream/buffer_test.go +++ b/purego/internal/stream/buffer_test.go @@ -2,6 +2,7 @@ package stream import ( "context" + "fmt" "sync" "sync/atomic" "testing" @@ -130,6 +131,36 @@ func TestBufferByteBackpressureAndRelease(t *testing.T) { } } +func TestBufferReservationReconciliationNeverOvercommits(t *testing.T) { + b := newBuffer[string](3, 10) + if err := b.reserve(context.Background(), 3); err != nil { + t.Fatalf("reserve: %v", err) + } + if err := b.reconcileReservation(3, 7); err != nil { + t.Fatalf("grow available reservation: %v", err) + } + if items, usedBytes := b.usage(); items != 1 || usedBytes != 7 { + t.Fatalf("grown usage = (%d,%d), want (1,7)", items, usedBytes) + } + if err := b.appendUnits(0, "first", 1, 7); err != nil { + t.Fatalf("append grown reservation: %v", err) + } + + if err := b.reserve(context.Background(), 3); err != nil { + t.Fatalf("reserve remaining bytes: %v", err) + } + if err := b.reconcileReservation(3, 4); err == nil { + t.Fatal("underestimated reservation overcommitted byte limit") + } + if items, usedBytes := b.usage(); items != 2 || usedBytes != 10 { + t.Fatalf("usage after rejected growth = (%d,%d), want (2,10)", items, usedBytes) + } + b.release(3) + if items, usedBytes := b.usage(); items != 1 || usedBytes != 7 { + t.Fatalf("usage after rollback = (%d,%d), want (1,7)", items, usedBytes) + } +} + func TestBufferCapacityWaitersAreFIFO(t *testing.T) { b := newBuffer[encodedMsg](3, 3) if err := b.reserve(context.Background(), 3); err != nil { @@ -210,6 +241,86 @@ func TestBufferRecoveryDoesNotDoubleChargeBytes(t *testing.T) { } } +func TestBufferRecoverySlicesOutsideMutex(t *testing.T) { + b := newBuffer[string](3, 30) + if err := b.reserve(context.Background(), 10); err != nil { + t.Fatalf("reserve first: %v", err) + } + if err := b.appendUnits(0, "full", 5, 10); err != nil { + t.Fatalf("append first: %v", err) + } + if _, err := b.next(context.Background()); err != nil { + t.Fatalf("next first: %v", err) + } + if _, _, err := b.acknowledge(AckResolution{ + FullyAcknowledgedOffset: -1, + PartialOffset: 0, + PartialUnits: 2, + AcknowledgedUnits: 2, + }); err != nil { + t.Fatalf("partial acknowledge: %v", err) + } + + sliceStarted := make(chan struct{}) + releaseSlice := make(chan struct{}) + requeueDone := make(chan error, 1) + go func() { + requeueDone <- b.requeueWithSlicer(func( + payload string, + acknowledgedPrefix uint64, + ) (string, error) { + close(sliceStarted) + <-releaseSlice + if payload != "full" || acknowledgedPrefix != 2 { + return "", fmt.Errorf( + "slice input = (%q,%d), want (full,2)", + payload, + acknowledgedPrefix, + ) + } + return "suffix", nil + }) + }() + <-sliceStarted + + // Queue admission must not wait for the decode/slice/re-encode transform. + admitted := make(chan error, 1) + go func() { + if err := b.reserve(context.Background(), 10); err != nil { + admitted <- err + return + } + admitted <- b.appendUnits(1, "later", 1, 10) + }() + select { + case err := <-admitted: + if err != nil { + t.Fatalf("concurrent admission: %v", err) + } + case <-time.After(time.Second): + t.Fatal("concurrent admission blocked behind recovery slicer") + } + + close(releaseSlice) + if err := <-requeueDone; err != nil { + t.Fatalf("requeueWithSlicer: %v", err) + } + first, err := b.next(context.Background()) + if err != nil { + t.Fatalf("next sliced: %v", err) + } + if first.payload != "suffix" || first.units != 3 || first.ackedUnits != 0 { + t.Fatalf("sliced item = %+v", first) + } + second, err := b.next(context.Background()) + if err != nil { + t.Fatalf("next concurrent: %v", err) + } + if second.payload != "later" || second.offset != 1 { + t.Fatalf("concurrently admitted item = %+v", second) + } +} + func TestBufferDrainMakesReservationRollbackIdempotent(t *testing.T) { t.Run("release", func(t *testing.T) { b := newBuffer[encodedMsg](4, 10) diff --git a/purego/internal/stream/core.go b/purego/internal/stream/core.go index 0196d3d2..5cd8c068 100644 --- a/purego/internal/stream/core.go +++ b/purego/internal/stream/core.go @@ -249,10 +249,11 @@ func (w *watermark) waitFor(ctx context.Context, target int64) error { // CoreStream is the protocol-agnostic ingestion core. It owns the buffer, // sender goroutine, receiver goroutine, ack watermark, and the supervisor // that reconnects on failure. It is generic over the wire request/response -// types (Req/Resp): proto/JSON instantiate it over EphemeralStream, Arrow over -// Flight. The three specialization points — encoder, ackModel, and the -// wireStream returned by opener — are injected, so this core is written once -// and never names a concrete proto type. +// types (Req/Resp): proto/JSON instantiate it over EphemeralStream. The +// specialization points — encoder, ackModel (optionally also resolvingAckModel +// for record-count durability), and the wireStream returned by opener +// (optionally also submissionReceiptStream for multi-frame sends) — are +// injected, so this core is written once and never names a concrete proto type. // // The per-stream goroutines: // @@ -291,6 +292,10 @@ type CoreStream[Req, Resp any] struct { offsetExhausted atomic.Bool // lastEnqueued is the highest queued offset. lastEnqueued atomic.Int64 + // durableProgress advances for full or partial acknowledgment progress. It + // lets a partially acknowledged multi-unit item reset the recovery budget + // even though the public logical watermark cannot advance yet. + durableProgress atomic.Uint64 // done is closed when the supervisor exits (terminal state). done chan struct{} @@ -497,6 +502,75 @@ func (cs *CoreStream[Req, Resp]) enqueueEncoded( ctx context.Context, weight int64, encodeFn func() (Req, error), +) (int64, error) { + return cs.enqueuePayload(ctx, weight, 0, encodeFn) +} + +// EnqueuePayload admits an already-built protocol payload. units is the number +// of durability units represented by the logical item and retainedBytes is its +// conservative buffer-memory charge. This exported/internal extension point is +// for typed protocol wrappers whose input is not []byte. +func (cs *CoreStream[Req, Resp]) EnqueuePayload( + ctx context.Context, + payload Req, + units uint64, + retainedBytes int64, +) (int64, error) { + if units == 0 { + return -1, fmt.Errorf("stream: payload must contain at least one durability unit") + } + if reported := cs.enc.unitCount(payload); reported != units { + return -1, fmt.Errorf( + "stream: payload reports %d durability units, caller supplied %d", + reported, units, + ) + } + return cs.enqueuePayload(ctx, retainedBytes, units, func() (Req, error) { + return payload, nil + }) +} + +// EnqueuePayloadBuilder reserves one item and a conservative byte estimate +// before invoking build. It is the typed protocol path for inputs whose +// materialization is expensive (for example a serialize-and-compress step). +// +// build must return the payload, its durability-unit count, and its actual +// retained byte charge. The reservation is reconciled before the payload is +// published. Underestimates can consume only byte capacity that is immediately +// available; they never wait while retaining a built payload and an item slot. +func (cs *CoreStream[Req, Resp]) EnqueuePayloadBuilder( + ctx context.Context, + estimatedRetainedBytes int64, + build func() (payload Req, units uint64, retainedBytes int64, err error), +) (int64, error) { + if build == nil { + return -1, fmt.Errorf("stream: payload builder is required") + } + return cs.enqueuePayloadReserved( + ctx, + estimatedRetainedBytes, + func() (Req, uint64, int64, error) { + return build() + }, + ) +} + +func (cs *CoreStream[Req, Resp]) enqueuePayload( + ctx context.Context, + weight int64, + explicitUnits uint64, + build func() (Req, error), +) (int64, error) { + return cs.enqueuePayloadReserved(ctx, weight, func() (Req, uint64, int64, error) { + payload, err := build() + return payload, explicitUnits, weight, err + }) +} + +func (cs *CoreStream[Req, Resp]) enqueuePayloadReserved( + ctx context.Context, + estimatedWeight int64, + build func() (Req, uint64, int64, error), ) (int64, error) { if cs.isClosed() { if err := cs.terminalErr(); err != nil { @@ -513,15 +587,36 @@ func (cs *CoreStream[Req, Resp]) enqueueEncoded( if cs.offsetExhausted.Load() { return -1, ErrOffsetExhausted } - if err := cs.buf.reserve(ctx, weight); err != nil { + if err := cs.buf.reserve(ctx, estimatedWeight); err != nil { return -1, err } - msg, err := encodeFn() + msg, explicitUnits, actualWeight, err := build() if err != nil { - cs.buf.release(weight) + cs.buf.release(estimatedWeight) return -1, err } + if err := cs.buf.reconcileReservation(estimatedWeight, actualWeight); err != nil { + cs.buf.release(estimatedWeight) + return -1, err + } + weight := actualWeight + units := explicitUnits + reportedUnits := cs.enc.unitCount(msg) + if units == 0 { + units = reportedUnits + } else if reportedUnits != units { + cs.buf.release(weight) + return -1, fmt.Errorf( + "stream: payload reports %d durability units, builder supplied %d", + reportedUnits, + units, + ) + } + if units == 0 { + cs.buf.release(weight) + return -1, fmt.Errorf("stream: encoded payload contains no durability units") + } if size := cs.enc.maxWireSize(msg); size > cs.cfg.MaxPayloadBytes { cs.buf.release(weight) return -1, fmt.Errorf("%w: encoded size %d exceeds MaxPayloadBytes=%d", @@ -541,7 +636,7 @@ func (cs *CoreStream[Req, Resp]) enqueueEncoded( } offset := cs.nextOffset cs.enc.stampOffset(msg, offset) - if err := cs.buf.append(offset, msg, weight); err != nil { + if err := cs.buf.appendUnits(offset, msg, units, weight); err != nil { cs.offsetMu.Unlock() return -1, err } @@ -614,10 +709,21 @@ func (cs *CoreStream[Req, Resp]) GetUnackedBatches() ([][][]byte, error) { items := cs.consolidateUnacked() out := make([][][]byte, 0, len(items)) for _, it := range items { + payload := it.payload + if it.ackedUnits > 0 { + var err error + payload, err = cs.enc.slice(payload, it.ackedUnits) + if err != nil { + return nil, fmt.Errorf( + "stream: recover unacknowledged logical offset %d: %w", + it.offset, err, + ) + } + } // Re-extract the raw record bytes from the encoded message so callers get // back the original record content. decode clones, so the retained payload // is never aliased. - out = append(out, cs.enc.decode(it.payload)) + out = append(out, cs.enc.decode(payload)) } return out, nil } @@ -744,10 +850,22 @@ func (cs *CoreStream[Req, Resp]) consolidateUnacked() []item[Req] { type sendEvent struct { logicalOffset int64 physicalOffset int64 + unitStart uint64 + unitEnd uint64 + receiptUnits uint64 completed bool err error } +type partialSubmissionFailure struct { + cause error + unitStart uint64 + unitEnd uint64 +} + +func (e *partialSubmissionFailure) Error() string { return e.cause.Error() } +func (e *partialSubmissionFailure) Unwrap() error { return e.cause } + // runOnce operates one transport stream until a worker exits. lifecycleCtx owns // the live connection after open; openingCtx bounds only this open attempt. // resetRecoveryBudget reports durable progress or a stable connection. @@ -782,9 +900,14 @@ func (cs *CoreStream[Req, Resp]) runOnce( cs.signalReady(nil) openedAt := time.Now() startAck := cs.wm.current() + startDurableProgress := cs.durableProgress.Load() - // Resend unacknowledged items on the new connection. - cs.buf.requeue() + // Resend unacknowledged items on the new connection. A record-count + // protocol slices any durably acknowledged prefix before replay. + if err := cs.buf.requeueWithSlicer(cs.enc.slice); err != nil { + stream.Close() + return wrapValidation(err), false + } // senderCtx controls only this connection's sender. senderCtx, cancelSender := context.WithCancel(lifecycleCtx) @@ -798,11 +921,20 @@ func (cs *CoreStream[Req, Resp]) runOnce( // cap 2: at most a {start, completed} pair for the one record in flight. sendEvents := make(chan sendEvent, 2) sendConsumed := make(chan struct{}, 1) + ackProgress := make(chan uint64, 1) go cs.sender(senderCtx, stream, senderExitCh, flightSignal, sendEvents, sendConsumed) go func() { defer close(receiverDone) - cs.receiver(stream, receiverExitCh, pauseCh, flightSignal, sendEvents, sendConsumed) + cs.receiver( + stream, + receiverExitCh, + pauseCh, + flightSignal, + sendEvents, + sendConsumed, + ackProgress, + ) }() var senderParked bool @@ -882,12 +1014,49 @@ waitLoop: default: } case errors.As(cause, &ps): - stream.Close() - if !senderExited { - <-senderExitCh - } - <-receiverDone + // A server-requested rotation is orderly: park the sender, half-close + // requests, and drain responses to EOF before reconnecting. This lets the + // service finish late acknowledgments and observe END_STREAM instead of + // an abrupt cancellation. gracefulTeardown still hard-aborts if the + // shared drain budget expires. + cs.gracefulTeardown( + stream, senderExited, senderExitCh, receiverDone, + ) default: + // A multi-frame Send can fail after earlier frames were accepted. Keep + // Recv alive briefly so a server ACK that follows the Send error can + // establish durability for that submitted prefix before recovery slices + // and replays the remainder. + var partial *partialSubmissionFailure + waitedForReceiver := false + if !receiverReported && errors.As(cause, &partial) { + waitedForReceiver = true + timer := time.NewTimer(cs.cfg.DrainTimeout) + waiting := true + for waiting { + select { + case acknowledged := <-ackProgress: + if acknowledged > partial.unitStart && + acknowledged <= partial.unitEnd { + waiting = false + } + case receiverErr := <-receiverExitCh: + receiverReported = true + if receiverErr != nil && !errors.Is(receiverErr, context.Canceled) { + cause = receiverErr + } + waiting = false + case <-timer.C: + waiting = false + } + } + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + } // gRPC reports a server-side abort to Send as an opaque io.EOF and carries // the real status (auth, schema, protocol) on Recv, so an EOF cause is not // authoritative yet. Let the receiver report first, bounded by DrainTimeout: @@ -895,7 +1064,7 @@ waitLoop: // leaving recovery to retry a permanent rejection and skip credential // invalidation. A failed send does not end the receiver, so it is still // reading and the server's status is what unblocks it. - if !receiverReported && errors.Is(cause, io.EOF) { + if !receiverReported && !waitedForReceiver && errors.Is(cause, io.EOF) { timer := time.NewTimer(cs.cfg.DrainTimeout) select { case receiverErr := <-receiverExitCh: @@ -914,8 +1083,23 @@ waitLoop: <-senderExitCh } <-receiverDone + // The prefix-ACK exit can reach here without consulting the receiver, + // and Close unblocks a Recv holding the server's status. Only a + // definitive rejection wins; a retryable one is an artifact of the abort. + if !receiverReported { + select { + case receiverErr := <-receiverExitCh: + if receiverErr != nil && + !errors.Is(receiverErr, context.Canceled) && + !isRetryable(receiverErr) { + cause = receiverErr + } + default: + } + } } resetRecoveryBudget = cs.wm.current() > startAck || + cs.durableProgress.Load() > startDurableProgress || time.Since(openedAt) >= cs.cfg.RecoveryResetAfter return cause, resetRecoveryBudget } @@ -986,6 +1170,7 @@ func (cs *CoreStream[Req, Resp]) sender( } physicalOffset := int64(0) physicalExhausted := false + submittedUnits := uint64(0) for { if physicalExhausted { errCh <- fmt.Errorf("stream: physical offset space exhausted") @@ -996,9 +1181,15 @@ func (cs *CoreStream[Req, Resp]) sender( errCh <- nil // ctx cancelled or buffer closed — clean exit return } + if it.units > ^uint64(0)-submittedUnits { + errCh <- fmt.Errorf("stream: submitted durability-unit space exhausted") + return + } + unitEnd := submittedUnits + it.units cs.enc.stampOffset(it.payload, physicalOffset) if !publish(sendEvent{ logicalOffset: it.offset, physicalOffset: physicalOffset, + unitStart: submittedUnits, unitEnd: unitEnd, }) { errCh <- nil return @@ -1008,15 +1199,62 @@ func (cs *CoreStream[Req, Resp]) sender( case flightSignal <- struct{}{}: default: } - err = stream.Send(it.payload) + var receipt SubmissionReceipt + if receiptStream, ok := stream.(submissionReceiptStream[Req]); ok { + receipt, err = receiptStream.SendWithReceipt(it.payload) + } else { + err = stream.Send(it.payload) + if err == nil { + receipt.SubmittedUnits = it.units + } + } + // A receipt over the payload size, or a short receipt from a successful + // Send, is a transport defect: a replay would hit the same defect. + if receipt.SubmittedUnits > it.units { + err = wrapValidation(fmt.Errorf( + "stream: send offset %d reported %d submitted units for %d-unit payload", + it.offset, + receipt.SubmittedUnits, + it.units, + )) + } else if err == nil && receipt.SubmittedUnits != it.units { + err = wrapValidation(fmt.Errorf( + "stream: send offset %d completed with %d of %d units submitted", + it.offset, + receipt.SubmittedUnits, + it.units, + )) + } sendEvents <- sendEvent{ logicalOffset: it.offset, physicalOffset: physicalOffset, + unitStart: submittedUnits, + unitEnd: unitEnd, + receiptUnits: receipt.SubmittedUnits, completed: true, err: err, } + // The receiver must consume the receipt before the supervisor tears down + // a failed multi-frame send, or an ACK already received for an earlier + // frame could be lost. + select { + case <-sendConsumed: + case <-senderCtx.Done(): + if err == nil { + errCh <- nil + return + } + } if err != nil { - errCh <- fmt.Errorf("stream: send offset %d: %w", it.offset, err) + sendErr := err + if receipt.SubmittedUnits > 0 { + sendErr = &partialSubmissionFailure{ + cause: err, + unitStart: submittedUnits, + unitEnd: submittedUnits + receipt.SubmittedUnits, + } + } + errCh <- fmt.Errorf("stream: send offset %d: %w", it.offset, sendErr) return } if physicalOffset == math.MaxInt64 { @@ -1024,12 +1262,7 @@ func (cs *CoreStream[Req, Resp]) sender( } else { physicalOffset++ } - select { - case <-sendConsumed: - case <-senderCtx.Done(): - errCh <- nil - return - } + submittedUnits = unitEnd } } @@ -1041,6 +1274,7 @@ func (cs *CoreStream[Req, Resp]) receiver( flightSignal <-chan struct{}, sendEvents <-chan sendEvent, sendConsumed chan<- struct{}, + ackProgress chan uint64, ) { type recvResult struct { resp Resp @@ -1065,10 +1299,12 @@ func (cs *CoreStream[Req, Resp]) receiver( } }() var stopRecvOnce sync.Once - stopRecv := func() { + stopRecv := func(abort bool) { stopRecvOnce.Do(func() { close(recvStop) - stream.Close() + if abort { + stream.Close() + } }) <-recvDone } @@ -1076,12 +1312,34 @@ func (cs *CoreStream[Req, Resp]) receiver( lackTimer := time.NewTimer(cs.cfg.LackOfAckTimeout) lackTimer.Stop() lackTimerArmed := false + var lackDeadline time.Time + _, usesRecordCountAcks := cs.ackMdl.(resolvingAckModel[Resp]) armLackTimer := func() { - if lackTimerArmed || cs.buf.inFlight() == 0 { + if !usesRecordCountAcks { + if lackTimerArmed || cs.buf.inFlight() == 0 { + return + } + lackTimer.Reset(cs.cfg.LackOfAckTimeout) + lackTimerArmed = true + return + } + deadline, ok := cs.buf.oldestInFlightDeadline(cs.cfg.LackOfAckTimeout) + if !ok { return } - lackTimer.Reset(cs.cfg.LackOfAckTimeout) + if lackTimerArmed && deadline.Equal(lackDeadline) { + return + } + if lackTimerArmed { + lackTimer.Stop() + } + wait := time.Until(deadline) + if wait < 0 { + wait = 0 + } + lackTimer.Reset(wait) lackTimerArmed = true + lackDeadline = deadline } disarmLackTimer := func() { if !lackTimerArmed { @@ -1090,17 +1348,29 @@ func (cs *CoreStream[Req, Resp]) receiver( // Go 1.23+ guarantees that Stop prevents stale timer values. lackTimer.Stop() lackTimerArmed = false + lackDeadline = time.Time{} } - // syncLackTimer realigns the ack-silence budget with the in-flight set. Only - // durable progress restarts it: a stale or duplicate cumulative ack must not - // extend the budget, or a server repeating one offset could postpone recovery - // indefinitely while later offsets stay unacknowledged. + // syncLackTimer preserves the proto/JSON ack-silence behavior while + // record-count protocols use the absolute deadline of their oldest pending + // logical item. A partial, stale, or duplicate row ACK cannot move it. syncLackTimer := func(progressed bool) { - if cs.buf.inFlight() == 0 { + if !usesRecordCountAcks { + if cs.buf.inFlight() == 0 { + disarmLackTimer() + return + } + if progressed { + disarmLackTimer() + } + armLackTimer() + return + } + deadline, ok := cs.buf.oldestInFlightDeadline(cs.cfg.LackOfAckTimeout) + if !ok { disarmLackTimer() return } - if progressed { + if lackTimerArmed && !deadline.Equal(lackDeadline) { disarmLackTimer() } armLackTimer() @@ -1117,120 +1387,283 @@ func (cs *CoreStream[Req, Resp]) receiver( } defer stopPauseTimer() - lastAckedPhysical := int64(-1) - sendingPhysical := int64(-1) - sendingLogical := int64(-1) - pendingPhysicalAck := int64(-1) - sentBasePhysical := int64(0) - var sentLogical []int64 - - applyLogicalAck := func(offset int64) error { - current := cs.wm.current() - if offset <= current { + lastAckedUnits := uint64(0) + submittedUnits := uint64(0) + nextPhysical := int64(0) + var submitted []SubmittedRange + var sending *SubmittedRange + var pendingAckUnits uint64 + pendingAck := false + + ackState := func(includeSending bool) AckState { + state := AckState{ + Ranges: submitted, + AcknowledgedUnits: lastAckedUnits, + SubmittedUnits: submittedUnits, + } + if includeSending && sending != nil { + state.Active = *sending + state.HasActive = true + state.SubmittedUnits = sending.UnitEnd + } + return state + } + // resolveAck maps one server response onto logical durability progress. A + // protocol-supplied model is re-resolved against the core's own view of the + // submitted ranges and must agree, since a hook that claimed progress those + // ranges do not support would corrupt the buffer. The built-in offset model + // already resolves through ResolveAcknowledgedUnits, so cross-checking it + // would compare a pure function against itself on every ack. + resolveAck := func(resp Resp, legacyOffset int64, state AckState) (AckResolution, error) { + model, ok := cs.ackMdl.(resolvingAckModel[Resp]) + if !ok { + return resolveOffsetAck(legacyOffset, state) + } + resolution, err := model.resolve(resp, state) + if err != nil { + return AckResolution{}, err + } + canonical, err := ResolveAcknowledgedUnits(resolution.AcknowledgedUnits, state) + if err != nil { + return AckResolution{}, err + } + if resolution != canonical { + return AckResolution{}, fmt.Errorf( + "stream: acknowledgment model returned inconsistent resolution: got %+v, want %+v", + resolution, canonical, + ) + } + return resolution, nil + } + applyResolution := func(resolution AckResolution) error { + if resolution.AcknowledgedUnits <= lastAckedUnits { syncLackTimer(false) return nil } - highest, ok := cs.buf.highestInFlight() - if !ok || offset > highest { - return fmt.Errorf( - "stream: server ack offset %d exceeds highest in-flight offset %d", - offset, highest, - ) + discarded, progressed, err := cs.buf.acknowledge(resolution) + if err != nil { + return err + } + if progressed { + cs.durableProgress.Add(1) + } + lastAckedUnits = resolution.AcknowledgedUnits + if progressed { + // Keep only the latest cumulative progress. runOnce uses this to + // bound late-ACK draining after a receipt-bearing Send failure. + select { + case ackProgress <- lastAckedUnits: + default: + select { + case <-ackProgress: + default: + } + select { + case ackProgress <- lastAckedUnits: + default: + } + } + } + if resolution.FullyAcknowledgedOffset >= 0 { + prune := 0 + for prune < len(submitted) && + submitted[prune].LogicalOffset <= resolution.FullyAcknowledgedOffset { + prune++ + } + submitted = submitted[prune:] + if len(submitted) == 0 { + submitted = nil + } } - discarded := cs.buf.discardThrough(offset) if discarded.count > 0 { cs.wm.advance(discarded.last) cs.dispatcher.enqueueAcks(discarded.first, discarded.last) } - syncLackTimer(discarded.count > 0) + syncLackTimer(progressed) if cs.buf.inFlight() == 0 && pauseState != nil { stopPauseTimer() return *pauseState } return nil } - applyPhysicalAck := func(offset int64) error { - if offset <= lastAckedPhysical { - syncLackTimer(false) - return nil - } - index := offset - sentBasePhysical - if index < 0 || index >= int64(len(sentLogical)) { - return fmt.Errorf( - "stream: server ack offset %d exceeds highest completed physical offset %d", - offset, sentBasePhysical+int64(len(sentLogical))-1, - ) - } - logical := sentLogical[index] - if err := applyLogicalAck(logical); err != nil { + applyAck := func(resp Resp, legacyOffset int64) error { + state := ackState(true) + resolution, err := resolveAck(resp, legacyOffset, state) + if err != nil { return err } - lastAckedPhysical = offset - sentLogical = sentLogical[index+1:] - sentBasePhysical = offset + 1 - if len(sentLogical) == 0 { - sentLogical = nil + if sending != nil && resolution.AcknowledgedUnits > submittedUnits { + if !pendingAck || resolution.AcknowledgedUnits > pendingAckUnits { + pendingAck = true + pendingAckUnits = resolution.AcknowledgedUnits + } + if submittedUnits <= lastAckedUnits { + syncLackTimer(false) + return nil + } + completed, err := ResolveAcknowledgedUnits( + submittedUnits, ackState(false), + ) + if err != nil { + return err + } + return applyResolution(completed) } - return nil + return applyResolution(resolution) } handleSendEvent := func(event sendEvent) error { if !event.completed { - sendingPhysical = event.physicalOffset - sendingLogical = event.logicalOffset + if sending != nil { + return fmt.Errorf("stream: overlapping Send operations") + } + if event.physicalOffset != nextPhysical { + return fmt.Errorf( + "stream: physical send offset %d is not contiguous after %d", + event.physicalOffset, nextPhysical-1, + ) + } + if event.unitStart != submittedUnits || event.unitEnd <= event.unitStart { + return fmt.Errorf( + "stream: submitted unit range [%d,%d) is not contiguous after %d", + event.unitStart, event.unitEnd, submittedUnits, + ) + } + sending = &SubmittedRange{ + WireOffset: event.physicalOffset, + LogicalOffset: event.logicalOffset, + UnitStart: event.unitStart, + UnitEnd: event.unitEnd, + ItemUnitEnd: event.unitEnd, + } return nil } + // rejectSubmission releases the sender and clears the active-send state + // before reporting a protocol violation. The completion event has + // already been consumed, so leaving sending or pendingAck set would + // make resolvePendingOnExit wait for an event that can never arrive. + rejectSubmission := func(err error) error { + sending = nil + pendingAck = false + pendingAckUnits = 0 + sendConsumed <- struct{}{} + return err + } + if sending == nil || + event.physicalOffset != sending.WireOffset || + event.logicalOffset != sending.LogicalOffset || + event.unitStart != sending.UnitStart || + event.unitEnd != sending.UnitEnd { + return rejectSubmission(fmt.Errorf( + "stream: Send completion does not match active submission", + )) + } + submittedCount := event.unitEnd - event.unitStart + if event.receiptUnits > submittedCount { + return rejectSubmission(wrapValidation(fmt.Errorf( + "stream: submission receipt %d exceeds active range size %d", + event.receiptUnits, + submittedCount, + ))) + } + if event.err == nil && event.receiptUnits != submittedCount { + return rejectSubmission(wrapValidation(fmt.Errorf( + "stream: completed Send submitted %d of %d units", + event.receiptUnits, + submittedCount, + ))) + } + var result error + submissionLimit := submittedUnits + event.receiptUnits if event.err == nil { - expected := sentBasePhysical + int64(len(sentLogical)) - if event.physicalOffset != expected { - result = fmt.Errorf( - "stream: physical send offset %d is not contiguous after %d", - event.physicalOffset, expected-1, - ) - } else { - sentLogical = append(sentLogical, event.logicalOffset) + submitted = append(submitted, *sending) + submittedUnits = sending.UnitEnd + submissionLimit = submittedUnits + if nextPhysical < math.MaxInt64 { + nextPhysical++ } + } else if event.receiptUnits > 0 { + partial := *sending + partial.UnitEnd = submissionLimit + submitted = append(submitted, partial) + submittedUnits = submissionLimit } - sendingPhysical = -1 - sendingLogical = -1 - var ackErr error - if event.err == nil && pendingPhysicalAck >= 0 { - ackErr = applyPhysicalAck(pendingPhysicalAck) - } - pendingPhysicalAck = -1 - if result == nil { - result = ackErr + if pendingAck { + // A failed Send leaves the core unsure how much of the item reached + // the server, so an ack covering units it cannot account for is + // unusable rather than a server protocol violation. Keep the part + // the submitted ranges do support and let the retryable send failure + // drive recovery. + if pendingAckUnits > submittedUnits { + pendingAckUnits = submittedUnits + } + if pendingAckUnits > lastAckedUnits { + resolution, err := ResolveAcknowledgedUnits( + pendingAckUnits, ackState(false), + ) + if err != nil { + result = &invalidAcknowledgment{cause: err} + } else { + result = applyResolution(resolution) + } + } } + sending = nil + pendingAck = false + pendingAckUnits = 0 sendConsumed <- struct{}{} return result } - resolvePendingOnExit := func() error { - if pendingPhysicalAck < 0 || sendingPhysical < 0 { + resolvePendingOnExit := func(abort bool) error { + if !pendingAck || sending == nil { return nil } - if sendingLogical < 0 { - return fmt.Errorf("stream: missing logical offset for active send") + if abort { + stream.Close() + return handleSendEvent(<-sendEvents) + } + // Preserve an orderly pause when the active Send is about to complete. + // A stuck Send cannot hold rotation forever: fall back to a hard abort + // after the same bounded drain budget used by lifecycle teardown. + timer := time.NewTimer(cs.cfg.DrainTimeout) + defer timer.Stop() + select { + case event := <-sendEvents: + return handleSendEvent(event) + case <-timer.C: + stream.Close() + return handleSendEvent(<-sendEvents) } - stream.Close() - return handleSendEvent(<-sendEvents) } handleTerminal := func(err error) { var ps pauseSignal - if err != nil && !errors.As(err, &ps) { + isPause := errors.As(err, &ps) + if err != nil && !isPause { // Real receiver failures are authoritative. Publish before any // operation that can close the transport and unblock Send. errCh <- err - stopRecv() + // An ACK can arrive while a multi-frame Send is still active. Once + // the receiver reports a real failure, aborting the transport makes + // that Send return its authoritative submission receipt. Reconcile + // the receipt and apply the buffered ACK before recovery snapshots + // the buffer, but never replace the receiver's error with a receipt + // or acknowledgment reconciliation error. + _ = resolvePendingOnExit(true) + stopRecv(true) return } - if pendingErr := resolvePendingOnExit(); pendingErr != nil { - if err == nil || errors.As(err, &ps) { + if pendingErr := resolvePendingOnExit(!isPause); pendingErr != nil { + if err == nil || isPause { err = pendingErr + // Reconciling the buffered ack can itself complete a pending + // rotation, so the replacement decides whether this exit is + // still an orderly pause rather than a hard failure. + isPause = errors.As(err, &ps) } } // Publish the resolved clean/pause outcome before final pump teardown. errCh <- err - stopRecv() + stopRecv(err != nil && !isPause) } for { @@ -1251,10 +1684,11 @@ func (cs *CoreStream[Req, Resp]) receiver( } continue case <-flightSignal: - armLackTimer() + syncLackTimer(false) continue case <-lackC: lackTimerArmed = false + lackDeadline = time.Time{} if pauseState != nil { continue } @@ -1285,19 +1719,44 @@ func (cs *CoreStream[Req, Resp]) receiver( return } - kind, offset, pause := cs.ackMdl.classify(r.resp) - if kind == unknownResponse || kind == malformedResponse { + classified := cs.ackMdl.classify(r.resp) + if classified.failure == unknownResponse || + classified.failure == malformedResponse { // A response the ack model can't interpret (unrecognized type, or an // ack missing/with a negative offset) is a protocol violation. Tear the // stream down rather than silently dropping it; the supervisor decides // whether to reconnect. handleTerminal reaps the pump's next Recv. - handleTerminal(fmt.Errorf("stream: unusable server response (kind %d)", kind)) + handleTerminal(fmt.Errorf( + "stream: unusable server response (failure %d)", + classified.failure, + )) return } - if kind == pauseResponse { + if classified.hasAck { + for { + select { + case event := <-sendEvents: + if err := handleSendEvent(event); err != nil { + handleTerminal(err) + return + } + default: + goto sendsDrained + } + } + sendsDrained: + if err := applyAck(r.resp, classified.legacyOffset); err != nil { + handleTerminal(err) + return + } + } + + // Apply inline durability progress before parking the sender or deciding + // that rotation can begin immediately. + if classified.pause != nil { if pauseState == nil { - pauseCopy := pause - wait := cs.effectivePauseWait(pause.duration) + pauseCopy := *classified.pause + wait := cs.effectivePauseWait(pauseCopy.duration) pauseCopy.resumeAt = time.Now().Add(wait) pauseState = &pauseCopy select { @@ -1318,33 +1777,6 @@ func (cs *CoreStream[Req, Resp]) receiver( // Drain late acks during the pause. continue } - - if kind == ackResponse { - for { - select { - case event := <-sendEvents: - if err := handleSendEvent(event); err != nil { - handleTerminal(err) - return - } - default: - goto sendsDrained - } - } - sendsDrained: - if sendingPhysical >= 0 && offset == sendingPhysical { - if offset > pendingPhysicalAck { - pendingPhysicalAck = offset - } - continue - } - if err := applyPhysicalAck(offset); err != nil { - handleTerminal(err) - return - } - } - // Non-ack, non-pause kinds already returned above; the next Recv is - // already outstanding. } } diff --git a/purego/internal/stream/core_test.go b/purego/internal/stream/core_test.go index 1f2a8e72..55430099 100644 --- a/purego/internal/stream/core_test.go +++ b/purego/internal/stream/core_test.go @@ -1122,6 +1122,64 @@ func TestCoreStreamRejectsAckWhenSendFails(t *testing.T) { } } +// TestCoreStreamRecoversAfterAckDuringFailedSend covers a server that +// acknowledges a record while the client's Send for it is still outstanding and +// then aborts the stream. The unusable acknowledgment must be discarded and the +// retryable send failure must drive recovery: the core does not know how much of +// that Send reached the server, which is not a server protocol violation. +func TestCoreStreamRecoversAfterAckDuringFailedSend(t *testing.T) { + blocked := newControlledSendRPC() + replay := newFakeRPC() + opener := &mixedOpener{streams: []wireStream[encodedMsg, ephemeralResp]{ + transport.NewFakeStreamForTesting(blocked), + transport.NewFakeStreamForTesting(replay), + }} + + cfg := testConfig() + cfg.RecoveryRetries = 2 + cs := newCoreForTest(testParams(), cfg, opener, nil) + t.Cleanup(func() { cs.Terminate() }) + + offset, err := cs.Ingest(context.Background(), []byte(`{}`)) + if err != nil { + t.Fatalf("Ingest: %v", err) + } + <-blocked.started + blocked.ack(offset) + blocked.result <- errors.New("send failed") + + // The record was never confirmed submitted, so it must be replayed rather + // than the stream dying with a non-retryable protocol error. + waitCondition(t, func() bool { return len(replay.sends) > 0 }, 3*time.Second) + replay.ack(0) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := cs.WaitForOffset(ctx, offset); err != nil { + t.Fatalf("WaitForOffset after recovery: %v (terminal=%v)", err, cs.terminalErr()) + } +} + +// mixedOpener hands out pre-built connections of differing fake types in order. +type mixedOpener struct { + mu sync.Mutex + streams []wireStream[encodedMsg, ephemeralResp] + idx int +} + +func (o *mixedOpener) Open( + context.Context, transport.StreamParams, +) (wireStream[encodedMsg, ephemeralResp], error) { + o.mu.Lock() + defer o.mu.Unlock() + if o.idx >= len(o.streams) { + return nil, fmt.Errorf("mixedOpener: no connection left") + } + stream := o.streams[o.idx] + o.idx++ + return stream, nil +} + // TestCoreStreamMalformedAckTearsDownAndRecovers verifies that an // uninterpretable server response (an ack missing its offset) tears the stream // down instead of being silently ignored, and the supervisor reconnects. @@ -1326,6 +1384,18 @@ func TestIsRetryable(t *testing.T) { {name: "closed", err: errClosed, want: false}, {name: "classified retryable", err: fmt.Errorf("wrapped: %w", &classifiedError{retryable: true}), want: true}, {name: "classified terminal", err: &classifiedError{retryable: false}, want: false}, + // An acknowledgment the submitted ranges cannot support is a protocol + // violation: reconnecting would replay into the same disagreement. + { + name: "invalid acknowledgment", + err: &invalidAcknowledgment{cause: errors.New("ack of 9 units")}, + want: false, + }, + { + name: "invalid acknowledgment when wrapped", + err: fmt.Errorf("recv: %w", &invalidAcknowledgment{cause: errors.New("bad ack")}), + want: false, + }, {name: "open budget", err: &openBudgetExceeded{cause: context.DeadlineExceeded}, want: true}, {name: "context deadline", err: context.DeadlineExceeded, want: false}, {name: "ordinary transport", err: errors.New("connection reset"), want: true}, diff --git a/purego/internal/stream/durability_model_test.go b/purego/internal/stream/durability_model_test.go new file mode 100644 index 00000000..28a694a6 --- /dev/null +++ b/purego/internal/stream/durability_model_test.go @@ -0,0 +1,1243 @@ +package stream + +import ( + "context" + "errors" + "fmt" + "io" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/databricks/zerobus-sdk/purego/internal/transport" +) + +// ---- fake record-count protocol -------------------------------------------- +// +// The proto/JSON protocols are atomic: one logical item is always one +// durability unit, so they exercise none of the partial-acknowledgment, +// payload-slicing, or submission-receipt behaviour the core now supports. This +// fake protocol makes a single row the durability unit, so a server can +// acknowledge a prefix of a multi-row item and a transport can submit a prefix +// of one logical Send. + +// rowPayload is one logical item whose durability unit is a single row. +type rowPayload struct { + offset int64 + rows []string +} + +// rowResponse is the fake server response. It carries a cumulative row count +// for the record-count model and a connection-local wire offset for the atomic +// model, so one transport can drive both acknowledgment shapes. +type rowResponse struct { + ackedRows uint64 + wireOffset int64 + pause time.Duration + hasAck bool + hasPause bool + unknown bool +} + +func rowEncoderHooks() EncoderHooks[*rowPayload] { + return EncoderHooks[*rowPayload]{ + EncodeRecord: func(record []byte) (*rowPayload, error) { + return &rowPayload{rows: []string{string(record)}}, nil + }, + EncodeBatch: func(records [][]byte) (*rowPayload, error) { + if len(records) == 0 { + return nil, fmt.Errorf("row batch must not be empty") + } + rows := make([]string, len(records)) + for i, record := range records { + rows[i] = string(record) + } + return &rowPayload{rows: rows}, nil + }, + StampOffset: func(payload *rowPayload, offset int64) { payload.offset = offset }, + UnitCount: func(payload *rowPayload) uint64 { return uint64(len(payload.rows)) }, + Slice: func(payload *rowPayload, acknowledgedPrefix uint64) (*rowPayload, error) { + if acknowledgedPrefix >= uint64(len(payload.rows)) { + return nil, fmt.Errorf( + "acknowledged prefix %d covers all %d rows", + acknowledgedPrefix, len(payload.rows), + ) + } + return &rowPayload{ + offset: payload.offset, + rows: slices.Clone(payload.rows[acknowledgedPrefix:]), + }, nil + }, + Decode: func(payload *rowPayload) [][]byte { + out := make([][]byte, len(payload.rows)) + for i, row := range payload.rows { + out[i] = []byte(row) + } + return out + }, + MaxWireSize: func(payload *rowPayload) int { + size := 0 + for _, row := range payload.rows { + size += len(row) + } + return size + }, + RetainedSize: func(rawBytes, recordCount int) int64 { + return int64(rawBytes + recordCount) + }, + } +} + +func rowAckHooks() AckModelHooks[*rowResponse] { + return AckModelHooks[*rowResponse]{ + Classify: func(resp *rowResponse) ResponseClassification { + if resp == nil || resp.unknown { + return ResponseClassification{Status: ResponseUnknown} + } + return ResponseClassification{ + Status: ResponseOK, + HasAck: resp.hasAck, + HasPause: resp.hasPause, + PauseDuration: resp.pause, + } + }, + Resolve: func(resp *rowResponse, state AckState) (AckResolution, error) { + return ResolveAcknowledgedUnits(resp.ackedRows, state) + }, + } +} + +// atomicRowAckHooks omits Resolve, the configuration an atomic protocol uses: +// the core resolves the hook-supplied connection-local wire offset itself +// instead of asking the protocol to translate a unit count. +func atomicRowAckHooks() AckModelHooks[*rowResponse] { + hooks := rowAckHooks() + hooks.Classify = func(resp *rowResponse) ResponseClassification { + if resp == nil || resp.unknown { + return ResponseClassification{Status: ResponseUnknown} + } + return ResponseClassification{ + Status: ResponseOK, + HasAck: resp.hasAck, + LegacyOffset: resp.wireOffset, + HasPause: resp.hasPause, + PauseDuration: resp.pause, + } + } + hooks.Resolve = nil + return hooks +} + +// rowWire is an in-process transport for the fake row protocol. +type rowWire struct { + sends chan *rowPayload + resps chan *rowResponse + delivered chan struct{} + closeOnce sync.Once + closed atomic.Bool +} + +func newRowWire() *rowWire { + return &rowWire{ + sends: make(chan *rowPayload, 64), + resps: make(chan *rowResponse, 64), + delivered: make(chan struct{}, 64), + } +} + +func (w *rowWire) ServerID() string { return "row-wire" } + +func (w *rowWire) Send(payload *rowPayload) error { + if w.closed.Load() { + return io.EOF + } + select { + case w.sends <- payload: + return nil + default: + return fmt.Errorf("rowWire: sends channel full") + } +} + +func (w *rowWire) Recv() (*rowResponse, error) { + resp, ok := <-w.resps + if !ok { + return nil, io.EOF + } + // Signals that the receiver has taken this response, so a test can order a + // later event after it. + select { + case w.delivered <- struct{}{}: + default: + } + return resp, nil +} + +func (w *rowWire) CloseSend() error { + w.shutdown() + return nil +} + +func (w *rowWire) Close() { w.shutdown() } + +func (w *rowWire) shutdown() { + w.closeOnce.Do(func() { + w.closed.Store(true) + close(w.resps) + }) +} + +// ackRows publishes a cumulative connection-local row acknowledgment. +func (w *rowWire) ackRows(n uint64) { + if w.closed.Load() { + return + } + w.resps <- &rowResponse{hasAck: true, ackedRows: n} +} + +// ackWireOffset publishes a cumulative connection-local wire offset ack. +func (w *rowWire) ackWireOffset(offset int64) { + if w.closed.Load() { + return + } + w.resps <- &rowResponse{hasAck: true, wireOffset: offset} +} + +// nextSend waits for the transport to receive one logical payload. +func (w *rowWire) nextSend(t *testing.T) *rowPayload { + t.Helper() + select { + case payload := <-w.sends: + return payload + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for a payload to reach the transport") + return nil + } +} + +// partialRowWire submits only a prefix of the first logical Send and then +// fails, modelling a multi-frame transport whose later frame is rejected. +type partialRowWire struct { + *rowWire + submitRows uint64 + failWith error + failed atomic.Bool +} + +func (w *partialRowWire) SendWithReceipt(payload *rowPayload) (SubmissionReceipt, error) { + if w.failed.CompareAndSwap(false, true) { + prefix := &rowPayload{ + offset: payload.offset, + rows: slices.Clone(payload.rows[:w.submitRows]), + } + if err := w.rowWire.Send(prefix); err != nil { + return SubmissionReceipt{}, err + } + return SubmissionReceipt{SubmittedUnits: w.submitRows}, w.failWith + } + if err := w.rowWire.Send(payload); err != nil { + return SubmissionReceipt{}, err + } + return SubmissionReceipt{SubmittedUnits: uint64(len(payload.rows))}, nil +} + +// overReportingRowWire holds one Send open until released, then claims more +// submitted units than the payload contains — the shape of a buggy protocol +// adapter, which must fail the stream rather than wedge it. +type overReportingRowWire struct { + *rowWire + started chan struct{} + release chan struct{} + startOnce sync.Once +} + +func newOverReportingRowWire() *overReportingRowWire { + return &overReportingRowWire{ + rowWire: newRowWire(), + started: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (w *overReportingRowWire) SendWithReceipt( + payload *rowPayload, +) (SubmissionReceipt, error) { + w.startOnce.Do(func() { close(w.started) }) + <-w.release + if err := w.rowWire.Send(payload); err != nil { + return SubmissionReceipt{}, err + } + return SubmissionReceipt{SubmittedUnits: uint64(len(payload.rows)) + 5}, nil +} + +// shortSuccessRowWire submits every row but reports a prefix receipt with a nil +// error, so accepting it would strand the rows it leaves out. +type shortSuccessRowWire struct { + *rowWire + reportRows uint64 +} + +func (w *shortSuccessRowWire) SendWithReceipt( + payload *rowPayload, +) (SubmissionReceipt, error) { + if err := w.rowWire.Send(payload); err != nil { + return SubmissionReceipt{}, err + } + return SubmissionReceipt{SubmittedUnits: w.reportRows}, nil +} + +// authRejectingRowWire delivers a definitive server rejection from the Recv that +// teardown unblocks, the way gRPC carries a status after Send has already failed. +type authRejectingRowWire struct { + *partialRowWire +} + +func (w *authRejectingRowWire) Recv() (*rowResponse, error) { + resp, err := w.partialRowWire.Recv() + if errors.Is(err, io.EOF) { + return nil, status.Error(codes.Unauthenticated, "expired credentials") + } + return resp, err +} + +// rowOpener hands out pre-built connections in order. +type rowOpener struct { + mu sync.Mutex + wires []WireStream[*rowPayload, *rowResponse] + opened int +} + +func newRowOpener(wires ...WireStream[*rowPayload, *rowResponse]) *rowOpener { + return &rowOpener{wires: wires} +} + +// count reports how many connections have been handed out. +func (o *rowOpener) count() int { + o.mu.Lock() + defer o.mu.Unlock() + return o.opened +} + +func (o *rowOpener) open( + context.Context, StreamParams, +) (WireStream[*rowPayload, *rowResponse], error) { + o.mu.Lock() + defer o.mu.Unlock() + if o.opened >= len(o.wires) { + return nil, fmt.Errorf("rowOpener: no connection left") + } + wire := o.wires[o.opened] + o.opened++ + return wire, nil +} + +func newRowStream( + t *testing.T, + cfg Config, + opener *rowOpener, +) *CoreStream[*rowPayload, *rowResponse] { + t.Helper() + return newRowStreamWithAcks(t, cfg, opener, rowAckHooks()) +} + +func newRowStreamWithAcks( + t *testing.T, + cfg Config, + opener *rowOpener, + acks AckModelHooks[*rowResponse], +) *CoreStream[*rowPayload, *rowResponse] { + t.Helper() + return newRowStreamWithParams(t, testParams(), cfg, opener, acks) +} + +func newRowStreamWithParams( + t *testing.T, + params StreamParams, + cfg Config, + opener *rowOpener, + acks AckModelHooks[*rowResponse], +) *CoreStream[*rowPayload, *rowResponse] { + t.Helper() + cs, err := NewCoreStreamWithHooks[*rowPayload, *rowResponse]( + context.Background(), + params, + cfg, + opener.open, + rowEncoderHooks(), + acks, + nil, + ) + if err != nil { + t.Fatalf("NewCoreStreamWithHooks: %v", err) + } + t.Cleanup(func() { cs.Terminate() }) + return cs +} + +func rowConfig() Config { + cfg := testConfig() + cfg.DrainTimeout = 2 * time.Second + return cfg +} + +// ---- ResolveAcknowledgedUnits ---------------------------------------------- + +func TestResolveAcknowledgedUnits(t *testing.T) { + twoItems := []SubmittedRange{ + {LogicalOffset: 0, UnitStart: 0, UnitEnd: 3, ItemUnitEnd: 3}, + {LogicalOffset: 1, UnitStart: 3, UnitEnd: 5, ItemUnitEnd: 5}, + } + + tests := []struct { + name string + ackedUnits uint64 + state AckState + want AckResolution + wantErr string + }{ + { + name: "no new progress", + ackedUnits: 2, + state: AckState{ + Ranges: twoItems, AcknowledgedUnits: 2, SubmittedUnits: 5, + }, + want: AckResolution{ + AcknowledgedUnits: 2, FullyAcknowledgedOffset: -1, PartialOffset: -1, + }, + }, + { + // A watermark that advances but lands below every submitted range + // resolves to no offset, so accepting it would raise the connection's + // ack watermark without any durable progress behind it. + name: "advancing ack below the first submitted range", + ackedUnits: 2, + state: AckState{ + SubmittedUnits: 6, + Ranges: []SubmittedRange{ + {LogicalOffset: 0, UnitStart: 3, UnitEnd: 6, ItemUnitEnd: 6}, + }, + }, + wantErr: "does not intersect an unacknowledged submitted range", + }, + { + // The head range is exactly consumed and the next one has not + // started, which is ordinary progress rather than a gap. + name: "ack landing on a later range boundary", + ackedUnits: 3, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + want: AckResolution{ + AcknowledgedUnits: 3, FullyAcknowledgedOffset: 0, PartialOffset: -1, + }, + }, + { + // The still-sending item resolves like any other range, so an ack + // arriving mid-send can retire it without the core having to + // combine it into Ranges first. + name: "active range resolves fully", + ackedUnits: 7, + state: AckState{ + Ranges: twoItems, + Active: SubmittedRange{LogicalOffset: 2, UnitStart: 5, UnitEnd: 7, ItemUnitEnd: 7}, + HasActive: true, + SubmittedUnits: 7, + }, + want: AckResolution{ + AcknowledgedUnits: 7, FullyAcknowledgedOffset: 2, PartialOffset: -1, + }, + }, + { + name: "active range resolves partially", + ackedUnits: 6, + state: AckState{ + Ranges: twoItems, + Active: SubmittedRange{LogicalOffset: 2, UnitStart: 5, UnitEnd: 7, ItemUnitEnd: 7}, + HasActive: true, + SubmittedUnits: 7, + }, + want: AckResolution{ + AcknowledgedUnits: 6, FullyAcknowledgedOffset: 1, + PartialOffset: 2, PartialUnits: 1, + }, + }, + { + // Contiguity is enforced across the Ranges/Active seam too, so a + // gap there is caught rather than silently resolved. + name: "active range not contiguous with ranges", + ackedUnits: 7, + state: AckState{ + Ranges: twoItems, + Active: SubmittedRange{LogicalOffset: 2, UnitStart: 6, UnitEnd: 8, ItemUnitEnd: 8}, + HasActive: true, + SubmittedUnits: 8, + }, + wantErr: "submitted unit range starts at 6 after 5", + }, + { + // With no Ranges behind it the active item is the whole window. + name: "active range alone", + ackedUnits: 2, + state: AckState{ + Active: SubmittedRange{LogicalOffset: 9, UnitStart: 0, UnitEnd: 4, ItemUnitEnd: 4}, + HasActive: true, + SubmittedUnits: 4, + }, + want: AckResolution{ + AcknowledgedUnits: 2, FullyAcknowledgedOffset: -1, + PartialOffset: 9, PartialUnits: 2, + }, + }, + { + name: "partial prefix of first item", + ackedUnits: 2, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + want: AckResolution{ + AcknowledgedUnits: 2, FullyAcknowledgedOffset: -1, + PartialOffset: 0, PartialUnits: 2, + }, + }, + { + name: "exactly one full item", + ackedUnits: 3, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + want: AckResolution{ + AcknowledgedUnits: 3, FullyAcknowledgedOffset: 0, PartialOffset: -1, + }, + }, + { + name: "full item plus partial next", + ackedUnits: 4, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + want: AckResolution{ + AcknowledgedUnits: 4, FullyAcknowledgedOffset: 0, + PartialOffset: 1, PartialUnits: 1, + }, + }, + { + name: "all submitted items", + ackedUnits: 5, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + want: AckResolution{ + AcknowledgedUnits: 5, FullyAcknowledgedOffset: 1, PartialOffset: -1, + }, + }, + { + name: "zero ItemUnitEnd defaults to UnitEnd", + ackedUnits: 2, + state: AckState{ + Ranges: []SubmittedRange{{LogicalOffset: 7, UnitStart: 0, UnitEnd: 2}}, + SubmittedUnits: 2, + }, + want: AckResolution{ + AcknowledgedUnits: 2, FullyAcknowledgedOffset: 7, PartialOffset: -1, + }, + }, + { + // A multi-frame Send that failed after submitting a prefix leaves + // UnitEnd short of ItemUnitEnd. Acking the whole submitted prefix + // makes the item partially, not fully, durable. + name: "submitted prefix of an incomplete item", + ackedUnits: 2, + state: AckState{ + Ranges: []SubmittedRange{ + {LogicalOffset: 4, UnitStart: 0, UnitEnd: 2, ItemUnitEnd: 5}, + }, + SubmittedUnits: 2, + }, + want: AckResolution{ + AcknowledgedUnits: 2, FullyAcknowledgedOffset: -1, + PartialOffset: 4, PartialUnits: 2, + }, + }, + { + name: "ack beyond submitted units", + ackedUnits: 6, + state: AckState{Ranges: twoItems, SubmittedUnits: 5}, + wantErr: "only 5 units were submitted", + }, + { + name: "watermark beyond submitted units", + ackedUnits: 1, + state: AckState{ + Ranges: twoItems, AcknowledgedUnits: 9, SubmittedUnits: 5, + }, + wantErr: "watermark 9 exceeds submitted units 5", + }, + { + name: "empty submitted range", + ackedUnits: 1, + state: AckState{ + Ranges: []SubmittedRange{{LogicalOffset: 0, UnitStart: 2, UnitEnd: 2}}, + SubmittedUnits: 4, + }, + wantErr: "invalid submitted range", + }, + { + name: "item end precedes submitted end", + ackedUnits: 1, + state: AckState{ + Ranges: []SubmittedRange{ + {LogicalOffset: 0, UnitStart: 0, UnitEnd: 4, ItemUnitEnd: 2}, + }, + SubmittedUnits: 4, + }, + wantErr: "precedes submitted range end", + }, + { + name: "non-contiguous ranges", + ackedUnits: 4, + state: AckState{ + Ranges: []SubmittedRange{ + {LogicalOffset: 0, UnitStart: 0, UnitEnd: 2, ItemUnitEnd: 2}, + {LogicalOffset: 1, UnitStart: 3, UnitEnd: 5, ItemUnitEnd: 5}, + }, + SubmittedUnits: 5, + }, + wantErr: "starts at 3 after 2", + }, + { + name: "ack intersects no submitted range", + ackedUnits: 3, + state: AckState{SubmittedUnits: 5}, + wantErr: "does not intersect", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := ResolveAcknowledgedUnits(tc.ackedUnits, tc.state) + if tc.wantErr != "" { + if err == nil { + t.Fatalf("want error containing %q, got resolution %+v", tc.wantErr, got) + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("error = %q, want it to contain %q", err, tc.wantErr) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tc.want { + t.Fatalf("resolution = %+v, want %+v", got, tc.want) + } + }) + } +} + +// ---- generic hook seam ------------------------------------------------------ + +func TestNewCoreStreamWithHooksRejectsIncompleteHooks(t *testing.T) { + fullEncoder := rowEncoderHooks() + fullAcks := rowAckHooks() + opener := newRowOpener(newRowWire()) + + tests := []struct { + name string + open OpenFunc[*rowPayload, *rowResponse] + enc EncoderHooks[*rowPayload] + acks AckModelHooks[*rowResponse] + wantErr string + }{ + { + name: "missing open", open: nil, enc: fullEncoder, acks: fullAcks, + wantErr: "protocol Open hook is required", + }, + { + name: "missing encoder hook", open: opener.open, + enc: EncoderHooks[*rowPayload]{EncodeRecord: fullEncoder.EncodeRecord}, + acks: fullAcks, + wantErr: "all encoder hooks are required", + }, + { + name: "missing classify", open: opener.open, enc: fullEncoder, + acks: AckModelHooks[*rowResponse]{}, + wantErr: "Classify hook is required", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cs, err := NewCoreStreamWithHooks[*rowPayload, *rowResponse]( + context.Background(), testParams(), rowConfig(), + tc.open, tc.enc, tc.acks, nil, + ) + if err == nil { + cs.Terminate() + t.Fatalf("want error containing %q, got a stream", tc.wantErr) + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("error = %q, want it to contain %q", err, tc.wantErr) + } + }) + } +} + +// TestHookAckModelClassifyTranslatesStatuses covers the adapter that turns a +// protocol's ResponseClassification into the core's internal classification. +// Only the OK+ack shape is reached by the end-to-end tests, so the remaining +// branches — including the guard against a response carrying no signal at all — +// are pinned here. +func TestHookAckModelClassifyTranslatesStatuses(t *testing.T) { + tests := []struct { + name string + in ResponseClassification + wantFailure responseFailure + wantAck bool + wantOffset int64 + wantPause bool + wantDuration time.Duration + }{ + { + name: "unknown", + in: ResponseClassification{Status: ResponseUnknown}, + wantFailure: unknownResponse, + }, + { + name: "malformed", + in: ResponseClassification{Status: ResponseMalformed}, + wantFailure: malformedResponse, + }, + { + name: "ack only", + in: ResponseClassification{Status: ResponseOK, HasAck: true, LegacyOffset: 7}, + wantAck: true, + wantOffset: 7, + }, + { + name: "pause only", + in: ResponseClassification{ + Status: ResponseOK, HasPause: true, PauseDuration: 3 * time.Second, + }, + wantPause: true, + wantDuration: 3 * time.Second, + }, + { + name: "ack and pause together", + in: ResponseClassification{ + Status: ResponseOK, HasAck: true, LegacyOffset: 2, + HasPause: true, PauseDuration: time.Second, + }, + wantAck: true, wantOffset: 2, + wantPause: true, wantDuration: time.Second, + }, + { + name: "ok but no signal", + in: ResponseClassification{Status: ResponseOK}, + wantFailure: unknownResponse, + }, + { + name: "unrecognized status", + in: ResponseClassification{Status: ResponseStatus(99)}, + wantFailure: unknownResponse, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + model := hookAckModel[*rowResponse]{ + hooks: AckModelHooks[*rowResponse]{ + Classify: func(*rowResponse) ResponseClassification { return tc.in }, + }, + } + got := model.classify(nil) + if got.failure != tc.wantFailure { + t.Fatalf("failure = %v, want %v", got.failure, tc.wantFailure) + } + if got.hasAck != tc.wantAck { + t.Fatalf("hasAck = %v, want %v", got.hasAck, tc.wantAck) + } + if got.hasAck && got.legacyOffset != tc.wantOffset { + t.Fatalf("legacyOffset = %d, want %d", got.legacyOffset, tc.wantOffset) + } + if (got.pause != nil) != tc.wantPause { + t.Fatalf("pause = %v, want pause %v", got.pause, tc.wantPause) + } + if got.pause != nil && got.pause.duration != tc.wantDuration { + t.Fatalf("pause duration = %v, want %v", got.pause.duration, tc.wantDuration) + } + }) + } +} + +// TestHookProtocolWithoutResolveUsesWireOffsets covers the other half of the +// acknowledgment seam: a protocol that omits Resolve is atomic, so the core +// resolves its wire offsets itself and uses the connection-wide ack-silence +// budget rather than the per-item deadline. +func TestHookProtocolWithoutResolveUsesWireOffsets(t *testing.T) { + wire := newRowWire() + cs := newRowStreamWithAcks(t, rowConfig(), newRowOpener(wire), atomicRowAckHooks()) + + for _, record := range []string{"a", "b", "c"} { + if _, err := cs.Ingest(context.Background(), []byte(record)); err != nil { + t.Fatalf("Ingest %q: %v", record, err) + } + wire.nextSend(t) + } + + // Wire offset 1 makes the first two items durable, not the third. + wire.ackWireOffset(1) + if err := cs.WaitForOffset(context.Background(), 1); err != nil { + t.Fatalf("WaitForOffset(1): %v", err) + } + flushCtx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + if err := cs.Flush(flushCtx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Flush before final ack = %v, want DeadlineExceeded", err) + } + + wire.ackWireOffset(2) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after final ack: %v", err) + } + if got := cs.buf.inFlight(); got != 0 { + t.Fatalf("in-flight items after full ack = %d, want 0", got) + } +} + +// TestHookProtocolIngestAndAckByRows covers the ordinary path: a multi-row item +// becomes durable only once every one of its rows is acknowledged. +func TestHookProtocolIngestAndAckByRows(t *testing.T) { + wire := newRowWire() + cs := newRowStream(t, rowConfig(), newRowOpener(wire)) + + offset, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("a"), []byte("b"), []byte("c"), + }) + if err != nil { + t.Fatalf("IngestBatch: %v", err) + } + if offset != 0 { + t.Fatalf("offset = %d, want 0", offset) + } + if sent := wire.nextSend(t); len(sent.rows) != 3 { + t.Fatalf("transport received %d rows, want 3", len(sent.rows)) + } + + // A prefix acknowledgment must not advance the logical watermark. + wire.ackRows(2) + flushCtx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + if err := cs.Flush(flushCtx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Flush after partial ack = %v, want DeadlineExceeded", err) + } + + wire.ackRows(3) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after full ack: %v", err) + } + + // A single record is just a one-unit item on the same connection, so its + // rows continue the cumulative count. + single, err := cs.Ingest(context.Background(), []byte("d")) + if err != nil { + t.Fatalf("Ingest: %v", err) + } + if single != 1 { + t.Fatalf("single-record offset = %d, want 1", single) + } + if sent := wire.nextSend(t); !slices.Equal(sent.rows, []string{"d"}) { + t.Fatalf("transport received %v, want [d]", sent.rows) + } + wire.ackRows(4) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after single record: %v", err) + } +} + +// TestHookProtocolPartialAckReplaysOnlyRemainder is the core guarantee of the +// record-count model: rows the server already made durable are not resent. +func TestHookProtocolPartialAckReplaysOnlyRemainder(t *testing.T) { + first := newRowWire() + second := newRowWire() + cs := newRowStream(t, rowConfig(), newRowOpener(first, second)) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), []byte("r2"), []byte("r3"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + if sent := first.nextSend(t); len(sent.rows) != 4 { + t.Fatalf("first connection received %d rows, want 4", len(sent.rows)) + } + + first.ackRows(2) + // Drop the connection so the supervisor recovers onto the second wire. + first.shutdown() + + replayed := second.nextSend(t) + if got := replayed.rows; !slices.Equal(got, []string{"r2", "r3"}) { + t.Fatalf("replayed rows = %v, want [r2 r3]", got) + } + + second.ackRows(2) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after replay: %v", err) + } +} + +// TestHookProtocolSubmissionReceiptPreservesPrefixAck covers a multi-frame Send +// that fails after earlier frames were accepted: an acknowledgment for the +// submitted prefix stays authoritative, so recovery replays only the rest. +func TestHookProtocolSubmissionReceiptPreservesPrefixAck(t *testing.T) { + failing := &partialRowWire{ + rowWire: newRowWire(), + submitRows: 2, + failWith: fmt.Errorf("frame rejected"), + } + second := newRowWire() + cs := newRowStream(t, rowConfig(), newRowOpener(failing, second)) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), []byte("r2"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + + submitted := failing.nextSend(t) + if got := submitted.rows; !slices.Equal(got, []string{"r0", "r1"}) { + t.Fatalf("submitted prefix = %v, want [r0 r1]", got) + } + // The server acknowledges the frames that did land before the Send failed. + failing.ackRows(2) + + replayed := second.nextSend(t) + if got := replayed.rows; !slices.Equal(got, []string{"r2"}) { + t.Fatalf("replayed rows = %v, want [r2]", got) + } + + second.ackRows(1) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after receipt-based recovery: %v", err) + } +} + +// TestHookProtocolBadReceiptDoesNotWedgeTeardown covers a transport that +// reports an impossible submission receipt while an acknowledgment is buffered +// against the still-active Send. Rejecting the receipt must also release that +// buffered state, or teardown waits forever for a completion event that was +// already consumed. +func TestHookProtocolBadReceiptDoesNotWedgeTeardown(t *testing.T) { + wire := newOverReportingRowWire() + cfg := rowConfig() + cfg.Recovery = RecoveryDisabled + cs := newRowStreamWithAcks(t, cfg, newRowOpener(wire), rowAckHooks()) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + + // Buffer an ack against the in-progress Send: the sender is parked inside + // SendWithReceipt, so the receiver cannot see a completion event first. + <-wire.started + wire.ackRows(1) + <-wire.delivered + close(wire.release) + + done := make(chan error, 1) + go func() { done <- cs.Terminate() }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("Terminate hung after an invalid submission receipt") + } +} + +// TestHookProtocolImpossibleReceiptIsTerminal covers the receipt shapes no +// correct transport produces. Both are deterministic, so the spare connection +// offered here must go unused. +func TestHookProtocolImpossibleReceiptIsTerminal(t *testing.T) { + tests := []struct { + name string + wire func() WireStream[*rowPayload, *rowResponse] + release func(WireStream[*rowPayload, *rowResponse]) + wantErr string + }{ + { + name: "receipt exceeds payload", + wire: func() WireStream[*rowPayload, *rowResponse] { + return newOverReportingRowWire() + }, + release: func(w WireStream[*rowPayload, *rowResponse]) { + over := w.(*overReportingRowWire) + <-over.started + close(over.release) + }, + // The receiver rejects first, so its wording reaches the caller. + wantErr: "exceeds active range size", + }, + { + name: "successful send reports a short receipt", + wire: func() WireStream[*rowPayload, *rowResponse] { + return &shortSuccessRowWire{rowWire: newRowWire(), reportRows: 1} + }, + release: func(WireStream[*rowPayload, *rowResponse]) {}, + wantErr: "of 2 units submitted", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + first := tc.wire() + second := newRowWire() + opener := newRowOpener(first, second) + cfg := rowConfig() + cfg.Recovery = RecoveryEnabled + cfg.RecoveryRetries = 3 + cfg.RecoveryBackoff = time.Millisecond + cs := newRowStreamWithAcks(t, cfg, opener, rowAckHooks()) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + tc.release(first) + + waitCondition(t, cs.IsClosed, 5*time.Second) + err := cs.terminalErr() + if err == nil { + t.Fatal("stream closed without a terminal error") + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("terminal error = %v, want it to mention %q", err, tc.wantErr) + } + if isRetryable(err) { + t.Errorf("terminal error %v is retryable, so recovery replays it", err) + } + if got := opener.count(); got != 1 { + t.Errorf("opened %d connections, want 1: payload was replayed", got) + } + if got := len(second.sends); got != 0 { + t.Errorf("second connection received %d sends, want 0", got) + } + }) + } +} + +// TestHookProtocolPrefixAckPreservesTerminalRecvStatus covers a prefix ACK that +// ends the durability wait before the receiver is consulted. The rejection that +// teardown then unblocks must still win over the send error. +func TestHookProtocolPrefixAckPreservesTerminalRecvStatus(t *testing.T) { + wire := &authRejectingRowWire{ + partialRowWire: &partialRowWire{ + rowWire: newRowWire(), + submitRows: 1, + // gRPC reports a server-side abort to Send as an opaque io.EOF. + failWith: io.EOF, + }, + } + second := newRowWire() + opener := newRowOpener(wire, second) + provider := &countingHeadersProvider{} + params := testParams() + params.HeadersProvider = provider + cfg := rowConfig() + cfg.Recovery = RecoveryEnabled + cfg.RecoveryRetries = 3 + cfg.RecoveryBackoff = time.Millisecond + cs := newRowStreamWithParams(t, params, cfg, opener, rowAckHooks()) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + // End the wait on durable progress rather than a receiver exit. + wire.nextSend(t) + wire.ackRows(1) + + waitCondition(t, cs.IsClosed, 5*time.Second) + err := cs.terminalErr() + if !transport.IsAuthRejection(err) { + t.Fatalf("terminal error = %v, want the server's rejection", err) + } + if got := provider.invalidations.Load(); got != 1 { + t.Errorf("Invalidate calls = %d, want 1", got) + } + if got := opener.count(); got != 1 { + t.Errorf("opened %d connections, want 1: the rejection was retried", got) + } +} + +// TestHookProtocolPromotedItemGetsFreshAckBudget covers an item promoted to head +// by a full ack of its predecessor. Ordered acks mean it could not have been made +// durable sooner, so it must start a fresh budget on promotion instead of +// inheriting the deadline it was given on entering flight. +func TestHookProtocolPromotedItemGetsFreshAckBudget(t *testing.T) { + const timeout = time.Second + first := newRowWire() + second := newRowWire() + opener := newRowOpener(first, second) + cfg := rowConfig() + cfg.Recovery = RecoveryEnabled + cfg.RecoveryRetries = 3 + cfg.RecoveryBackoff = time.Millisecond + cfg.LackOfAckTimeout = timeout + cs := newRowStreamWithAcks(t, cfg, opener, rowAckHooks()) + + // A and B enter flight together, so they share one original deadline. + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("a0"), []byte("a1"), + }); err != nil { + t.Fatalf("IngestBatch(A): %v", err) + } + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("b0"), + }); err != nil { + t.Fatalf("IngestBatch(B): %v", err) + } + first.nextSend(t) + first.nextSend(t) + + // Partial progress carries A past that shared deadline. + time.Sleep(timeout / 2) + first.ackRows(1) + + // Complete A on a row boundary once B's original deadline has passed. + time.Sleep(timeout * 7 / 10) + first.ackRows(2) + + // Leave room for an unrefreshed deadline to fire before B is acked. A + // promotion that keeps the stale deadline reconnects here rather than + // closing, so the replay is what the assertion has to catch. + time.Sleep(timeout / 5) + if got := opener.count(); got != 1 { + t.Fatalf("opened %d connections, want 1: promoting B replayed it", got) + } + + first.ackRows(3) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush after promoting B to head: %v", err) + } + if got := opener.count(); got != 1 { + t.Errorf("opened %d connections, want 1: B was replayed", got) + } + if got := len(second.sends); got != 0 { + t.Errorf("second connection received %d sends, want 0", got) + } +} + +// TestHookProtocolGetUnackedSlicesAcknowledgedPrefix checks that recovering +// unacknowledged work after teardown excludes rows already made durable. +func TestHookProtocolGetUnackedSlicesAcknowledgedPrefix(t *testing.T) { + wire := newRowWire() + cfg := rowConfig() + cfg.Recovery = RecoveryDisabled + cs := newRowStream(t, cfg, newRowOpener(wire)) + + if _, err := cs.IngestBatch(context.Background(), [][]byte{ + []byte("r0"), []byte("r1"), []byte("r2"), + }); err != nil { + t.Fatalf("IngestBatch: %v", err) + } + wire.nextSend(t) + wire.ackRows(1) + waitCondition(t, func() bool { + unacked, err := cs.GetUnackedBatches() + return err == nil && len(unacked) > 0 + }, 3*time.Second) + + unacked, err := cs.GetUnackedBatches() + if err != nil { + t.Fatalf("GetUnackedBatches: %v", err) + } + if len(unacked) != 1 { + t.Fatalf("unacked groups = %d, want 1", len(unacked)) + } + got := make([]string, len(unacked[0])) + for i, record := range unacked[0] { + got[i] = string(record) + } + if !slices.Equal(got, []string{"r1", "r2"}) { + t.Fatalf("unacked rows = %v, want [r1 r2]", got) + } +} + +// ---- typed payload admission ------------------------------------------------ + +func TestEnqueuePayloadValidatesUnitCount(t *testing.T) { + wire := newRowWire() + cs := newRowStream(t, rowConfig(), newRowOpener(wire)) + + if _, err := cs.EnqueuePayload( + context.Background(), &rowPayload{rows: []string{"a"}}, 0, 8, + ); err == nil { + t.Fatal("zero-unit payload was admitted") + } + if _, err := cs.EnqueuePayload( + context.Background(), &rowPayload{rows: []string{"a", "b"}}, 5, 8, + ); err == nil { + t.Fatal("payload with a mismatched unit count was admitted") + } + + offset, err := cs.EnqueuePayload( + context.Background(), &rowPayload{rows: []string{"a", "b"}}, 2, 8, + ) + if err != nil { + t.Fatalf("EnqueuePayload: %v", err) + } + if sent := wire.nextSend(t); len(sent.rows) != 2 || sent.offset != offset { + t.Fatalf("transport received %+v, want 2 rows at offset %d", sent, offset) + } + wire.ackRows(2) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush: %v", err) + } +} + +func TestEnqueuePayloadBuilderReconcilesReservation(t *testing.T) { + wire := newRowWire() + cs := newRowStream(t, rowConfig(), newRowOpener(wire)) + + if _, err := cs.EnqueuePayloadBuilder(context.Background(), 16, nil); err == nil { + t.Fatal("nil builder was accepted") + } + + buildErr := fmt.Errorf("materialization failed") + if _, err := cs.EnqueuePayloadBuilder( + context.Background(), 16, + func() (*rowPayload, uint64, int64, error) { return nil, 0, 0, buildErr }, + ); !errors.Is(err, buildErr) { + t.Fatalf("builder error = %v, want %v", err, buildErr) + } + + // A builder that reports a different unit count than the payload carries is + // a protocol bug and must not reach the buffer. + if _, err := cs.EnqueuePayloadBuilder( + context.Background(), 16, + func() (*rowPayload, uint64, int64, error) { + return &rowPayload{rows: []string{"a"}}, 4, 8, nil + }, + ); err == nil { + t.Fatal("builder with a mismatched unit count was admitted") + } + + // The rejected attempts must not have leaked buffer capacity. + if items, bytes := cs.buf.usage(); items != 0 || bytes != 0 { + t.Fatalf("buffer usage after failed builds = (%d,%d), want (0,0)", items, bytes) + } + + if _, err := cs.EnqueuePayloadBuilder( + context.Background(), 16, + func() (*rowPayload, uint64, int64, error) { + return &rowPayload{rows: []string{"a", "b", "c"}}, 3, 6, nil + }, + ); err != nil { + t.Fatalf("EnqueuePayloadBuilder: %v", err) + } + if sent := wire.nextSend(t); len(sent.rows) != 3 { + t.Fatalf("transport received %d rows, want 3", len(sent.rows)) + } + wire.ackRows(3) + if err := cs.Flush(context.Background()); err != nil { + t.Fatalf("Flush: %v", err) + } +} diff --git a/purego/internal/stream/encoder.go b/purego/internal/stream/encoder.go index db8db020..f9be7158 100644 --- a/purego/internal/stream/encoder.go +++ b/purego/internal/stream/encoder.go @@ -12,16 +12,17 @@ import ( // encodedMsg is a wire-ready EphemeralStream ingest request, built once at // Ingest time and held in the buffer until the sender transmits it. It is the -// Req type the proto/JSON core is instantiated with; the Arrow path will use a -// Flight frame instead. Encoding is eager so the buffer never retains live user -// objects. +// Req type the proto/JSON core is instantiated with; another protocol +// instantiates the core over its own payload type. Encoding is eager so the +// buffer never retains live user objects. type encodedMsg = *zerobuspb.EphemeralStreamRequest // encoder turns user records into offset-independent wire messages and recovers // them again for GetUnacked. It is the send-side per-encoding seam. The core is -// generic over the wire message type Req and -// never names a concrete proto type — proto/JSON supply encoder[encodedMsg]; -// Arrow will supply encoder[flightFrame]. +// generic over the wire message type Req and never names a concrete proto type. +// proto/JSON supply encoder[encodedMsg] directly; a protocol defined outside +// this package supplies EncoderHooks, which hookEncoder adapts to this +// interface. // // The sender replaces the encoded offset with a connection-local wire offset. type encoder[Req any] interface { @@ -34,11 +35,19 @@ type encoder[Req any] interface { encodeBatch(records [][]byte) (Req, error) // stampOffset assigns the connection-local wire offset. stampOffset(msg Req, offset int64) + // unitCount reports the number of protocol durability units in msg. + // Atomic protocols return one even when msg contains a record batch. + unitCount(msg Req) uint64 + // slice removes an acknowledged durability-unit prefix. It is called only + // for a partially acknowledged item during recovery or GetUnacked. + slice(msg Req, acknowledgedPrefix uint64) (Req, error) // decode recovers the raw record bytes from a wire message so GetUnacked can // return original content. A single-record message yields one entry; a batch // yields all of its records so no unacked record is silently dropped. decode(msg Req) [][]byte - // maxWireSize reports an upper bound across every offset stamp. + // maxWireSize reports an upper bound for any one transport frame produced by + // Send, across every offset stamp. A wire stream may expand one logical msg + // into multiple frames behind its single Send call. maxWireSize(msg Req) int // retainedSize estimates the heap retained by an encoded message before it // is built. It includes raw bytes, per-record containers, framing, and @@ -99,6 +108,18 @@ func (protoEncoder) encodeBatch(records [][]byte) (encodedMsg, error) { func (protoEncoder) decode(msg encodedMsg) [][]byte { return extractEphemeralRecords(msg) } +func (protoEncoder) unitCount(encodedMsg) uint64 { return 1 } + +func (protoEncoder) slice(msg encodedMsg, acknowledgedPrefix uint64) (encodedMsg, error) { + if acknowledgedPrefix == 0 { + return msg, nil + } + return nil, fmt.Errorf( + "stream: proto payload is atomic and cannot drop %d acknowledged units", + acknowledgedPrefix, + ) +} + func (protoEncoder) maxWireSize(msg encodedMsg) int { if msg == nil { return 0 @@ -153,6 +174,18 @@ func (jsonEncoder) encodeBatch(records [][]byte) (encodedMsg, error) { func (jsonEncoder) decode(msg encodedMsg) [][]byte { return extractEphemeralRecords(msg) } +func (jsonEncoder) unitCount(encodedMsg) uint64 { return 1 } + +func (jsonEncoder) slice(msg encodedMsg, acknowledgedPrefix uint64) (encodedMsg, error) { + if acknowledgedPrefix == 0 { + return msg, nil + } + return nil, fmt.Errorf( + "stream: JSON payload is atomic and cannot drop %d acknowledged units", + acknowledgedPrefix, + ) +} + func (jsonEncoder) maxWireSize(msg encodedMsg) int { if msg == nil { return 0 @@ -201,6 +234,56 @@ func newEncoder(rt zerobuspb.RecordType) (encoder[encodedMsg], error) { } } +// EncoderHooks adapts protocol functions into the generic encoding seam. It is +// exported within the internal package boundary for protocol implementations +// whose payload type is defined outside stream. +type EncoderHooks[Req any] struct { + EncodeRecord func(record []byte) (Req, error) + EncodeBatch func(records [][]byte) (Req, error) + StampOffset func(msg Req, offset int64) + UnitCount func(msg Req) uint64 + Slice func(msg Req, acknowledgedPrefix uint64) (Req, error) + Decode func(msg Req) [][]byte + MaxWireSize func(msg Req) int + RetainedSize func(rawBytes, recordCount int) int64 +} + +type hookEncoder[Req any] struct { + hooks EncoderHooks[Req] +} + +func (e hookEncoder[Req]) encode(record []byte) (Req, error) { + return e.hooks.EncodeRecord(record) +} + +func (e hookEncoder[Req]) encodeBatch(records [][]byte) (Req, error) { + return e.hooks.EncodeBatch(records) +} + +func (e hookEncoder[Req]) stampOffset(msg Req, offset int64) { + e.hooks.StampOffset(msg, offset) +} + +func (e hookEncoder[Req]) unitCount(msg Req) uint64 { + return e.hooks.UnitCount(msg) +} + +func (e hookEncoder[Req]) slice(msg Req, acknowledgedPrefix uint64) (Req, error) { + return e.hooks.Slice(msg, acknowledgedPrefix) +} + +func (e hookEncoder[Req]) decode(msg Req) [][]byte { + return e.hooks.Decode(msg) +} + +func (e hookEncoder[Req]) maxWireSize(msg Req) int { + return e.hooks.MaxWireSize(msg) +} + +func (e hookEncoder[Req]) retainedSize(rawBytes, recordCount int) int64 { + return e.hooks.RetainedSize(rawBytes, recordCount) +} + // extractEphemeralRecords recovers the raw record bytes from an EphemeralStream // wire message. A single-record message yields one entry; a batch yields all of // its records. Shared by the proto and JSON encoders' decode. diff --git a/purego/internal/stream/encoder_test.go b/purego/internal/stream/encoder_test.go index a957cdc3..c5f7514e 100644 --- a/purego/internal/stream/encoder_test.go +++ b/purego/internal/stream/encoder_test.go @@ -253,3 +253,42 @@ func TestEncoderReusesOffsetPointer(t *testing.T) { t.Fatalf("offset = %d, want 2", got) } } + +// TestAtomicEncoderUnitCountAndSlice pins the durability-unit contract for the +// proto and JSON encoders. slice runs on every reconnect via requeueWithSlicer, +// so its identity case is on the recovery path for the shipping protocols. +func TestAtomicEncoderUnitCountAndSlice(t *testing.T) { + encoders := map[string]encoder[encodedMsg]{ + "proto": protoEncoder{}, + "json": jsonEncoder{}, + } + for name, enc := range encoders { + t.Run(name, func(t *testing.T) { + single, err := enc.encode([]byte("record")) + if err != nil { + t.Fatalf("encode: %v", err) + } + batch, err := enc.encodeBatch([][]byte{[]byte("a"), []byte("b")}) + if err != nil { + t.Fatalf("encodeBatch: %v", err) + } + // A batch is atomic to the server, so it is one unit despite + // carrying several records. + for label, msg := range map[string]encodedMsg{"single": single, "batch": batch} { + if got := enc.unitCount(msg); got != 1 { + t.Fatalf("%s unitCount = %d, want 1", label, got) + } + got, err := enc.slice(msg, 0) + if err != nil { + t.Fatalf("%s slice(0): %v", label, err) + } + if got != msg { + t.Fatalf("%s slice(0) must return the payload unchanged", label) + } + if _, err := enc.slice(msg, 1); err == nil { + t.Fatalf("%s accepted a non-zero acknowledged prefix", label) + } + } + }) + } +} diff --git a/purego/internal/stream/export_test.go b/purego/internal/stream/export_test.go index 80dc6c9e..6608346f 100644 --- a/purego/internal/stream/export_test.go +++ b/purego/internal/stream/export_test.go @@ -17,6 +17,25 @@ func (b *buffer[Req]) enqueue(ctx context.Context, offset int64, msg Req) error return b.append(offset, msg, 1) } +// append adds a single-unit item, the shape every atomic-protocol test uses. +func (b *buffer[Req]) append(offset int64, msg Req, weight int64) error { + return b.appendUnits(offset, msg, 1, weight) +} + +// requeue replays the in-flight set without a payload slicer. +func (b *buffer[Req]) requeue() { + _ = b.requeueWithSlicer(nil) +} + +// discardThrough acknowledges every in-flight item up to and including offset. +func (b *buffer[Req]) discardThrough(offset int64) discardResult { + result, _, _ := b.acknowledge(AckResolution{ + FullyAcknowledgedOffset: offset, + PartialOffset: -1, + }) + return result +} + func (b *buffer[Req]) usage() (int, int64) { b.mu.Lock() defer b.mu.Unlock() diff --git a/purego/internal/stream/helpers_test.go b/purego/internal/stream/helpers_test.go index 8f565966..1b1f5c3a 100644 --- a/purego/internal/stream/helpers_test.go +++ b/purego/internal/stream/helpers_test.go @@ -665,7 +665,7 @@ type blockingAckModel struct { once sync.Once } -func (m *blockingAckModel) classify(resp ephemeralResp) (respKind, int64, pauseSignal) { +func (m *blockingAckModel) classify(resp ephemeralResp) responseClassification { m.once.Do(func() { close(m.entered) }) <-m.release return (offsetAckModel{}).classify(resp) diff --git a/purego/internal/stream/wirestream.go b/purego/internal/stream/wirestream.go index fb718d81..dcaa9412 100644 --- a/purego/internal/stream/wirestream.go +++ b/purego/internal/stream/wirestream.go @@ -2,16 +2,19 @@ package stream import ( "context" + "fmt" "github.com/databricks/zerobus-sdk/purego/internal/transport" ) -// wireStream abstracts an open bidirectional transport stream. -// The core uses one sender goroutine. -type wireStream[Req, Resp any] interface { +// WireStream abstracts an open bidirectional transport stream. Send accepts one +// logical item; an implementation may emit multiple transport frames before it +// returns. The core uses one sender goroutine. +type WireStream[Req, Resp any] interface { // ServerID returns the identifier assigned when this connection opened. ServerID() string - // Send writes one request to the server. + // Send writes one logical request to the server. It returns only after every + // transport frame for that request has been submitted. Send(req Req) error // Recv returns io.EOF when the server ends the stream. Recv() (Resp, error) @@ -22,11 +25,49 @@ type wireStream[Req, Resp any] interface { Close() } +// SubmissionReceipt reports how much of one logical request was successfully +// submitted before SendWithReceipt returned. It lets a multi-frame protocol +// preserve authoritative acknowledgments for earlier frames when a later frame +// fails to send. +type SubmissionReceipt struct { + SubmittedUnits uint64 +} + +// submissionReceiptStream is an optional extension implemented by transports +// that expand one logical request into multiple independently submitted frames. +// Plain WireStream implementations retain the atomic Send behavior. +type submissionReceiptStream[Req any] interface { + SendWithReceipt(req Req) (SubmissionReceipt, error) +} + +// wireStream is the package-local spelling of WireStream, kept so the core and +// its tests are not churned by the exported name. +type wireStream[Req, Resp any] interface { + WireStream[Req, Resp] +} + // opener creates transport streams. type opener[Req, Resp any] interface { Open(ctx context.Context, p StreamParams) (wireStream[Req, Resp], error) } +// OpenFunc opens one protocol transport connection. +type OpenFunc[Req, Resp any] func( + ctx context.Context, + p StreamParams, +) (WireStream[Req, Resp], error) + +type hookOpener[Req, Resp any] struct { + open OpenFunc[Req, Resp] +} + +func (o hookOpener[Req, Resp]) Open( + ctx context.Context, + p StreamParams, +) (wireStream[Req, Resp], error) { + return o.open(ctx, p) +} + // StreamParams aliases the transport parameters. type StreamParams = transport.StreamParams @@ -68,3 +109,42 @@ func NewProtoJSONStream( openingCtx, params, cfg, NewEphemeralOpener(conn), enc, ackMdl, callback, ), nil } + +// NewCoreStreamWithHooks constructs a generic stream from exported/internal +// protocol hooks. It is intended for protocol adapters in sibling internal +// packages; user-facing SDK constructors should wrap it. +func NewCoreStreamWithHooks[Req, Resp any]( + openingCtx context.Context, + params StreamParams, + cfg Config, + open OpenFunc[Req, Resp], + enc EncoderHooks[Req], + acks AckModelHooks[Resp], + callback AckCallback, +) (*CoreStream[Req, Resp], error) { + if open == nil { + return nil, fmt.Errorf("stream: protocol Open hook is required") + } + if enc.EncodeRecord == nil || enc.EncodeBatch == nil || + enc.StampOffset == nil || enc.UnitCount == nil || enc.Slice == nil || + enc.Decode == nil || enc.MaxWireSize == nil || enc.RetainedSize == nil { + return nil, fmt.Errorf("stream: all encoder hooks are required") + } + if acks.Classify == nil { + return nil, fmt.Errorf("stream: acknowledgment Classify hook is required") + } + + var ackMdl ackModel[Resp] = hookAckModel[Resp]{hooks: acks} + if acks.Resolve != nil { + ackMdl = resolvingHookAckModel[Resp]{hookAckModel: hookAckModel[Resp]{hooks: acks}} + } + return NewCoreStream[Req, Resp]( + openingCtx, + params, + cfg, + hookOpener[Req, Resp]{open: open}, + hookEncoder[Req]{hooks: enc}, + ackMdl, + callback, + ), nil +}