Handle nil control connection when reconnecting (#142)

Author: lukasz-antoniak

CHANGELOG/CHANGELOG-2.3.md

@@ -4,6 +4,10 @@ Changelog for the ZDM Proxy, new PRs should update the `unreleased` section.
 
 When cutting a new release, update the `unreleased` heading to the tag being generated and date, like `## vX.Y.Z - YYYY-MM-DD` and create a new placeholder section for `unreleased` entries.
 
+## v2.3.3 - 2025-04-29
+
+* [#142](https://github.com/datastax/zdm-proxy/pull/142): Handle nil control connection when reconnecting
+
 ## v2.3.2 - 2025-04-14
 
 * [#139](https://github.com/datastax/zdm-proxy/pull/139): Ignore forwarding CQL requests for DSE Insights Client to target cluster

RELEASE_NOTES.md

@@ -6,13 +6,21 @@ Build artifacts are available at [Docker Hub](https://hub.docker.com/repository/
 
 For additional details on the changes included in a specific release, see the associated CHANGELOG-x.x.md file.
 
+## v2.3.3 - 2025-04-29
+
+Fix race condition that caused a crash while reconnecting the control connection
+
+[Changelog](CHANGELOG/CHANGELOG-2.3.md#v233---2024-04-29)
+
 ## v2.3.2 - 2025-04-14
 
 Ignore forwarding CQL requests for DSE Insights Client to target cluster
 
 Upgrade software dependencies to resolve vulnerabilities:
 - GoLang to 1.24.2
 
+[Changelog](CHANGELOG/CHANGELOG-2.3.md#v232---2024-04-14)
+
 ## v2.3.1 - 2024-11-08
 
 Upgrade software dependencies to resolve vulnerabilities:

integration-tests/connect_test.go

@@ -46,6 +46,31 @@ func TestGoCqlConnect(t *testing.T) {
 	require.Equal(t, "fake", iter.Columns()[0].Name)
 }
 
+func TestCannotConnectWithoutControlConnection(t *testing.T) {
+	c := setup.NewTestConfig("", "")
+	testSetup, err := setup.NewSimulacronTestSetupWithSessionAndNodesAndConfig(t, true, false, 1, c, nil)
+	require.Nil(t, err)
+	defer testSetup.Cleanup()
+
+	// try to force a scenario where the control connection is reconnecting (cqlConn is nil)
+	// when a new client handler is being created
+	// controlConn.Open() triggers a cqlConn.Close() on the existing connection (and making it nil) before it opens a new one
+	go func() {
+		_, err := testSetup.Proxy.GetOriginControlConn().Open(false, context.Background())
+		if err != nil {
+			t.Logf("err while opening cc in the background: %v", err)
+		}
+	}()
+
+	for i := 0; i < 1000; i++ {
+		// connect to proxy as a "client"
+		client := cqlClient.NewCqlClient("127.0.0.1:14002", nil)
+		conn, err := client.ConnectAndInit(context.Background(), primitive.ProtocolVersion4, 0)
+		require.Nil(t, err)
+		_ = conn.Close()
+	}
+}
+
 // Simulacron-based test to make sure that we can handle invalid protocol error and downgrade
 // used protocol on control connection. ORIGIN and TARGET are using the same C* version
 func TestControlConnectionProtocolVersionNegotiation(t *testing.T) {

proxy/launch.go

@@ -13,7 +13,7 @@ import (
 )
 
 // TODO: to be managed externally
-const ZdmVersionString = "2.3.2"
+const ZdmVersionString = "2.3.3"
 
 var displayVersion = flag.Bool("version", false, "display the ZDM proxy version and exit")
 var configFile = flag.String("config", "", "specify path to ZDM configuration file")

proxy/pkg/zdmproxy/clienthandler.go

@@ -159,6 +159,15 @@ func NewClientHandler(
 		}
 	}
 
+	originCCProtoVer, err := originControlConn.LoadProtoVersion()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load protocol version from %v: %w", common.ClusterTypeOrigin, err)
+	}
+	targetCCProtoVer, err := targetControlConn.LoadProtoVersion()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load protocol version from %v: %w", common.ClusterTypeTarget, err)
+	}
+
 	nodeMetrics, err := metricHandler.GetNodeMetrics(originEndpointId, targetEndpointId, asyncEndpointId)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create node metrics: %w", err)
@@ -169,8 +178,6 @@ func NewClientHandler(
 	requestsDoneCtx, requestsDoneCancelFn := context.WithCancel(context.Background())
 
 	// Initialize stream id processors to manage the ids sent to the clusters
-	originCCProtoVer := originControlConn.cqlConn.GetProtocolVersion()
-	targetCCProtoVer := targetControlConn.cqlConn.GetProtocolVersion()
 	// Calculate maximum number of stream IDs. Take the oldest protocol version negotiated between two clusters
 	// and apply limit defined in proxy configuration. If origin or target cluster are still running protocol V2,
 	// we will limit maximum number of stream IDs to 127 on both clusters. Logic is based on Java driver version 3.x.

proxy/pkg/zdmproxy/controlconn.go

@@ -55,6 +55,7 @@ type ControlConn struct {
 	protocolEventSubscribers map[ProtocolEventObserver]interface{}
 	authEnabled              *atomic.Value
 	metricsHandler           *metrics.MetricHandler
+	controlConnProtoVersion  *atomic.Value
 }
 
 const ProxyVirtualRack = "rack0"
@@ -102,6 +103,7 @@ func NewControlConn(ctx context.Context, defaultPort int, connConfig ConnectionC
 		protocolEventSubscribers: map[ProtocolEventObserver]interface{}{},
 		authEnabled:              authEnabled,
 		metricsHandler:           metricsHandler,
+		controlConnProtoVersion:  &atomic.Value{},
 	}
 }
 
@@ -378,7 +380,7 @@ func (cc *ControlConn) connAndNegotiateProtoVer(endpoint Endpoint, initialProtoV
 				cc.connConfig.GetClusterType(), endpoint.GetEndpointIdentifier(), err)
 			return nil, err
 		}
-		newConn := NewCqlConnection(endpoint, tcpConn, cc.username, cc.password, ccReadTimeout, ccWriteTimeout, cc.conf, protoVer)
+		newConn := NewCqlConnection(cc, endpoint, tcpConn, cc.username, cc.password, ccReadTimeout, ccWriteTimeout, cc.conf, protoVer)
 		err = newConn.InitializeContext(protoVer, ctx)
 		var respErr *ResponseError
 		if err != nil && errors.As(err, &respErr) && respErr.IsProtocolError() && strings.Contains(err.Error(), "Invalid or unsupported protocol version") {
@@ -717,6 +719,26 @@ func (cc *ControlConn) RemoveObserver(observer ProtocolEventObserver) {
 	delete(cc.protocolEventSubscribers, observer)
 }
 
+func (cc *ControlConn) StoreProtoVersion(protoVersion primitive.ProtocolVersion) {
+	cc.controlConnProtoVersion.Store(protoVersion)
+}
+
+func (cc *ControlConn) LoadProtoVersion() (primitive.ProtocolVersion, error) {
+	var protoVersion primitive.ProtocolVersion
+	protoVersionAny := cc.controlConnProtoVersion.Load()
+	if protoVersionAny == nil {
+		return protoVersion, errors.New("protocol version could not be retrieved")
+	}
+
+	var ok bool
+	protoVersion, ok = protoVersionAny.(primitive.ProtocolVersion)
+	if ok {
+		return protoVersion, nil
+	}
+
+	panic("invalid type for protocol version")
+}
+
 func computeAssignedHosts(index int, count int, orderedHosts []*Host) []*Host {
 	i := 0
 	assignedHosts := make([]*Host, 0)

proxy/pkg/zdmproxy/cqlconn.go

@@ -47,6 +47,7 @@ type CqlConnection interface {
 
 // Not thread safe
 type cqlConn struct {
+	controlConn           *ControlConn
 	readTimeout           time.Duration
 	writeTimeout          time.Duration
 	endpoint              Endpoint
@@ -82,12 +83,13 @@ func (c *cqlConn) String() string {
 }
 
 func NewCqlConnection(
-	endpoint Endpoint, conn net.Conn,
+	controlConn *ControlConn, endpoint Endpoint, conn net.Conn,
 	username string, password string,
 	readTimeout time.Duration, writeTimeout time.Duration,
 	conf *config.Config, protoVer primitive.ProtocolVersion) CqlConnection {
 	ctx, cFn := context.WithCancel(context.Background())
 	cqlConn := &cqlConn{
+		controlConn:  controlConn,
 		readTimeout:  readTimeout,
 		writeTimeout: writeTimeout,
 		endpoint:     endpoint,
@@ -257,6 +259,7 @@ func (c *cqlConn) InitializeContext(version primitive.ProtocolVersion, ctx conte
 	}
 
 	c.protocolVersion.Store(version)
+	c.controlConn.StoreProtoVersion(version)
 	c.initialized = true
 	c.authEnabled = authEnabled
 	return nil