diff --git a/pg/cancel.go b/pg/cancel.go new file mode 100644 index 000000000..7e2073df3 --- /dev/null +++ b/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 +} diff --git a/pg/cancel_test.go b/pg/cancel_test.go new file mode 100644 index 000000000..7512109ac --- /dev/null +++ b/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") + } +} diff --git a/pg/message/message.go b/pg/message/message.go index b7de6683f..caf8214d3 100644 --- a/pg/message/message.go +++ b/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 +} diff --git a/pg/protocol.go b/pg/protocol.go index 8a50baf05..55ed8b750 100644 --- a/pg/protocol.go +++ b/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() diff --git a/pg/server.go b/pg/server.go index 55d374005..9cde6b374 100644 --- a/pg/server.go +++ b/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. diff --git a/pg/server_test.go b/pg/server_test.go index c2f2e897e..af8537ce8 100644 --- a/pg/server_test.go +++ b/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)) + } + }) } diff --git a/server/pg.go b/server/pg.go index d29168ea0..43fb22454 100644 --- a/server/pg.go +++ b/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), }, } }