From 6b0c2f5a4ca13a8ea2f35cb906db8005f3ed768f Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Sun, 28 Jun 2026 18:03:16 +0700 Subject: [PATCH] fix(runtime): initialize resources per service --- README.md | 23 +++----- adapter/adapter.go | 37 +++++++++++++ adapter/config_test.go | 53 ++++++++++++++++++ adapter/couchbase.go | 51 ++++++++++++++---- adapter/couchbase_initialization.go | 77 ++++++++++++++++++++++++++ adapter/couchbase_test.go | 45 ++++++++++++++++ adapter/mongodb.go | 19 ++++++- adapter/mysql.go | 26 ++++++++- adapter/postgresql.go | 28 +++++++++- adapter/redis.go | 13 ++++- adapter/valkey.go | 9 +++- config.example.yml | 3 ++ main.go | 24 +-------- runner.go | 84 +++++++++++++++++++++-------- runner_test.go | 81 ++++++++++++++++++++++++++++ 15 files changed, 501 insertions(+), 72 deletions(-) create mode 100644 adapter/couchbase_initialization.go create mode 100644 adapter/couchbase_test.go create mode 100644 runner_test.go diff --git a/README.md b/README.md index 05a6add..22ef53f 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ services: | `postgresql` | `github.com/lib/pq` | `url` | | `mysql` | `github.com/go-sql-driver/mysql` | `dsn` | | `mongodb` | `go.mongodb.org/mongo-driver/v2` | `uri`, `database`, `collection` | -| `couchbase` | `github.com/couchbase/gocb/v2` | `connection_string`, `username`, `password`, `bucket_name`, `scope_name`, `collection_name` | +| `couchbase` | `github.com/couchbase/gocb/v2` | `connection_string`, `username`, `password`, `bucket_name`, `scope_name`, `collection_name`, optional `ready_timeout`, optional `bucket_ram_quota_mb` | ## Quick start (Compose) @@ -86,22 +86,15 @@ go run . ## How it works -On every tick the chosen adapter performs the cheapest write that proves the cluster is alive. `counter_key` selects the key/doc ID and defaults to `counter`. +On startup each adapter initializes the minimum resource it owns, then every tick performs the cheapest write that proves the cluster is alive. `counter_key` selects the key/doc ID and defaults to `counter`. -- **Redis/Valkey** — `INCR key` -- **PostgreSQL** — `UPDATE keepalive SET value = value + 1 WHERE key = $1 RETURNING value` -- **MySQL** — `UPDATE` + `SELECT` by key inside a transaction -- **MongoDB** — `FindOneAndUpdate({_id: key}, {$inc: {count: 1}}, upsert)` -- **Couchbase** — `GET key` -> `++` -> `UPSERT key` +- **Redis/Valkey** — initialize with `SETNX key 0`, then `INCR key` +- **PostgreSQL** — `CREATE TABLE IF NOT EXISTS keepalive`, seed `key`, then `UPDATE ... RETURNING` +- **MySQL** — `CREATE TABLE IF NOT EXISTS keepalive`, seed `key`, then `UPDATE` + `SELECT` +- **MongoDB** — upsert `{_id: key, count: 0}` on connect, then `FindOneAndUpdate({_id: key}, {$inc: {count: 1}}, upsert)` +- **Couchbase** — optionally create the bucket when `bucket_ram_quota_mb` is set, create configured scope/collection when missing, insert `key = 0` if missing, then `GET key` -> `++` -> `UPSERT key` -The PostgreSQL and MySQL adapters expect a table: - -```sql -CREATE TABLE keepalive (key TEXT PRIMARY KEY, value BIGINT NOT NULL DEFAULT 0); -INSERT INTO keepalive (key, value) VALUES ('counter', 0); -``` - -Seed the value with your configured `counter_key` when it is not `counter`. MySQL uses backticked identifiers — see `adapter/mysql.go`. +Each configured service starts independently. If one service cannot connect, it logs the error and retries without stopping other services in the same deployment. ## Adding a new adapter diff --git a/adapter/adapter.go b/adapter/adapter.go index 677c0bf..a392839 100644 --- a/adapter/adapter.go +++ b/adapter/adapter.go @@ -9,6 +9,9 @@ package adapter import ( "context" "fmt" + "strconv" + "strings" + "time" ) type Adapter interface { @@ -34,6 +37,40 @@ func (c Config) Optional(name, def string) string { return def } +func (c Config) OptionalDuration(name string, def time.Duration) (time.Duration, error) { + value := strings.TrimSpace(c[name]) + if value == "" { + return def, nil + } + if d, err := time.ParseDuration(value); err == nil { + if d <= 0 { + return 0, fmt.Errorf("config %s must be greater than zero", name) + } + return d, nil + } + seconds, err := strconv.Atoi(value) + if err != nil { + return 0, fmt.Errorf("config %s must be a duration like 30s or an integer number of seconds", name) + } + d := time.Duration(seconds) * time.Second + if d <= 0 { + return 0, fmt.Errorf("config %s must be greater than zero", name) + } + return d, nil +} + +func (c Config) OptionalUint64(name string, def uint64) (uint64, error) { + value := strings.TrimSpace(c[name]) + if value == "" { + return def, nil + } + out, err := strconv.ParseUint(value, 10, 64) + if err != nil { + return 0, fmt.Errorf("config %s must be an unsigned integer", name) + } + return out, nil +} + type Factory func(Config) (Adapter, error) var Registry = map[string]Factory{} diff --git a/adapter/config_test.go b/adapter/config_test.go index c7c195a..ecd277d 100644 --- a/adapter/config_test.go +++ b/adapter/config_test.go @@ -3,6 +3,7 @@ package adapter import ( "strings" "testing" + "time" ) func TestConfigRequiredReturnsConfiguredValue(t *testing.T) { @@ -46,3 +47,55 @@ func TestConfigOptionalUsesDefault(t *testing.T) { t.Fatalf("Optional() = %q, want counter", got) } } + +func TestConfigOptionalDurationParsesDuration(t *testing.T) { + cfg := Config{"ready_timeout": "45s"} + + got, err := cfg.OptionalDuration("ready_timeout", time.Second) + if err != nil { + t.Fatalf("OptionalDuration returned error: %v", err) + } + if got != 45*time.Second { + t.Fatalf("OptionalDuration() = %s, want 45s", got) + } +} + +func TestConfigOptionalDurationParsesSeconds(t *testing.T) { + cfg := Config{"ready_timeout": "30"} + + got, err := cfg.OptionalDuration("ready_timeout", time.Second) + if err != nil { + t.Fatalf("OptionalDuration returned error: %v", err) + } + if got != 30*time.Second { + t.Fatalf("OptionalDuration() = %s, want 30s", got) + } +} + +func TestConfigOptionalDurationRejectsNonPositive(t *testing.T) { + cfg := Config{"ready_timeout": "0s"} + + if _, err := cfg.OptionalDuration("ready_timeout", time.Second); err == nil { + t.Fatal("OptionalDuration returned nil error") + } +} + +func TestConfigOptionalUint64ParsesValue(t *testing.T) { + cfg := Config{"bucket_ram_quota_mb": "128"} + + got, err := cfg.OptionalUint64("bucket_ram_quota_mb", 0) + if err != nil { + t.Fatalf("OptionalUint64 returned error: %v", err) + } + if got != 128 { + t.Fatalf("OptionalUint64() = %d, want 128", got) + } +} + +func TestConfigOptionalUint64RejectsInvalidValue(t *testing.T) { + cfg := Config{"bucket_ram_quota_mb": "-1"} + + if _, err := cfg.OptionalUint64("bucket_ram_quota_mb", 0); err == nil { + t.Fatal("OptionalUint64 returned nil error") + } +} diff --git a/adapter/couchbase.go b/adapter/couchbase.go index e354eab..601af66 100644 --- a/adapter/couchbase.go +++ b/adapter/couchbase.go @@ -33,14 +33,24 @@ func init() { if err != nil { return nil, err } + readyTimeout, err := cfg.OptionalDuration("ready_timeout", defaultCouchbaseReadyTimeout) + if err != nil { + return nil, err + } + bucketRAMQuotaMB, err := cfg.OptionalUint64("bucket_ram_quota_mb", 0) + if err != nil { + return nil, err + } return &couchbaseAdapter{ - conn: conn, - user: user, - pass: pass, - bucket: bucket, - scope: scope, - collName: collName, - docID: cfg.Optional("counter_key", "counter"), + conn: conn, + user: user, + pass: pass, + bucket: bucket, + scope: scope, + collName: collName, + docID: cfg.Optional("counter_key", "counter"), + readyTimeout: readyTimeout, + bucketRAMQuotaMB: bucketRAMQuotaMB, }, nil } } @@ -55,9 +65,12 @@ type couchbaseAdapter struct { scope string collName string docID string + + readyTimeout time.Duration + bucketRAMQuotaMB uint64 } -func (a *couchbaseAdapter) Connect(_ context.Context) error { +func (a *couchbaseAdapter) Connect(ctx context.Context) error { opts := gocb.ClusterOptions{ Authenticator: gocb.PasswordAuthenticator{Username: a.user, Password: a.pass}, } @@ -68,12 +81,29 @@ func (a *couchbaseAdapter) Connect(_ context.Context) error { if err != nil { return err } + connected := false + defer func() { + if !connected { + cluster.Close(nil) + } + }() + + if err := a.ensureBucket(ctx, cluster); err != nil { + return err + } b := cluster.Bucket(a.bucket) - if err := b.WaitUntilReady(5*time.Second, nil); err != nil { + if err := b.WaitUntilReady(a.readyTimeout, &gocb.WaitUntilReadyOptions{Context: ctx}); err != nil { + return err + } + if err := a.ensureScopeAndCollection(ctx, b); err != nil { return err } a.cluster = cluster a.coll = b.Scope(a.scope).Collection(a.collName) + if err := a.ensureDocument(ctx); err != nil { + return err + } + connected = true return nil } @@ -98,5 +128,8 @@ func (a *couchbaseAdapter) Increment(_ context.Context) (int64, error) { } func (a *couchbaseAdapter) Close(_ context.Context) error { + if a.cluster == nil { + return nil + } return a.cluster.Close(nil) } diff --git a/adapter/couchbase_initialization.go b/adapter/couchbase_initialization.go new file mode 100644 index 0000000..66557c1 --- /dev/null +++ b/adapter/couchbase_initialization.go @@ -0,0 +1,77 @@ +package adapter + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/couchbase/gocb/v2" +) + +const defaultCouchbaseReadyTimeout = 30 * time.Second + +func (a *couchbaseAdapter) ensureBucket(ctx context.Context, cluster *gocb.Cluster) error { + if a.bucketRAMQuotaMB == 0 { + return nil + } + err := cluster.Buckets().CreateBucket(gocb.CreateBucketSettings{ + BucketSettings: gocb.BucketSettings{ + Name: a.bucket, + RAMQuotaMB: a.bucketRAMQuotaMB, + BucketType: gocb.CouchbaseBucketType, + }, + }, &gocb.CreateBucketOptions{Context: ctx, Timeout: a.readyTimeout}) + if err != nil && !errors.Is(err, gocb.ErrBucketExists) { + return fmt.Errorf("create bucket %q: %w", a.bucket, err) + } + return nil +} + +func (a *couchbaseAdapter) ensureScopeAndCollection(ctx context.Context, bucket *gocb.Bucket) error { + manager := bucket.Collections() + if a.scope != "_default" { + err := manager.CreateScope(a.scope, &gocb.CreateScopeOptions{Context: ctx, Timeout: a.readyTimeout}) + if err != nil && !errors.Is(err, gocb.ErrScopeExists) { + return fmt.Errorf("create scope %q: %w", a.scope, err) + } + } + if a.scope == "_default" && a.collName == "_default" { + return nil + } + err := manager.CreateCollection(gocb.CollectionSpec{Name: a.collName, ScopeName: a.scope}, &gocb.CreateCollectionOptions{Context: ctx, Timeout: a.readyTimeout}) + if err != nil && !errors.Is(err, gocb.ErrCollectionExists) { + return fmt.Errorf("create collection %q.%q: %w", a.scope, a.collName, err) + } + return nil +} + +func (a *couchbaseAdapter) ensureDocument(ctx context.Context) error { + deadline := time.Now().Add(a.readyTimeout) + for { + _, err := a.coll.Insert(a.docID, uint64(0), &gocb.InsertOptions{Context: ctx, Timeout: a.readyTimeout}) + if err == nil || errors.Is(err, gocb.ErrDocumentExists) { + return nil + } + if !errors.Is(err, gocb.ErrScopeNotFound) && !errors.Is(err, gocb.ErrCollectionNotFound) { + return err + } + if time.Now().After(deadline) { + return fmt.Errorf("wait for collection %q.%q: %w", a.scope, a.collName, err) + } + if !sleepContext(ctx, time.Second) { + return ctx.Err() + } + } +} + +func sleepContext(ctx context.Context, d time.Duration) bool { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/adapter/couchbase_test.go b/adapter/couchbase_test.go new file mode 100644 index 0000000..cc8dc91 --- /dev/null +++ b/adapter/couchbase_test.go @@ -0,0 +1,45 @@ +package adapter + +import ( + "testing" + "time" +) + +func TestCouchbaseFactoryParsesInitializationOptions(t *testing.T) { + a, err := New("couchbase", Config{ + "connection_string": "couchbases://cb.example.com", + "username": "user", + "password": "pass", + "bucket_name": "keepalive", + "scope_name": "scope", + "collection_name": "collection", + "ready_timeout": "45s", + "bucket_ram_quota_mb": "128", + }) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + + got := a.(*couchbaseAdapter) + if got.readyTimeout != 45*time.Second { + t.Fatalf("readyTimeout = %s, want 45s", got.readyTimeout) + } + if got.bucketRAMQuotaMB != 128 { + t.Fatalf("bucketRAMQuotaMB = %d, want 128", got.bucketRAMQuotaMB) + } +} + +func TestCouchbaseFactoryRejectsInvalidReadyTimeout(t *testing.T) { + _, err := New("couchbase", Config{ + "connection_string": "couchbases://cb.example.com", + "username": "user", + "password": "pass", + "bucket_name": "keepalive", + "scope_name": "scope", + "collection_name": "collection", + "ready_timeout": "0s", + }) + if err == nil { + t.Fatal("New returned nil error") + } +} diff --git a/adapter/mongodb.go b/adapter/mongodb.go index 0edc4c9..ed8db05 100644 --- a/adapter/mongodb.go +++ b/adapter/mongodb.go @@ -48,7 +48,21 @@ func (a *mongoAdapter) Connect(ctx context.Context) error { } a.client = client a.coll = client.Database(a.dbName).Collection(a.collName) - return client.Ping(ctx, nil) + if err := client.Ping(ctx, nil); err != nil { + client.Disconnect(ctx) + return err + } + if err := a.ensureInitialized(ctx); err != nil { + client.Disconnect(ctx) + return err + } + return nil +} + +func (a *mongoAdapter) ensureInitialized(ctx context.Context) error { + update := bson.M{"$setOnInsert": bson.M{"count": int64(0)}} + _, err := a.coll.UpdateOne(ctx, bson.M{"_id": a.docID}, update, options.UpdateOne().SetUpsert(true)) + return err } func (a *mongoAdapter) Increment(ctx context.Context) (int64, error) { @@ -65,5 +79,8 @@ func (a *mongoAdapter) Increment(ctx context.Context) (int64, error) { } func (a *mongoAdapter) Close(ctx context.Context) error { + if a.client == nil { + return nil + } return a.client.Disconnect(ctx) } diff --git a/adapter/mysql.go b/adapter/mysql.go index 42c938b..94799cc 100644 --- a/adapter/mysql.go +++ b/adapter/mysql.go @@ -36,7 +36,28 @@ func (a *mysqlAdapter) Connect(ctx context.Context) error { db.SetMaxOpenConns(10) db.SetMaxIdleConns(10) a.db = db - return a.db.PingContext(ctx) + if err := a.db.PingContext(ctx); err != nil { + a.db.Close() + return err + } + if err := a.ensureInitialized(ctx); err != nil { + a.db.Close() + return err + } + return nil +} + +func (a *mysqlAdapter) ensureInitialized(ctx context.Context) error { + if _, err := a.db.ExecContext(ctx, + "CREATE TABLE IF NOT EXISTS `keepalive` (`key` VARCHAR(255) PRIMARY KEY, `value` BIGINT NOT NULL DEFAULT 0)", + ); err != nil { + return err + } + _, err := a.db.ExecContext(ctx, + "INSERT IGNORE INTO `keepalive` (`key`, `value`) VALUES (?, 0)", + a.key, + ) + return err } func (a *mysqlAdapter) Increment(ctx context.Context) (int64, error) { @@ -63,5 +84,8 @@ func (a *mysqlAdapter) Increment(ctx context.Context) (int64, error) { } func (a *mysqlAdapter) Close(_ context.Context) error { + if a.db == nil { + return nil + } return a.db.Close() } diff --git a/adapter/postgresql.go b/adapter/postgresql.go index dec4c14..104ea19 100644 --- a/adapter/postgresql.go +++ b/adapter/postgresql.go @@ -33,7 +33,30 @@ func (a *postgresAdapter) Connect(ctx context.Context) error { return err } a.db = db - return a.db.PingContext(ctx) + if err := a.db.PingContext(ctx); err != nil { + a.db.Close() + return err + } + if err := a.ensureInitialized(ctx); err != nil { + a.db.Close() + return err + } + return nil +} + +func (a *postgresAdapter) ensureInitialized(ctx context.Context) error { + if _, err := a.db.ExecContext(ctx, ` +CREATE TABLE IF NOT EXISTS keepalive ( + key TEXT PRIMARY KEY, + value BIGINT NOT NULL DEFAULT 0 +)`); err != nil { + return err + } + _, err := a.db.ExecContext(ctx, + `INSERT INTO keepalive (key, value) VALUES ($1, 0) ON CONFLICT (key) DO NOTHING`, + a.key, + ) + return err } func (a *postgresAdapter) Increment(ctx context.Context) (int64, error) { @@ -53,5 +76,8 @@ func (a *postgresAdapter) Increment(ctx context.Context) (int64, error) { } func (a *postgresAdapter) Close(_ context.Context) error { + if a.db == nil { + return nil + } return a.db.Close() } diff --git a/adapter/redis.go b/adapter/redis.go index 030aee3..6a53de0 100644 --- a/adapter/redis.go +++ b/adapter/redis.go @@ -31,7 +31,15 @@ func (a *redisAdapter) Connect(ctx context.Context) error { return err } a.client = redis.NewClient(opt) - return a.client.Ping(ctx).Err() + if err := a.client.Ping(ctx).Err(); err != nil { + a.client.Close() + return err + } + if err := a.client.SetNX(ctx, a.key, 0, 0).Err(); err != nil { + a.client.Close() + return err + } + return nil } func (a *redisAdapter) Increment(ctx context.Context) (int64, error) { @@ -39,5 +47,8 @@ func (a *redisAdapter) Increment(ctx context.Context) (int64, error) { } func (a *redisAdapter) Close(_ context.Context) error { + if a.client == nil { + return nil + } return a.client.Close() } diff --git a/adapter/valkey.go b/adapter/valkey.go index d7133f0..99dec2c 100644 --- a/adapter/valkey.go +++ b/adapter/valkey.go @@ -25,7 +25,7 @@ type valkeyAdapter struct { key string } -func (a *valkeyAdapter) Connect(_ context.Context) error { +func (a *valkeyAdapter) Connect(ctx context.Context) error { opt, err := valkey.ParseURL(a.url) if err != nil { return err @@ -35,6 +35,10 @@ func (a *valkeyAdapter) Connect(_ context.Context) error { return err } a.client = client + if err := a.client.Do(ctx, a.client.B().Setnx().Key(a.key).Value("0").Build()).Error(); err != nil { + a.client.Close() + return err + } return nil } @@ -43,6 +47,9 @@ func (a *valkeyAdapter) Increment(ctx context.Context) (int64, error) { } func (a *valkeyAdapter) Close(_ context.Context) error { + if a.client == nil { + return nil + } a.client.Close() return nil } diff --git a/config.example.yml b/config.example.yml index 6066b96..b42c381 100644 --- a/config.example.yml +++ b/config.example.yml @@ -25,3 +25,6 @@ services: bucket_name: keepalive scope_name: _default collection_name: _default + ready_timeout: 30s + # Optional. Set only when this user can create a missing bucket. + # bucket_ram_quota_mb: 128 diff --git a/main.go b/main.go index a4dc2a9..676a120 100644 --- a/main.go +++ b/main.go @@ -7,8 +7,6 @@ import ( "os/signal" "sync" "syscall" - - "github.com/tiennm99/keepalive/adapter" ) func main() { @@ -23,26 +21,9 @@ func main() { } ctx, cancel := context.WithCancel(context.Background()) - running := make([]runningService, 0, len(services)) - for _, svcConfig := range services { - a, err := adapter.New(svcConfig.AdapterType, svcConfig.Config) - if err != nil { - cancel() - closeServices(running) - log.Fatalf("[%s] init adapter: %v", svcConfig.Name, err) - } - if err := a.Connect(ctx); err != nil { - cancel() - closeServices(running) - log.Fatalf("[%s] connect: %v", svcConfig.Name, err) - } - running = append(running, runningService{config: svcConfig, adapter: a}) - log.Printf("[%s] keepalive: %s every %s", svcConfig.Name, svcConfig.AdapterType, svcConfig.Interval) - } - var wg sync.WaitGroup - for _, svc := range running { - runService(ctx, &wg, svc) + for _, svcConfig := range services { + runService(ctx, &wg, svcConfig) } sigCh := make(chan os.Signal, 1) @@ -51,5 +32,4 @@ func main() { cancel() wg.Wait() - closeServices(running) } diff --git a/runner.go b/runner.go index dbd6b10..44f7070 100644 --- a/runner.go +++ b/runner.go @@ -9,43 +9,85 @@ import ( "github.com/tiennm99/keepalive/adapter" ) +var reconnectDelay = 10 * time.Second + type runningService struct { config serviceConfig adapter adapter.Adapter } -func runService(ctx context.Context, wg *sync.WaitGroup, svc runningService) { +func runService(ctx context.Context, wg *sync.WaitGroup, config serviceConfig) { wg.Add(1) go func() { defer wg.Done() - ticker := time.NewTicker(svc.config.Interval) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): + a, err := adapter.New(config.AdapterType, config.Config) + if err != nil { + log.Printf("[%s] init adapter: %v", config.Name, err) return - case <-ticker.C: - tickCtx, cancel := context.WithTimeout(ctx, 3*time.Second) - count, err := svc.adapter.Increment(tickCtx) - cancel() - if err != nil { - log.Printf("[%s] increment: %v", svc.config.Name, err) - continue - } - log.Printf("[%s] counter: %d", svc.config.Name, count) } + + if err := a.Connect(ctx); err != nil { + if ctx.Err() != nil { + closeService(ctx, config.Name, a) + return + } + log.Printf("[%s] connect: %v", config.Name, err) + closeService(ctx, config.Name, a) + if !waitContext(ctx, reconnectDelay) { + return + } + continue + } + + log.Printf("[%s] keepalive: %s every %s", config.Name, config.AdapterType, config.Interval) + runConnectedService(ctx, runningService{config: config, adapter: a}) + closeService(ctx, config.Name, a) + return } }() } -func closeServices(services []runningService) { - for _, svc := range services { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - if err := svc.adapter.Close(shutdownCtx); err != nil { - log.Printf("[%s] close: %v", svc.config.Name, err) +func runConnectedService(ctx context.Context, svc runningService) { + ticker := time.NewTicker(svc.config.Interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + tickCtx, cancel := context.WithTimeout(ctx, 3*time.Second) + count, err := svc.adapter.Increment(tickCtx) + cancel() + if err != nil { + log.Printf("[%s] increment: %v", svc.config.Name, err) + continue + } + log.Printf("[%s] counter: %d", svc.config.Name, count) } - cancel() + } +} + +func closeService(_ context.Context, name string, a adapter.Adapter) { + if a == nil { + return + } + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := a.Close(shutdownCtx); err != nil { + log.Printf("[%s] close: %v", name, err) + } +} + +func waitContext(ctx context.Context, d time.Duration) bool { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true } } diff --git a/runner_test.go b/runner_test.go new file mode 100644 index 0000000..164306b --- /dev/null +++ b/runner_test.go @@ -0,0 +1,81 @@ +package main + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/tiennm99/keepalive/adapter" +) + +type retryConnectAdapter struct { + connects *atomic.Int32 + connected chan<- struct{} + once *sync.Once +} + +func (a *retryConnectAdapter) Connect(context.Context) error { + if a.connects.Add(1) == 1 { + return errors.New("connect failed") + } + a.once.Do(func() { close(a.connected) }) + return nil +} + +func (a *retryConnectAdapter) Increment(context.Context) (int64, error) { + return 0, nil +} + +func (a *retryConnectAdapter) Close(context.Context) error { + return nil +} + +func TestRunServiceRetriesConnectFailure(t *testing.T) { + oldReconnectDelay := reconnectDelay + reconnectDelay = 10 * time.Millisecond + defer func() { reconnectDelay = oldReconnectDelay }() + + var connects atomic.Int32 + connected := make(chan struct{}) + var once sync.Once + adapter.Registry["retry-test"] = func(adapter.Config) (adapter.Adapter, error) { + return &retryConnectAdapter{connects: &connects, connected: connected, once: &once}, nil + } + defer delete(adapter.Registry, "retry-test") + + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + runService(ctx, &wg, serviceConfig{ + Name: "retry-test", + AdapterType: "retry-test", + Interval: time.Hour, + Config: adapter.Config{}, + }) + + select { + case <-connected: + case <-time.After(time.Second): + cancel() + t.Fatal("service did not retry and connect") + } + + cancel() + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("service did not stop after context cancellation") + } + + if got := connects.Load(); got != 2 { + t.Fatalf("connect attempts = %d, want 2", got) + } +}