Refactor awaitSchemaAgreement logic

conn.go

@@ -1815,7 +1815,7 @@ func (c *Conn) querySystemLocal(ctx context.Context) *Iter {
 	return c.query(ctx, "SELECT * FROM system.local WHERE key='local'"+usingClause)
 }
 
-func (c *Conn) awaitSchemaAgreement(ctx context.Context) (err error) {
+func (c *Conn) awaitSchemaAgreement(ctx context.Context) error {
 	usingClause := ""
 	if c.session.control != nil {
 		usingClause = c.session.usingTimeoutClause
@@ -1827,20 +1827,35 @@ func (c *Conn) awaitSchemaAgreement(ctx context.Context) (err error) {
 
 	endDeadline := time.Now().Add(c.session.cfg.MaxWaitSchemaAgreement)
 
-	for time.Now().Before(endDeadline) {
+	var err error
+	ticker := time.NewTicker(200 * time.Millisecond) // Create a ticker that ticks every 200ms
+	defer ticker.Stop()
+
+	waitForNextTick := func() error {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+			return nil
+		}
+	}
+
+	getSchemaAgreement := func() error {
 		iter := c.querySystemPeers(ctx, c.host.version)
 
 		versions = make(map[string]struct{})
 
-		rows, err := iter.SliceMap()
+		var rows []map[string]interface{}
+		rows, err = iter.SliceMap()
 		if err != nil {
-			goto cont
+			return err
 		}
 
 		for _, row := range rows {
-			host, err := c.session.hostInfoFromMap(row, &HostInfo{connectAddress: c.host.ConnectAddress(), port: c.session.cfg.Port})
+			var host *HostInfo
+			host, err = hostInfoFromMap(row, &HostInfo{connectAddress: c.host.ConnectAddress(), port: c.session.cfg.Port}, c.session.cfg.translateAddressPort)
 			if err != nil {
-				goto cont
+				return err
 			}
 			if !isValidPeer(host) || host.schemaVersion == "" {
 				c.logger.Printf("invalid peer or peer with empty schema_version: peer=%q", host)
@@ -1851,7 +1866,7 @@ func (c *Conn) awaitSchemaAgreement(ctx context.Context) (err error) {
 		}
 
 		if err = iter.Close(); err != nil {
-			goto cont
+			return err
 		}
 
 		iter = c.query(ctx, localSchemas)
@@ -1861,32 +1876,34 @@ func (c *Conn) awaitSchemaAgreement(ctx context.Context) (err error) {
 		}
 
 		if err = iter.Close(); err != nil {
-			goto cont
+			return err
 		}
 
-		if len(versions) <= 1 {
-			return nil
-		}
+		if len(versions) > 1 {
+			schemas := make([]string, 0, len(versions))
+			for schema := range versions {
+				schemas = append(schemas, schema)
+			}
 
-	cont:
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-time.After(200 * time.Millisecond):
+			return &ErrSchemaMismatch{schemas: schemas}
 		}
-	}
 
-	if err != nil {
-		return err
+		return nil
 	}
 
-	schemas := make([]string, 0, len(versions))
-	for schema := range versions {
-		schemas = append(schemas, schema)
+	for time.Now().Before(endDeadline) {
+		err = getSchemaAgreement()
+
+		if err == ErrConnectionClosed || err == nil {
+			return err
+		}
+
+		if tickerErr := waitForNextTick(); tickerErr != nil {
+			return tickerErr
+		}
 	}
 
-	// not exported
-	return fmt.Errorf("gocql: cluster schema versions not consistent: %+v", schemas)
+	return err
 }
 
 var (
@@ -1896,3 +1913,11 @@ var (
 	ErrConnectionClosed  = errors.New("gocql: connection closed waiting for response")
 	ErrNoStreams         = errors.New("gocql: no streams available on connection")
 )
+
+type ErrSchemaMismatch struct {
+	schemas []string
+}
+
+func (e *ErrSchemaMismatch) Error() string {
+	return fmt.Sprintf("gocql: cluster schema versions not consistent: %+v", e.schemas)
+}

host_source.go

@@ -658,7 +658,7 @@ func checkSystemSchema(control *controlConn) (bool, error) {
 
 // Given a map that represents a row from either system.local or system.peers
 // return as much information as we can in *HostInfo
-func (s *Session) hostInfoFromMap(row map[string]interface{}, host *HostInfo) (*HostInfo, error) {
+func hostInfoFromMap(row map[string]interface{}, host *HostInfo, translateAddressPort func(addr net.IP, port int) (net.IP, int)) (*HostInfo, error) {
 	const assertErrorMsg = "Assertion failed for %s"
 	var ok bool
 
@@ -771,7 +771,7 @@ func (s *Session) hostInfoFromMap(row map[string]interface{}, host *HostInfo) (*
 	}
 
 	host.untranslatedConnectAddress = host.ConnectAddress()
-	ip, port := s.cfg.translateAddressPort(host.untranslatedConnectAddress, host.port)
+	ip, port := translateAddressPort(host.untranslatedConnectAddress, host.port)
 	host.connectAddress = ip
 	host.port = port
 
@@ -789,7 +789,7 @@ func (s *Session) hostInfoFromIter(iter *Iter, connectAddress net.IP, defaultPor
 		return nil, errors.New("query returned 0 rows")
 	}
 
-	host, err := s.hostInfoFromMap(rows[0], &HostInfo{connectAddress: connectAddress, port: defaultPort})
+	host, err := hostInfoFromMap(rows[0], &HostInfo{connectAddress: connectAddress, port: defaultPort}, s.cfg.translateAddressPort)
 	if err != nil {
 		return nil, err
 	}
@@ -840,7 +840,7 @@ func (r *ringDescriber) getClusterPeerInfo(localHost *HostInfo) ([]*HostInfo, er
 
 	for _, row := range rows {
 		// extract all available info about the peer
-		host, err := r.session.hostInfoFromMap(row, &HostInfo{port: r.session.cfg.Port})
+		host, err := hostInfoFromMap(row, &HostInfo{port: r.session.cfg.Port}, r.session.cfg.translateAddressPort)
 		if err != nil {
 			return nil, err
 		} else if !isValidPeer(host) {
@@ -913,7 +913,7 @@ func (r *ringDescriber) getHostInfo(hostID UUID) (*HostInfo, error) {
 			}
 
 			for _, row := range rows {
-				h, err := r.session.hostInfoFromMap(row, &HostInfo{port: r.session.cfg.Port})
+				h, err := hostInfoFromMap(row, &HostInfo{port: r.session.cfg.Port}, r.session.cfg.translateAddressPort)
 				if err != nil {
 					return nil, err
 				}

integration_test.go

@@ -6,6 +6,7 @@ package gocql
 // This file groups integration tests where Cassandra has to be set up with some special integration variables
 import (
 	"context"
+	"errors"
 	"reflect"
 	"testing"
 	"time"
@@ -233,3 +234,24 @@ func TestSessionAwaitSchemaAgreement(t *testing.T) {
 		t.Fatalf("expected session.AwaitSchemaAgreement to not return an error but got '%v'", err)
 	}
 }
+
+func TestSessionAwaitSchemaAgreementSessionClosed(t *testing.T) {
+	session := createSession(t)
+	session.Close()
+
+	if err := session.AwaitSchemaAgreement(context.Background()); !errors.Is(err, ErrConnectionClosed) {
+		t.Fatalf("expected session.AwaitSchemaAgreement to return ErrConnectionClosed but got '%v'", err)
+	}
+
+}
+
+func TestSessionAwaitSchemaAgreementContextCanceled(t *testing.T) {
+	session := createSession(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	if err := session.AwaitSchemaAgreement(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected session.AwaitSchemaAgreement to return 'context canceled' but got '%v'", err)
+	}
+
+}