// Tests of the round-robin connection selection policy implementation
func TestRoundRobinConnPolicy(t *testing.T) {
policy := RoundRobinConnPolicy()()
conn0 := &Conn{streams: streams.New(1)}
conn1 := &Conn{streams: streams.New(1)}
conn := []*Conn{
conn0,
conn1,
}
policy.SetConns(conn)
if actual := policy.Pick(nil); actual != conn0 {
t.Error("Expected conn1")
}
if actual := policy.Pick(nil); actual != conn1 {
t.Error("Expected conn0")
}
if actual := policy.Pick(nil); actual != conn0 {
t.Error("Expected conn1")
}
}
// Connect establishes a connection to a Cassandra node.
func Connect(host *HostInfo, addr string, cfg *ConnConfig,
errorHandler ConnErrorHandler, session *Session) (*Conn, error) {
var (
err error
conn net.Conn
)
dialer := &net.Dialer{
Timeout: cfg.Timeout,
}
if cfg.tlsConfig != nil {
// the TLS config is safe to be reused by connections but it must not
// be modified after being used.
conn, err = tls.DialWithDialer(dialer, "tcp", addr, cfg.tlsConfig)
} else {
conn, err = dialer.Dial("tcp", addr)
}
if err != nil {
return nil, err
}
// going to default to proto 2
if cfg.ProtoVersion < protoVersion1 || cfg.ProtoVersion > protoVersion4 {
log.Printf("unsupported protocol version: %d using 2\n", cfg.ProtoVersion)
cfg.ProtoVersion = 2
}
headerSize := 8
if cfg.ProtoVersion > protoVersion2 {
headerSize = 9
}
c := &Conn{
conn: conn,
r: bufio.NewReader(conn),
cfg: cfg,
calls: make(map[int]*callReq),
timeout: cfg.Timeout,
version: uint8(cfg.ProtoVersion),
addr: conn.RemoteAddr().String(),
errorHandler: errorHandler,
compressor: cfg.Compressor,
auth: cfg.Authenticator,
headerBuf: make([]byte, headerSize),
quit: make(chan struct{}),
session: session,
streams: streams.New(cfg.ProtoVersion),
host: host,
}
if cfg.Keepalive > 0 {
c.setKeepalive(cfg.Keepalive)
}
go c.serve()
if err := c.startup(); err != nil {
conn.Close()
return nil, err
}
c.started = true
return c, nil
}
// Connect establishes a connection to a Cassandra node.
func Connect(host *HostInfo, cfg *ConnConfig, errorHandler ConnErrorHandler, session *Session) (*Conn, error) {
// TODO(zariel): remove these
if host == nil {
panic("host is nil")
} else if len(host.Peer()) == 0 {
panic("host missing peer ip address")
} else if host.Port() == 0 {
panic("host missing port")
}
var (
err error
conn net.Conn
)
dialer := &net.Dialer{
Timeout: cfg.Timeout,
}
// TODO(zariel): handle ipv6 zone
translatedPeer, translatedPort := session.cfg.translateAddressPort(host.Peer(), host.Port())
addr := (&net.TCPAddr{IP: translatedPeer, Port: translatedPort}).String()
//addr := (&net.TCPAddr{IP: host.Peer(), Port: host.Port()}).String()
if cfg.tlsConfig != nil {
// the TLS config is safe to be reused by connections but it must not
// be modified after being used.
conn, err = tls.DialWithDialer(dialer, "tcp", addr, cfg.tlsConfig)
} else {
conn, err = dialer.Dial("tcp", addr)
}
if err != nil {
return nil, err
}
c := &Conn{
conn: conn,
r: bufio.NewReader(conn),
cfg: cfg,
calls: make(map[int]*callReq),
timeout: cfg.Timeout,
version: uint8(cfg.ProtoVersion),
addr: conn.RemoteAddr().String(),
errorHandler: errorHandler,
compressor: cfg.Compressor,
auth: cfg.Authenticator,
quit: make(chan struct{}),
session: session,
streams: streams.New(cfg.ProtoVersion),
host: host,
}
if cfg.Keepalive > 0 {
c.setKeepalive(cfg.Keepalive)
}
var (
ctx context.Context
cancel func()
)
if c.timeout > 0 {
ctx, cancel = context.WithTimeout(context.Background(), c.timeout)
} else {
ctx, cancel = context.WithCancel(context.Background())
}
defer cancel()
frameTicker := make(chan struct{}, 1)
startupErr := make(chan error)
go func() {
for range frameTicker {
err := c.recv()
if err != nil {
select {
case startupErr <- err:
case <-ctx.Done():
}
return
}
}
}()
go func() {
defer close(frameTicker)
err := c.startup(ctx, frameTicker)
select {
case startupErr <- err:
case <-ctx.Done():
}
}()
select {
case err := <-startupErr:
if err != nil {
c.Close()
return nil, err
}
case <-ctx.Done():
c.Close()
return nil, errors.New("gocql: no response to connection startup within timeout")
}
go c.serve()
return c, nil
}
