From 57b85f4f8dae611d346f6ca6059d5d08de49d8a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= Date: Tue, 9 Jul 2024 16:15:25 +0200 Subject: [PATCH] Optimisation for sonar --- relay/client/client.go | 101 ++++++++++++++++++++++++----------------- 1 file changed, 60 insertions(+), 41 deletions(-) diff --git a/relay/client/client.go b/relay/client/client.go index 9c8bd0365..35c338e34 100644 --- a/relay/client/client.go +++ b/relay/client/client.go @@ -297,7 +297,7 @@ func (c *Client) readLoop(relayConn net.Conn) { c.log.Debugf("failed to read message from relay server: %s", errExit) } c.mu.Unlock() - goto Exit + break } msgType, err := messages.DetermineServerMsgType(buf[:n]) @@ -306,55 +306,74 @@ func (c *Client) readLoop(relayConn net.Conn) { continue } - switch msgType { - case messages.MsgTypeHealthCheck: - msg := messages.MarshalHealthcheck() - _, wErr := c.relayConn.Write(msg) - if wErr != nil { - if c.serviceIsRunning && !internallyStoppedFlag.isSet() { - c.log.Errorf("failed to send heartbeat: %s", wErr) - } - } - hc.Heartbeat() - case messages.MsgTypeTransport: - peerID, payload, err := messages.UnmarshalTransportMsg(buf[:n]) - if err != nil { - if c.serviceIsRunning && !internallyStoppedFlag.isSet() { - c.log.Errorf("failed to parse transport message: %v", err) - } - continue - } - stringID := messages.HashIDToString(peerID) - - c.mu.Lock() - if !c.serviceIsRunning { - c.mu.Unlock() - goto Exit - } - container, ok := c.conns[stringID] - c.mu.Unlock() - if !ok { - c.log.Errorf("peer not found: %s", stringID) - continue - } - - container.writeMsg(Msg{ - bufPool: c.bufPool, - bufPtr: bufPtr, - Payload: payload}) - case messages.MsgTypeClose: - log.Debugf("relay connection close by server") - goto Exit + if !c.handleMsg(msgType, buf[:n], bufPtr, hc, internallyStoppedFlag) { + break } } -Exit: hc.Stop() c.notifyDisconnected() c.wgReadLoop.Done() _ = c.close(false) } +func (c *Client) handleMsg(msgType messages.MsgType, buf []byte, bufPtr *[]byte, hc *healthcheck.Receiver, internallyStoppedFlag *internalStopFlag) (continueLoop bool) { + switch msgType { + case messages.MsgTypeHealthCheck: + msg := messages.MarshalHealthcheck() + _, wErr := c.relayConn.Write(msg) + if wErr != nil { + if c.serviceIsRunning && !internallyStoppedFlag.isSet() { + c.log.Errorf("failed to send heartbeat: %s", wErr) + } + } + hc.Heartbeat() + c.bufPool.Put(bufPtr) + case messages.MsgTypeTransport: + return c.handleTrasnportMsg(buf, bufPtr, internallyStoppedFlag) + case messages.MsgTypeClose: + log.Debugf("relay connection close by server") + c.bufPool.Put(bufPtr) + return false + } + + return true +} + +func (c *Client) handleTrasnportMsg(buf []byte, bufPtr *[]byte, internallyStoppedFlag *internalStopFlag) bool { + peerID, payload, err := messages.UnmarshalTransportMsg(buf) + if err != nil { + if c.serviceIsRunning && !internallyStoppedFlag.isSet() { + c.log.Errorf("failed to parse transport message: %v", err) + } + + c.bufPool.Put(bufPtr) + return true + } + stringID := messages.HashIDToString(peerID) + + c.mu.Lock() + if !c.serviceIsRunning { + c.mu.Unlock() + c.bufPool.Put(bufPtr) + return false + } + container, ok := c.conns[stringID] + c.mu.Unlock() + if !ok { + c.log.Errorf("peer not found: %s", stringID) + c.bufPool.Put(bufPtr) + return true + } + msg := Msg{ + bufPool: c.bufPool, + bufPtr: bufPtr, + Payload: payload, + } + container.writeMsg(msg) + return true +} + // todo check by reference too, the id is not enought because the id come from the outer conn func (c *Client) writeTo(id string, dstID []byte, payload []byte) (int, error) { c.mu.Lock()