Merge pull request #760 from niaow/pg-local-cancel

Add support for local cancellation to postgres endpoint

pg/cancel.go

@@ -0,0 +1,137 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pg
+
+import (
+	"context"
+	"encoding/binary"
+	"io"
+	"sync"
+
+	"github.com/pkg/errors"
+)
+
+// ErrCancelledMissingConnection is an error triggered by cancelling a connection that does not exist.
+var ErrCancelledMissingConnection = errors.New("cancelled connection does not exist")
+
+// CancellationToken is a value used to identify a backend for cancellation.
+type CancellationToken struct {
+	PID, Key int32
+}
+
+// CancellationManager manages postgres connection cancellation.
+type CancellationManager interface {
+	// Token acquires a new cancellation token.
+	// The returned channel is sent to every time the connection is cancelled.
+	// The connection may be cancelled an unlimited number of times.
+	Token() (<-chan struct{}, context.CancelFunc, CancellationToken, error)
+
+	// Cancel sends a cancellation notification to the connection with the associated token.
+	// If the token is not associated with a connection, this returns ErrCancelledMissingConnection.
+	Cancel(CancellationToken) error
+}
+
+// NewLocalCancellationManager creates an in-memory CancellationManager using randomly generated tokens.
+// The provided reader is expected to be secure (e.g. crypto/rand.Reader).
+func NewLocalCancellationManager(rand io.Reader) CancellationManager {
+	return &localCancellationManager{
+		rand:        rand,
+		connections: make(map[CancellationToken]chan<- struct{}),
+	}
+}
+
+type localCancellationManager struct {
+	mu          sync.RWMutex
+	rand        io.Reader
+	connections map[CancellationToken]chan<- struct{}
+}
+
+func (c *localCancellationManager) Token() (<-chan struct{}, context.CancelFunc, CancellationToken, error) {
+	notify := make(chan struct{}, 1)
+
+gen:
+	token, err := c.generateToken()
+	if err != nil {
+		return nil, nil, CancellationToken{}, err
+	}
+	cancel := c.registerToken(token, notify)
+	if cancel == nil {
+		goto gen
+	}
+
+	return notify, cancel, token, nil
+}
+
+func (c *localCancellationManager) generateToken() (CancellationToken, error) {
+	var data [8]byte
+	for {
+		var n int
+		for n < 8 {
+			nn, err := c.rand.Read(data[n:])
+			if err != nil {
+				return CancellationToken{}, errors.Wrap(err, "generating a cancellation token")
+			}
+			n += nn
+		}
+
+		pid := int32(binary.LittleEndian.Uint32(data[:4]))
+		if pid < 0 {
+			continue
+		}
+		key := int32(binary.LittleEndian.Uint32(data[4:]))
+		if key < 0 {
+			continue
+		}
+
+		return CancellationToken{PID: pid, Key: key}, nil
+	}
+}
+
+func (c *localCancellationManager) registerToken(token CancellationToken, notify chan<- struct{}) context.CancelFunc {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if _, collision := c.connections[token]; collision {
+		return nil
+	}
+
+	c.connections[token] = notify
+
+	return func() {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+
+		delete(c.connections, token)
+
+		close(notify)
+	}
+}
+
+func (c *localCancellationManager) Cancel(token CancellationToken) error {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	ch := c.connections[token]
+	if ch == nil {
+		return ErrCancelledMissingConnection
+	}
+
+	select {
+	case ch <- struct{}{}:
+	default:
+	}
+
+	return nil
+}

pg/cancel_test.go

