fix(runtime): initialize resources per service

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
 

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{}

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")
+	}
+}

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)
 }

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
+	}
+}

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")
+	}
+}

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)
 }

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()
 }

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()
 }

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()
 }

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
 }

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

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)
 }

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
 	}
 }

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)
+	}
+}