Fix proxy close logic (#82)

`Accept()` can return `proxy.test: error: accept tcp [::]:9042: use of closed network connection` when closing. This case needs to be handled specially when the proxy is closed properly.

Author: mpenick

proxy/proxy.go

@@ -44,6 +44,8 @@ var (
 	encodedZeroValue, _ = proxycore.EncodeType(datatype.Int, primitive.ProtocolVersion4, 0)
 )
 
+var ErrProxyClosed = errors.New("proxy closed")
+
 const preparedIdSize = 16
 
 type Config struct {
@@ -76,6 +78,8 @@ type Proxy struct {
 	clientIdGen         uint64
 	lb                  proxycore.LoadBalancer
 	systemLocalValues   map[string]message.Column
+	closed              chan struct{}
+	closingMu           sync.Mutex
 }
 
 func (p *Proxy) OnEvent(event proxycore.Event) {
@@ -114,6 +118,7 @@ func NewProxy(ctx context.Context, config Config) *Proxy {
 		ctx:    ctx,
 		config: config,
 		logger: proxycore.GetOrCreateNopLogger(config.Logger),
+		closed: make(chan struct{}),
 	}
 }
 
@@ -195,13 +200,25 @@ func (p *Proxy) Serve() error {
 	for {
 		conn, err := p.listener.AcceptTCP()
 		if err != nil {
-			return err
+			select {
+			case <-p.closed:
+				return ErrProxyClosed
+			default:
+				return err
+			}
 		}
 		p.handle(conn)
 	}
 }
 
-func (p *Proxy) Shutdown() error {
+func (p *Proxy) Close() error {
+	p.closingMu.Lock()
+	defer p.closingMu.Unlock()
+	select {
+	case <-p.closed:
+	default:
+		close(p.closed)
+	}
 	return p.listener.Close()
 }
 

proxy/proxy_retries_test.go

@@ -382,7 +382,7 @@ func testProxyRetry(t *testing.T, query message.Message, response message.Error)
 	})
 	defer func() {
 		cluster.Shutdown()
-		_ = proxy.Shutdown()
+		_ = proxy.Close()
 	}()
 
 	cl := connectTestClient(t, ctx)

proxy/proxy_test.go

@@ -74,7 +74,7 @@ func TestProxy_ListenAndServe(t *testing.T) {
 	})
 	defer func() {
 		cluster.Shutdown()
-		_ = proxy.Shutdown()
+		_ = proxy.Close()
 	}()
 
 	cl := connectTestClient(t, ctx)
@@ -138,7 +138,7 @@ func TestProxy_Unprepared(t *testing.T) {
 	})
 	defer func() {
 		cluster.Shutdown()
-		_ = proxy.Shutdown()
+		_ = proxy.Close()
 	}()
 
 	cl := connectTestClient(t, ctx)
@@ -174,7 +174,7 @@ func TestProxy_UseKeyspace(t *testing.T) {
 	cluster, proxy := setupProxyTest(t, ctx, 1, nil)
 	defer func() {
 		cluster.Shutdown()
-		_ = proxy.Shutdown()
+		_ = proxy.Close()
 	}()
 
 	cl := connectTestClient(t, ctx)
@@ -195,7 +195,7 @@ func TestProxy_NegotiateProtocolV5(t *testing.T) {
 	cluster, proxy := setupProxyTest(t, ctx, 1, nil)
 	defer func() {
 		cluster.Shutdown()
-		_ = proxy.Shutdown()
+		_ = proxy.Close()
 	}()
 
 	cl, err := proxycore.ConnectClient(ctx, proxycore.NewEndpoint(testProxyContactPoint), proxycore.ClientConnConfig{})

proxy/run.go

@@ -269,14 +269,14 @@ func listenAndServe(p *Proxy, ctx context.Context, logger *zap.Logger) (err erro
 		case <-ctx.Done():
 			logger.Debug("proxy interrupted/killed")
 			_ = server.Close()
-			_ = p.Shutdown()
+			_ = p.Close()
 		}
 	}()
 
 	go func() {
 		defer wg.Done()
 		err = p.Serve()
-		if err != nil {
+		if err != nil && err != ErrProxyClosed {
 			ch <- err
 		}
 	}()