-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathserver.go
More file actions
840 lines (718 loc) · 24.1 KB
/
Copy pathserver.go
File metadata and controls
840 lines (718 loc) · 24.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
package gqm
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"os"
"os/signal"
"runtime"
"strings"
"sync"
"syscall"
"time"
"github.com/benedict-erwin/gqm/monitor"
"github.com/redis/go-redis/v9"
)
// AuthUser represents a user for the monitoring API.
type AuthUser struct {
Username string
PasswordHash string // bcrypt hash
Role string // "admin" or "viewer"; defaults to "admin" if empty
}
// AuthAPIKey represents an API key for programmatic access.
type AuthAPIKey struct {
Name string
Key string
Role string // "admin" or "viewer"; defaults to "admin" if empty
}
// ServerOption configures a Server.
type ServerOption func(*serverConfig)
type serverConfig struct {
redisOpts []RedisOption
globalTimeout time.Duration
gracePeriod time.Duration
shutdownTimeout time.Duration
logger *slog.Logger
// Phase 5: config-driven fields
defaultTimezone string // fallback timezone for cron entries (IANA)
schedulerEnabled bool // whether to start the scheduler goroutine
schedulerPollInterval time.Duration // poll interval for scheduled/cron jobs
logLevel string // auto-create logger if no WithLogger (debug/info/warn/error)
catchAllPool string // pool name with job_types: ["*"]
// Phase 6: monitoring
apiEnabled bool
apiAddr string
dashEnabled bool
dashPathPrefix string
dashCustomDir string
authEnabled bool
authSessionTTL int // seconds, default 86400
authUsers []AuthUser // loaded from config
apiKeys []AuthAPIKey
apiRateLimit int // requests/second per IP; 0 = default (100), -1 = disabled
// trustProxy allows the client-supplied X-Forwarded-Proto header to decide
// whether a connection counts as HTTPS; cookieSecure states it outright.
trustProxy bool
cookieSecure bool
// Retention windows in seconds, applied to terminal jobs that carry no
// per-job override. -1 retains permanently; 0 deletes immediately.
//
// Pointers so that the zero value of serverConfig means "use the built-in
// default" rather than 0, which would mean "delete every terminal job
// immediately" — the most destructive possible reading of an unset field.
resultTTL *int
failureTTL *int
}
// resultRetention resolves the retention window for a job that completed
// successfully: per-job override first, then the server setting, then the
// built-in default.
func (c *serverConfig) resultRetention(override *int) int {
return retentionTTL(override, retentionTTL(c.resultTTL, defaultResultTTLSeconds))
}
// failureRetention resolves the retention window for a job that reached a
// failure terminal state, with the same precedence as resultRetention.
func (c *serverConfig) failureRetention(override *int) int {
return retentionTTL(override, retentionTTL(c.failureTTL, defaultFailureTTLSeconds))
}
// WithServerRedis sets the Redis address for the server.
func WithServerRedis(addr string) ServerOption {
return func(cfg *serverConfig) {
cfg.redisOpts = append(cfg.redisOpts, WithRedisAddr(addr))
}
}
// WithServerRedisOpts sets Redis options for the server.
func WithServerRedisOpts(opts ...RedisOption) ServerOption {
return func(cfg *serverConfig) {
cfg.redisOpts = append(cfg.redisOpts, opts...)
}
}
// WithServerRedisClient injects a pre-configured *redis.Client for the
// server. This enables Redis Sentinel or any custom go-redis setup.
// Connection options (WithServerRedis, etc.) are ignored when this is
// used; only the key prefix (WithPrefix via WithServerRedisOpts) still
// applies.
func WithServerRedisClient(rdb *redis.Client) ServerOption {
return func(cfg *serverConfig) {
cfg.redisOpts = append(cfg.redisOpts, WithRedisClient(rdb))
}
}
// WithGlobalTimeout sets the global default job timeout.
// Must be > 0; the global timeout cannot be disabled.
func WithGlobalTimeout(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
if d > 0 {
cfg.globalTimeout = d
}
}
}
// WithGracePeriod sets the default grace period after context cancellation.
func WithGracePeriod(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
if d > 0 {
cfg.gracePeriod = d
}
}
}
// WithResultTTL sets how long completed jobs are retained server-wide.
// Individual jobs can override it with the ResultTTL enqueue option.
//
// A negative duration means TTLPermanent (retain forever). A zero duration
// deletes each job hash as soon as it completes. Note that retaining jobs
// permanently leaves no way to reclaim their memory.
func WithResultTTL(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
cfg.resultTTL = ttlSeconds(d)
}
}
// WithFailureTTL sets how long jobs that reached a failure terminal state
// (dead-lettered, canceled, or stopped) are retained server-wide. Individual
// jobs can override it with the FailureTTL enqueue option.
//
// Follows the same convention as WithResultTTL.
func WithFailureTTL(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
cfg.failureTTL = ttlSeconds(d)
}
}
// WithShutdownTimeout sets the maximum wait time during graceful shutdown.
func WithShutdownTimeout(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
if d > 0 {
cfg.shutdownTimeout = d
}
}
}
// WithLogger sets a custom slog.Logger.
func WithLogger(l *slog.Logger) ServerOption {
return func(cfg *serverConfig) { cfg.logger = l }
}
// WithDefaultTimezone sets the fallback timezone (IANA name) for cron entries
// that don't specify their own timezone. Defaults to UTC.
func WithDefaultTimezone(tz string) ServerOption {
return func(cfg *serverConfig) { cfg.defaultTimezone = tz }
}
// WithSchedulerEnabled controls whether the scheduler goroutine is started.
// Defaults to true. Set to false for worker-only instances.
func WithSchedulerEnabled(enabled bool) ServerOption {
return func(cfg *serverConfig) { cfg.schedulerEnabled = enabled }
}
// WithSchedulerPollInterval sets the poll interval for the scheduler engine.
// Defaults to 1s.
func WithSchedulerPollInterval(d time.Duration) ServerOption {
return func(cfg *serverConfig) {
if d > 0 {
cfg.schedulerPollInterval = d
}
}
}
// WithLogLevel sets the log level for the auto-created logger.
// Only takes effect if no WithLogger() is provided.
// Valid values: "debug", "info", "warn", "error".
func WithLogLevel(level string) ServerOption {
return func(cfg *serverConfig) { cfg.logLevel = level }
}
// WithAPI enables the HTTP monitoring API on the given address.
func WithAPI(enabled bool, addr string) ServerOption {
return func(cfg *serverConfig) {
cfg.apiEnabled = enabled
if addr != "" {
cfg.apiAddr = addr
}
}
}
// WithDashboard enables the web dashboard.
func WithDashboard(enabled bool) ServerOption {
return func(cfg *serverConfig) { cfg.dashEnabled = enabled }
}
// WithDashboardDir sets a custom directory for serving dashboard assets.
func WithDashboardDir(dir string) ServerOption {
return func(cfg *serverConfig) { cfg.dashCustomDir = dir }
}
// WithDashboardPathPrefix sets the URL prefix for the dashboard.
func WithDashboardPathPrefix(prefix string) ServerOption {
return func(cfg *serverConfig) { cfg.dashPathPrefix = prefix }
}
// WithAuthEnabled enables authentication for the monitoring API.
func WithAuthEnabled(enabled bool) ServerOption {
return func(cfg *serverConfig) { cfg.authEnabled = enabled }
}
// WithTrustProxy allows the X-Forwarded-Proto header to decide whether a
// request counts as HTTPS when marking the session cookie Secure.
//
// The header comes from the client, so this is only safe when a proxy in front
// of the server sets it and strips any incoming value. Leaving it off is not
// dangerous on its own — see WithCookieSecure for the case that is.
func WithTrustProxy(trust bool) ServerOption {
return func(cfg *serverConfig) { cfg.trustProxy = trust }
}
// WithCookieSecure marks the session cookie Secure regardless of how the
// connection looks to the server.
//
// Use this behind a TLS terminating proxy. Without it the server sees plain
// HTTP on the proxy-to-app hop, issues the cookie without Secure and with only
// SameSite=Lax, and the browser will then send the session token over plain
// HTTP. Stating the deployment fact is more reliable than inferring it from a
// header the proxy may not send.
func WithCookieSecure(secure bool) ServerOption {
return func(cfg *serverConfig) { cfg.cookieSecure = secure }
}
// WithAuthUsers sets the users for the monitoring API.
func WithAuthUsers(users []AuthUser) ServerOption {
return func(cfg *serverConfig) { cfg.authUsers = users }
}
// WithAPIKeys sets the API keys for programmatic access.
func WithAPIKeys(keys []AuthAPIKey) ServerOption {
return func(cfg *serverConfig) { cfg.apiKeys = keys }
}
// Server manages worker pools and processes jobs.
type Server struct {
cfg *serverConfig
rc *RedisClient
scripts *scriptRegistry
handlers map[string]Handler
handlerConfigs map[string]*handlerConfig
middlewares []MiddlewareFunc
pools []*pool
logger *slog.Logger
// jobTypePool tracks which pool each job type is assigned to.
// Key: job type, Value: pool name.
// Populated by Pool() (explicit) and Handle() with Workers() (implicit).
jobTypePool map[string]string
// poolNames tracks registered pool names to detect duplicates.
poolNames map[string]bool
// cronEntries holds registered cron entries indexed by ID.
cronEntries map[string]*CronEntry
// cronMu protects concurrent access to CronEntry fields (Enabled, UpdatedAt)
// between the scheduler goroutine (evalCron reader) and HTTP handlers
// (setCronEnabled writer). The map itself is only mutated pre-Start via
// Schedule(), so only field-level access needs synchronization at runtime.
cronMu sync.Mutex
// mon holds the HTTP monitoring server (nil if API disabled).
mon *monitor.Monitor
// serverID uniquely identifies this server instance (hostname:pid).
serverID string
startedAt time.Time
mu sync.RWMutex
running bool
stopCh chan struct{}
stopOnce sync.Once
}
// NewServer creates a new Server with the given options.
func NewServer(opts ...ServerOption) (*Server, error) {
cfg := &serverConfig{
globalTimeout: defaultGlobalTimeout,
gracePeriod: defaultGracePeriod,
shutdownTimeout: defaultShutdownTimeout,
schedulerEnabled: true,
schedulerPollInterval: defaultSchedulerPollInterval,
authSessionTTL: 86400, // 24 hours
}
for _, opt := range opts {
opt(cfg)
}
if cfg.logger == nil {
cfg.logger = newLoggerFromLevel(cfg.logLevel)
}
rc, err := NewRedisClient(cfg.redisOpts...)
if err != nil {
return nil, fmt.Errorf("creating server redis client: %w", err)
}
sr := newScriptRegistry()
if err := sr.load(); err != nil {
rc.Close()
return nil, fmt.Errorf("loading lua scripts: %w", err)
}
hostname, _ := os.Hostname()
serverID := fmt.Sprintf("%s-%d", hostname, os.Getpid())
s := &Server{
cfg: cfg,
rc: rc,
scripts: sr,
handlers: make(map[string]Handler),
handlerConfigs: make(map[string]*handlerConfig),
jobTypePool: make(map[string]string),
poolNames: make(map[string]bool),
cronEntries: make(map[string]*CronEntry),
logger: cfg.logger,
serverID: serverID,
stopCh: make(chan struct{}),
}
// Initialize monitor if API is enabled
if cfg.apiEnabled {
monCfg := monitor.Config{
APIAddr: cfg.apiAddr,
AuthEnabled: cfg.authEnabled,
AuthSessionTTL: cfg.authSessionTTL,
DashEnabled: cfg.dashEnabled,
DashPathPrefix: cfg.dashPathPrefix,
DashCustomDir: cfg.dashCustomDir,
RateLimit: cfg.apiRateLimit,
TrustProxy: cfg.trustProxy,
CookieSecure: cfg.cookieSecure,
}
for _, u := range cfg.authUsers {
monCfg.AuthUsers = append(monCfg.AuthUsers, monitor.AuthUser{
Username: u.Username,
PasswordHash: u.PasswordHash,
Role: u.Role,
})
}
for _, k := range cfg.apiKeys {
monCfg.APIKeys = append(monCfg.APIKeys, monitor.AuthAPIKey{
Name: k.Name,
Key: k.Key,
Role: k.Role,
})
}
s.mon = monitor.New(rc.Unwrap(), rc.Prefix(), cfg.logger, monCfg, s)
}
return s, nil
}
// Handle registers a handler for a job type.
func (s *Server) Handle(jobType string, handler Handler, opts ...HandleOption) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.running {
return fmt.Errorf("cannot register handler while server is running")
}
if _, exists := s.handlers[jobType]; exists {
return fmt.Errorf("%w: %s", ErrDuplicateHandler, jobType)
}
hcfg := &handlerConfig{}
for _, opt := range opts {
opt(hcfg)
}
s.handlers[jobType] = handler
s.handlerConfigs[jobType] = hcfg
if hcfg.workers > 0 {
// Check for conflict: job type already assigned to another pool
if existingPool, ok := s.jobTypePool[jobType]; ok {
return fmt.Errorf("%w: job type %q already assigned to pool %q, cannot create implicit pool",
ErrJobTypeConflict, jobType, existingPool)
}
// Implicit pool: dedicated pool + queue per job type
pcfg := newDefaultPoolConfig(jobType, jobType)
pcfg.concurrency = hcfg.workers
pcfg.gracePeriod = s.cfg.gracePeriod
s.pools = append(s.pools, newPool(pcfg, s))
s.jobTypePool[jobType] = jobType
s.poolNames[jobType] = true
}
return nil
}
// Use registers middleware that wraps all handlers. Middleware is applied
// in the order registered: Use(a, b) executes as a → b → handler.
//
// Must be called before Start(). Middleware is applied to handlers once
// at startup, so Use() and Handle() can be called in any order.
func (s *Server) Use(mws ...MiddlewareFunc) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.running {
return fmt.Errorf("cannot register middleware while server is running")
}
for i, mw := range mws {
if mw == nil {
return fmt.Errorf("middleware at index %d is nil", i)
}
}
s.middlewares = append(s.middlewares, mws...)
return nil
}
// Pool registers an explicit worker pool configuration (Layer 3).
// This allows grouping multiple job types into a single pool with shared
// concurrency, custom queues, dequeue strategy, and retry policy.
//
// Must be called before Start(). Returns an error if the pool name is
// already registered or if the configuration is invalid.
func (s *Server) Pool(cfg PoolConfig) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.running {
return fmt.Errorf("cannot register pool while server is running")
}
if cfg.Name == "" {
return fmt.Errorf("pool name must not be empty")
}
if s.poolNames[cfg.Name] {
return fmt.Errorf("%w: %s", ErrDuplicatePool, cfg.Name)
}
// Check for conflict: job type already assigned to another pool
for _, jt := range cfg.JobTypes {
if existingPool, ok := s.jobTypePool[jt]; ok {
return fmt.Errorf("%w: job type %q already assigned to pool %q, cannot assign to %q",
ErrJobTypeConflict, jt, existingPool, cfg.Name)
}
}
pcfg := cfg.toInternal(s.cfg.gracePeriod)
s.pools = append(s.pools, newPool(pcfg, s))
s.poolNames[cfg.Name] = true
// Track job type → pool assignments
for _, jt := range cfg.JobTypes {
s.jobTypePool[jt] = cfg.Name
}
return nil
}
// Schedule registers a cron entry for recurring job scheduling.
// Must be called before Start(). The entry's cron expression is parsed and
// validated. Duplicate IDs are rejected with ErrDuplicateCronEntry.
func (s *Server) Schedule(entry CronEntry) error {
s.mu.Lock()
defer s.mu.Unlock()
if s.running {
return fmt.Errorf("cannot register cron entry while server is running")
}
entry.applyDefaults()
if err := entry.validate(); err != nil {
return err
}
if _, exists := s.cronEntries[entry.ID]; exists {
return fmt.Errorf("%w: %s", ErrDuplicateCronEntry, entry.ID)
}
now := time.Now()
entry.CreatedAt = now.Unix()
entry.UpdatedAt = now.Unix()
s.cronEntries[entry.ID] = &entry
return nil
}
// Start begins processing jobs. It blocks until the server is stopped
// via signal (SIGTERM/SIGINT) or the context is cancelled.
// The server is single-use: after Stop or Start returns, create a new Server
// instance instead of calling Start again.
func (s *Server) Start(ctx context.Context) error {
s.mu.Lock()
if s.running {
s.mu.Unlock()
return fmt.Errorf("server already running")
}
s.running = true
s.mu.Unlock()
if err := s.rc.Ping(ctx); err != nil {
return fmt.Errorf("redis connection check: %w", err)
}
// The acknowledgement is checked here rather than inside warnIfUnprotected,
// so that function stays independently testable.
if !unprotectedRedisAcknowledged.Load() {
s.rc.warnIfUnprotected(os.Stderr, s.cfg.authEnabled)
}
// Ensure a default pool exists for handlers without Workers()
s.ensureDefaultPool()
// Apply middleware chain to all handlers (once).
// Wrap in reverse order so Use(a, b) executes as a → b → handler.
// Clear middlewares after application to prevent double-wrapping.
if len(s.middlewares) > 0 {
for jobType, handler := range s.handlers {
wrapped := handler
for i := len(s.middlewares) - 1; i >= 0; i-- {
wrapped = s.middlewares[i](wrapped)
}
s.handlers[jobType] = wrapped
}
s.middlewares = nil
}
// Save cron entries to Redis
if err := s.saveCronEntries(ctx); err != nil {
return fmt.Errorf("saving cron entries: %w", err)
}
s.startedAt = time.Now()
s.logger.Info("server starting",
"server_id", s.serverID,
"pools", len(s.pools),
"handlers", len(s.handlers),
"cron_entries", len(s.cronEntries),
)
// Create a cancellable context for all pools
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// Signal handling
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
defer signal.Stop(sigCh)
// Start all pools and optionally the scheduler
var wg sync.WaitGroup
// Scheduler engine (handles retry/delayed jobs + cron evaluation)
if s.cfg.schedulerEnabled {
sched := newSchedulerEngine(s)
wg.Add(1)
go func() {
defer wg.Done()
sched.run(ctx)
}()
}
// Server-level heartbeat
wg.Add(1)
go func() {
defer wg.Done()
s.serverHeartbeatLoop(ctx)
}()
for _, p := range s.pools {
wg.Add(1)
go func(p *pool) {
defer wg.Done()
p.run(ctx)
}(p)
}
// Start HTTP monitor server if configured
if s.mon != nil {
go func() {
if err := s.mon.Start(); err != nil {
s.logger.Error("monitor server error", "error", err)
}
}()
}
// Wait for shutdown signal, context cancellation, or Stop() call.
select {
case sig := <-sigCh:
s.logger.Info("received signal, initiating shutdown", "signal", sig)
cancel()
case <-ctx.Done():
s.logger.Info("context cancelled, initiating shutdown")
case <-s.stopCh:
s.logger.Info("Stop() called, initiating shutdown")
cancel()
}
// Graceful shutdown: wait for pools to finish within timeout
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
s.logger.Info("all pools stopped gracefully")
case <-time.After(s.cfg.shutdownTimeout):
s.logger.Warn("shutdown timeout reached, waiting for goroutines before closing Redis",
"timeout", s.cfg.shutdownTimeout)
// Wait for goroutines to finish even after timeout, so they don't
// use a closed Redis connection. Hard limit capped at 10s to avoid
// doubling the configured shutdown timeout.
hardLimit := s.cfg.shutdownTimeout
if hardLimit > 10*time.Second {
hardLimit = 10 * time.Second
}
hardTimeout := time.NewTimer(hardLimit)
select {
case <-done:
hardTimeout.Stop()
case <-hardTimeout.C:
s.logger.Error("hard shutdown timeout reached, closing Redis with goroutines still running")
}
}
// Stop monitor server before closing Redis
if s.mon != nil {
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := s.mon.Stop(shutdownCtx); err != nil {
s.logger.Error("monitor server shutdown error", "error", err)
}
shutdownCancel()
}
// Remove server registration from Redis
s.cleanupServerHeartbeat()
s.mu.Lock()
s.running = false
s.mu.Unlock()
return s.rc.Close()
}
// Stop signals the server to stop.
func (s *Server) Stop() {
s.stopOnce.Do(func() {
close(s.stopCh)
})
}
// defaultServerHeartbeatInterval is the interval between server-level heartbeat updates.
const defaultServerHeartbeatInterval = 10 * time.Second
// serverHeartbeatLoop periodically updates the server registration in Redis.
func (s *Server) serverHeartbeatLoop(ctx context.Context) {
s.sendServerHeartbeat(ctx)
ticker := time.NewTicker(defaultServerHeartbeatInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
s.sendServerHeartbeat(ctx)
}
}
}
// sendServerHeartbeat writes server info to Redis.
func (s *Server) sendServerHeartbeat(ctx context.Context) {
serverKey := s.rc.Key("server", s.serverID)
poolNames := make([]string, 0, len(s.pools))
for _, p := range s.pools {
poolNames = append(poolNames, p.cfg.name)
}
pipe := s.rc.rdb.Pipeline()
pipe.HSet(ctx, serverKey,
"id", s.serverID,
"hostname", s.hostname(),
"pid", os.Getpid(),
"go_version", runtime.Version(),
"started_at", s.startedAt.Unix(),
"last_heartbeat", time.Now().Unix(),
"pools", strings.Join(poolNames, ","),
"num_pools", len(s.pools),
"concurrency_total", s.totalConcurrency(),
"status", "active",
)
pipe.SAdd(ctx, s.rc.Key("servers"), s.serverID)
pipe.Expire(ctx, serverKey, 3*defaultServerHeartbeatInterval)
if _, err := pipe.Exec(ctx); err != nil {
if ctx.Err() == nil {
s.logger.Error("server heartbeat update failed", "error", err)
}
}
}
// cleanupServerHeartbeat removes the server registration on graceful shutdown.
func (s *Server) cleanupServerHeartbeat() {
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
pipe := s.rc.rdb.Pipeline()
pipe.Del(ctx, s.rc.Key("server", s.serverID))
pipe.SRem(ctx, s.rc.Key("servers"), s.serverID)
if _, err := pipe.Exec(ctx); err != nil {
s.logger.Error("server heartbeat cleanup failed", "error", err)
}
}
// hostname returns the hostname or "unknown" on error.
func (s *Server) hostname() string {
h, err := os.Hostname()
if err != nil {
return "unknown"
}
return h
}
// totalConcurrency returns the sum of concurrency across all pools.
func (s *Server) totalConcurrency() int {
total := 0
for _, p := range s.pools {
total += p.cfg.concurrency
}
return total
}
// ensureDefaultPool creates a "default" pool for handlers not assigned to any pool.
// If a catch-all pool (job_types: ["*"]) is configured, unassigned handlers are
// routed to that pool instead of creating a new default pool.
func (s *Server) ensureDefaultPool() {
var unassigned []string
for jt := range s.handlers {
if _, assigned := s.jobTypePool[jt]; !assigned {
unassigned = append(unassigned, jt)
}
}
if len(unassigned) == 0 {
return
}
// If a catch-all pool is configured, assign unassigned handlers to it.
if s.cfg.catchAllPool != "" {
for _, jt := range unassigned {
s.jobTypePool[jt] = s.cfg.catchAllPool
}
s.logger.Info("assigned unassigned handlers to catch-all pool",
"pool", s.cfg.catchAllPool,
"handlers", unassigned,
)
return
}
pcfg := newDefaultPoolConfig("default", "default")
pcfg.gracePeriod = s.cfg.gracePeriod
s.pools = append(s.pools, newPool(pcfg, s))
s.logger.Info("created default pool",
"concurrency", pcfg.concurrency,
"unassigned_handlers", unassigned,
)
}
// saveCronEntries persists registered cron entries to Redis.
func (s *Server) saveCronEntries(ctx context.Context) error {
if len(s.cronEntries) == 0 {
return nil
}
entriesKey := s.rc.Key("cron", "entries")
pipe := s.rc.rdb.Pipeline()
for id, entry := range s.cronEntries {
data, err := json.Marshal(entry)
if err != nil {
return fmt.Errorf("marshaling cron entry %q: %w", id, err)
}
pipe.HSet(ctx, entriesKey, id, string(data))
}
if _, err := pipe.Exec(ctx); err != nil {
return fmt.Errorf("writing cron entries to redis: %w", err)
}
s.logger.Info("cron entries saved to redis", "count", len(s.cronEntries))
return nil
}
// resolveTimeout returns the effective timeout for a job.
func (s *Server) resolveTimeout(job *Job, pcfg *poolConfig) time.Duration {
// Job-level timeout (highest priority)
if job.Timeout > 0 {
return time.Duration(job.Timeout) * time.Second
}
// Pool-level timeout
if pcfg.jobTimeout > 0 {
return pcfg.jobTimeout
}
// Global default (always set, cannot be disabled)
return s.cfg.globalTimeout
}