Merge pull request #396 from dolthub/andy/fix-poll-for-closed-connection

/server: Unwrap `netutil.ConnWithTimeouts` in pollForClosedConnection

server/handler.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/dolthub/vitess/go/mysql"
+	"github.com/dolthub/vitess/go/netutil"
 	"github.com/dolthub/vitess/go/sqltypes"
 	"github.com/dolthub/vitess/go/vt/proto/query"
 	"github.com/dolthub/vitess/go/vt/sqlparser"
@@ -452,7 +453,7 @@ func (h *Handler) errorWrappedDoQuery(
 // the supplied error channel if it has. Meant to be run in a separate goroutine from the query handler routine.
 // Returns immediately on platforms that can't support TCP socket checks.
 func (h *Handler) pollForClosedConnection(c *mysql.Conn, errChan chan error, quit chan struct{}, query string) {
-	tcpConn, ok := c.Conn.(*net.TCPConn)
+	tcpConn, ok := maybeGetTCPConn(c.Conn)
 	if !ok {
 		logrus.Debug("Connection checker exiting, connection isn't TCP")
 		return
@@ -496,6 +497,20 @@ func (h *Handler) pollForClosedConnection(c *mysql.Conn, errChan chan error, qui
 	}
 }
 
+func maybeGetTCPConn(conn net.Conn) (*net.TCPConn, bool) {
+	wrap, ok := conn.(netutil.ConnWithTimeouts)
+	if ok {
+		conn = wrap.Conn
+	}
+
+	tcp, ok := conn.(*net.TCPConn)
+	if ok {
+		return tcp, true
+	}
+
+	return nil, false
+}
+
 func isSessionAutocommit(ctx *sql.Context) (bool, error) {
 	autoCommitSessionVar, err := ctx.GetSessionVariable(ctx, sql.AutoCommitSessionVar)
 	if err != nil {