add graceful shutdown, add reconnect on connection loss

internal/foundry/foundry.go

@@ -1,154 +1,200 @@
 package foundry
 
 import (
+	"context"
 	"errors"
+	"log/slog"
+	"sync"
 	"time"
 
-	"gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/actions"
 	"gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/transport"
 	"gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/types"
 )
 
 var (
-	ListenIsDone = errors.New("Listen for websocket data in foundry is stopped")
+	ListenIsDone    = errors.New("Listen for websocket data in foundry is stopped")
+	ChannelIsClosed = errors.New("Channel is closed")
+
+	CloseTimeoutExceed = errors.New("Timeout of websocket close is exceed")
 )
 
 type FoundryApi struct {
 	//TODO: make check of admin's authentication
-	Transport   *transport.FoundryTransport
-	IsReadReady bool
+	transport   *transport.FoundryTransport
+	IsAvailable bool
 
+	Logger *slog.Logger
 	Status types.FoundryStatus
+	wg     sync.WaitGroup
 }
 
 func NewFoundry() *FoundryApi {
 	return &FoundryApi{}
 }
 
-func (foundry *FoundryApi) StartListen() error {
-	foundry.Transport.ReadChan = *types.InitWsChannels()
-
-	wsChannels := &foundry.Transport.ReadChan
-	defer foundry.Transport.ReadChan.Close()
-	go foundry.ListenAndServeWS()
+func (foundry *FoundryApi) SetTransport(tr *transport.FoundryTransport) {
+	foundry.transport = tr
+}
 
+func (foundry *FoundryApi) ListenAndServeWS() error {
 	var err error
+
+	foundry.transport.ReadChan = *types.InitWsChannels()
+	defer foundry.transport.ReadChan.Close()
+
+	foundry.background(foundry.transport.ListenWebSocket)
+
+	err = foundry.ServeWebSocket()
+	if err != nil {
+		foundry.IsAvailable = false
+	}
+	return err
+}
+
+func (foundry *FoundryApi) ServeWebSocket() error {
+	var err error
+	wsChannels := &foundry.transport.ReadChan
+
 	for {
 		select {
-		// case msg, ok := <-foundry.ws.channels.Msg():
-		// 	if !ok {
-		// 		time.Sleep(5 * time.Microsecond)
-		// 		continue
-		// 	}
-		// 	var foundryState *json_model.FoundryState
-		// 	foundryState, err = json_model.ParseSetupModel(msg.ToByteSlice())
-		// 	if err != nil {
-		// 		return nil
-		// 	}
-		// 	foundry.models.FoundryState.Insert(foundryState)
+		case message, ok := <-wsChannels.Msg():
+			if !ok {
+				foundry.transport.Logger.Debug("Read channel is closed", "type", types.WebSocketCode)
+				return ChannelIsClosed
+			}
+			foundry.transport.Logger.Debug("Data has been received\n", "msg", message.Code, "type", types.WebSocketCode) //, "msg", message.MsgJson)
+
+			switch message.Code {
+			case types.RespPingCode, types.RespSessionDataCode:
+				err := foundry.transport.SendOnlyCodeRequest(message.Code)
+				if err != nil {
+					wsChannels.Err() <- &types.FoundryError{Direction: types.WriterCode, Err: err, Type: types.WebSocketCode}
+					continue
+				}
+			case types.RespServerChangeCode:
+				foundry.IsAvailable = true
+				// data, err := actions.SplitToTypeAndData([]byte(message.MsgJson))
+				// if err != nil {
+				// 	wsChannels.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode}
+				// 	continue
+				// }
+
+				// if data == nil {
+				// 	continue
+				// }
+
+				// err = data.Action(foundry.transport, &foundry.Status)
+				// if err != nil {
+				// 	wsChannels.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode}
+				// 	continue
+				// }
+			case types.RespDataCode:
+				foundry.transport.ExchangeChan.Msgs[message.Id] = make(chan []byte, 1)
+				foundry.transport.ExchangeChan.Msgs[message.Id] <- []byte(message.MsgJson)
+				go foundry.transport.CloseMsgChannel(message.Id, 5*time.Second)
+			default:
+			}
 		case err = <-wsChannels.Err():
 			var foundryErr *types.FoundryError
 			if errors.As(err, &foundryErr) {
 				if foundryErr.IsFatal {
-					foundry.Transport.CloseWebSocketConn()
+					foundry.transport.CloseWebSocketConn()
 					return foundryErr
 				}
-				foundry.Transport.Logger.Warn("Got error when listening or served", "err", err.Error(), "type", foundryErr.Type, "direction", foundryErr.Direction)
+				foundry.transport.Logger.Warn("Got error when listening or served", "err", err.Error(), "type", foundryErr.Type, "direction", foundryErr.Direction)
 			}
 		case <-wsChannels.Done():
-			foundry.Transport.CloseWebSocketConn()
+			foundry.transport.CloseWebSocketConn()
 			return ListenIsDone
 		}
 	}
 }
 
-func (foundry *FoundryApi) ServeWebSocket() {
-	for {
-		select {
-		case message, ok := <-foundry.Transport.ReadChan.Msg():
-			if !ok {
-				foundry.Transport.Logger.Debug("Read channel is closed", "type", types.WebSocketCode)
-				return
-			}
-			foundry.Transport.Logger.Debug("Data has been received\n", "msg", message.Code, "type", types.WebSocketCode) //, "msg", message.MsgJson)
-
-			switch message.Code {
-			case types.RespPingCode, types.RespSessionDataCode:
-				err := foundry.Transport.SendOnlyCodeRequest(message.Code)
-				if err != nil {
-					foundry.IsReadReady = false
-					foundry.Transport.ReadChan.Err() <- &types.FoundryError{Direction: types.WriterCode, Err: err, Type: types.WebSocketCode}
-					continue
-				}
-			case types.RespServerChangeCode:
-				foundry.IsReadReady = true
-				data, err := actions.SplitToTypeAndData([]byte(message.MsgJson))
-				if err != nil {
-					foundry.Transport.ReadChan.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode}
-					continue
-				}
-
-				if data == nil {
-					continue
-				}
-
-				err = data.Action(foundry.Transport, &foundry.Status)
-				if err != nil {
-					foundry.Transport.ReadChan.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode}
-					continue
-				}
-			case types.RespDataCode:
-				foundry.Transport.ExchangeChan.Msgs[message.Id] = make(chan []byte, 1)
-				foundry.Transport.ExchangeChan.Msgs[message.Id] <- []byte(message.MsgJson)
-				go foundry.Transport.CloseMsgChannel(message.Id, 5*time.Second)
-			default:
-			}
-		case <-foundry.Transport.ReadChan.Done():
-			return
-		}
-	}
-}
-
-func (foundry *FoundryApi) ListenAndServeWS() {
-	go foundry.Transport.ListenWebSocket()
-	foundry.ServeWebSocket()
-}
-
 func (foundry *FoundryApi) HandleWSRequest(msgType string) ([]byte, error) {
-	msg := types.NewWsMessageByPage(msgType, foundry.Transport.CurrWsId)
+	msg := types.NewWsMessageByPage(msgType, foundry.transport.CurrWsId)
 
-	return foundry.Transport.HandleWebsocketRequest(msg)
+	return foundry.transport.HandleWebsocketRequest(msg)
 }
 
