fix: implement more robust retries for Watch

Author: miparnisari

internal/client/client.go

@@ -44,6 +44,7 @@ const (
 	defaultRetryJitterFraction     = 0.5
 	bulkImportRoute                = "/authzed.api.v1.ExperimentalService/BulkImportRelationships"
 	importBulkRoute                = "/authzed.api.v1.PermissionsService/ImportBulkRelationships"
+	watchRoute                     = "/authzed.api.v1.WatchService/Watch"
 )
 
 // NewClient defines an (overridable) means of creating a new client.
@@ -232,7 +233,7 @@ func DialOptsFromFlags(cmd *cobra.Command, token storage.Token) ([]grpc.DialOpti
 
 	streamInterceptors := []grpc.StreamClientInterceptor{
 		zgrpcutil.StreamLogDispatchTrailers,
-		selector.StreamClientInterceptor(retry.StreamClientInterceptor(retryOpts...), selector.MatchFunc(isNoneOf(bulkImportRoute, importBulkRoute))),
+		selector.StreamClientInterceptor(retry.StreamClientInterceptor(retryOpts...), selector.MatchFunc(isNoneOf(bulkImportRoute, importBulkRoute, watchRoute))),
 	}
 
 	if !cobrautil.MustGetBool(cmd, "skip-version-check") {

internal/client/client_test.go

@@ -133,35 +133,41 @@ func TestGetCurrentTokenWithCLIOverrideWithoutSecretFile(t *testing.T) {
 	require.Equal(&bTrue, token.Insecure)
 }
 
-type fakeSchemaServer struct {
+type fakeServer struct {
 	v1.UnimplementedSchemaServiceServer
 	v1.UnimplementedExperimentalServiceServer
+	v1.UnimplementedWatchServiceServer
 	v1.UnimplementedPermissionsServiceServer
 	testFunc func()
 }
 
-func (fss *fakeSchemaServer) ReadSchema(_ context.Context, _ *v1.ReadSchemaRequest) (*v1.ReadSchemaResponse, error) {
+func (fss *fakeServer) ReadSchema(_ context.Context, _ *v1.ReadSchemaRequest) (*v1.ReadSchemaResponse, error) {
 	fss.testFunc()
 	return nil, status.Error(codes.Unavailable, "")
 }
 
-func (fss *fakeSchemaServer) BulkImportRelationships(grpc.ClientStreamingServer[v1.BulkImportRelationshipsRequest, v1.BulkImportRelationshipsResponse]) error {
+func (fss *fakeServer) BulkImportRelationships(grpc.ClientStreamingServer[v1.BulkImportRelationshipsRequest, v1.BulkImportRelationshipsResponse]) error {
 	fss.testFunc()
 	return status.Errorf(codes.Aborted, "")
 }
 
-func (fss *fakeSchemaServer) ImportBulkRelationships(grpc.ClientStreamingServer[v1.ImportBulkRelationshipsRequest, v1.ImportBulkRelationshipsResponse]) error {
+func (fss *fakeServer) ImportBulkRelationships(grpc.ClientStreamingServer[v1.ImportBulkRelationshipsRequest, v1.ImportBulkRelationshipsResponse]) error {
 	fss.testFunc()
 	return status.Errorf(codes.Aborted, "")
 }
 
+func (fss *fakeServer) Watch(*v1.WatchRequest, grpc.ServerStreamingServer[v1.WatchResponse]) error {
+	fss.testFunc()
+	return status.Errorf(codes.Unavailable, "")
+}
+
 func TestRetries(t *testing.T) {
 	ctx := t.Context()
 	var callCount uint
 	lis := bufconn.Listen(1024 * 1024)
 	s := grpc.NewServer()
 
-	fakeServer := &fakeSchemaServer{testFunc: func() {
+	fakeServer := &fakeServer{testFunc: func() {
 		callCount++
 	}}
 	v1.RegisterSchemaServiceServer(s, fakeServer)
@@ -190,22 +196,25 @@ func TestRetries(t *testing.T) {
 	c, err := authzed.NewClient("passthrough://bufnet", dialOpts...)
 	require.NoError(t, err)
 
-	_, err = c.ReadSchema(ctx, &v1.ReadSchemaRequest{})
-	grpcutil.RequireStatus(t, codes.Unavailable, err)
-	require.Equal(t, retries, callCount)
+	t.Run("read_schema", func(t *testing.T) {
+		_, err = c.ReadSchema(ctx, &v1.ReadSchemaRequest{})
+		grpcutil.RequireStatus(t, codes.Unavailable, err)
+		require.Equal(t, retries, callCount)
+	})
 }
 
-func TestDoesNotRetryBackupRestore(t *testing.T) {
+func TestDoesNotRetry(t *testing.T) {
 	ctx := t.Context()
 	var callCount uint
 	lis := bufconn.Listen(1024 * 1024)
 	s := grpc.NewServer()
 
-	fakeServer := &fakeSchemaServer{testFunc: func() {
+	fakeServer := &fakeServer{testFunc: func() {
 		callCount++
 	}}
 	v1.RegisterPermissionsServiceServer(s, fakeServer)
 	v1.RegisterExperimentalServiceServer(s, fakeServer)
+	v1.RegisterWatchServiceServer(s, fakeServer)
 
 	go func() {
 		_ = s.Serve(lis)
@@ -231,20 +240,34 @@ func TestDoesNotRetryBackupRestore(t *testing.T) {
 	c, err := authzed.NewClientWithExperimentalAPIs("passthrough://bufnet", dialOpts...)
 	require.NoError(t, err)
 
-	ibc, err := c.ImportBulkRelationships(ctx)
-	require.NoError(t, err)
-	err = ibc.SendMsg(&v1.ImportBulkRelationshipsRequest{})
-	require.NoError(t, err)
-	_, err = ibc.CloseAndRecv()
-	grpcutil.RequireStatus(t, codes.Aborted, err)
-	require.Equal(t, uint(1), callCount)
+	t.Run("import_bulk", func(t *testing.T) {
+		ibc, err := c.ImportBulkRelationships(ctx)
+		require.NoError(t, err)
+		err = ibc.SendMsg(&v1.ImportBulkRelationshipsRequest{})
+		require.NoError(t, err)
+		_, err = ibc.CloseAndRecv()
+		grpcutil.RequireStatus(t, codes.Aborted, err)
+		require.Equal(t, uint(1), callCount)
+	})
 
-	callCount = 0
-	bic, err := c.BulkImportRelationships(ctx)
-	require.NoError(t, err)
-	err = bic.SendMsg(&v1.BulkImportRelationshipsRequest{})
-	require.NoError(t, err)
-	_, err = bic.CloseAndRecv()
-	grpcutil.RequireStatus(t, codes.Aborted, err)
-	require.Equal(t, uint(1), callCount)
+	t.Run("bulk_import", func(t *testing.T) {
+		callCount = 0
+		bic, err := c.BulkImportRelationships(ctx)
+		require.NoError(t, err)
+		err = bic.SendMsg(&v1.BulkImportRelationshipsRequest{})
+		require.NoError(t, err)
+		_, err = bic.CloseAndRecv()
+		grpcutil.RequireStatus(t, codes.Aborted, err)
+		require.Equal(t, uint(1), callCount)
+	})
+
+	t.Run("watch", func(t *testing.T) {
+		callCount = 0
+		watchReq, err := c.Watch(ctx, &v1.WatchRequest{})
+		require.NoError(t, err)
+		resp, err := watchReq.Recv()
+		require.Nil(t, resp)
+		grpcutil.RequireStatus(t, codes.Unavailable, err)
+		require.Equal(t, uint(1), callCount)
+	})
 }

internal/commands/watch.go

@@ -74,26 +74,27 @@ func watchCmdFunc(cmd *cobra.Command, _ []string) error {
 		relFilters = append(relFilters, relFilter)
 	}
 
-	req := &v1.WatchRequest{
-		OptionalObjectTypes:         watchObjectTypes,
-		OptionalRelationshipFilters: relFilters,
-	}
-	if watchRevision != "" {
-		req.OptionalStartCursor = &v1.ZedToken{Token: watchRevision}
-	}
-
 	ctx, cancel := context.WithCancel(cmd.Context())
 	defer cancel()
 
 	signalctx, interruptCancel := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
 	defer interruptCancel()
 
-	watchStream, err := cli.Watch(ctx, req)
-	if err != nil {
-		return err
+	req := &v1.WatchRequest{
+		OptionalObjectTypes:         watchObjectTypes,
+		OptionalRelationshipFilters: relFilters,
 	}
 
 	for {
+		if watchRevision != "" {
+			req.OptionalStartCursor = &v1.ZedToken{Token: watchRevision}
+		}
+
+		watchStream, err := cli.Watch(ctx, req)
+		if err != nil {
+			return err
+		}
+
 		select {
 		case <-signalctx.Done():
 			console.Errorf("stream interrupted after program termination\n")
@@ -104,7 +105,15 @@ func watchCmdFunc(cmd *cobra.Command, _ []string) error {
 		default:
 			resp, err := watchStream.Recv()
 			if err != nil {
-				return err
+				if !strings.Contains(err.Error(), "stream timeout") && !strings.Contains(err.Error(), "RST_STREAM closed stream") {
+					return err
+				}
+				// retry
+				continue
+			}
+
+			if resp.ChangesThrough != nil {
+				watchRevision = resp.ChangesThrough.Token
 			}
 
 			for _, update := range resp.Updates {