internal/grpcsync: add ScheduleAndWait to CallbackSerializer (#9162)

Adds `ScheduleAndWait` to `CallbackSerializer`. It schedules a callback
and blocks until the callback has run. If the serializer was already
closed before the callback could be scheduled, it returns
`ErrSerializerClosed`, so callers can tell whether the callback actually
ran.

This replaces the common pattern of calling `ScheduleOr` with a `done`
channel just to wait for a callback to finish.

Implements the design discussed in #8501.

Fixes #8501

RELEASE NOTES: none

Author: ulascansenturk

internal/grpcsync/callback_serializer.go

@@ -20,10 +20,15 @@ package grpcsync
 
 import (
 	"context"
+	"errors"
 
 	"google.golang.org/grpc/internal/buffer"
 )
 
+// ErrSerializerClosed is returned by ScheduleAndWait if the CallbackSerializer
+// was closed before the callback could be scheduled.
+var ErrSerializerClosed = errors.New("callback serializer is closed")
+
 // CallbackSerializer provides a mechanism to schedule callbacks in a
 // synchronized manner. It provides a FIFO guarantee on the order of execution
 // of scheduled callbacks. New callbacks can be scheduled by invoking the
@@ -77,6 +82,27 @@ func (cs *CallbackSerializer) ScheduleOr(f func(ctx context.Context), onFailure
 	}
 }
 
+// ScheduleAndWait schedules the provided callback function f to be executed in
+// the order it was added and blocks until f has run. If the context passed to
+// NewCallbackSerializer was canceled before this method is called, f is not run
+// and ScheduleAndWait returns ErrSerializerClosed.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) ScheduleAndWait(f func(ctx context.Context)) error {
+	done := make(chan struct{})
+	var err error
+	cs.ScheduleOr(func(ctx context.Context) {
+		f(ctx)
+		close(done)
+	}, func() {
+		err = ErrSerializerClosed
+		close(done)
+	})
+	<-done
+	return err
+}
+
 func (cs *CallbackSerializer) run(ctx context.Context) {
 	defer close(cs.done)
 

internal/grpcsync/callback_serializer_test.go

@@ -204,3 +204,44 @@ func (s) TestCallbackSerializer_Schedule_Close(t *testing.T) {
 	case <-done:
 	}
 }
+
+// TestCallbackSerializer_ScheduleAndWait verifies that ScheduleAndWait runs the
+// provided callback and blocks until it has finished executing, returning a nil
+// error.
+func (s) TestCallbackSerializer_ScheduleAndWait(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cs := NewCallbackSerializer(ctx)
+
+	executed := false
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		executed = true
+	}); err != nil {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want nil", err)
+	}
+
+	// The callback must have run by the time ScheduleAndWait returns. No
+	// synchronization is needed here because ScheduleAndWait is expected to
+	// block until the callback completes.
+	if !executed {
+		t.Fatal("ScheduleAndWait() returned before the callback was executed")
+	}
+}
+
+// TestCallbackSerializer_ScheduleAndWait_Closed verifies that ScheduleAndWait
+// does not run the callback and returns ErrSerializerClosed when the serializer
+// has already been closed.
+func (s) TestCallbackSerializer_ScheduleAndWait_Closed(t *testing.T) {
+	serializerCtx, serializerCancel := context.WithCancel(context.Background())
+	cs := NewCallbackSerializer(serializerCtx)
+
+	// Close the serializer and wait for it to finish shutting down.
+	serializerCancel()
+	<-cs.Done()
+
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		t.Fatal("Callback executed on a closed serializer")
+	}); err != ErrSerializerClosed {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want %v", err, ErrSerializerClosed)
+	}
+}