-func (foundry *FoundryApi) ConnectToWebSocket() error {
-	var err error
-	if !foundry.Transport.HasSessionId() {
-		err = foundry.Transport.InitSessionId()
-		if err != nil {
-			return err
-		}
-	}
+func (foundry *FoundryApi) Shutdown() error {
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
 
-	err = foundry.Transport.ConnectToFoundry()
-	if err != nil {
-		return err
-	}
+	foundry.transport.Logger.Info("Completing foundry background tasks")
 
-	err = foundry.Transport.InitWebSocketConnection()
-	if err != nil {
-		return err
-	}
+	foundryClosed := make(chan struct{})
+	go func() {
+		types.CloseChannel(foundry.transport.ReadChan.Done())
 
-	status, err := foundry.Transport.Http.GetStatus()
-	if err != nil {
-		return err
-	}
-	foundry.Status = *types.NewFoundryStatus(status)
+		foundry.wg.Wait()
+		foundry.transport.ExchangeChan.Close()
+		foundry.transport.ReadChan.Close()
 
-	err = foundry.StartListen()
-	if err != nil && err != ListenIsDone {
-		return err
+		types.CloseChannel(foundryClosed)
+	}()
+
+	select {
+	case <-foundryClosed:
+		foundry.transport.Logger.Info("Stopped foundry server")
+	case <-ctx.Done():
+		return CloseTimeoutExceed
 	}
 
 	return nil
 }
+
+func (foundry *FoundryApi) ConnectToWebSocket() (bool, error) {
+	var err error
+	if !foundry.transport.HasSessionId() {
+		err = foundry.transport.InitSessionId()
+		if err != nil {
+			return false, err
+		}
+	}
+
+	err = foundry.transport.ConnectToFoundry()
+	if err != nil {
+		return false, err
+	}
+
+	err = foundry.transport.InitWebSocketConnection()
+	if err != nil {
+		return false, err
+	}
+
+	status, err := foundry.transport.Http.GetStatus()
+	if err != nil {
+		return false, err
+	}
+	foundry.Status = *types.NewFoundryStatus(status)
+
+	err = foundry.ListenAndServeWS()
+	if err != nil && err != ListenIsDone {
+		return true, err
+	}
+
+	return true, nil
+}
+
+func (foundry *FoundryApi) StartListenFoundry() {
+	foundry.background(func() {
+		timeInterval := 1 * time.Second
+
+		for {
+			ok, err := foundry.ConnectToWebSocket()
+			if err != nil {
+				foundry.transport.Logger.Error("Error raised", "err", err.Error())
+
+				if ok {
+					timeInterval = 1 * time.Second
+				}
+				foundry.transport.Logger.Info("Trying to reconnect", "timer", timeInterval.String())
+
+				time.Sleep(timeInterval)
+				timeInterval = min(timeInterval*2, 15*time.Second)
+				continue
+			}
+			return
+		}
+	})
+}

internal/foundry/helpers.go

@@ -0,0 +1,18 @@
+package foundry
+
+import "fmt"
+
+func (f *FoundryApi) background(fn func()) {
+	f.wg.Add(1)
+	go func() {
+		defer f.wg.Done()
+
+		defer func() {
+			if err := recover(); err != nil {
+				f.transport.Logger.Error(fmt.Sprintf("%s", err))
+			}
+		}()
+
+		fn()
+	}()
+}

internal/foundry/transport/websocket.go

@@ -80,7 +80,12 @@ func (tr *FoundryTransport) ListenWebSocket() {
 	for {
 		_, message, err := tr.WsConn.ReadMessage()
 		if err != nil {
-			tr.ReadChan.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode}
+			select {
+			case <-tr.ReadChan.Done():
+				return
+			default:
+				tr.ReadChan.Err() <- &types.FoundryError{Direction: types.ReaderCode, Err: err, Type: types.WebSocketCode, IsFatal: true}
+			}
 			return
 		}
 		data, err := types.ParseWsRespMessage(string(message), types.RequestCodes)
@@ -105,7 +110,7 @@ func (tr *FoundryTransport) GetJsonData(stateType string) ([]byte, error) {
 	return tr.HandleWebsocketRequest(msg)
 }
 
-func (tr *FoundryTransport) CloseWebSocketConn() {
+func (tr *FoundryTransport) CloseWebSocketConn() error {
 	tr.Logger.Debug("Websocket has been closed")
-	tr.WsConn.Close()
+	return tr.WsConn.Close()
 }

internal/foundry/types/channels.go

@@ -8,6 +8,27 @@ type ExchangeChannels struct {
 	ProgressMutex sync.Mutex
 }
 
+func CloseMapChannel[K comparable, V any](channel map[K]chan V) {
+	for _, v := range channel {
+		select {
+		case _, ok := <-v:
+			if ok {
+				close(v)
+			}
+		default:
+			close(v)
+		}
+	}
+	for k := range channel {
+		delete(channel, k)
+	}
+}
+
+func (channels *ExchangeChannels) Close() {
+	CloseMapChannel(channels.Msgs)
+	CloseMapChannel(channels.ProgressMsg)
+}
+
 type ReadChannels struct {
 	done chan struct{}
 	msg  chan *WsMessage
@@ -26,10 +47,21 @@ func (channels ReadChannels) Done() chan struct{} {
 	return channels.done
 }
 
+func CloseChannel[V any](channel chan V) {
+	select {
+	case _, ok := <-channel:
+		if ok {
+			close(channel)
+		}
+	default:
+		close(channel)
+	}
+}
+
 func (channels *ReadChannels) Close() {
-	close(channels.done)
-	close(channels.err)
-	close(channels.msg)
+	CloseChannel(channels.done)
+	CloseChannel(channels.err)
+	CloseChannel(channels.msg)
 }
 
 func InitWsChannels() *ReadChannels {