Introduced ability to provide custom Dial functions.

Author: niemeyer

cluster.go

@@ -51,10 +51,16 @@ type mongoCluster struct {
 	direct       bool
 	cachedIndex  map[string]bool
 	sync         chan bool
+	dial         dialer
 }
 
-func newCluster(userSeeds []string, direct bool) *mongoCluster {
-	cluster := &mongoCluster{userSeeds: userSeeds, references: 1, direct: direct}
+func newCluster(userSeeds []string, direct bool, dial dialer) *mongoCluster {
+	cluster := &mongoCluster{
+		userSeeds: userSeeds,
+		references: 1,
+		direct: direct,
+		dial: dial,
+	}
 	cluster.serverSynced.L = cluster.RWMutex.RLocker()
 	cluster.sync = make(chan bool, 1)
 	go cluster.syncServersLoop()
@@ -329,7 +335,7 @@ func (cluster *mongoCluster) syncServersIteration(direct bool) {
 		go func() {
 			defer wg.Done()
 
-			server, err := newServer(addr, cluster.sync)
+			server, err := newServer(addr, cluster.sync, cluster.dial)
 			if err != nil {
 				log("SYNC Failed to start sync of ", addr, ": ", err.Error())
 				return

cluster_test.go

@@ -27,10 +27,12 @@
 package mgo_test
 
 import (
+	"fmt"
 	"io"
 	. "launchpad.net/gocheck"
 	"labix.org/v2/mgo"
 	"labix.org/v2/mgo/bson"
+	"net"
 	"strings"
 	"time"
 )
@@ -1043,3 +1045,38 @@ func (s *S) TestSetModeEventualIterBug(c *C) {
 	c.Assert(iter.Err(), Equals, nil)
 	c.Assert(i, Equals, N)
 }
+
+func (s *S) TestCustomDial(c *C) {
+	dials := make(chan bool, 16)
+	dial := func(addr net.Addr) (net.Conn, error) {
+		tcpaddr, ok := addr.(*net.TCPAddr)
+		if !ok {
+			return nil, fmt.Errorf("unexpected address type: %T", addr)
+		}
+		dials <- true
+		return net.DialTCP("tcp", nil, tcpaddr)
+	}
+	info := mgo.DialInfo{
+		Addrs: []string{"localhost:40012"},
+		Dial: dial,
+	}
+
+	// Use hostname here rather than IP, to make things trickier.
+	session, err := mgo.DialWithInfo(&info)
+	c.Assert(err, IsNil)
+	defer session.Close()
+
+	const N = 3
+	for i := 0; i < N; i++ {
+		select {
+		case <-dials:
+		case <-time.After(5 * time.Second):
+			c.Fatalf("expected %d dials, got %d", N, i)
+		}
+	}
+	select {
+	case <-dials:
+		c.Fatalf("got more dials than expected")
+	case <-time.After(100 * time.Millisecond):
+	}
+}

server.go

@@ -46,9 +46,12 @@ type mongoServer struct {
 	closed        bool
 	master        bool
 	sync          chan bool
+	dial          dialer
 }
 
-func newServer(addr string, sync chan bool) (server *mongoServer, err error) {
+type dialer func(addr net.Addr) (net.Conn, error)
+
+func newServer(addr string, sync chan bool, dial dialer) (server *mongoServer, err error) {
 	tcpaddr, err := net.ResolveTCPAddr("tcp", addr)
 	if err != nil {
 		log("Failed to resolve ", addr, ": ", err.Error())
@@ -64,6 +67,7 @@ func newServer(addr string, sync chan bool) (server *mongoServer, err error) {
 		ResolvedAddr: resolvedAddr,
 		tcpaddr:      tcpaddr,
 		sync:         sync,
+		dial:         dial,
 	}
 	return
 }
@@ -115,10 +119,17 @@ func (server *mongoServer) Connect() (*mongoSocket, error) {
 	addr := server.Addr
 	tcpaddr := server.tcpaddr
 	master := server.master
+	dial := server.dial
 	server.RUnlock()
 
 	log("Establishing new connection to ", addr, "...")
-	conn, err := net.DialTCP("tcp", nil, tcpaddr)
+	var conn net.Conn
+	var err error
+	if dial == nil {
+		conn, err = net.DialTCP("tcp", nil, tcpaddr)
+	} else {
+		conn, err = dial(tcpaddr)
+	}
 	if err != nil {
 		log("Connection to ", addr, " failed: ", err.Error())
 		return nil, err

session.go

@@ -33,6 +33,7 @@ import (
 	"fmt"
 	"labix.org/v2/mgo/bson"
 	"math"
+	"net"
 	"reflect"
 	"runtime"
 	"sort"
@@ -244,11 +245,16 @@ type DialInfo struct {
 	// or the "admin" database otherwise. See the Session.Login method too.
 	Username string
 	Password string
+
+	// Dial optionally specifies the dial function for creating connections.
+	// At the moment addr will have type *net.TCPAddr, but other types may
+	// be provided in the future, so check and fail if necessary.
+	Dial func(addr net.Addr) (net.Conn, error)
 }
 
 // DialWithInfo establishes a new session to the cluster identified by info.
 func DialWithInfo(info *DialInfo) (*Session, error) {
-	cluster := newCluster(info.Addrs, info.Direct)
+	cluster := newCluster(info.Addrs, info.Direct, info.Dial)
 	session := newSession(Eventual, cluster, info.Timeout)
 	session.defaultdb = info.Database
 	if session.defaultdb == "" {

socket.go

@@ -38,7 +38,7 @@ type replyFunc func(err error, reply *replyOp, docNum int, docData []byte)
 type mongoSocket struct {
 	sync.Mutex
 	server        *mongoServer // nil when cached
-	conn          *net.TCPConn
+	conn          net.Conn
 	addr          string // For debugging only.
 	nextRequestId uint32
 	replyFuncs    map[uint32]replyFunc
@@ -97,7 +97,7 @@ type requestInfo struct {
 	replyFunc replyFunc
 }
 
-func newSocket(server *mongoServer, conn *net.TCPConn) *mongoSocket {
+func newSocket(server *mongoServer, conn net.Conn) *mongoSocket {
 	socket := &mongoSocket{conn: conn, addr: server.Addr}
 	socket.gotNonce.L = &socket.Mutex
 	socket.replyFuncs = make(map[uint32]replyFunc)
@@ -371,7 +371,7 @@ func (socket *mongoSocket) Query(ops ...interface{}) (err error) {
 	return err
 }
 
-func fill(r *net.TCPConn, b []byte) error {
+func fill(r net.Conn, b []byte) error {
 	l := len(b)
 	n, err := r.Read(b)
 	for n != l && err == nil {