fix(coordination): prevent spurious reconnect after session start

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed coordination sessions spuriously reconnecting immediately after creation, which could cause non-idempotent operations to fail with `operation status is unknown`
+
 ## v3.149.0
 * Added `balancers.PreferPrimaryPile` and `balancers.PreferPrimaryPileWithFallback` endpoint selection policies
 

internal/coordination/session.go

@@ -442,8 +442,10 @@ func (s *session) receiveLoop( //nolint:funlen
 
 			return
 		case *Ydb_Coordination.SessionResponse_SessionStarted_:
-			sessionStarted <- message.GetSessionStarted()
+			// mainLoop may enter the keep-alive loop as soon as the notification is received.
+			// Publish it only after the timestamp used by that loop has been initialized.
 			s.updateLastGoodResponseTime()
+			sessionStarted <- message.GetSessionStarted()
 		case *Ydb_Coordination.SessionResponse_SessionStopped_:
 			sessionStopped <- message.GetSessionStopped()
 			s.cancel()

internal/coordination/session_test.go

@@ -1,11 +1,72 @@
 package coordination
 
 import (
+	"context"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Coordination"
+	"go.uber.org/mock/gomock"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+func TestReceiveLoopUpdatesLastGoodResponseTimeBeforeSessionStarted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	client := NewMockCoordinationService_SessionClient(ctrl)
+	streamCtx, cancelStream := context.WithCancel(context.Background())
+	startedResponse := &Ydb_Coordination.SessionResponse_SessionStarted{
+		SessionId: 42,
+	}
+	recvReturned := make(chan struct{})
+
+	firstRecv := client.EXPECT().Recv().DoAndReturn(func() (*Ydb_Coordination.SessionResponse, error) {
+		close(recvReturned)
+
+		return &Ydb_Coordination.SessionResponse{
+			Response: &Ydb_Coordination.SessionResponse_SessionStarted_{
+				SessionStarted: startedResponse,
+			},
+		}, nil
+	})
+	secondRecv := client.EXPECT().Recv().DoAndReturn(func() (*Ydb_Coordination.SessionResponse, error) {
+		<-streamCtx.Done()
+
+		return nil, streamCtx.Err()
+	})
+	gomock.InOrder(firstRecv.Call, secondRecv.Call)
+
+	s := &session{
+		trace: &trace.Coordination{},
+	}
+	sessionStarted := make(chan *Ydb_Coordination.SessionResponse_SessionStarted, 1)
+	sessionStopped := make(chan *Ydb_Coordination.SessionResponse_SessionStopped, 1)
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	s.mutex.Lock()
+	go s.receiveLoop(&wg, client, cancelStream, sessionStarted, sessionStopped)
+
+	<-recvReturned
+	var start *Ydb_Coordination.SessionResponse_SessionStarted
+	select {
+	case start = <-sessionStarted:
+	case <-time.After(100 * time.Millisecond): //nolint:mnd
+	}
+	publishedBeforeTimestamp := start != nil
+	s.mutex.Unlock()
+	if start == nil {
+		start = <-sessionStarted
+	}
+
+	cancelStream()
+	wg.Wait()
+	require.False(t, publishedBeforeTimestamp, "session started was published before the keep-alive timestamp update")
+	require.Same(t, startedResponse, start)
+	require.False(t, s.getLastGoodResponseTime().IsZero())
+}
+
 func TestNewProtectionKey(t *testing.T) {
 	key1 := newProtectionKey()
 	require.NotNil(t, key1)