drpcmigrate: avoid potential double close of default listener

Change-Id: I700a43f306b000e1540558c0f302ca751b38ea05
storj · Sep 22, 2022 · 04b0ec3 · 04b0ec3
1 parent f85532f
commit 04b0ec3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 16 deletions.
diff --git a/drpcmigrate/mux.go b/drpcmigrate/mux.go
@@ -95,8 +95,10 @@ func (m *ListenMux) Run(ctx context.Context) error {
 	for _, lis := range m.routes {
 		<-lis.done
 	}
-	m.def.err = Closed
-	close(m.def.done)
+
+	_ = m.def.Close()
+	<-m.def.done
+
 	return m.err
 }
 

diff --git a/drpcmigrate/mux_test.go b/drpcmigrate/mux_test.go
@@ -94,28 +94,18 @@ func TestMuxLoopClose(t *testing.T) {
 
 	expectedErrs := make(chan error, 3)
 	muxErr := make(chan error, 1)
-	go func() {
-		expectedErrs <- run(lis1)
-	}()
-	go func() {
-		expectedErrs <- run(lis2)
+	go func() { expectedErrs <- run(lis1) }()
+	go func() { expectedErrs <- run(lis2) }()
+	go func() { expectedErrs <- run(mux.Default()) }()
 
-	}()
-	go func() {
-		expectedErrs <- run(mux.Default())
-	}()
-
-	go func() {
-		muxErr <- mux.Run(ctx)
-	}()
+	go func() { muxErr <- mux.Run(ctx) }()
 
 	for i := 0; i < 3; i++ {
 		select {
 		case <-processed:
 		case <-timeout.C:
 			t.Fatal("test is timed out")
 		}
-
 	}
 
 	// stopping the mux