Переглянути джерело

config: make dialer an interface (#1398)

Chris Bannister 5 роки тому
батько
коміт
68212b7a2c
2 змінених файлів з 12 додано та 6 видалено
  1. 6 1
      cluster.go
  2. 6 5
      conn.go

+ 6 - 1
cluster.go

@@ -5,6 +5,7 @@
 package gocql
 
 import (
+	"context"
 	"errors"
 	"net"
 	"time"
@@ -146,12 +147,16 @@ type ClusterConfig struct {
 
 	// Dialer will be used to establish all connections created for this Cluster.
 	// If not provided, a default dialer configured with ConnectTimeout will be used.
-	Dialer *net.Dialer
+	Dialer Dialer
 
 	// internal config for testing
 	disableControlConn bool
 }
 
+type Dialer interface {
+	DialContext(ctx context.Context, network, addr string) (net.Conn, error)
+}
+
 // NewCluster generates a new config for the default cluster implementation.
 //
 // The supplied hosts are used to initially connect to the cluster then the rest of

+ 6 - 5
conn.go

@@ -99,7 +99,7 @@ type ConnConfig struct {
 	CQLVersion     string
 	Timeout        time.Duration
 	ConnectTimeout time.Duration
-	Dialer         *net.Dialer
+	Dialer         Dialer
 	Compressor     Compressor
 	Authenticator  Authenticator
 	AuthProvider   func(h *HostInfo) (Authenticator, error)
@@ -203,14 +203,15 @@ func (s *Session) dialWithoutObserver(ctx context.Context, host *HostInfo, cfg *
 
 	dialer := cfg.Dialer
 	if dialer == nil {
-		dialer = &net.Dialer{
+		d := &net.Dialer{
 			Timeout: cfg.ConnectTimeout,
 		}
+		if cfg.Keepalive > 0 {
+			d.KeepAlive = cfg.Keepalive
+		}
+		dialer = d
 	}
 
-	if cfg.Keepalive > 0 {
-		dialer.KeepAlive = cfg.Keepalive
-	}
 
 	conn, err := dialer.DialContext(ctx, "tcp", host.HostnameAndPort())
 	if err != nil {