Fix closing RouteGroup logic

skycoin · Nov 25, 2019 · 29768d1 · 29768d1
1 parent 7898035
commit 29768d1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/pkg/router/route_group.go b/pkg/router/route_group.go
@@ -124,6 +124,15 @@ func (r *RouteGroup) Read(p []byte) (n int, err error) {
 		return 0, nil
 	}
 
+	r.mu.Lock()
+	if r.readBuf.Len() > 0 {
+		data, err := r.readBuf.Read(p)
+		r.mu.Unlock()
+
+		return data, err
+	}
+	r.mu.Unlock()
+
 	data, ok := <-r.readCh
 	if !ok {
 		return 0, io.ErrClosedPipe
@@ -205,7 +214,7 @@ func (r *RouteGroup) Close() error {
 
 	r.once.Do(func() {
 		close(r.done)
-		// close(r.readCh) // TODO: fix panics and uncomment
+		close(r.readCh)
 	})
 
 	return nil

diff --git a/pkg/router/route_group_test.go b/pkg/router/route_group_test.go
@@ -386,12 +386,16 @@ func pushPackets(ctx context.Context, t *testing.T, from *transport.Manager, to
 		select {
 		case <-ctx.Done():
 			return
+		case <-to.done:
+			return
 		default:
 			packet, err := from.ReadPacket()
 			assert.NoError(t, err)
 			select {
 			case <-ctx.Done():
 				return
+			case <-to.done:
+				return
 			case to.readCh <- packet.Payload():
 			}
 		}