Fix router tests

pkg/router/route_group.go

@@ -136,7 +136,6 @@ func (rg *RouteGroup) Read(p []byte) (n int, err error) {
 	}
 
 	if rg.readDeadline.Closed() {
-		rg.logger.Infoln("TIMEOUT ERROR?")
 		return 0, timeoutError{}
 	}
 
@@ -147,6 +146,91 @@ func (rg *RouteGroup) Read(p []byte) (n int, err error) {
 	return rg.read(p)
 }
 
+// Write writes payload to a RouteGroup
+// For the first version, only the first ForwardRule (fwd[0]) is used for writing.
+func (rg *RouteGroup) Write(p []byte) (n int, err error) {
+	if rg.isClosed() {
+		return 0, io.ErrClosedPipe
+	}
+
+	if rg.writeDeadline.Closed() {
+		return 0, timeoutError{}
+	}
+
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	rg.mu.Lock()
+	tp, err := rg.tp()
+	if err != nil {
+		rg.mu.Unlock()
+		return 0, err
+	}
+
+	rule, err := rg.rule()
+	if err != nil {
+		rg.mu.Unlock()
+		return 0, err
+	}
+	// we don't need to keep holding mutex from this point on
+	rg.mu.Unlock()
+
+	return rg.write(p, tp, rule)
+}
+
+// Close closes a RouteGroup.
+func (rg *RouteGroup) Close() error {
+	if rg.isClosed() {
+		return io.ErrClosedPipe
+	}
+
+	if rg.isRemoteClosed() {
+		// remote already closed, everything is cleaned up,
+		// we just need to close signal channel at this point
+		close(rg.closed)
+		return nil
+	}
+
+	atomic.StoreInt32(&rg.closeInitiated, 1)
+
+	rg.mu.Lock()
+	defer rg.mu.Unlock()
+
+	return rg.close(routing.CloseRequested)
+}
+
+// LocalAddr returns destination address of underlying RouteDescriptor.
+func (rg *RouteGroup) LocalAddr() net.Addr {
+	return rg.desc.Dst()
+}
+
+// RemoteAddr returns source address of underlying RouteDescriptor.
+func (rg *RouteGroup) RemoteAddr() net.Addr {
+	return rg.desc.Src()
+}
+
+// SetDeadline sets both read and write deadlines.
+func (rg *RouteGroup) SetDeadline(t time.Time) error {
+	if err := rg.SetReadDeadline(t); err != nil {
+		return err
+	}
+
+	return rg.SetWriteDeadline(t)
+}
+
+// SetReadDeadline sets read deadline.
+func (rg *RouteGroup) SetReadDeadline(t time.Time) error {
+	rg.readDeadline.Set(t)
+	return nil
+}
+
+// SetWriteDeadline sets write deadline.
+func (rg *RouteGroup) SetWriteDeadline(t time.Time) error {
+	rg.writeDeadline.Set(t)
+	return nil
+}
+
 // read reads incoming data. It tries to fetch the data from the internal buffer.
 // If buffer is empty it blocks on receiving from the data channel
 func (rg *RouteGroup) read(p []byte) (int, error) {
@@ -178,35 +262,8 @@ func (rg *RouteGroup) read(p []byte) (int, error) {
 	}
 }
 
