From aaec9b0fe2d01bd3ff761345ba47bb94a9b0dc59 Mon Sep 17 00:00:00 2001
From: Erik Dubbelboer <erik@dubbelboer.com>
Date: Sat, 2 Feb 2019 23:40:05 +0100
Subject: [PATCH] Make InmemoryListener.Dial return when the connection is
 accepted

This makes InmemoryListener deterministic which makes our tests much
less flacky under high load or when GOMAXPROCS=1
---
 fasthttputil/inmemory_listener.go | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/fasthttputil/inmemory_listener.go b/fasthttputil/inmemory_listener.go
index 80aca3f..1b1a5f3 100644
--- a/fasthttputil/inmemory_listener.go
+++ b/fasthttputil/inmemory_listener.go
@@ -13,13 +13,18 @@ import (
 type InmemoryListener struct {
 	lock   sync.Mutex
 	closed bool
-	conns  chan net.Conn
+	conns  chan acceptConn
+}
+
+type acceptConn struct {
+	conn     net.Conn
+	accepted chan struct{}
 }
 
 // NewInmemoryListener returns new in-memory dialer<->net.Listener.
 func NewInmemoryListener() *InmemoryListener {
 	return &InmemoryListener{
-		conns: make(chan net.Conn, 1024),
+		conns: make(chan acceptConn, 1024),
 	}
 }
 
@@ -33,7 +38,8 @@ func (ln *InmemoryListener) Accept() (net.Conn, error) {
 	if !ok {
 		return nil, fmt.Errorf("InmemoryListener is already closed: use of closed network connection")
 	}
-	return c, nil
+	close(c.accepted)
+	return c.conn, nil
 }
 
 // Close implements net.Listener's Close.
@@ -59,8 +65,9 @@ func (ln *InmemoryListener) Addr() net.Addr {
 	}
 }
 
-// Dial creates new client<->server connection, enqueues server side
-// of the connection to Accept and returns client side of the connection.
+// Dial creates new client<->server connection.
+// Just like a real Dial it only returns once the server
+// has accepted the connection.
 //
 // It is safe calling Dial from concurrently running goroutines.
 func (ln *InmemoryListener) Dial() (net.Conn, error) {
@@ -68,8 +75,11 @@ func (ln *InmemoryListener) Dial() (net.Conn, error) {
 	cConn := pc.Conn1()
 	sConn := pc.Conn2()
 	ln.lock.Lock()
+	accepted := make(chan struct{})
 	if !ln.closed {
-		ln.conns <- sConn
+		ln.conns <- acceptConn{sConn, accepted}
+		// Wait until the connection has been accepted.
+		<-accepted
 	} else {
 		sConn.Close()
 		cConn.Close()