Fix in client the close event

relay/client/client.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
@@ -29,20 +30,25 @@ type connContainer struct {
 
 type Client struct {
 	log           *log.Entry
+	ctx           context.Context
+	ctxCancel     context.CancelFunc
 	serverAddress string
 	hashedID      []byte
 
-	conns map[string]*connContainer
+	conns map[string]*connContainer // todo handle it in thread safe way
 
 	relayConn      net.Conn
 	relayConnState bool
 	mu             sync.Mutex
 }
 
-func NewClient(serverAddress, peerID string) *Client {
+func NewClient(ctx context.Context, serverAddress, peerID string) *Client {
+	ctx, ctxCancel := context.WithCancel(ctx)
 	hashedID, hashedStringId := messages.HashID(peerID)
 	return &Client{
 		log:           log.WithField("client_id", hashedStringId),
+		ctx:           ctx,
+		ctxCancel:     ctxCancel,
 		serverAddress: serverAddress,
 		hashedID:      hashedID,
 		conns:         make(map[string]*connContainer),
@@ -51,7 +57,11 @@ func NewClient(serverAddress, peerID string) *Client {
 
 func (c *Client) Connect() error {
 	c.mu.Lock()
-	defer c.mu.Unlock()
+	if c.relayConnState {
+		c.mu.Unlock()
+		return nil
+	}
+
 	conn, err := udp.Dial(c.serverAddress)
 	if err != nil {
 		return err
@@ -68,18 +78,39 @@ func (c *Client) Connect() error {
 		return err
 	}
 
-	err = c.relayConn.SetReadDeadline(time.Time{})
-	if err != nil {
-		log.Errorf("failed to reset read deadline: %s", err)
-		return err
-	}
-
 	c.relayConnState = true
-	go c.readLoop()
+	c.mu.Unlock()
+
+	go func() {
+		<-c.ctx.Done()
+		cErr := c.close()
+		if cErr != nil {
+			log.Errorf("failed to close relay connection: %s", cErr)
+		}
+	}()
+	// blocking function
+	c.readLoop()
+
+	c.mu.Lock()
+
+	// close all Conn types
+	for _, container := range c.conns {
+		close(container.messages)
+	}
+	c.conns = make(map[string]*connContainer)
+
+	c.mu.Unlock()
+
 	return nil
 }
 
 func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if !c.relayConnState {
+		return nil, fmt.Errorf("relay connection is not established")
+	}
+
 	hashedID, hashedStringID := messages.HashID(dstPeerID)
 	log.Infof("open connection to peer: %s", hashedStringID)
 	messageBuffer := make(chan Msg, 2)
@@ -93,6 +124,11 @@ func (c *Client) OpenConn(dstPeerID string) (net.Conn, error) {
 }
 
 func (c *Client) Close() error {
+	c.ctxCancel()
+	return c.close()
+}
+
+func (c *Client) close() error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -101,11 +137,20 @@ func (c *Client) Close() error {
 	}
 
 	c.relayConnState = false
+
 	err := c.relayConn.Close()
+
 	return err
 }
 
 func (c *Client) handShake() error {
+	defer func() {
+		err := c.relayConn.SetReadDeadline(time.Time{})
+		if err != nil {
+			log.Errorf("failed to reset read deadline: %s", err)
+		}
+	}()
+
 	msg, err := messages.MarshalHelloMsg(c.hashedID)
 	if err != nil {
 		log.Errorf("failed to marshal hello message: %s", err)
@@ -145,7 +190,7 @@ func (c *Client) handShake() error {
 
 func (c *Client) readLoop() {
 	defer func() {
-		c.log.Debugf("exit from read loop")
+		c.log.Tracef("exit from read loop")
 	}()
 	var errExit error
 	var n int

relay/client/dialer/ws/client_conn.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
 )
 
 type Conn struct {
@@ -51,6 +52,9 @@ func (c *Conn) SetDeadline(t time.Time) error {
 }
 
 func (c *Conn) Close() error {
-	_ = c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5))
+	err := c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5))
+	if err != nil {
+		log.Errorf("failed to close conn?: %s", err)
+	}
 	return c.Conn.Close()
 }

relay/client/manager.go

@@ -2,42 +2,74 @@ package client
 
 import (
 	"context"
+	"net"
 	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
 )
 
 type Manager struct {
 	ctx        context.Context
-	ctxCancel  context.CancelFunc
 	srvAddress string
 	peerID     string
 
-	wg sync.WaitGroup
+	reconnectTime time.Duration
 
-	clients      map[string]*Client
-	clientsMutex sync.RWMutex
+	mu     sync.Mutex
+	client *Client
 }
 
 func NewManager(ctx context.Context, serverAddress string, peerID string) *Manager {
-	ctx, cancel := context.WithCancel(ctx)
 	return &Manager{
-		ctx:        ctx,
-		ctxCancel:  cancel,
-		srvAddress: serverAddress,
-		peerID:     peerID,
-		clients:    make(map[string]*Client),
+		ctx:           ctx,
+		srvAddress:    serverAddress,
+		peerID:        peerID,
+		reconnectTime: 5 * time.Second,
 	}
 }
 
-func (m *Manager) Teardown() {
-	m.ctxCancel()
-	m.wg.Wait()
-}
-
-func (m *Manager) newSrvConnection(address string) {
-	if _, ok := m.clients[address]; ok {
+func (m *Manager) Serve() {
+	ok := m.mu.TryLock()
+	if !ok {
 		return
 	}
 
-	//	client := NewClient(address, m.peerID)
-	//err = client.Connect()
+	m.client = NewClient(m.ctx, m.srvAddress, m.peerID)
+
+	go func() {
+		defer m.mu.Unlock()
+
+		// todo this is not thread safe
+		for {
+			select {
+			case <-m.ctx.Done():
+				return
+			default:
+				m.connect()
+			}
+
+			select {
+			case <-m.ctx.Done():
+				return
+			case <-time.After(2 * time.Second): //timeout
+			}
+		}
+	}()
+}
+
+func (m *Manager) OpenConn(peerKey string) (net.Conn, error) {
+	// todo m.client nil check
+	return m.client.OpenConn(peerKey)
+}
+
+// connect is blocking
+func (m *Manager) connect() {
+	err := m.client.Connect()
+	if err != nil {
+		if m.ctx.Err() != nil {
+			return
+		}
+		log.Errorf("connection error with '%s': %s", m.srvAddress, err)
+	}
 }