TTL for discovery values

api/pkg/discovery/kv.go

@@ -3,14 +3,21 @@ package discovery
 import (
 	"encoding/json"
 	"errors"
+	"os"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/nats-io/nats.go"
 	"github.com/tech/sendico/pkg/mlogger"
 	"go.uber.org/zap"
 )
 
-const DefaultKVBucket = "discovery_registry"
+const (
+	DefaultKVBucket     = "discovery_registry"
+	kvTTLSecondsEnv     = "DISCOVERY_KV_TTL_SECONDS"
+	defaultKVTTLSeconds = 3600
+)
 
 type KVStore struct {
 	logger mlogger.Logger
@@ -29,6 +36,7 @@ func NewKVStore(logger mlogger.Logger, js nats.JetStreamContext, bucket string)
 	if bucket == "" {
 		bucket = DefaultKVBucket
 	}
+	ttl := kvTTL(logger)
 	kv, err := js.KeyValue(bucket)
 	if err != nil {
 		if errors.Is(err, nats.ErrBucketNotFound) {
@@ -36,14 +44,21 @@ func NewKVStore(logger mlogger.Logger, js nats.JetStreamContext, bucket string)
 				Bucket:      bucket,
 				Description: "service discovery registry",
 				History:     1,
+				TTL:         ttl,
 			})
 			if err == nil && logger != nil {
-				logger.Info("Discovery KV bucket created", zap.String("bucket", bucket))
+				fields := []zap.Field{zap.String("bucket", bucket)}
+				if ttl > 0 {
+					fields = append(fields, zap.Duration("ttl", ttl))
+				}
+				logger.Info("Discovery KV bucket created", fields...)
 			}
 		}
 		if err != nil {
 			return nil, err
 		}
+	} else if ttl > 0 {
+		ensureKVTTL(logger, js, kv, bucket, ttl)
 	}
 
 	return &KVStore{
@@ -53,6 +68,59 @@ func NewKVStore(logger mlogger.Logger, js nats.JetStreamContext, bucket string)
 	}, nil
 }
 
+func kvTTL(logger mlogger.Logger) time.Duration {
+	raw := strings.TrimSpace(os.Getenv(kvTTLSecondsEnv))
+	if raw == "" {
+		return time.Duration(defaultKVTTLSeconds) * time.Second
+	}
+	secs, err := strconv.Atoi(raw)
+	if err != nil {
+		if logger != nil {
+			logger.Warn("Invalid discovery KV TTL seconds, using default", zap.String("env", kvTTLSecondsEnv), zap.String("value", raw), zap.Int("default_seconds", defaultKVTTLSeconds))
+		}
+		return time.Duration(defaultKVTTLSeconds) * time.Second
+	}
+	if secs <= 0 {
+		return 0
+	}
+	return time.Duration(secs) * time.Second
+}
+
+func ensureKVTTL(logger mlogger.Logger, js nats.JetStreamContext, kv nats.KeyValue, bucket string, ttl time.Duration) {
+	if kv == nil || js == nil || ttl <= 0 {
+		return
+	}
+	status, err := kv.Status()
+	if err != nil {
+		if logger != nil {
+			logger.Warn("Failed to read discovery KV status", zap.String("bucket", bucket), zap.Error(err))
+		}
+		return
+	}
+	if status.TTL() == ttl {
+		return
+	}
+	stream := "KV_" + bucket
+	info, err := js.StreamInfo(stream)
+	if err != nil {
+		if logger != nil {
+			logger.Warn("Failed to read discovery KV stream info", zap.String("bucket", bucket), zap.String("stream", stream), zap.Error(err))
+		}
+		return
+	}
+	cfg := info.Config
+	cfg.MaxAge = ttl
+	if _, err := js.UpdateStream(&cfg); err != nil {
+		if logger != nil {
+			logger.Warn("Failed to update discovery KV TTL", zap.String("bucket", bucket), zap.Duration("ttl", ttl), zap.Error(err))
+		}
+		return
+	}
+	if logger != nil {
+		logger.Info("Discovery KV TTL updated", zap.String("bucket", bucket), zap.Duration("ttl", ttl))
+	}
+}
+
 func (s *KVStore) Put(entry RegistryEntry) error {
 	if s == nil || s.kv == nil {
 		return errors.New("discovery kv: not configured")