fix: implement more robust retries for Watch

Author: miparnisari

internal/client/client.go

@@ -44,6 +44,7 @@ const (
 	defaultRetryJitterFraction     = 0.5
 	importBulkRoute                = "/authzed.api.v1.PermissionsService/ImportBulkRelationships"
 	exportBulkRoute                = "/authzed.api.v1.PermissionsService/ExportBulkRelationships"
+	watchRoute                     = "/authzed.api.v1.WatchService/Watch"
 )
 
 // NewClient defines an (overridable) means of creating a new client.
@@ -235,7 +236,7 @@ func DialOptsFromFlags(cmd *cobra.Command, token storage.Token) ([]grpc.DialOpti
 		// retrying the bulk import in backup/restore logic is handled manually.
 		// retrying bulk export is also handled manually, because the default behavior is
 		// to start at the beginning of the stream, which produces duplicate values.
-		selector.StreamClientInterceptor(retry.StreamClientInterceptor(retryOpts...), selector.MatchFunc(isNoneOf(importBulkRoute, exportBulkRoute))),
+		selector.StreamClientInterceptor(retry.StreamClientInterceptor(retryOpts...), selector.MatchFunc(isNoneOf(importBulkRoute, exportBulkRoute, watchRoute))),
 	}
 
 	if !cobrautil.MustGetBool(cmd, "skip-version-check") {

internal/client/client_test.go

@@ -133,30 +133,36 @@ 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) 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)
@@ -185,22 +191,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)
@@ -226,20 +235,23 @@ 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.ImportBulkRelationships(ctx)
-	require.NoError(t, err)
-	err = bic.SendMsg(&v1.ImportBulkRelationshipsRequest{})
-	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

@@ -9,7 +9,10 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -74,26 +77,32 @@ 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,
+		OptionalUpdateKinds: []v1.WatchKind{
+			v1.WatchKind_WATCH_KIND_INCLUDE_CHECKPOINTS,
+			v1.WatchKind_WATCH_KIND_INCLUDE_SCHEMA_UPDATES,
+		},
 	}
 
 	for {
+		log.Trace().Msg("calling watch with token " + watchRevision)
+		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 +113,26 @@ func watchCmdFunc(cmd *cobra.Command, _ []string) error {
 		default:
 			resp, err := watchStream.Recv()
 			if err != nil {
-				return err
+				statusErr, ok := status.FromError(err)
+				if !ok || (ok && statusErr.Code() != codes.Unavailable) {
+					// cannot retry; return
+					return err
+				}
+
+				log.Trace().Err(err).Msg("will retry")
+				continue
+			}
+
+			if resp.ChangesThrough != nil {
+				watchRevision = resp.ChangesThrough.Token
+				log.Trace().Msg("updated watch revision to " + watchRevision)
+			}
+
+			if resp.SchemaUpdated {
+				console.Println("SCHEMA UPDATED")
+			}
+			if resp.IsCheckpoint {
+				console.Println("CHECKPOINT REACHED")
 			}
 
 			for _, update := range resp.Updates {

internal/commands/watch_test.go

@@ -1,10 +1,21 @@
 package commands
 
 import (
+	"context"
+	"net"
 	"reflect"
+	"sync"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+
+	"github.com/authzed/zed/internal/client"
+	zedtesting "github.com/authzed/zed/internal/testing"
 )
 
 func TestParseRelationshipFilter(t *testing.T) {
@@ -108,3 +119,91 @@ func TestParseRelationshipFilter(t *testing.T) {
 		}
 	}
 }
+
+type mockWatchServer struct {
+	v1.UnimplementedWatchServiceServer
+	sendOnce uint
+}
+
+func (mws *mockWatchServer) Watch(_ *v1.WatchRequest, stream grpc.ServerStreamingServer[v1.WatchResponse]) error {
+	update := &v1.RelationshipUpdate{
+		Operation: v1.RelationshipUpdate_OPERATION_CREATE,
+		Relationship: &v1.Relationship{
+			Resource: &v1.ObjectReference{
+				ObjectType: "document",
+				ObjectId:   "1",
+			},
+			Relation: "viewer",
+			Subject: &v1.SubjectReference{
+				Object: &v1.ObjectReference{
+					ObjectType: "user",
+					ObjectId:   "alice",
+				},
+			},
+		},
+	}
+
+	response := &v1.WatchResponse{
+		Updates:        []*v1.RelationshipUpdate{update},
+		ChangesThrough: &v1.ZedToken{Token: "revision1"},
+	}
+
+	if mws.sendOnce == 0 {
+		mws.sendOnce++
+		return stream.Send(response)
+	}
+
+	return nil
+}
+
+func TestWatchCmdFunc(t *testing.T) {
+	lis := bufconn.Listen(1024 * 1024)
+	s := grpc.NewServer()
+
+	mockServer := &mockWatchServer{}
+	v1.RegisterWatchServiceServer(s, mockServer)
+
+	go func() {
+		_ = s.Serve(lis)
+	}()
+	t.Cleanup(s.Stop)
+
+	conn, err := grpc.NewClient("passthrough://bufnet",
+		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
+			return lis.Dial()
+		}),
+		grpc.WithInsecure(), // nolint:staticcheck
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { conn.Close() })
+
+	client.NewClient = zedtesting.ClientFromConn(conn)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t,
+		zedtesting.StringFlag{FlagName: "log-level", FlagValue: "trace", Changed: true},
+	)
+	cmd.SetContext(ctx)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	watchErr := make(chan error, 1)
+
+	go func() {
+		defer wg.Done()
+		watchErr <- watchCmdFunc(cmd, []string{})
+		// NOTE: no assertions on the changes themselves, as the cmd just prints them out and we have no way of capturing that
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+
+	cancel()
+
+	wg.Wait()
+
+	err = <-watchErr
+	require.ErrorContains(t, err, "EOF")
+}