grpcsync: simplifications in CallbackSerializer

Author: dfawley

internal/grpcsync/callback_serializer.go

@@ -20,7 +20,6 @@ package grpcsync
 
 import (
 	"context"
-	"sync"
 
 	"google.golang.org/grpc/internal/buffer"
 )
@@ -38,8 +37,6 @@ type CallbackSerializer struct {
 	done chan struct{}
 
 	callbacks *buffer.Unbounded
-	closedMu  sync.Mutex
-	closed    bool
 }
 
 // NewCallbackSerializer returns a new CallbackSerializer instance. The provided
@@ -65,53 +62,34 @@ func NewCallbackSerializer(ctx context.Context) *CallbackSerializer {
 // callbacks to be executed by the serializer. It is not possible to add
 // callbacks once the context passed to NewCallbackSerializer is cancelled.
 func (cs *CallbackSerializer) Schedule(f func(ctx context.Context)) bool {
-	cs.closedMu.Lock()
-	defer cs.closedMu.Unlock()
-
-	if cs.closed {
-		return false
-	}
-	cs.callbacks.Put(f)
-	return true
+	return cs.callbacks.Put(f) == nil
 }
 
 func (cs *CallbackSerializer) run(ctx context.Context) {
-	var backlog []func(context.Context)
-
 	defer close(cs.done)
+
+	// TODO: when Go 1.21 is the oldest supported version, this loop and Close
+	// can be replaced with:
+	//
+	// context.AfterFunc(ctx, cs.callbacks.Close)
 	for ctx.Err() == nil {
 		select {
 		case <-ctx.Done():
 			// Do nothing here. Next iteration of the for loop will not happen,
 			// since ctx.Err() would be non-nil.
-		case callback := <-cs.callbacks.Get():
+		case cb := <-cs.callbacks.Get():
 			cs.callbacks.Load()
-			callback.(func(ctx context.Context))(ctx)
+			cb.(func(context.Context))(ctx)
 		}
 	}
 
-	// Fetch pending callbacks if any, and execute them before returning from
-	// this method and closing cs.done.
-	cs.closedMu.Lock()
-	cs.closed = true
-	backlog = cs.fetchPendingCallbacks()
+	// Close the buffer to prevent new callbacks from being added.
 	cs.callbacks.Close()
-	cs.closedMu.Unlock()
-	for _, b := range backlog {
-		b(ctx)
-	}
-}
 
-func (cs *CallbackSerializer) fetchPendingCallbacks() []func(context.Context) {
-	var backlog []func(context.Context)
-	for {
-		select {
-		case b := <-cs.callbacks.Get():
-			backlog = append(backlog, b.(func(context.Context)))
-			cs.callbacks.Load()
-		default:
-			return backlog
-		}
+	// Run all pending callbacks.
+	for cb := range cs.callbacks.Get() {
+		cs.callbacks.Load()
+		cb.(func(context.Context))(ctx)
 	}
 }