-// Write writes payload to a RouteGroup
-// For the first version, only the first ForwardRule (fwd[0]) is used for writing.
-func (rg *RouteGroup) Write(p []byte) (n int, err error) {
-	if rg.isClosed() {
-		return 0, io.ErrClosedPipe
-	}
-
-	if rg.writeDeadline.Closed() {
-		return 0, timeoutError{}
-	}
-
-	if len(p) == 0 {
-		return 0, nil
-	}
-
-	rg.mu.Lock()
-	defer rg.mu.Unlock()
-
-	tp, err := rg.tp()
-	if err != nil {
-		return 0, err
-	}
-
-	rule, err := rg.rule()
-	if err != nil {
-		return 0, err
-	}
-
-	packet := routing.MakeDataPacket(rule.KeyRouteID(), p)
+func (rg *RouteGroup) write(data []byte, tp *transport.ManagedTransport, rule routing.Rule) (int, error) {
+	packet := routing.MakeDataPacket(rule.KeyRouteID(), data)
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -223,20 +280,28 @@ func (rg *RouteGroup) Write(p []byte) (n int, err error) {
 
 		atomic.StoreInt64(&rg.lastSent, time.Now().UnixNano())
 
-		return len(p), nil
+		return len(data), nil
 	}
 }
 
 func (rg *RouteGroup) writePacketAsync(ctx context.Context, tp *transport.ManagedTransport, packet routing.Packet) chan error {
 	errCh := make(chan error)
 	go func() {
-		errCh <- tp.WritePacket(ctx, packet)
-		close(errCh)
+		defer close(errCh)
+		err := tp.WritePacket(ctx, packet)
+		select {
+		case <-ctx.Done():
+			return
+		case errCh <- err:
+			return
+		}
 	}()
 
 	return errCh
 }
 
+// rule fetches first available forward rule.
+// NOTE: not thread-safe.
 func (rg *RouteGroup) rule() (routing.Rule, error) {
 	if len(rg.fwd) == 0 {
 		return nil, ErrNoRules
@@ -247,6 +312,8 @@ func (rg *RouteGroup) rule() (routing.Rule, error) {
 	return rule, nil
 }
 
+// tp fetches first available transport.
+// NOTE: not thread-safe.
 func (rg *RouteGroup) tp() (*transport.ManagedTransport, error) {
 	if len(rg.tps) == 0 {
 		return nil, ErrNoTransports
@@ -261,58 +328,6 @@ func (rg *RouteGroup) tp() (*transport.ManagedTransport, error) {
 	return tp, nil
 }
 
-// Close closes a RouteGroup.
-func (rg *RouteGroup) Close() error {
-	if rg.isClosed() {
-		return io.ErrClosedPipe
-	}
-
-	if rg.isRemoteClosed() {
-		// remote already closed, everything is cleaned up,
-		// we just need to close signal channel at this point
-		close(rg.closed)
-		return nil
-	}
-
-	atomic.StoreInt32(&rg.closeInitiated, 1)
-
-	rg.mu.Lock()
-	defer rg.mu.Unlock()
-
-	return rg.close(routing.CloseRequested)
-}
-
-// LocalAddr returns destination address of underlying RouteDescriptor.
-func (rg *RouteGroup) LocalAddr() net.Addr {
-	return rg.desc.Dst()
-}
-
-// RemoteAddr returns source address of underlying RouteDescriptor.
-func (rg *RouteGroup) RemoteAddr() net.Addr {
-	return rg.desc.Src()
-}
-
-// SetDeadline sets both read and write deadlines.
-func (rg *RouteGroup) SetDeadline(t time.Time) error {
-	if err := rg.SetReadDeadline(t); err != nil {
-		return err
-	}
-
-	return rg.SetWriteDeadline(t)
-}
-
-// SetReadDeadline sets read deadline.
-func (rg *RouteGroup) SetReadDeadline(t time.Time) error {
-	rg.readDeadline.Set(t)
-	return nil
-}
-
-// SetWriteDeadline sets write deadline.
-func (rg *RouteGroup) SetWriteDeadline(t time.Time) error {
-	rg.writeDeadline.Set(t)
-	return nil
-}
-
 func (rg *RouteGroup) keepAliveLoop(interval time.Duration) {
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
@@ -354,7 +369,6 @@ func (rg *RouteGroup) sendKeepAlive() error {
 
 	packet := routing.MakeKeepAlivePacket(rule.KeyRouteID())
 	if err := tp.WritePacket(context.Background(), packet); err != nil {
-		rg.logger.WithError(err).Error("Failed to send keep-alive packet")
 		return err
 	}
 
@@ -432,7 +446,7 @@ func (rg *RouteGroup) broadcastClosePackets(code routing.CloseCode) {
 	for i := 0; i < len(rg.tps); i++ {
 		packet := routing.MakeClosePacket(rg.fwd[i].KeyRouteID(), code)
 		if err := rg.tps[i].WritePacket(context.Background(), packet); err != nil {
-			rg.logger.WithError(err).Error("Failed to send close packet")
+			rg.logger.WithError(err).Errorf("Failed to send close packet to %s", rg.tps[i].Remote())
 		}
 	}
 }