coordination: fix flaky TestCoordinationSemaphore without changing public API

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `coordination.Session.WaitConnected(ctx context.Context) error` to allow callers to wait until the session has fully reconnected after a `Reconnect()` call
+
 ## v3.135.10
 * Fixed the SDK's `database/sql` driver to consistently map session-invalidating YDB errors to `driver.ErrBadConn` where possible, so `database/sql` can detect and discard bad connections
 

coordination/client.go

@@ -124,6 +124,11 @@ type Session interface {
 	// Reconnect forcibly shuts down the underlying gRPC stream and initiates a new one. This method is highly unlikely
 	// to be of use in a typical application but is extremely useful for testing an API implementation.
 	Reconnect()
+
+	// WaitConnected blocks until the session has (re)connected to the server and the session handshake is complete, or
+	// until either the provided ctx is canceled or the session itself is closed, lost, or otherwise canceled. It is
+	// intended to be called after Reconnect to wait for the reconnect to finish before performing further operations.
+	WaitConnected(ctx context.Context) error
 }
 
 // Lease is the object which defines the rights of the session to the acquired semaphore. Lease is alive until its

coordination/example_test.go

@@ -114,6 +114,11 @@ func Example_semaphore() {
 	fmt.Printf("session 1 acquired semaphore 10\n")
 
 	s.Reconnect()
+	if err = s.WaitConnected(ctx); err != nil {
+		fmt.Printf("failed to wait for reconnect: %v", err)
+
+		return
+	}
 	fmt.Printf("session 1 reconnected\n")
 
 	desc, err := s.DescribeSemaphore(

internal/coordination/session.go

@@ -40,9 +40,12 @@ type (
 		controller        *conversation.Controller
 		sessionID         uint64
 
-		mutex                sync.Mutex // guards the field below
+		mutex                sync.Mutex // guards the fields below
 		lastGoodResponseTime time.Time
 		cancelStream         context.CancelFunc
+		// connectedChan is nil when the session is connected; it is set to a new unclosed channel by Reconnect and
+		// closed (and reset to nil) by mainLoop when the session reconnects successfully.
+		connectedChan chan struct{}
 
 		onCreate []func(s *session)
 		onClose  []func(s *session)
@@ -306,10 +309,16 @@ func (s *session) mainLoop(ctx context.Context, path string, sessionStartedChan
 			trace.CoordinationOnSessionStarted(s.trace, start.GetSessionId(), s.sessionID)
 			if s.sessionID == 0 {
 				s.sessionID = start.GetSessionId()
+				// Mark as connected before closing sessionStartedChan so that WaitConnected callers see the correct
+				// state even if Reconnect is called immediately after Session() returns.
+				s.markConnected()
 				close(sessionStartedChan)
 			} else if start.GetSessionId() != s.sessionID {
 				// Reconnect if the server response is invalid.
 				cancelStream()
+			} else {
+				// Reconnect has successfully restored the existing session.
+				s.markConnected()
 			}
 			close(startSending)
 		case <-time.After(s.sessionStartTimeout):
@@ -536,13 +545,55 @@ func (s *session) Close(ctx context.Context) error {
 	return nil
 }
 
+// markConnected signals that the session has successfully connected (or reconnected) to the server. It closes
+// the notification channel set by Reconnect and resets it to nil. It is safe to call even when no reconnect is
+// pending (i.e. connectedChan is already nil), in which case it is a no-op. The method acquires s.mutex
+// internally to synchronize with Reconnect.
+func (s *session) markConnected() {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if s.connectedChan != nil {
+		close(s.connectedChan)
+		s.connectedChan = nil
+	}
+}
+
 func (s *session) Reconnect() {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	if s.cancelStream != nil {
 		s.cancelStream()
 	}
+
+	// Create a new notification channel so that callers of WaitConnected will block until the next successful
+	// reconnect completes. If a channel already exists (i.e. a previous reconnect is still in progress), keep it so
+	// that all waiters are unblocked by the same event.
+	if s.connectedChan == nil {
+		s.connectedChan = make(chan struct{})
+	}
+}
+
+// WaitConnected blocks until the session has reconnected to the server following a Reconnect call, or until the
+// provided ctx is canceled or the session context is done.
+func (s *session) WaitConnected(ctx context.Context) error {
+	s.mutex.Lock()
+	ch := s.connectedChan
+	s.mutex.Unlock()
+
+	if ch == nil {
+		return nil
+	}
+
+	select {
+	case <-ch:
+		return nil
+	case <-s.ctx.Done():
+		return s.ctx.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func (s *session) SessionID() uint64 {

tests/integration/coordination_test.go

@@ -85,6 +85,9 @@ func TestCoordinationSemaphore(sourceTest *testing.T) {
 	fmt.Printf("session 1 acquired semaphore 10\n")
 
 	s.Reconnect()
+	if err = s.WaitConnected(ctx); err != nil {
+		t.Fatalf("failed to wait for reconnect: %v", err)
+	}
 	fmt.Printf("session 1 reconnected\n")
 
 	desc, err := s.DescribeSemaphore(