From ac72bd7bcb764da2695e343b28492bf4fc651c5a Mon Sep 17 00:00:00 2001 From: Vivek Patel Date: Thu, 22 Apr 2021 11:44:55 +0530 Subject: [PATCH] Improves sending pending prompts via websocket and adds config for log enabled for websocket --- app/boot_levels.go | 1 + automation/service/session.go | 36 ++++++++++++----------------------- pkg/options/websocket.gen.go | 1 + pkg/options/websocket.yaml | 4 ++++ websocket/config.go | 1 + websocket/encoder.go | 6 ++---- websocket/router.go | 4 +--- websocket/session.go | 14 ++++++++++++-- websocket/types.go | 5 +++-- websocket/websocket.go | 4 ++++ 10 files changed, 41 insertions(+), 35 deletions(-) diff --git a/app/boot_levels.go b/app/boot_levels.go index 436721b30..33adee462 100644 --- a/app/boot_levels.go +++ b/app/boot_levels.go @@ -320,6 +320,7 @@ func (app *CortezaApp) InitServices(ctx context.Context) (err error) { corredor.Service().SetRoleFinder(sysService.DefaultRole) app.WsServer = websocket.Websocket(&websocket.Config{ + LogEnabled: app.Opt.Websocket.LogEnabled, Timeout: app.Opt.Websocket.Timeout, PingTimeout: app.Opt.Websocket.PingTimeout, PingPeriod: app.Opt.Websocket.PingPeriod, diff --git a/automation/service/session.go b/automation/service/session.go index 7f789d60c..c89e7da6a 100644 --- a/automation/service/session.go +++ b/automation/service/session.go @@ -236,7 +236,6 @@ func (svc *session) spawn(g *wfexec.Graph, workflowID uint64, trace bool) (ses * func (svc *session) Watch(ctx context.Context) { gcTicker := time.NewTicker(time.Second) - promptTicker := time.NewTicker(time.Second) go func() { defer sentry.Recover() @@ -266,6 +265,18 @@ func (svc *session) Watch(ctx context.Context) { // @todo cleanup pool when sessions are complete case <-gcTicker.C: + // Sends pending prompt via websocket + ws := websocket.Session(ctx, nil, nil) + userIds := ws.GetActiveUserIDs() + for _, s := range svc.pool { + for _, u := range userIds { + pp := s.PendingPrompts(u) + if len(pp) > 0 { + _ = ws.Send(websocket.Message(websocket.StatusOK, s.WorkflowID, pp), u) + } + } + } + svc.gc() } } @@ -274,29 +285,6 @@ func (svc *session) Watch(ctx context.Context) { //svc.suspendAll(ctx) }() - // Prompt routine - go func() { - defer sentry.Recover() - defer promptTicker.Stop() - defer svc.log.Info("stopped") - - for { - select { - case <-ctx.Done(): - return - case <-promptTicker.C: - activeSessions := websocket.GetActiveSessions() - for _, s := range activeSessions { - var ctxr = context.Background() - pp := svc.PendingPrompts(auth.SetIdentityToContext(ctxr, s.User())) - if len(pp) > 0 { - _ = s.Send(websocket.Message(websocket.StatusOK, websocket.WorkflowApplication, pp), s.User().Identity()) - } - } - } - } - }() - svc.log.Debug("watcher initialized") } diff --git a/pkg/options/websocket.gen.go b/pkg/options/websocket.gen.go index 0f52c29c8..dd3c3c524 100644 --- a/pkg/options/websocket.gen.go +++ b/pkg/options/websocket.gen.go @@ -14,6 +14,7 @@ import ( type ( WebsocketOpt struct { + LogEnabled bool `env:"WEBSOCKET_LOG_ENABLED"` Timeout time.Duration `env:"WEBSOCKET_TIMEOUT"` PingTimeout time.Duration `env:"WEBSOCKET_PING_TIMEOUT"` PingPeriod time.Duration `env:"WEBSOCKET_PING_PERIOD"` diff --git a/pkg/options/websocket.yaml b/pkg/options/websocket.yaml index 167ce9610..5ecde67de 100644 --- a/pkg/options/websocket.yaml +++ b/pkg/options/websocket.yaml @@ -6,6 +6,10 @@ docs: description: A Websocket server emphasize the trigger events and actions. props: + - name: LogEnabled + type: bool + description: Enable extra logging for authentication flows + - name: Timeout type: time.Duration default: 15 * time.Second diff --git a/websocket/config.go b/websocket/config.go index 386189cf6..2121ea36a 100644 --- a/websocket/config.go +++ b/websocket/config.go @@ -6,6 +6,7 @@ import ( type ( Config struct { + LogEnabled bool Timeout time.Duration PingTimeout time.Duration PingPeriod time.Duration diff --git a/websocket/encoder.go b/websocket/encoder.go index 0db815a56..0377683da 100644 --- a/websocket/encoder.go +++ b/websocket/encoder.go @@ -7,14 +7,12 @@ import ( const ( StatusOK = "ok" StatusError = "error" - - WorkflowApplication = "Workflow" ) type ( message struct { Status string `json:"status"` - Application string `json:"application"` + Application uint64 `json:"application"` Data interface{} `json:"data"` } @@ -23,7 +21,7 @@ type ( } ) -func Message(status, application string, data interface{}) *message { +func Message(status string, application uint64, data interface{}) *message { return &message{ Status: status, Application: application, diff --git a/websocket/router.go b/websocket/router.go index 1d1dd5413..ae57b1a2e 100644 --- a/websocket/router.go +++ b/websocket/router.go @@ -10,7 +10,5 @@ import ( // If it's valid then we keep the connection open or close it func (ws *websocket) MountRoutes(r chi.Router) { // Initialize handlers & controllers. - r.Group(func(r chi.Router) { - r.Get("/", ws.Open) - }) + r.Get("/", ws.Open) } diff --git a/websocket/session.go b/websocket/session.go index d4b7b6ac5..2e5e5349e 100644 --- a/websocket/session.go +++ b/websocket/session.go @@ -310,6 +310,16 @@ func (s *session) sendBytes(p []byte) error { return nil } -func GetActiveSessions() map[uint64]*session { - return sessions +func (s *session) Walk(callback func(*session)) { + for _, sess := range sessions { + callback(sess) + } +} + +// GetActiveUserIDs return userIDs from active ws sessions +func (s *session) GetActiveUserIDs() (ids []uint64) { + s.Walk(func(sess *session) { + ids = append(ids, sess.user.Identity()) + }) + return } diff --git a/websocket/types.go b/websocket/types.go index fd85b8a91..3ee6a1db7 100644 --- a/websocket/types.go +++ b/websocket/types.go @@ -31,6 +31,7 @@ func (a *Auth) ParseWithClaims() (jwt.MapClaims, error) { } } -func Unmarshal(raw []byte) (p *Payload, err error) { - return p, json.Unmarshal(raw, p) +func Unmarshal(raw []byte) (*Payload, error) { + var p Payload + return &p, json.Unmarshal(raw, &p) } diff --git a/websocket/websocket.go b/websocket/websocket.go index 1849f5a60..55a5d00cf 100644 --- a/websocket/websocket.go +++ b/websocket/websocket.go @@ -27,6 +27,10 @@ type ( ) func Websocket(config *Config, logger *zap.Logger) *websocket { + if !config.LogEnabled { + logger = zap.NewNop() + } + return &websocket{ config: config, logger: logger,