@@ -0,0 +1,78 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pg
+
+import (
+	"crypto/rand"
+	"testing"
+)
+
+func TestCancel(t *testing.T) {
+	mgr := NewLocalCancellationManager(rand.Reader)
+	notify, cancel, token, err := mgr.Token()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case <-notify:
+		t.Fatal("unexpected cancellation")
+	default:
+	}
+
+	err = mgr.Cancel(token)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case <-notify:
+	default:
+		t.Fatal("cancellation not propogated")
+	}
+
+	select {
+	case <-notify:
+		t.Fatal("unexpected cancellation")
+	default:
+	}
+
+	err = mgr.Cancel(CancellationToken{PID: -1, Key: -1})
+	if err == nil {
+		t.Fatal("invalid cancellation completed")
+	}
+
+	select {
+	case <-notify:
+		t.Fatal("unexpected cancellation")
+	default:
+	}
+
+	cancel()
+
+	select {
+	case _, ok := <-notify:
+		if ok {
+			t.Fatal("unexpected cancellation")
+		}
+	default:
+		t.Fatal("expected cancellation channel to be closed")
+	}
+
+	err = mgr.Cancel(token)
+	if err == nil {
+		t.Fatal("invalid cancellation completed")
+	}
+}

pg/message/message.go

@@ -49,6 +49,9 @@ const (
 
 	// TypeSimpleQuery is a simple query request.
 	TypeSimpleQuery Type = 'Q'
+
+	// TypeBackendKeyData contains a cancellation key for the client to use later.
+	TypeBackendKeyData Type = 'K'
 )
 
 // AuthenticationOK is a message indicating that authentication has completed.
@@ -342,3 +345,23 @@ func (e *Encoder) NegotiateProtocolVersion(maxMinor int32, unrecognizedOptions .
 		Data: e.buf.Bytes(),
 	}, nil
 }
+
+// BackendKeyData encodes a Message with a cancellation key.
+func (e *Encoder) BackendKeyData(pid, key int32) (Message, error) {
+	e.buf.Reset()
+
+	err := e.i32(pid)
+	if err != nil {
+		return Message{}, err
+	}
+
+	err = e.i32(key)
+	if err != nil {
+		return Message{}, err
+	}
+
+	return Message{
+		Type: TypeBackendKeyData,
+		Data: e.buf.Bytes(),
+	}, nil
+}

pg/protocol.go

@@ -139,13 +139,7 @@ startup:
 	proto := Protocol(binary.BigEndian.Uint32(data))
 	data = data[4:]
 
-	// Handle special protocols.
-	switch proto {
-	case ProtocolCancel:
-		// TODO: send an actual postgres error message.
-		return errors.New("cancellation protocol not yet supported")
-
-	case ProtocolSSL:
+	if proto == ProtocolSSL {
 		if s.TLSConfig != nil {
 			// Upgrade the connection to TLS and renegotiate on the tunneled connection.
 			_, err = conn.Write([]byte{'S'})
@@ -177,8 +171,14 @@ startup:
 		return errors.Errorf("client at %s attempted to initiate an unsecured postgres conenction", conn.RemoteAddr())
 	}
 
-	// Handle regular postgres.
-	return s.handleStandard(ctx, proto, conn, data)
+	switch proto {
+	case ProtocolCancel:
+		// Handle cancellation.
+		return s.handleCancel(ctx, conn, data)
+	default:
+		// Handle regular postgres.
+		return s.handleStandard(ctx, proto, conn, data)
+	}
 }
 
 // parseParams parses a parameter list from a startup packet.
@@ -207,6 +207,33 @@ func parseParams(data []byte) (map[string]string, error) {
 	}
 }
 
+// handleCancel handles cancel request connections.
+func (s *Server) handleCancel(ctx context.Context, conn net.Conn, data []byte) error {
+	if len(data) != 8 {
+		return errors.New("malformed cancellation packet")
+	}
+
+	if s.CancellationManager == nil {
+		return errors.New("cancellation is not configured")
+	}
+
+	pid := int32(binary.BigEndian.Uint32(data[:4]))
+	key := int32(binary.BigEndian.Uint32(data[4:]))
+
+	err := s.CancellationManager.Cancel(CancellationToken{PID: pid, Key: key})
+	switch err {
+	case nil:
+	case ErrCancelledMissingConnection:
+		// This is usually not a real error (race condition in the protocol).
+		// This can happen if a client cancels a request and shuts down.
+		s.Logger.Debugf("client at %v sent a mismatched cancellation token (is a load balancer misconfigured?)", conn.RemoteAddr())
+	default:
+		return err
+	}
+
+	return nil
+}
+
 // handleStandard handles a connection in the standard postgres wire protocol.
 // The client is responsible for closing the connection when this finishes.
 func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Conn, data []byte) error {
@@ -296,6 +323,25 @@ func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Co
 		return errors.Wrap(err, "sending authentication confirmation")
 	}
 
