Restructure sending messages

Use an internal buffer in pgproto3.Frontend and pgproto3.Backend instead
of directly writing to the underlying net.Conn. This will allow tracing
messages as well as simplify pipeline mode.

pgproto3/backend.go

@@ -11,6 +11,8 @@ type Backend struct {
 	cr *chunkReader
 	w  io.Writer
 
+	wbuf []byte
+
 	// Frontend message flyweights
 	bind           Bind
 	cancelRequest  CancelRequest
@@ -47,10 +49,28 @@ func NewBackend(r io.Reader, w io.Writer) *Backend {
 	return &Backend{cr: cr, w: w}
 }
 
-// Send sends a message to the frontend.
-func (b *Backend) Send(msg BackendMessage) error {
-	_, err := b.w.Write(msg.Encode(nil))
-	return err
+// Send sends a message to the frontend (i.e. the client). The message is not guaranteed to be written until Flush is
+// called.
+func (b *Backend) Send(msg BackendMessage) {
+	b.wbuf = msg.Encode(b.wbuf)
+}
+
+// Flush writes any pending messages to the frontend (i.e. the client).
+func (b *Backend) Flush() error {
+	n, err := b.w.Write(b.wbuf)
+
+	const maxLen = 1024
+	if len(b.wbuf) > maxLen {
+		b.wbuf = make([]byte, 0, maxLen)
+	} else {
+		b.wbuf = b.wbuf[:0]
+	}
+
+	if err != nil {
+		return &writeError{err: err, safeToRetry: n == 0}
+	}
+
+	return nil
 }
 
 // ReceiveStartupMessage receives the initial connection message. This method is used of the normal Receive method

pgproto3/frontend.go

@@ -12,6 +12,8 @@ type Frontend struct {
 	cr *chunkReader
 	w  io.Writer
 
+	wbuf []byte
+
 	// Backend message flyweights
 	authenticationOk                AuthenticationOk
 	authenticationCleartextPassword AuthenticationCleartextPassword
@@ -56,10 +58,28 @@ func NewFrontend(r io.Reader, w io.Writer) *Frontend {
 	return &Frontend{cr: cr, w: w}
 }
 
-// Send sends a message to the backend.
-func (f *Frontend) Send(msg FrontendMessage) error {
-	_, err := f.w.Write(msg.Encode(nil))
-	return err
+// Send sends a message to the backend (i.e. the server). The message is not guaranteed to be written until Flush is
+// called.
+func (f *Frontend) Send(msg FrontendMessage) {
+	f.wbuf = msg.Encode(f.wbuf)
+}
+
+// Flush writes any pending messages to the backend (i.e. the server).
+func (f *Frontend) Flush() error {
+	n, err := f.w.Write(f.wbuf)
+
+	const maxLen = 1024
+	if len(f.wbuf) > maxLen {
+		f.wbuf = make([]byte, 0, maxLen)
+	} else {
+		f.wbuf = f.wbuf[:0]
+	}
+
+	if err != nil {
+		return &writeError{err: err, safeToRetry: n == 0}
+	}
+
+	return nil
 }
 
 func translateEOFtoErrUnexpectedEOF(err error) error {

pgproto3/pgproto3.go

@@ -17,11 +17,13 @@ type Message interface {
 	Encode(dst []byte) []byte
 }
 
+// FrontendMessage is a message sent by the frontend (i.e. the client).
 type FrontendMessage interface {
 	Message
 	Frontend() // no-op method to distinguish frontend from backend methods
 }
 
+// BackendMessage is a message sent by the backend (i.e. the server).
 type BackendMessage interface {
 	Message
 	Backend() // no-op method to distinguish frontend from backend methods
@@ -50,6 +52,23 @@ func (e *invalidMessageFormatErr) Error() string {
 	return fmt.Sprintf("%s body is invalid", e.messageType)
 }
 
+type writeError struct {
+	err         error
+	safeToRetry bool
+}
+
+func (e *writeError) Error() string {
+	return fmt.Sprintf("write failed: %s", e.err.Error())
+}
+
+func (e *writeError) SafeToRetry() bool {
+	return e.safeToRetry
+}
+
+func (e *writeError) Unwrap() error {
+	return e.err
+}
+
 // getValueFromJSON gets the value from a protocol message representation in JSON.
 func getValueFromJSON(v map[string]string) ([]byte, error) {
 	if v == nil {