fix(streamtimeout): exempt bulk export/import methods

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Caveats: compiled caveats (and their CEL environments) are now cached per schema version — hung off the stored schema (`ReadOnlyStoredSchema`) and rebuilt only when the schema changes — rather than rebuilt on every check, reducing check cost for schemas with many caveats (https://github.com/authzed/spicedb/pull/3166)
 
 ### Fixed
+- Bulk export and import streams are no longer cancelled by the streaming timeout while healthy (clients saw `RST_STREAM NO_ERROR`). The timeout now treats messages received from the client as activity, so it bounds inactivity rather than the total duration of a call: bulk import is no longer capped at `--streaming-api-response-delay-timeout` regardless of how actively the client is sending, and committing an import once the client half-closes is not bounded either. Bulk export, which sends one batch per datastore page and holds no transaction between pages, is exempt from the timeout entirely. All other streaming methods are unchanged. (https://github.com/authzed/spicedb/pull/3143)
 - Fixed a nil pointer dereference panic in `CheckBulkPermissions` that could occur under concurrent load when a tracing-enabled check shared a singleflight dispatch with a non-tracing bulk check. Debug-enabled checks are no longer singleflighted together with non-debug checks. (https://github.com/authzed/spicedb/pull/3174)
 - Fixed a nil pointer dereference panic in the Postgres FDW (https://github.com/authzed/spicedb/pull/3235)
 - CockroachDB: deletes performed by CockroachDB's row-level TTL job for expired relationships are no longer emitted as `DELETE` events by the Watch API. On CockroachDB ≥ 24.1, SpiceDB sets the `ttl_disable_changefeed_replication` storage parameter on the relationship tables at startup (if it lacks `ALTER TABLE` privileges, it logs a warning with the statement to run manually); on older versions a startup warning is logged and TTL deletes continue to be emitted. Note that the parameter affects any changefeed over these tables — external changefeeds that want TTL deletes can opt back in with `ignore_disable_changefeed_replication`. Delete-only transactions also no longer write an internal transaction-metadata marker row, reducing write amplification. (https://github.com/authzed/spicedb/pull/3210)

internal/middleware/streamtimeout/streamtimeout.go

@@ -12,41 +12,78 @@ import (
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
+// Option configures MustStreamServerInterceptor.
+type Option func(*options)
+
+type options struct {
+	exemptMethods map[string]struct{}
+}
+
+// WithExemptMethods exempts the given gRPC full method names from the streaming
+// timeout entirely. Use only for methods that can legitimately go quiet for
+// arbitrarily long without the stream being unhealthy, and whose handlers hold
+// no per-stream resources while quiet. Prefer the default behavior, which
+// bounds inactivity rather than total duration.
+//
+// Pass the generated <Service>_<Method>_FullMethodName constants rather than
+// string literals: a literal missing the leading slash never matches, and the
+// exemption silently becomes a no-op.
+func WithExemptMethods(methods ...string) Option {
+	return func(o *options) {
+		if o.exemptMethods == nil {
+			o.exemptMethods = make(map[string]struct{}, len(methods))
+		}
+		for _, m := range methods {
+			o.exemptMethods[m] = struct{}{}
+		}
+	}
+}
+
 // MustStreamServerInterceptor returns a new stream server interceptor that cancels the context
-// after a timeout if no new data has been received.
-func MustStreamServerInterceptor(timeout time.Duration) grpc.StreamServerInterceptor {
+// after a timeout if the stream has been inactive. Messages sent to the client and messages
+// received from it both count as activity, so the timeout bounds inactivity rather than the
+// total duration of the call. The timer is stopped once the stream stops making progress in
+// either direction, including when the client half-closes and RecvMsg returns io.EOF.
+func MustStreamServerInterceptor(timeout time.Duration, opts ...Option) grpc.StreamServerInterceptor {
 	if timeout <= 0 {
 		panic("timeout must be >= 0 for streaming timeout interceptor")
 	}
+	o := options{}
+	for _, opt := range opts {
+		opt(&o)
+	}
 
 	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		if _, exempt := o.exemptMethods[info.FullMethod]; exempt {
+			return handler(srv, stream)
+		}
 		ctx := stream.Context()
 		withCancel, internalCancelFn := context.WithCancelCause(ctx)
 		timer := time.AfterFunc(timeout, func() {
 			internalCancelFn(spiceerrors.WithCodeAndDetailsAsError(fmt.Errorf("operation took longer than allowed %v to complete", timeout), codes.DeadlineExceeded))
 		})
-		wrapper := &sendWrapper{stream, withCancel, timer, timeout}
+		wrapper := &activityWrapper{stream, withCancel, timer, timeout}
 		return handler(srv, wrapper)
 	}
 }
 
-type sendWrapper struct {
+type activityWrapper struct {
 	grpc.ServerStream
 
 	ctx     context.Context
 	timer   *time.Timer
 	timeout time.Duration
 }
 
-func (s *sendWrapper) Context() context.Context {
+func (s *activityWrapper) Context() context.Context {
 	return s.ctx
 }
 
-func (s *sendWrapper) SetTrailer(_ metadata.MD) {
+func (s *activityWrapper) SetTrailer(_ metadata.MD) {
 	s.timer.Stop()
 }
 
-func (s *sendWrapper) SendMsg(m any) error {
+func (s *activityWrapper) SendMsg(m any) error {
 	err := s.ServerStream.SendMsg(m)
 	if err != nil {
 		s.timer.Stop()
@@ -55,3 +92,18 @@ func (s *sendWrapper) SendMsg(m any) error {
 	}
 	return err
 }
+
+// RecvMsg counts a message received from the client as activity. Without this, a
+// client-streaming method such as ImportBulkRelationships would be bounded by the
+// total duration of the call rather than by inactivity, because its handler does
+// not send anything until SendAndClose. Once the client half-closes, RecvMsg
+// returns io.EOF and the timer stops, so committing the work is not bounded either.
+func (s *activityWrapper) RecvMsg(m any) error {
+	err := s.ServerStream.RecvMsg(m)
+	if err != nil {
+		s.timer.Stop()
+	} else {
+		s.timer.Reset(s.timeout)
+	}
+	return err
+}

internal/middleware/streamtimeout/streamtimeout_test.go

@@ -2,7 +2,9 @@ package streamtimeout
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
 	"testing"
 	"time"
 
@@ -87,3 +89,182 @@ func (s *testSuite) TestStreamTimeout() {
 		s.Require().LessOrEqual(maxCounter, int32(6), "stream was not properly canceled: %d", maxCounter)
 	}
 }
+
+type exemptTestServer struct {
+	testpb.UnimplementedTestServiceServer
+}
+
+func (t exemptTestServer) PingList(_ *testpb.PingListRequest, server testpb.TestService_PingListServer) error {
+	// Sleep well past the 50ms timeout before sending anything. If the
+	// interceptor exemption works, the handler receives the original stream
+	// whose context is not bound to the timer, so Context().Err() stays nil.
+	// If the exemption is broken, the interceptor wraps the stream with a
+	// cancelable context that fires at 50ms — Context().Err() then returns
+	// the DeadlineExceeded cause and this test fails (the smoking gun that
+	// distinguishes the fix from a no-op).
+	time.Sleep(150 * time.Millisecond)
+	if err := server.Context().Err(); err != nil {
+		return fmt.Errorf("expected uncanceled context on exempt method, got %w", err)
+	}
+	if err := server.Send(&testpb.PingListResponse{Counter: 1}); err != nil {
+		return err
+	}
+	return nil
+}
+
+type exemptTestSuite struct {
+	*testpb.InterceptorTestSuite
+}
+
+func TestStreamTimeoutExemptedMethods(t *testing.T) {
+	s := &exemptTestSuite{
+		InterceptorTestSuite: &testpb.InterceptorTestSuite{
+			TestService: &exemptTestServer{},
+			ServerOpts: []grpc.ServerOption{
+				grpc.StreamInterceptor(MustStreamServerInterceptor(
+					50*time.Millisecond,
+					WithExemptMethods(testpb.TestService_PingList_FullMethodName),
+				)),
+			},
+		},
+	}
+	suite.Run(t, s)
+}
+
+func (s *exemptTestSuite) TestExemptedMethodIsNotCanceled() {
+	stream, err := s.Client.PingList(s.SimpleCtx(), &testpb.PingListRequest{Value: "exempt"})
+	s.Require().NoError(err)
+
+	resp, err := stream.Recv()
+	s.Require().NoError(err, "exempt method must not be canceled by streamtimeout")
+	s.Require().Equal(int32(1), resp.Counter)
+
+	_, err = stream.Recv()
+	s.Require().ErrorContains(err, "EOF")
+}
+
+type recvActivityTestServer struct {
+	testpb.UnimplementedTestServiceServer
+}
+
+// PingClientStream mirrors the shape of ImportBulkRelationships: it consumes
+// batches from the client and does not send anything until SendAndClose.
+func (t recvActivityTestServer) PingClientStream(server testpb.TestService_PingClientStreamServer) error {
+	var received int32
+	for {
+		_, err := server.Recv()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return err
+		}
+		received++
+	}
+
+	// The client has half-closed, so nothing further will arrive to reset the timer.
+	// Committing the received work must not be bounded by the timeout either, so stay
+	// quiet well past it before responding.
+	time.Sleep(150 * time.Millisecond)
+	if err := server.Context().Err(); err != nil {
+		return fmt.Errorf("context canceled after client half-close: %w", err)
+	}
+
+	return server.SendAndClose(&testpb.PingClientStreamResponse{Counter: received})
+}
+
+type recvActivityTestSuite struct {
+	*testpb.InterceptorTestSuite
+}
+
+func TestStreamTimeoutResetByReceives(t *testing.T) {
+	s := &recvActivityTestSuite{
+		InterceptorTestSuite: &testpb.InterceptorTestSuite{
+			TestService: &recvActivityTestServer{},
+			ServerOpts: []grpc.ServerOption{
+				grpc.StreamInterceptor(MustStreamServerInterceptor(50 * time.Millisecond)),
+			},
+		},
+	}
+	suite.Run(t, s)
+}
+
+func (s *recvActivityTestSuite) TestReceivesResetTheTimeout() {
+	stream, err := s.Client.PingClientStream(s.SimpleCtx())
+	s.Require().NoError(err)
+
+	// Send batches spaced under the 50ms timeout but totaling well over it. Each
+	// receive has to reset the timer, otherwise the call is canceled partway through
+	// even though the client never stopped sending.
+	for range 5 {
+		s.Require().NoError(stream.Send(&testpb.PingClientStreamRequest{Value: "batch"}))
+		time.Sleep(30 * time.Millisecond)
+	}
+
+	resp, err := stream.CloseAndRecv()
+	s.Require().NoError(err, "client-streaming call must not be canceled while the client is sending")
+	s.Require().Equal(int32(5), resp.Counter)
+}
+
+type mixedTestServer struct {
+	testpb.UnimplementedTestServiceServer
+}
+
+func (t mixedTestServer) PingList(_ *testpb.PingListRequest, server testpb.TestService_PingListServer) error {
+	time.Sleep(150 * time.Millisecond)
+	if err := server.Context().Err(); err != nil {
+		return fmt.Errorf("exempt method was canceled: %w", err)
+	}
+	return server.Send(&testpb.PingListResponse{Counter: 1})
+}
+
+// PingStream is quiet in both directions and is not exempt, so the timer must fire.
+// It waits on the context rather than blocking in Recv, because canceling the
+// interceptor's derived context does not unblock a receive already in flight.
+func (t mixedTestServer) PingStream(server testpb.TestService_PingStreamServer) error {
+	select {
+	case <-server.Context().Done():
+		return server.Context().Err()
+	case <-time.After(time.Second):
+		return fmt.Errorf("non-exempt method was not canceled by streamtimeout")
+	}
+}
+
+type mixedTestSuite struct {
+	*testpb.InterceptorTestSuite
+}
+
+// TestStreamTimeoutMixedExemptions guards against an exemption that applies too
+// broadly: one interceptor exempts PingList while PingStream must still time out.
+func TestStreamTimeoutMixedExemptions(t *testing.T) {
+	s := &mixedTestSuite{
+		InterceptorTestSuite: &testpb.InterceptorTestSuite{
+			TestService: &mixedTestServer{},
+			ServerOpts: []grpc.ServerOption{
+				grpc.StreamInterceptor(MustStreamServerInterceptor(
+					50*time.Millisecond,
+					WithExemptMethods(testpb.TestService_PingList_FullMethodName),
+				)),
+			},
+		},
+	}
+	suite.Run(t, s)
+}
+
+func (s *mixedTestSuite) TestExemptMethodSurvives() {
+	stream, err := s.Client.PingList(s.SimpleCtx(), &testpb.PingListRequest{Value: "exempt"})
+	s.Require().NoError(err)
+
+	resp, err := stream.Recv()
+	s.Require().NoError(err)
+	s.Require().Equal(int32(1), resp.Counter)
+}
+
+func (s *mixedTestSuite) TestNonExemptMethodStillTimesOut() {
+	stream, err := s.Client.PingStream(s.SimpleCtx())
+	s.Require().NoError(err)
+
+	_, err = stream.Recv()
+	s.Require().Error(err, "non-exempt method must still be canceled by streamtimeout")
+	s.Require().ErrorContains(err, "context canceled")
+}

internal/services/v1/experimental.go

@@ -112,7 +112,16 @@ func NewExperimentalServer(dispatch dispatch.Dispatcher, permServerConfig Permis
 				grpcvalidate.StreamServerInterceptor(validator),
 				handwrittenvalidation.StreamServerInterceptor,
 				usagemetrics.StreamServerInterceptor(),
-				streamtimeout.MustStreamServerInterceptor(config.StreamReadTimeout),
+				streamtimeout.MustStreamServerInterceptor(
+					config.StreamReadTimeout,
+					// See NewPermissionsServer: bulk export goes quiet while fetching each
+					// datastore page and holds nothing between them, while bulk import is
+					// left bounded because it holds an open write transaction and its
+					// receives from the client count as activity.
+					streamtimeout.WithExemptMethods(
+						v1.ExperimentalService_BulkExportRelationships_FullMethodName,
+					),
+				),
 				perfinsights.StreamServerInterceptor(permServerConfig.PerformanceInsightMetricsEnabled),
 			),
 		},

internal/services/v1/relationships.go

@@ -185,7 +185,21 @@ func NewPermissionsServer(
 				grpcvalidate.StreamServerInterceptor(validator),
 				handwrittenvalidation.StreamServerInterceptor,
 				usagemetrics.StreamServerInterceptor(),
-				streamtimeout.MustStreamServerInterceptor(configWithDefaults.StreamingAPITimeout),
+				streamtimeout.MustStreamServerInterceptor(
+					configWithDefaults.StreamingAPITimeout,
+					// Bulk export sends one batch per datastore page and sends nothing
+					// while fetching the next one, so the quiet stretches between sends
+					// scale with the size of the data being exported rather than with
+					// the health of the stream. Its handler holds no transaction and no
+					// connection between pages, so exempting it costs nothing.
+					// Bulk import is deliberately not exempt: it holds an open write
+					// transaction for the life of the stream, and receiving a batch from
+					// the client counts as activity, so the timeout still bounds a client
+					// that stops sending.
+					streamtimeout.WithExemptMethods(
+						v1.PermissionsService_ExportBulkRelationships_FullMethodName,
+					),
+				),
 				perfinsights.StreamServerInterceptor(configWithDefaults.PerformanceInsightMetricsEnabled),
 			),
 		},

pkg/cmd/serve.go

@@ -116,7 +116,7 @@ func RegisterServeFlags(cmd *cobra.Command, config *server.Config) error {
 	apiFlags.Uint16Var(&config.MaximumUpdatesPerWrite, "write-relationships-max-updates-per-call", 1000, "maximum number of updates allowed for WriteRelationships calls")
 	apiFlags.IntVar(&config.MaxCaveatContextSize, "max-caveat-context-size", 4096, "maximum allowed size of request caveat context in bytes. A value of zero or less means no limit")
 	apiFlags.IntVar(&config.MaxRelationshipContextSize, "max-relationship-context-size", 25000, "maximum allowed size of the context to be stored in a relationship")
-	apiFlags.DurationVar(&config.StreamingAPITimeout, "streaming-api-response-delay-timeout", 30*time.Second, "maximum time that streaming APIs (LookupSubjects, LookupResources, ReadRelationships and ExportBulkRelationships) can be allowed to run but no response be sent to the client before the stream times out")
+	apiFlags.DurationVar(&config.StreamingAPITimeout, "streaming-api-response-delay-timeout", 30*time.Second, "maximum time that streaming APIs (LookupSubjects, LookupResources, ReadRelationships and ImportBulkRelationships) can be allowed to run but no data be sent to or received from the client before the stream times out; ExportBulkRelationships is exempt")
 	apiFlags.DurationVar(&config.WatchHeartbeat, "watch-api-heartbeat", 1*time.Second, "heartbeat time on the watch in the API. 0 means to default to the datastore's minimum.")
 	apiFlags.Uint32Var(&config.MaxReadRelationshipsLimit, "max-read-relationships-limit", 1000, "maximum number of relationships that can be read in a single request")
 	apiFlags.Uint32Var(&config.MaxDeleteRelationshipsLimit, "max-delete-relationships-limit", 1000, "maximum number of relationships that can be deleted in a single request")