Improve listener handling and fix tests (#92)

* Improve listener handling and fix tests

While improving the close handling of the listener in `Proxy` I noticed
that this could be generalized to handle multiple listeners (like
`http.Server`). This could be quite useful in the future for supporting
different multi-proxy features in a more user-friendly and resource
efficient way.  So `Proxy`, the type, now allows multiple listeners in a
single proxy. This hasn't been exposed in an external feature yet.

This also improves clean up for servers, and their listeners, used
during testing and now tests generate unique, unbound ports for each
test to avoid sporadic errors I was seeing in CI.

* Fix comment

Author: mpenick

proxy/proxy.go

@@ -31,7 +31,6 @@ import (
 	"strconv"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/datastax/cql-proxy/parser"
@@ -50,6 +49,8 @@ var (
 )
 
 var ErrProxyClosed = errors.New("proxy closed")
+var ErrProxyAlreadyConnected = errors.New("proxy already connected")
+var ErrProxyNotConnected = errors.New("proxy not connected")
 
 const preparedIdSize = 16
 
@@ -83,18 +84,19 @@ type Proxy struct {
 	ctx                 context.Context
 	config              Config
 	logger              *zap.Logger
-	listener            *net.TCPListener
 	cluster             *proxycore.Cluster
 	sessions            [primitive.ProtocolVersionDse2 + 1]sync.Map // Cache sessions per protocol version
-	sessMu              sync.Mutex
-	schemaEventClients  sync.Map
+	mu                  sync.Mutex
+	isConnected         bool
+	isClosing           bool
+	clients             map[*client]struct{}
+	listeners           map[*net.Listener]struct{}
+	eventClients        sync.Map
 	preparedCache       proxycore.PreparedCache
 	preparedIdempotence sync.Map
-	clientIdGen         uint64
 	lb                  proxycore.LoadBalancer
 	systemLocalValues   map[string]message.Column
 	closed              chan struct{}
-	closingMu           sync.Mutex
 	localNode           *node
 	nodes               []*node
 }
@@ -109,16 +111,15 @@ func (p *Proxy) OnEvent(event proxycore.Event) {
 	switch evt := event.(type) {
 	case *proxycore.SchemaChangeEvent:
 		frm := frame.NewFrame(p.cluster.NegotiatedVersion, -1, evt.Message)
-		p.schemaEventClients.Range(func(key, value interface{}) bool {
-			cl := value.(*client)
+		p.eventClients.Range(func(key, _ interface{}) bool {
+			cl := key.(*client)
 			err := cl.conn.Write(proxycore.SenderFunc(func(writer io.Writer) error {
 				return codec.EncodeFrame(frm, writer)
 			}))
 			cl.conn.LocalAddr()
 			if err != nil {
 				p.logger.Error("unable to send schema change event",
 					zap.Stringer("client", cl.conn.RemoteAddr()),
-					zap.Uint64("id", cl.id),
 					zap.Error(err))
 				_ = cl.conn.Close()
 			}
@@ -138,24 +139,24 @@ func NewProxy(ctx context.Context, config Config) *Proxy {
 		config.RetryPolicy = NewDefaultRetryPolicy()
 	}
 	return &Proxy{
-		ctx:    ctx,
-		config: config,
-		logger: proxycore.GetOrCreateNopLogger(config.Logger),
-		closed: make(chan struct{}),
+		ctx:       ctx,
+		config:    config,
+		logger:    proxycore.GetOrCreateNopLogger(config.Logger),
+		clients:   make(map[*client]struct{}),
+		listeners: make(map[*net.Listener]struct{}),
+		closed:    make(chan struct{}),
 	}
 }
 
-func (p *Proxy) ListenAndServe(address string) error {
-	err := p.Listen(address)
-	if err != nil {
-		return err
+func (p *Proxy) Connect() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.isConnected {
+		return ErrProxyAlreadyConnected
 	}
-	return p.Serve()
-}
 
-func (p *Proxy) Listen(address string) error {
 	var err error
-
 	p.preparedCache, err = getOrCreateDefaultPreparedCache(p.config.PreparedCache)
 	if err != nil {
 		return fmt.Errorf("unable to create prepared cache %w", err)
@@ -210,23 +211,23 @@ func (p *Proxy) Listen(address string) error {
 
 	p.sessions[p.cluster.NegotiatedVersion].Store("", sess) // No keyspace
 
-	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
-	if err != nil {
-		return err
-	}
-	p.listener, err = net.ListenTCP("tcp", tcpAddr)
-	if err != nil {
-		return err
-	}
-
-	p.logger.Info("proxy is listening", zap.Stringer("address", p.listener.Addr()))
-
+	p.isConnected = true
 	return nil
 }
 
-func (p *Proxy) Serve() error {
+// Serve the proxy using the specified listener. It can be called multiple times with different listeners allowing
+// them to share the same backend clusters.
+func (p *Proxy) Serve(l net.Listener) (err error) {
+	l = &closeOnceListener{Listener: l}
+	defer l.Close()
+
+	if err = p.addListener(&l); err != nil {
+		return err
+	}
+	defer p.removeListener(&l)
+
 	for {
-		conn, err := p.listener.AcceptTCP()
+		conn, err := l.Accept()
 		if err != nil {
 			select {
 			case <-p.closed:
@@ -239,15 +240,45 @@ func (p *Proxy) Serve() error {
 	}
 }
 
+func (p *Proxy) addListener(l *net.Listener) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.isClosing {
+		return ErrProxyClosed
+	}
+	if !p.isConnected {
+		return ErrProxyNotConnected
+	}
+	p.listeners[l] = struct{}{}
+	return nil
+}
+
+func (p *Proxy) removeListener(l *net.Listener) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	delete(p.listeners, l)
+}
+
 func (p *Proxy) Close() error {
-	p.closingMu.Lock()
-	defer p.closingMu.Unlock()
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	select {
 	case <-p.closed:
 	default:
 		close(p.closed)
 	}
-	return p.listener.Close()
+	var err error
+	for l := range p.listeners {
+		if closeErr := (*l).Close(); closeErr != nil && err == nil {
+			err = closeErr
+		}
+	}
+	for cl := range p.clients {
+		_ = cl.conn.Close()
+		p.eventClients.Delete(cl)
+		delete(p.clients, cl)
+	}
+	return err
 }
 
 func (p *Proxy) Ready() bool {
@@ -258,28 +289,29 @@ func (p *Proxy) OutageDuration() time.Duration {
 	return p.cluster.OutageDuration()
 }
 
-func (p *Proxy) handle(conn *net.TCPConn) {
-	if err := conn.SetKeepAlive(false); err != nil {
-		p.logger.Warn("failed to disable keepalive on connection", zap.Error(err))
-	}
-
-	if err := conn.SetNoDelay(true); err != nil {
-		p.logger.Warn("failed to set TCP_NODELAY on connection", zap.Error(err))
+func (p *Proxy) handle(conn net.Conn) {
+	if tcpConn, ok := conn.(*net.TCPConn); ok {
+		if err := tcpConn.SetKeepAlive(false); err != nil {
+			p.logger.Warn("failed to disable keepalive on connection", zap.Error(err))
+		}
+		if err := tcpConn.SetNoDelay(true); err != nil {
+			p.logger.Warn("failed to set TCP_NODELAY on connection", zap.Error(err))
+		}
 	}
 
 	cl := &client{
 		ctx:                 p.ctx,
 		proxy:               p,
-		id:                  atomic.AddUint64(&p.clientIdGen, 1),
 		preparedSystemQuery: make(map[[preparedIdSize]byte]interface{}),
 	}
+	p.addClient(cl)
 	cl.conn = proxycore.NewConn(conn, cl)
 	cl.conn.Start()
 }
 
 func (p *Proxy) maybeCreateSession(version primitive.ProtocolVersion, keyspace string) (*proxycore.Session, error) {
-	p.sessMu.Lock()
-	defer p.sessMu.Unlock()
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	if cachedSession, ok := p.sessions[version].Load(keyspace); ok {
 		return cachedSession.(*proxycore.Session), nil
 	} else {
@@ -463,12 +495,30 @@ func (p *Proxy) maybeStorePreparedIdempotence(raw *frame.RawFrame, msg message.M
 	}
 }
 
+func (p *Proxy) addClient(cl *client) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.clients[cl] = struct{}{}
+}
+
+func (p *Proxy) registerForEvents(cl *client) {
+	p.eventClients.Store(cl, struct{}{})
+}
+
+func (p *Proxy) removeClient(cl *client) {
+	p.eventClients.Delete(cl)
+
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	delete(p.clients, cl)
+
+}
+
 type client struct {
 	ctx                 context.Context
 	proxy               *Proxy
 	conn                *proxycore.Conn
 	keyspace            string
-	id                  uint64
 	preparedSystemQuery map[[16]byte]interface{}
 }
 
@@ -505,7 +555,7 @@ func (c *client) Receive(reader io.Reader) error {
 	case *message.Register:
 		for _, t := range msg.EventTypes {
 			if t == primitive.EventTypeSchemaChange {
-				c.proxy.schemaEventClients.Store(c.id, c)
+				c.proxy.registerForEvents(c)
 			}
 		}
 		c.send(raw.Header, &message.Ready{})
@@ -746,7 +796,7 @@ func (c *client) send(hdr *frame.Header, msg message.Message) {
 }
 
 func (c *client) Closing(_ error) {
-	c.proxy.schemaEventClients.Delete(c.id)
+	c.proxy.removeClient(c)
 }
 
 func getOrCreateDefaultPreparedCache(cache proxycore.PreparedCache) (proxycore.PreparedCache, error) {
@@ -818,3 +868,17 @@ func compareIPAddr(a *net.IPAddr, b *net.IPAddr) int {
 
 	return 0
 }
+
+// Wrap the listener so that if it's closed in the serve loop it doesn't race with proxy Close()
+type closeOnceListener struct {
+	net.Listener
+	once     sync.Once
+	closeErr error
+}
+
+func (oc *closeOnceListener) Close() error {
+	oc.once.Do(oc.close)
+	return oc.closeErr
+}
+
+func (oc *closeOnceListener) close() { oc.closeErr = oc.Listener.Close() }

proxy/proxy_retries_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/datastax/go-cassandra-native-protocol/message"
 	"github.com/datastax/go-cassandra-native-protocol/primitive"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 const idempotentQuery = "INSERT INTO test.test (k, v) VALUES ('a', 123e4567-e89b-12d3-a456-426614174000)"
@@ -195,7 +196,7 @@ func TestProxy_Retries(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		numNodesTried, retryCount, err := testProxyRetry(t, &message.Query{Query: tt.query}, tt.response)
+		numNodesTried, retryCount, err := testProxyRetry(t, &message.Query{Query: tt.query}, tt.response, tt.msg)
 		assert.Error(t, err, tt.msg)
 		assert.IsType(t, err, &proxycore.CqlError{}, tt.msg)
 		assert.Equal(t, tt.numNodesTried, numNodesTried, tt.msg)
@@ -228,7 +229,7 @@ func TestProxy_PreparedRetries(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		numNodesTried, retryCount, err := testProxyRetry(t, tt.execute, tt.response)
+		numNodesTried, retryCount, err := testProxyRetry(t, tt.execute, tt.response, tt.msg)
 		assert.Error(t, err, tt.msg)
 		assert.IsType(t, err, &proxycore.CqlError{}, tt.msg)
 		assert.Equal(t, tt.numNodesTried, numNodesTried, tt.msg)
@@ -307,23 +308,21 @@ func TestProxy_BatchRetries(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		numNodesTried, retryCount, err := testProxyRetry(t, tt.batch, tt.response)
+		numNodesTried, retryCount, err := testProxyRetry(t, tt.batch, tt.response, tt.msg)
 		assert.Error(t, err, tt.msg)
 		assert.IsType(t, err, &proxycore.CqlError{}, tt.msg)
 		assert.Equal(t, tt.numNodesTried, numNodesTried, tt.msg)
 		assert.Equal(t, tt.retryCount, retryCount, tt.msg)
 	}
 }
 
-func testProxyRetry(t *testing.T, query message.Message, response message.Error) (numNodesTried, retryCount int, responseError error) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
+func testProxyRetry(t *testing.T, query message.Message, response message.Error, testMessage string) (numNodesTried, retryCount int, responseError error) {
 	var mu sync.Mutex
 	tried := make(map[string]int)
 	prepared := make(map[[16]byte]string)
 
-	cluster, proxy := setupProxyTest(t, ctx, 3, proxycore.MockRequestHandlers{
+	ctx, cancel := context.WithCancel(context.Background())
+	tester, proxyContactPoint, err := setupProxyTest(ctx, 3, proxycore.MockRequestHandlers{
 		primitive.OpCodeQuery: func(cl *proxycore.MockClient, frm *frame.Frame) message.Message {
 			if msg := cl.InterceptQuery(frm.Header, frm.Body.Message.(*message.Query)); msg != nil {
 				return msg
@@ -381,13 +380,14 @@ func testProxyRetry(t *testing.T, query message.Message, response message.Error)
 		},
 	})
 	defer func() {
-		cluster.Shutdown()
-		_ = proxy.Close()
+		cancel()
+		tester.shutdown()
 	}()
+	require.NoError(t, err, testMessage)
 
-	cl := connectTestClient(t, ctx)
+	cl := connectTestClient(t, ctx, proxyContactPoint)
 
-	_, err := cl.Query(ctx, primitive.ProtocolVersion4, query)
+	_, err = cl.Query(ctx, primitive.ProtocolVersion4, query)
 
 	if cqlErr, ok := err.(*proxycore.CqlError); ok {
 		if unprepared, ok := cqlErr.Message.(*message.Unprepared); ok {