[tor-commits] [snowflake/master] Wait until all goroutines finish before shutdown

[cohosh]

commit effc6675448a3a6e62d2784557d67e2c46e376d5
Author: Cecylia Bocovich <coh...@torproject.org>
Date:   Fri Dec 4 10:50:00 2020 -0500

    Wait until all goroutines finish before shutdown
---
 client/snowflake.go | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/client/snowflake.go b/client/snowflake.go
index a1a679e..55addc1 100644
--- a/client/snowflake.go
+++ b/client/snowflake.go
@@ -12,6 +12,7 @@ import (
        "os/signal"
        "path/filepath"
        "strings"
+       "sync"
        "syscall"
        "time"
 
@@ -27,7 +28,7 @@ const (
 )
 
 // Accept local SOCKS connections and pass them to the handler.
-func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan 
struct{}) {
+func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan 
struct{}, wg sync.WaitGroup) {
        defer ln.Close()
        for {
                conn, err := ln.AcceptSocks()
@@ -40,6 +41,8 @@ func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, 
shutdown chan struc
                }
                log.Printf("SOCKS accepted: %v", conn.Req)
                go func() {
+                       wg.Add(1)
+                       defer wg.Done()
                        defer conn.Close()
 
                        err := conn.Grant(&net.TCPAddr{IP: net.IPv4zero, Port: 
0})
@@ -173,6 +176,7 @@ func main() {
        }
        listeners := make([]net.Listener, 0)
        shutdown := make(chan struct{})
+       var wg sync.WaitGroup
        for _, methodName := range ptInfo.MethodNames {
                switch methodName {
                case "snowflake":
@@ -183,7 +187,7 @@ func main() {
                                break
                        }
                        log.Printf("Started SOCKS listener at %v.", ln.Addr())
-                       go socksAcceptLoop(ln, dialer, shutdown)
+                       go socksAcceptLoop(ln, dialer, shutdown, wg)
                        pt.Cmethod(methodName, ln.Version(), ln.Addr())
                        listeners = append(listeners, ln)
                default:
@@ -216,6 +220,7 @@ func main() {
                ln.Close()
        }
        close(shutdown)
+       wg.Wait()
        log.Println("snowflake is done.")
 }
 

_______________________________________________
tor-commits mailing list
tor-commits@lists.torproject.org
https://lists.torproject.org/cgi-bin/mailman/listinfo/tor-commits