diff --git a/.examples.env b/.examples.env index 5c02770..6859dea 100644 --- a/.examples.env +++ b/.examples.env @@ -1,3 +1,4 @@ PORT=6756 LOG_LEVEL=DEBUG -MODE=DEV \ No newline at end of file +MODE=DEV +HEARTBEAT=5 \ No newline at end of file diff --git a/hub/internal/app/app.go b/hub/internal/app/app.go index 4f7eb4d..18687fe 100644 --- a/hub/internal/app/app.go +++ b/hub/internal/app/app.go @@ -59,9 +59,9 @@ func (a *App) Run() { a.log.Info().Msg("migrations applied successfully") hubStore := store.NewHubStore(DBConn) - hubService := hub_service.NewHubService(hubStore, a.log) + hubService := hub_service.NewHubService(hubStore, a.cfg.Heartbeat, a.log) statusNotifier := notifier.NewStatusNotifier() - connManger := connection_manager.NewConnectionManager(hubStore, statusNotifier, a.log) + connManger := connection_manager.NewConnectionManager(hubStore, statusNotifier, a.cfg.Heartbeat*1000, a.log) a.log.Info().Msg("starting hub service") diff --git a/hub/internal/service/connection_manager/manager.go b/hub/internal/service/connection_manager/manager.go index 13acc81..5a5480d 100644 --- a/hub/internal/service/connection_manager/manager.go +++ b/hub/internal/service/connection_manager/manager.go @@ -8,17 +8,16 @@ import ( "google.golang.org/grpc/metadata" ) -const heartbeatTimeoutMS = 6000 - type ConnectionManager struct { - heartbeat heartbeatStore - log zerolog.Logger - status statusNotifier - agentConnStore *AgentConnStore + heartbeat heartbeatStore + log zerolog.Logger + status statusNotifier + agentConnStore *AgentConnStore + heartbeatTimeout int } -func NewConnectionManager(heartbeat heartbeatStore, status statusNotifier, logger zerolog.Logger) *ConnectionManager { - return &ConnectionManager{heartbeat: heartbeat, log: logger, status: status, agentConnStore: NewAgentConnStore()} +func NewConnectionManager(heartbeat heartbeatStore, status statusNotifier, heartbeatTimeoutMS int, logger zerolog.Logger) *ConnectionManager { + return &ConnectionManager{heartbeat: heartbeat, log: logger, status: status, agentConnStore: NewAgentConnStore(), heartbeatTimeout: heartbeatTimeoutMS + 600} } func (c *ConnectionManager) NewConnection(stream StreamConn) error { @@ -32,7 +31,7 @@ func (c *ConnectionManager) NewConnection(stream StreamConn) error { status := c.status.New(AgentID) - agent := newAgentConnection(AgentID, stream, c.heartbeat, status, heartbeatTimeoutMS, c.log) + agent := newAgentConnection(AgentID, stream, c.heartbeat, status, c.heartbeatTimeout, c.log) c.agentConnStore.Add(AgentID, agent) defer c.agentConnStore.Delete(AgentID) diff --git a/hub/internal/service/connection_manager/manager_test.go b/hub/internal/service/connection_manager/manager_test.go index 61a16c6..b346130 100644 --- a/hub/internal/service/connection_manager/manager_test.go +++ b/hub/internal/service/connection_manager/manager_test.go @@ -24,7 +24,7 @@ func newConnectionManagerTestHarness(t *testing.T) *connectionManagerTestHarness heartbeat := &heartBeatMock{doneCh: make(chan struct{}, 2)} status := &statusNotifierMock{agentIDCh: make(chan string, 1)} - manager := NewConnectionManager(heartbeat, status, zerolog.New(nil)) + manager := NewConnectionManager(heartbeat, status, 10000, zerolog.New(nil)) return &connectionManagerTestHarness{manager: manager, status: status, heartbeat: heartbeat} } diff --git a/hub/internal/service/hub_service/hub.go b/hub/internal/service/hub_service/hub.go index 00d7181..f9ae961 100644 --- a/hub/internal/service/hub_service/hub.go +++ b/hub/internal/service/hub_service/hub.go @@ -12,8 +12,6 @@ import ( "github.com/rs/zerolog" ) -const HEARTBEAT = 5 - type Store interface { NewAgent(ctx context.Context, agent domainHub.CreateAgentModel) error GetAgentByAgentID(ctx context.Context, AgentID string) (domainHub.AgentModel, error) @@ -21,12 +19,13 @@ type Store interface { } type HubService struct { - store Store - log zerolog.Logger + store Store + log zerolog.Logger + heartbeatTimeout int } -func NewHubService(store Store, logger zerolog.Logger) *HubService { - return &HubService{log: logger, store: store} +func NewHubService(store Store, heartbeatTimeout int, logger zerolog.Logger) *HubService { + return &HubService{log: logger, store: store, heartbeatTimeout: heartbeatTimeout} } func (h *HubService) RegisterAgent(ctx context.Context, data domain.RegisterAgentRequest) (domain.RegisterAgentResponse, error) { @@ -47,7 +46,7 @@ func (h *HubService) RegisterAgent(ctx context.Context, data domain.RegisterAgen return domain.RegisterAgentResponse{}, fmt.Errorf("update agent in db: %w", err) } h.log.Debug().Str("agentId", agent.AgentID).Msg("agent updated successfully") - return domain.RegisterAgentResponse{AgentID: agent.AgentID, Heartbeat: HEARTBEAT}, nil + return domain.RegisterAgentResponse{AgentID: agent.AgentID, Heartbeat: h.heartbeatTimeout}, nil } AgentID, err := hasher.MakeID(data.Host, data.AgentName) @@ -62,5 +61,5 @@ func (h *HubService) RegisterAgent(ctx context.Context, data domain.RegisterAgen if err := h.store.NewAgent(ctx, agentStore); err != nil { return domain.RegisterAgentResponse{}, fmt.Errorf("insert new agent: %w", err) } - return domain.RegisterAgentResponse{AgentID: AgentID, Heartbeat: HEARTBEAT}, nil + return domain.RegisterAgentResponse{AgentID: AgentID, Heartbeat: h.heartbeatTimeout}, nil } diff --git a/shared/config/config.go b/shared/config/config.go index d7667c7..565787c 100644 --- a/shared/config/config.go +++ b/shared/config/config.go @@ -9,9 +9,10 @@ import ( ) type Config struct { - LogLevel string `env:"LOG_LEVEL" env-default:"INFO"` - Mode string `env:"MODE" env-default:"DEV"` - Port int `env:"PORT" env-default:"9000"` + LogLevel string `env:"LOG_LEVEL" env-default:"INFO"` + Mode string `env:"MODE" env-default:"DEV"` + Port int `env:"PORT" env-default:"9000"` + Heartbeat int `env:"HEARTBEAT" env-default:"5"` } func NewConfig() (*Config, error) {