From a465e6a657da3a9b12825bd4fa7f569921f23f97 Mon Sep 17 00:00:00 2001 From: lbenedar Date: Fri, 10 Apr 2026 13:59:16 +0300 Subject: [PATCH] add graceful shutdown, add reconnect on connection loss --- Makefile | 2 +- cmd/api/main.go | 45 +--- cmd/api/server.go | 70 +++++ internal/foundry/foundry.go | 250 +++++++++++------- internal/foundry/helpers.go | 18 ++ internal/foundry/transport/websocket.go | 11 +- .../types/{ws_channels.go => channels.go} | 38 ++- .../types/{ws_message.go => message.go} | 0 8 files changed, 286 insertions(+), 148 deletions(-) create mode 100644 cmd/api/server.go create mode 100644 internal/foundry/helpers.go rename internal/foundry/types/{ws_channels.go => channels.go} (53%) rename internal/foundry/types/{ws_message.go => message.go} (100%) diff --git a/Makefile b/Makefile index 7bd71fd..82608e0 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ -include .env api/run: - go run ./cmd/api -foundry_pass=${FOUNDRY_PASS} -foundry_host=${FOUNDRY_HOST} -port=${LISTEN_PORT} -log_level=${LOG_LEVEL} -db-dsn=${DB_DSN} + go run ./cmd/api -foundry_pass=${FOUNDRY_PASS} -foundry_host=${FOUNDRY_HOST} -port=${LISTEN_PORT} -env=development -log_level=${LOG_LEVEL} -db-dsn=${DB_DSN} db/migration/new: @echo 'Creating migration files for ${name}...' diff --git a/cmd/api/main.go b/cmd/api/main.go index 2024b26..31e3ae7 100644 --- a/cmd/api/main.go +++ b/cmd/api/main.go @@ -6,7 +6,6 @@ import ( "flag" "log" "log/slog" - "net/http" "os" "time" @@ -28,6 +27,7 @@ const ( type config struct { mode service_mode port string + env string db struct { dsn string @@ -72,6 +72,7 @@ func openDB(cfg config) (*sql.DB, error) { // TODO: change default values func (app *application) parseFlags() *requests.FoundryHttpRequest { flag.StringVar(&app.cfg.port, "port", ":9090", "Service port") + flag.StringVar(&app.cfg.env, "env", "development", "Enivornment (development|staging|production)") flag.StringVar(&app.cfg.db.dsn, "db-dsn", "file:db/default.db?cache=shared", "PostgreSQL DSN") flag.IntVar(&app.cfg.db.maxOpenConns, "db-max-open-conns", 25, "PostgreSQL max open connections") @@ -91,7 +92,6 @@ func (app *application) parseFlags() *requests.FoundryHttpRequest { flag.Parse() app.cfg.mode = service_mode(mode) - // app.foundryApp.SetConfig(&foundryConfig) logLevel, ok := StringToLogLevel[logLevelStr] if !ok { @@ -105,19 +105,7 @@ func (app *application) parseFlags() *requests.FoundryHttpRequest { return &foundryHttpData } -// TODO: Handle application exit if wrong data -func (app *application) StartListenFoundry() { - var err error - - for { - err = app.foundryApp.ConnectToWebSocket() - if err != nil { - app.slogger.Error("Error raised", "err", err.Error()) - return - } - } -} - +// TODO: Handle grace application exit func main() { app := application{foundryApp: &foundry.FoundryApi{}} @@ -133,29 +121,8 @@ func main() { defer db.Close() app.slogger.Info("Database connection pool established") - app.foundryApp.Transport = transport.NewFoundryTransport(db, app.slogger, foundryHttpData) - go app.StartListenFoundry() + app.foundryApp.SetTransport(transport.NewFoundryTransport(db, app.slogger, foundryHttpData)) + go app.foundryApp.StartListenFoundry() - errLog := slog.NewLogLogger(app.slogger.Handler(), slog.LevelError) - server := &http.Server{ - Addr: app.cfg.port, - Handler: app.routes(), - ErrorLog: errLog, - // TLSConfig: tlsConfig, - IdleTimeout: time.Minute, - ReadTimeout: 5 * time.Second, - WriteTimeout: 10 * time.Second, - } - - switch app.cfg.mode { - case API_MODE: - app.slogger.Info("Server is started", "port", app.cfg.port) - server.ListenAndServe() - case DISCORD_BOT_MODE: - app.slogger.Warn("TODO: implement this mode") - case TG_BOT_MODE: - app.slogger.Warn("TODO: implement this mode") - default: - app.slogger.Error("Wrong Mode. Should be one of: 'api', 'discord', 'tg'", "mode", app.cfg.mode) - } + app.serve() } diff --git a/cmd/api/server.go b/cmd/api/server.go new file mode 100644 index 0000000..f376959 --- /dev/null +++ b/cmd/api/server.go @@ -0,0 +1,70 @@ +package main + +import ( + "context" + "errors" + "log/slog" + "net/http" + "os" + "os/signal" + "syscall" + "time" +) + +func (app *application) serve() error { + srv := &http.Server{ + Addr: app.cfg.port, + Handler: app.routes(), + ErrorLog: slog.NewLogLogger(app.slogger.Handler(), slog.LevelError), + // TLSConfig: tlsConfig, + IdleTimeout: time.Minute, + ReadTimeout: 5 * time.Second, + WriteTimeout: 10 * time.Second, + } + + shutdownError := make(chan error) + go func() { + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + s := <-quit + + app.slogger.Info("Caught signal", "signal", s.String()) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + err := srv.Shutdown(ctx) + if err != nil { + shutdownError <- err + } + + err = app.foundryApp.Shutdown() + if err != nil { + shutdownError <- err + } + + shutdownError <- nil + }() + + switch app.cfg.mode { + case API_MODE: + app.slogger.Info("Starting server", "addr", srv.Addr, "env", app.cfg.env) + err := srv.ListenAndServe() + if !errors.Is(err, http.ErrServerClosed) { + return err + } + + app.slogger.Info("Stopped server", "addr", srv.Addr) + err = <-shutdownError + if err != nil { + return err + } + case DISCORD_BOT_MODE: + app.slogger.Warn("TODO: implement this mode") + case TG_BOT_MODE: + app.slogger.Warn("TODO: implement this mode") + default: + app.slogger.Error("Wrong Mode. Should be one of: 'api', 'discord', 'tg'", "mode", app.cfg.mode) + } + return nil +} diff --git a/internal/foundry/foundry.go b/internal/foundry/foundry.go index 7510c9e..51ccc50 100644 --- a/internal/foundry/foundry.go +++ b/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 + } + }) +} diff --git a/internal/foundry/helpers.go b/internal/foundry/helpers.go new file mode 100644 index 0000000..00fdb0c --- /dev/null +++ b/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() + }() +} diff --git a/internal/foundry/transport/websocket.go b/internal/foundry/transport/websocket.go index f376bce..c849872 100644 --- a/internal/foundry/transport/websocket.go +++ b/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() } diff --git a/internal/foundry/types/ws_channels.go b/internal/foundry/types/channels.go similarity index 53% rename from internal/foundry/types/ws_channels.go rename to internal/foundry/types/channels.go index 432d898..005fba7 100644 --- a/internal/foundry/types/ws_channels.go +++ b/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 { diff --git a/internal/foundry/types/ws_message.go b/internal/foundry/types/message.go similarity index 100% rename from internal/foundry/types/ws_message.go rename to internal/foundry/types/message.go