+	var cancelNotify <-chan struct{}
+	if s.CancellationManager != nil {
+		notify, cancel, token, err := s.CancellationManager.Token()
+		if err != nil {
+			return errors.Wrap(err, "setting up cancellation")
+		}
+		defer cancel()
+
+		msg, err := encoder.BackendKeyData(token.PID, token.Key)
+		if err != nil {
+			return errors.Wrap(err, "encoding cancellation key data")
+		}
+		err = w.WriteMessage(msg)
+		if err != nil {
+			return errors.Wrap(err, "sending cancellation key data")
+		}
+		cancelNotify = notify
+	}
+
 	var queryReady bool
 	for {
 		if !queryReady {
@@ -368,54 +414,12 @@ func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Co
 			// Parse the query message (a null-terminated string).
 			query := SimpleQuery(strings.TrimSuffix(string(msg.Data), "\x00"))
 
-			// Set up a result writer.
-			// SELECT is used as a default tag, which seems to be handled decently by clients.
-			// The encoder is intentionally not used because its buffer may be huge.
-			qwriter := &queryResultWriter{
-				w:   w,
-				te:  s.TypeEngine,
-				tag: "SELECT",
+			// Execute the query.
+			err := s.handleQuery(w, query, cancelNotify)
+			if err != nil {
+				return err
 			}
 
-			// Dispatch the query handler.
-			// TODO: cancellation (requires crazy internode logic and fake process IDs)
-			// This is not the connection context, since we want the request to finish safely before connection shutdown.
-			qerr := s.QueryHandler.HandleQuery(context.Background(), qwriter, query)
-			if qerr != nil {
-				// There was an error in processing the query.
-				// Send the error back to the client and keep going.
-				s.Logger.Debugf("failed to execute query %q: %v", query, qerr)
-				msg, err = encoder.GoError(qerr)
-				if err != nil {
-					return errors.Wrap(err, "failed to send query error to client")
-				}
-				err = w.WriteMessage(msg)
-				if err != nil {
-					return errors.Wrap(err, "failed to send query error to client")
-				}
-			} else {
-				if !qwriter.wroteHeaders {
-					// The handler did not write headers.
-					// Write back an empty set of headers.
-					err = qwriter.WriteHeader()
-					if err != nil {
-						return errors.Wrap(err, "sending empty column headers")
-					}
-				}
-				// The query completed normally.
-				// Notify the client of completion.
-				msg, err = encoder.CommandComplete(qwriter.tag)
-				if err != nil {
-					return errors.Wrap(err, "sending command completion notification")
-				}
-				err = w.WriteMessage(msg)
-				if err != nil {
-					return errors.Wrap(err, "sending command completion notification")
-				}
-			}
-
-			// The data will be flushed after we write back the "ready for query" state.
-
 		default:
 			// The message is not supported yet.
 			// Send an error.
@@ -449,6 +453,90 @@ func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Co
 	}
 }
 
+// handleQuery processes a single query on a connection.
+func (s *Server) handleQuery(w message.Writer, query Query, cancelNotify <-chan struct{}) error {
+	// Configure cancellation.
+	// This is not the connection context, since we want the request to finish safely before connection shutdown.
+	ctx := context.Background()
+	if cancelNotify != nil {
+		defer func() {
+			// Flush any cancel notifications.
+			// This works on a best-effort basis.
+			// It is still entirely possible that the cancel notification may be delivered to the next request.
+			// Regardless of what we do, we either get false positives or false negatives.
+			// This code chooses false positives.
+			for len(cancelNotify) > 0 {
+				<-cancelNotify
+			}
+		}()
+
+		var wg sync.WaitGroup
+		defer wg.Add(1)
+
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
+		defer cancel()
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			select {
+			case <-ctx.Done():
+			case <-cancelNotify:
+				cancel()
+			}
+		}()
+	}
+
+	// Set up a result writer.
+	// SELECT is used as a default tag, which seems to be handled decently by clients.
+	// The encoder is intentionally not re-used because its buffer may be huge.
+	qwriter := &queryResultWriter{
+		w:   w,
+		te:  s.TypeEngine,
+		tag: "SELECT",
+	}
+
+	// Dispatch the query handler.
+	qerr := s.QueryHandler.HandleQuery(ctx, qwriter, query)
+	if qerr != nil {
+		// There was an error in processing the query.
+		// Send the error back to the client and keep going.
+		s.Logger.Debugf("failed to execute query %q: %v", query, qerr)
+		msg, err := qwriter.enc.GoError(qerr)
+		if err != nil {
+			return errors.Wrap(err, "failed to send query error to client")
+		}
+		err = w.WriteMessage(msg)
+		if err != nil {
+			return errors.Wrap(err, "failed to send query error to client")
+		}
+	} else {
+		if !qwriter.wroteHeaders {
+			// The handler did not write headers.
+			// Write back an empty set of headers.
+			err := qwriter.WriteHeader()
+			if err != nil {
+				return errors.Wrap(err, "sending empty column headers")
+			}
+		}
+		// The query completed normally.
+		// Notify the client of completion.
+		msg, err := qwriter.enc.CommandComplete(qwriter.tag)
+		if err != nil {
+			return errors.Wrap(err, "sending command completion notification")
+		}
+		err = w.WriteMessage(msg)
+		if err != nil {
+			return errors.Wrap(err, "sending command completion notification")
+		}
+	}
+
+	// The data will be flushed after we write back the "ready for query" state.
+	return nil
+}
+
 func (s *Server) handleShutdown(conn net.Conn, w message.Writer, encoder *message.Encoder, notice ...message.NoticeField) error {
 	var wg sync.WaitGroup
 	defer wg.Wait()

pg/server.go

@@ -55,6 +55,10 @@ type Server struct {
 
 	// Logger is the logger to use for error conditions and state changes.
 	Logger logger.Logger
+
+	// CancellationManager is the cancellation manager to use.
+	// If this is not set, no cancellations will be applied.
+	CancellationManager CancellationManager
 }
 
 // ServeConn serves a single connection.

pg/server_test.go

@@ -16,10 +16,12 @@ package pg_test
 
 import (
 	"context"
+	"crypto/rand"
 	"fmt"
 	"net"
 	"os/exec"
 	"strconv"
+	"syscall"
 	"testing"
 	"time"
 
@@ -173,46 +175,100 @@ func TestPSQLQuery(t *testing.T) {
 		t.Fatalf("searching for psql: %v", err)
 	}
 
-	server := &pg.Server{
-		QueryHandler: pgtest.HandlerFunc(func(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
-			err := w.WriteHeader(pg.ColumnInfo{
-				Name: "field",
-				Type: pg.TypeCharoid,
-			})
-			if err != nil {
-				return err
+	t.Run("Query", func(t *testing.T) {
+		server := &pg.Server{
+			QueryHandler: pgtest.HandlerFunc(func(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
+				err := w.WriteHeader(pg.ColumnInfo{
+					Name: "field",
+					Type: pg.TypeCharoid,
+				})
+				if err != nil {
+					return err
+				}
+
+				err = w.WriteRowText("h")
+				if err != nil {
+					return err
+				}
+
+				err = w.WriteRowText("xyzzy")
+				if err != nil {
+					return err
+				}
+
+				return nil
+			}),
+			TypeEngine:     pg.PrimitiveTypeEngine{},
+			StartupTimeout: time.Second,
+			Logger:         logger.NopLogger,
+		}
+		addr, shutdown, err := pgtest.ServeTCP(":0", server)
+		if err != nil {
+			t.Fatalf("starting postgres server: %v", err)
+		}
+		defer shutdown.Finish(t, "postgres server")
+
+		tcpAddr := addr.(*net.TCPAddr)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		cmd := exec.CommandContext(ctx, "psql", "-h", tcpAddr.IP.String(), "-p", strconv.Itoa(tcpAddr.Port), "-c", "test query")
+		data, err := cmd.CombinedOutput()
+		if err != nil {
+			t.Fatalf("psql failed: %v", string(data))
+		}
+	})
+
+	t.Run("Cancel", func(t *testing.T) {
+		var term func() error
+		var qerr error
+		var qdone bool
+		defer func() {
+			if qerr != nil {
+				t.Fatal(qerr)
 			}
-
-			err = w.WriteRowText("h")
-			if err != nil {
-				return err
+			if !qdone {
+				t.Fatal("query not done")
 			}
+		}()
 
-			err = w.WriteRowText("xyzzy")
-			if err != nil {
-				return err
-			}
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
 
-			return nil
-		}),
-		TypeEngine:     pg.PrimitiveTypeEngine{},
-		StartupTimeout: time.Second,
-		Logger:         logger.NopLogger,
-	}
-	addr, shutdown, err := pgtest.ServeTCP(":0", server)
-	if err != nil {
-		t.Fatalf("starting postgres server: %v", err)
-	}
-	defer shutdown.Finish(t, "postgres server")
+		server := &pg.Server{
+			QueryHandler: pgtest.HandlerFunc(func(qctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
+				defer func() { qdone = true }()
+				qerr = term()
+				if qerr != nil {
+					return qerr
+				}
+				select {
+				case <-qctx.Done():
+				case <-ctx.Done():
+					qerr = ctx.Err()
+					return qerr
+				}
+				return nil
+			}),
+			TypeEngine:          pg.PrimitiveTypeEngine{},
+			StartupTimeout:      time.Second,
+			Logger:              logger.NopLogger,
+			CancellationManager: pg.NewLocalCancellationManager(rand.Reader),
+		}
+		addr, shutdown, err := pgtest.ServeTCP(":0", server)
+		if err != nil {
+			t.Fatalf("starting postgres server: %v", err)
+		}
+		defer shutdown.Finish(t, "postgres server")
 
-	tcpAddr := addr.(*net.TCPAddr)
+		tcpAddr := addr.(*net.TCPAddr)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-
-	cmd := exec.CommandContext(ctx, "psql", "-h", tcpAddr.IP.String(), "-p", strconv.Itoa(tcpAddr.Port), "-c", "test query")
-	data, err := cmd.CombinedOutput()
-	if err != nil {
-		t.Fatalf("psql failed: %v", string(data))
-	}
+		cmd := exec.CommandContext(ctx, "psql", "-h", tcpAddr.IP.String(), "-p", strconv.Itoa(tcpAddr.Port), "-c", "test query")
+		term = func() error { return cmd.Process.Signal(syscall.SIGINT) }
+		data, err := cmd.CombinedOutput()
+		if err != nil {
+			t.Fatalf("psql failed: %v", string(data))
+		}
+	})
 }

server/pg.go

@@ -16,6 +16,7 @@ package server
 
 import (
 	"context"
+	"crypto/rand"
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
@@ -57,6 +58,9 @@ func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) *
 			MaxStartupSize: 8 * 1024 * 1024,
 			Logger:         logger,
 			TLSConfig:      tls,
+
+			// This is somewhat limited right now: it does not work with load balancers.
+			CancellationManager: pg.NewLocalCancellationManager(rand.Reader),
 		},
 	}
 }