i hope the programming gods can forgive me after that one

lib/gat/gatling/client/client.go

@@ -22,13 +22,14 @@ import (
 	"net"
 	"reflect"
 	"strings"
+	"sync"
 )
 
 // / client state, one per client
 type Client struct {
 	conn net.Conn
 	r    *bufio.Reader
-	wr   io.Writer
+	wr   *bufio.Writer
 
 	recv chan protocol.Packet
 
@@ -62,6 +63,8 @@ type Client struct {
 	state       rune
 
 	log zlog.Logger
+
+	mu sync.Mutex
 }
 
 func NewClient(
@@ -76,7 +79,7 @@ func NewClient(
 	c := &Client{
 		conn:       conn,
 		r:          bufio.NewReader(conn),
-		wr:         conn,
+		wr:         bufio.NewWriter(conn),
 		recv:       make(chan protocol.Packet),
 		addr:       conn.RemoteAddr(),
 		pid:        int32(pid.Int64()),
@@ -129,6 +132,10 @@ func (c *Client) Accept(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
+			err = c.wr.Flush()
+			if err != nil {
+				return err
+			}
 			startup = new(protocol.StartupMessage)
 			err = startup.Read(c.r)
 			if err != nil {
@@ -139,6 +146,10 @@ func (c *Client) Accept(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
+			err = c.wr.Flush()
+			if err != nil {
+				return err
+			}
 			//TODO: we need to do an ssl handshake here.
 			var cert tls.Certificate
 			cert, err = tls.LoadX509KeyPair(c.conf.General.TlsCertificate, c.conf.General.TlsPrivateKey)
@@ -151,7 +162,7 @@ func (c *Client) Accept(ctx context.Context) error {
 			}
 			c.conn = tls.Server(c.conn, cfg)
 			c.r = bufio.NewReader(c.conn)
-			c.wr = c.conn
+			c.wr = bufio.NewWriter(c.conn)
 			err = startup.Read(c.r)
 			if err != nil {
 				return err
@@ -200,7 +211,11 @@ func (c *Client) Accept(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	_, err = pkt.Write(c.wr)
+	err = c.Send(pkt)
+	if err != nil {
+		return err
+	}
+	err = c.Flush()
 	if err != nil {
 		return err
 	}
@@ -263,7 +278,7 @@ func (c *Client) Accept(ctx context.Context) error {
 
 	authOk := new(protocol.Authentication)
 	authOk.Fields.Code = 0
-	_, err = authOk.Write(c.wr)
+	err = c.Send(authOk)
 	if err != nil {
 		return err
 	}
@@ -271,7 +286,7 @@ func (c *Client) Accept(ctx context.Context) error {
 	//
 	info := c.server.GetServerInfo()
 	for _, inf := range info {
-		_, err = inf.Write(c.wr)
+		err = c.Send(inf)
 		if err != nil {
 			return err
 		}
@@ -279,19 +294,23 @@ func (c *Client) Accept(ctx context.Context) error {
 	backendKeyData := new(protocol.BackendKeyData)
 	backendKeyData.Fields.ProcessID = c.pid
 	backendKeyData.Fields.SecretKey = c.secret_key
-	_, err = backendKeyData.Write(c.wr)
+	err = c.Send(backendKeyData)
 	if err != nil {
 		return err
 	}
 	readyForQuery := new(protocol.ReadyForQuery)
 	readyForQuery.Fields.Status = byte('I')
-	_, err = readyForQuery.Write(c.wr)
+	err = c.Send(readyForQuery)
 	if err != nil {
 		return err
 	}
 	go c.recvLoop()
 	open := true
 	for open {
+		err = c.Flush()
+		if err != nil {
+			return err
+		}
 		open, err = c.tick(ctx)
 		if !open {
 			break
@@ -348,7 +367,7 @@ func (c *Client) handle_cancel(ctx context.Context, p *protocol.StartupMessage)
 func (c *Client) tick(ctx context.Context) (bool, error) {
 	var rsp protocol.Packet
 	select {
-	case rsp = <-c.Recv():
+	case rsp = <-c.recv:
 	case <-ctx.Done():
 		return false, ctx.Err()
 	}
@@ -492,11 +511,19 @@ func (c *Client) GetPortal(name string) *protocol.Bind {
 }
 
 func (c *Client) Send(pkt protocol.Packet) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	//log.Printf("sent packet(%s) %+v", reflect.TypeOf(pkt), pkt)
 	_, err := pkt.Write(c.wr)
 	return err
 }
 
+func (c *Client) Flush() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.wr.Flush()
+}
+
 func (c *Client) Recv() <-chan protocol.Packet {
 	return c.recv
 }

lib/gat/gatling/conn_pool/server/server.go

@@ -3,7 +3,6 @@ package server
 import (
 	"bufio"
 	"fmt"
-	"io"
 	"net"
 	"reflect"
 	"time"
@@ -28,7 +27,7 @@ type Server struct {
 	remote net.Addr
 	conn   net.Conn
 	r      *bufio.Reader
-	wr     io.Writer
+	wr     *bufio.Writer
 
 	server_info []*protocol.ParameterStatus
 
@@ -72,13 +71,9 @@ func Dial(ctx context.Context,
 	if err != nil {
 		return nil, err
 	}
-	err = s.conn.(*net.TCPConn).SetNoDelay(false)
-	if err != nil {
-		return nil, err
-	}
 	s.remote = s.conn.RemoteAddr()
 	s.r = bufio.NewReader(s.conn)
-	s.wr = s.conn
+	s.wr = bufio.NewWriter(s.conn)
 	s.user = *user
 	s.db = db
 
@@ -127,7 +122,7 @@ func (s *Server) startup(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	return nil
+	return s.flush()
 }
 
 func (s *Server) connect(ctx context.Context) error {
@@ -164,17 +159,20 @@ func (s *Server) connect(ctx context.Context) error {
 				if err != nil {
 					return err
 				}
-				func() {
+
 				rsp := new(protocol.AuthenticationResponse)
 				buf := bufpool.Get(len(scrm.Name()) + 1 + 4 + len(bts))
 				buf.Reset()
-					defer bufpool.Put(buf)
 				_, _ = protocol.WriteString(buf, scrm.Name())
 				_, _ = protocol.WriteInt32(buf, int32(len(bts)))
 				buf.Write(bts)
 				rsp.Fields.Data = buf.Bytes()
 				err = s.writePacket(rsp)
-				}()
+				bufpool.Put(buf)
+				if err != nil {
+					return err
+				}
+				err = s.flush()
 				if err != nil {
 					return err
 				}
@@ -189,6 +187,10 @@ func (s *Server) connect(ctx context.Context) error {
 				if err != nil {
 					return err
 				}
+				err = s.flush()
+				if err != nil {
+					return err
+				}
 			case 12: // SASL_FINAL
 				s.log.Debug().Str("method", "scram256").Msg("sasl final")
 				var done bool
@@ -244,6 +246,10 @@ func (s *Server) writePacket(pkt protocol.Packet) error {
 	return err
 }
 
+func (s *Server) flush() error {
+	return s.wr.Flush()
+}
+
 func (s *Server) readPacket() (protocol.Packet, error) {
 	return protocol.ReadBackend(s.r)
 }
@@ -311,6 +317,7 @@ func (s *Server) destructPreparedStatement(name string) {
 	query := new(protocol.Query)
 	query.Fields.Query = fmt.Sprintf("DEALLOCATE \"%s\"", name)
 	_ = s.writePacket(query)
+	_ = s.flush()
 	// await server ready
 	for {
 		r, _ := s.readPacket()
@@ -358,6 +365,10 @@ func (s *Server) Describe(client gat.Client, d *protocol.Describe) error {
 	if err != nil {
 		return err
 	}
+	err = s.flush()
+	if err != nil {
+		return err
+	}
 
 	return s.forwardTo(client, func(pkt protocol.Packet) (forward bool, finish bool, err error) {
 		//log.Println("forward packet(%s) %+v", reflect.TypeOf(pkt), pkt)
@@ -386,6 +397,10 @@ func (s *Server) Execute(client gat.Client, e *protocol.Execute) error {
 	if err != nil {
 		return err
 	}
+	err = s.flush()
+	if err != nil {
+		return err
+	}
 
 	return s.forwardTo(client, func(pkt protocol.Packet) (forward bool, finish bool, err error) {
 		//log.Println("forward packet(%s) %+v", reflect.TypeOf(pkt), pkt)
@@ -408,6 +423,10 @@ func (s *Server) SimpleQuery(ctx context.Context, client gat.Client, query strin
 	if err != nil {
 		return err
 	}
+	err = s.flush()
+	if err != nil {
+		return err
+	}
 
 	// this function seems wild but it has to be the way it is so we read the whole response, even if the
 	// client fails midway
@@ -435,6 +454,10 @@ func (s *Server) Transaction(ctx context.Context, client gat.Client, query strin
 	if err != nil {
 		return err
 	}
+	err = s.flush()
+	if err != nil {
+		return err
+	}
 	return s.forwardTo(client, func(pkt protocol.Packet) (forward bool, finish bool, err error) {
 		//log.Printf("got server pkt pkt(%s): %+v ", reflect.TypeOf(pkt), pkt)
 		switch p := pkt.(type) {
@@ -444,6 +467,8 @@ func (s *Server) Transaction(ctx context.Context, client gat.Client, query strin
 				// send to client and wait for next query
 				err = client.Send(pkt)
 
+				if err == nil {
+					err = client.Flush()
 					if err == nil {
 						select {
 						case r := <-client.Recv():
@@ -452,6 +477,7 @@ func (s *Server) Transaction(ctx context.Context, client gat.Client, query strin
 							case *protocol.Query:
 								//forward to server
 								_ = s.writePacket(r)
+								_ = s.flush()
 							default:
 								err = fmt.Errorf("expected an error in transaction state but got something else")
 							}
@@ -459,11 +485,13 @@ func (s *Server) Transaction(ctx context.Context, client gat.Client, query strin
 							err = ctx.Err()
 						}
 					}
+				}
 
 				if err != nil {
 					end := new(protocol.Query)
 					end.Fields.Query = "END;"
 					_ = s.writePacket(end)
+					_ = s.flush()
 				}
 			} else {
 				finish = true
@@ -479,6 +507,10 @@ func (s *Server) Transaction(ctx context.Context, client gat.Client, query strin
 }
 
 func (s *Server) CopyIn(ctx context.Context, client gat.Client) error {
+	err := client.Flush()
+	if err != nil {
+		return err
+	}
 	for {
 		// detect a disconneted /hanging client by waiting 30 seoncds, else timeout
 		// otherwise, just keep reading packets until a done or error is received
@@ -489,10 +521,11 @@ func (s *Server) CopyIn(ctx context.Context, client gat.Client) error {
 		case pkt = <-client.Recv():
 		case <-cctx.Done():
 			_ = s.writePacket(new(protocol.CopyFail))
+			_ = s.flush()
 			return cctx.Err()
 		}
 		cancel()
-		err := s.writePacket(pkt)
+		err = s.writePacket(pkt)
 		if err != nil {
 			return err
 		}
@@ -500,7 +533,7 @@ func (s *Server) CopyIn(ctx context.Context, client gat.Client) error {
 		switch pkt.(type) {
 		case *protocol.CopyDone, *protocol.CopyFail:
 			// don't error on copyfail because the client is the one that errored, it already knows
-			return nil
+			return s.flush()
 		}
 	}
 }
@@ -510,6 +543,10 @@ func (s *Server) CallFunction(client gat.Client, payload *protocol.FunctionCall)
 	if err != nil {
 		return err
 	}
+	err = s.flush()
+	if err != nil {
+		return err
+	}
 	// read responses
 	return s.forwardTo(client, func(pkt protocol.Packet) (forward bool, finish bool, err error) {
 		switch pkt.(type) {

lib/gat/gatling/gatling.go

@@ -134,11 +134,6 @@ func (g *Gatling) ListenAndServe(ctx context.Context) error {
 
 // TODO: TLS
 func (g *Gatling) handleConnection(ctx context.Context, c net.Conn) error {
-	err := c.(*net.TCPConn).SetNoDelay(false)
-	if err != nil {
-		return err
-	}
-
 	cl := client.NewClient(g, g.c, c, false)
 
 	func() {
@@ -152,10 +147,11 @@ func (g *Gatling) handleConnection(ctx context.Context, c net.Conn) error {
 		delete(g.clients, cl.Id())
 	}()
 
-	err = cl.Accept(ctx)
+	err := cl.Accept(ctx)
 	if err != nil {
 		log.Println("err in connection:", err.Error())
 		_ = cl.Send(pg_error.IntoPacket(err))
+		_ = cl.Flush()
 	}
 	_ = c.Close()
 	return nil

lib/gat/interfaces.go

@@ -19,6 +19,7 @@ type Client interface {
 	SetCurrentConn(conn Connection)
 
 	Send(pkt protocol.Packet) error
+	Flush() error
 	Recv() <-chan protocol.Packet
 }