From 835738495b0fd828204eb6251a19f7dc01a5de59 Mon Sep 17 00:00:00 2001
From: Michael Nikitochkin <michael.nikitochkin@shopify.com>
Date: Fri, 17 Sep 2021 14:19:37 +0200
Subject: [PATCH] Extract proxy.go chagnes to #322

---
 proxy.go | 85 +++++++++++++++++++++++---------------------------------
 1 file changed, 34 insertions(+), 51 deletions(-)

diff --git a/proxy.go b/proxy.go
index f3c83b2ab..a32f50e52 100644
--- a/proxy.go
+++ b/proxy.go
@@ -24,8 +24,7 @@ type Proxy struct {
 	Upstream string `json:"upstream"`
 	Enabled  bool   `json:"enabled"`
 
-	listener net.Listener
-	started  chan error
+	started chan error
 
 	tomb        tomb.Tomb
 	connections ConnectionList
@@ -89,64 +88,48 @@ func (proxy *Proxy) Stop() {
 	stop(proxy)
 }
 
-func (p *Proxy) listen() (net.Listener, error) {
-	listener, err := net.Listen("tcp", p.Listen)
-	if err != nil {
-		p.started <- err
-		return nil, err
-	}
-	p.listener = listener
-	p.Listen = listener.Addr().String()
-	p.started <- nil
-
-	logrus.WithFields(logrus.Fields{
-		"name":     p.Name,
-		"proxy":    p.Listen,
-		"upstream": p.Upstream,
-	}).Info("Started proxy")
-
-	return p.listener, nil
-}
-
-func (p *Proxy) close() {
-	// Unblock ln.Accept()
-	err := p.listener.Close()
-	if err != nil {
-		logrus.WithFields(logrus.Fields{
-			"proxy":  p.Name,
-			"listen": p.Listen,
-			"err":    err,
-		}).Warn("Attempted to close an already closed proxy server")
-	}
-}
-
-// This channel is to kill the blocking Accept() call below by closing the
-// net.Listener.
-func (p *Proxy) freeBlocker(acceptTomb *tomb.Tomb) {
-	<-p.tomb.Dying()
-
-	// Notify ln.Accept() that the shutdown was safe
-	acceptTomb.Killf("Shutting down from stop()")
-
-	p.close()
-
-	// Wait for the accept loop to finish processing
-	acceptTomb.Wait()
-	p.tomb.Done()
-}
-
 // server runs the Proxy server, accepting new clients and creating Links to
 // connect them to upstreams.
 func (proxy *Proxy) server() {
-	ln, err := proxy.listen()
+	ln, err := net.Listen("tcp", proxy.Listen)
 	if err != nil {
+		proxy.started <- err
 		return
 	}
 
-	acceptTomb := &tomb.Tomb{}
+	proxy.Listen = ln.Addr().String()
+	proxy.started <- nil
+
+	logrus.WithFields(logrus.Fields{
+		"name":     proxy.Name,
+		"proxy":    proxy.Listen,
+		"upstream": proxy.Upstream,
+	}).Info("Started proxy")
+
+	acceptTomb := tomb.Tomb{}
 	defer acceptTomb.Done()
 
-	go proxy.freeBlocker(acceptTomb)
+	// This channel is to kill the blocking Accept() call below by closing the
+	// net.Listener.
+	go func() {
+		<-proxy.tomb.Dying()
+
+		// Notify ln.Accept() that the shutdown was safe
+		acceptTomb.Killf("Shutting down from stop()")
+		// Unblock ln.Accept()
+		err := ln.Close()
+		if err != nil {
+			logrus.WithFields(logrus.Fields{
+				"proxy":  proxy.Name,
+				"listen": proxy.Listen,
+				"err":    err,
+			}).Warn("Attempted to close an already closed proxy server")
+		}
+
+		// Wait for the accept loop to finish processing
+		acceptTomb.Wait()
+		proxy.tomb.Done()
+	}()
 
 	for {
 		client, err := ln.Accept()