pgwire: implement cancellation protocol

pkg/sql/cluster_wide_id.go

@@ -53,3 +53,9 @@ func BytesToClusterWideID(b []byte) ClusterWideID {
 func (id ClusterWideID) GetNodeID() int32 {
 	return int32(0xFFFFFFFF & id.Lo)
 }
+
+// GetPGWireCancelInfo extracts the timestamp of the cluster wide ID as 2
+// int32s, for use by the pgwire cancellation protocol.
+func (id ClusterWideID) GetPGWireCancelInfo() (int32, int32) {
+	return int32(id.Hi >> 32), int32(0xFFFFFFFF & id.Hi)
+}

pkg/sql/conn_executor.go

@@ -461,6 +461,13 @@ func (h ConnectionHandler) GetParamStatus(ctx context.Context, varName string) s
 	return defVal
 }
 
+// GetStatusParam retrieves the configured value of the session
+// variable identified by varName. This is used for the initial
+// message sent to a client during a session set-up.
+func (h ConnectionHandler) GetSessionID() ClusterWideID {
+	return h.ex.sessionID
+}
+
 // ServeConn serves a client connection by reading commands from the stmtBuf
 // embedded in the ConnHandler.
 //
@@ -655,6 +662,7 @@ func (s *Server) newConnExecutor(
 		settings:          s.cfg.Settings,
 	}
 	ex.extraTxnState.txnRewindPos = -1
+	ex.sessionID = ex.generateID()
 	ex.mu.ActiveQueries = make(map[ClusterWideID]*queryMeta)
 	ex.machine = fsm.MakeMachine(TxnStateTransitions, stateNoTxn{}, &ex.state)
 
@@ -1282,7 +1290,6 @@ func (ex *connExecutor) run(
 	ex.ctxHolder.connCtx = ctx
 	ex.onCancelSession = onCancel
 
-	ex.sessionID = ex.generateID()
 	ex.server.cfg.SessionRegistry.register(ex.sessionID, ex)
 	ex.planner.extendedEvalCtx.setSessionID(ex.sessionID)
 	defer ex.server.cfg.SessionRegistry.deregister(ex.sessionID)
@@ -2209,6 +2216,15 @@ func (ex *connExecutor) cancelQuery(queryID ClusterWideID) bool {
 	return false
 }
 
+// cancelCurrentQueries is part of the registrySession interface.
+func (ex *connExecutor) cancelCurrentQueries() {
+	ex.mu.Lock()
+	defer ex.mu.Unlock()
+	for _, queryMeta := range ex.mu.ActiveQueries {
+		queryMeta.cancel()
+	}
+}
+
 // cancelSession is part of the registrySession interface.
 func (ex *connExecutor) cancelSession() {
 	if ex.onCancelSession == nil {

pkg/sql/exec_util.go

@@ -1098,36 +1098,52 @@ type SessionArgs struct {
 // Use register() and deregister() to modify this registry.
 type SessionRegistry struct {
 	syncutil.Mutex
-	sessions map[ClusterWideID]registrySession
+	sessions         map[ClusterWideID]registrySession
+	sessionsByPGWire map[uint64]registrySession
 }
 
 // NewSessionRegistry creates a new SessionRegistry with an empty set
 // of sessions.
 func NewSessionRegistry() *SessionRegistry {
-	return &SessionRegistry{sessions: make(map[ClusterWideID]registrySession)}
+	return &SessionRegistry{
+		sessions:         make(map[ClusterWideID]registrySession),
+		sessionsByPGWire: make(map[uint64]registrySession),
+	}
 }
 
 func (r *SessionRegistry) register(id ClusterWideID, s registrySession) {
 	r.Lock()
 	r.sessions[id] = s
+	r.sessionsByPGWire[id.Hi] = s
 	r.Unlock()
 }
 
 func (r *SessionRegistry) deregister(id ClusterWideID) {
 	r.Lock()
 	delete(r.sessions, id)
+	delete(r.sessionsByPGWire, id.Hi)
 	r.Unlock()
 }
 
 type registrySession interface {
 	user() string
 	cancelQuery(queryID ClusterWideID) bool
+	cancelCurrentQueries()
 	cancelSession()
 	// serialize serializes a Session into a serverpb.Session
 	// that can be served over RPC.
 	serialize() serverpb.Session
 }
 
+// CancelQueryByPGWire looks up the associated query in the session registry and cancels it.
+func (r *SessionRegistry) CancelQueryByPGWire(code, secret uint32) {
+	r.Lock()
+	defer r.Unlock()
+	if session, ok := r.sessionsByPGWire[(uint64(code)<<32)|uint64(secret)]; ok {
+		session.cancelCurrentQueries()
+	}
+}
+
 // CancelQuery looks up the associated query in the session registry and cancels it.
 func (r *SessionRegistry) CancelQuery(queryIDStr string, username string) (bool, error) {
 	queryID, err := StringToClusterWideID(queryIDStr)

pkg/sql/pgwire/conn.go

@@ -658,6 +658,18 @@ func (c *conn) sendInitialConnData(
 		return sql.ConnectionHandler{}, err
 	}
 
+	// Send our cancel code and secret to the client, so they can cancel the
+	// connection. The cancel code and secret are both int32's. The cancel code
+	// is supposed to be the connection's PID, but since we don't have one of
+	// those, we'll just generate a random int64 and use it instead.
+	code, secret := connHandler.GetSessionID().GetPGWireCancelInfo()
+	c.msgBuilder.initMsg(pgwirebase.ServerMsgBackendKeyData)
+	c.msgBuilder.putInt32(code)
+	c.msgBuilder.putInt32(secret)
+	if err := c.msgBuilder.finishMsg(c.conn); err != nil {
+		return sql.ConnectionHandler{}, err
+	}
+
 	// An initial readyForQuery message is part of the handshake.
 	c.msgBuilder.initMsg(pgwirebase.ServerMsgReady)
 	c.msgBuilder.writeByte(byte(sql.IdleTxnBlock))

pkg/sql/pgwire/pgwirebase/msg.go

@@ -37,6 +37,7 @@ const (
 	ClientMsgTerminate   ClientMessageType = 'X'
 
 	ServerMsgAuth                 ServerMessageType = 'R'
+	ServerMsgBackendKeyData       ServerMessageType = 'K'
 	ServerMsgBindComplete         ServerMessageType = '2'
 	ServerMsgCommandComplete      ServerMessageType = 'C'
 	ServerMsgCloseComplete        ServerMessageType = '3'

pkg/sql/pgwire/server.go

@@ -542,8 +542,15 @@ func (s *Server) ServeConn(ctx context.Context, conn net.Conn, socketType Socket
 		// If the client is really issuing a cancel request, close the door
 		// in their face (we don't support it yet). Make a note of that use
 		// in telemetry.
-		telemetry.Inc(sqltelemetry.CancelRequestCounter)
-		_ = conn.Close()
+		code, err := buf.GetUint32()
+		if err != nil {
+			return err
+		}
+		secret, err := buf.GetUint32()
+		if err != nil {
+			return err
+		}
+		s.execCfg.SessionRegistry.CancelQueryByPGWire(code, secret)
 		return nil
 
 	case version30: