MonocoEditor-With-Lsp-Backend/backend/internal/cluster/redis_registry.go

package cluster

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

type RedisRegistryConfig struct {
	Addr              string
	Password          string
	DB                int
	KeyPrefix         string
	InstanceID        string
	InstanceEndpoint  string
	SessionTTL        time.Duration
	InstanceTTL       time.Duration
	HeartbeatInterval time.Duration
}

type RedisRegistry struct {
	client *redis.Client

	keyPrefix         string
	instanceID        string
	instanceEndpoint  string
	sessionTTL        time.Duration
	instanceTTL       time.Duration
	heartbeatInterval time.Duration

	stopCh   chan struct{}
	stopOnce sync.Once
}

var claimSessionScript = redis.NewScript(`
local sessionKey = KEYS[1]
local owner = ARGV[1]
local now = ARGV[2]
local ttl = tonumber(ARGV[3])

local existing = redis.call('HGET', sessionKey, 'owner')
if (not existing) or existing == owner then
  redis.call('HSET', sessionKey, 'owner', owner, 'updatedAt', now)
  redis.call('PEXPIRE', sessionKey, ttl)
  return owner
end

redis.call('PEXPIRE', sessionKey, ttl)
return existing
`)

var releaseSessionScript = redis.NewScript(`
local sessionKey = KEYS[1]
local owner = ARGV[1]
local existing = redis.call('HGET', sessionKey, 'owner')
if existing == owner then
  redis.call('DEL', sessionKey)
  return 1
end
return 0
`)

func NewRedisRegistry(ctx context.Context, cfg RedisRegistryConfig) (*RedisRegistry, error) {
	if strings.TrimSpace(cfg.Addr) == "" {
		return nil, errors.New("redis addr is required")
	}
	if strings.TrimSpace(cfg.KeyPrefix) == "" {
		cfg.KeyPrefix = "lsp-gateway"
	}
	if cfg.SessionTTL <= 0 {
		cfg.SessionTTL = 20 * time.Minute
	}
	if cfg.InstanceTTL <= 0 {
		cfg.InstanceTTL = 30 * time.Second
	}
	if cfg.HeartbeatInterval <= 0 {
		cfg.HeartbeatInterval = 10 * time.Second
	}
	if cfg.InstanceID == "" {
		return nil, errors.New("instance id is required")
	}

	client := redis.NewClient(&redis.Options{
		Addr:     cfg.Addr,
		Password: cfg.Password,
		DB:       cfg.DB,
	})

	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("redis ping failed: %w", err)
	}

	registry := &RedisRegistry{
		client:            client,
		keyPrefix:         cfg.KeyPrefix,
		instanceID:        cfg.InstanceID,
		instanceEndpoint:  cfg.InstanceEndpoint,
		sessionTTL:        cfg.SessionTTL,
		instanceTTL:       cfg.InstanceTTL,
		heartbeatInterval: cfg.HeartbeatInterval,
		stopCh:            make(chan struct{}),
	}
	if err := registry.refreshInstance(ctx); err != nil {
		_ = client.Close()
		return nil, err
	}
	go registry.heartbeatLoop()
	return registry, nil
}

func (r *RedisRegistry) ClaimSession(
	ctx context.Context,
	language string,
	sessionID string,
) (ownerID string, ownerEndpoint string, err error) {
	key := r.sessionKey(language, sessionID)
	now := time.Now().UTC().Format(time.RFC3339Nano)
	raw, err := claimSessionScript.Run(
		ctx,
		r.client,
		[]string{key},
		r.instanceID,
		now,
		r.sessionTTL.Milliseconds(),
	).Result()
	if err != nil {
		return "", "", fmt.Errorf("claim session failed: %w", err)
	}

	owner := fmt.Sprint(raw)
	if owner == "" {
		return "", "", errors.New("claim session returned empty owner")
	}
	if owner == r.instanceID {
		return owner, r.instanceEndpoint, nil
	}

	endpoint, err := r.resolveInstanceEndpoint(ctx, owner)
	if err != nil {
		return owner, "", err
	}
	return owner, endpoint, nil
}

func (r *RedisRegistry) ReleaseSession(ctx context.Context, language, sessionID string) error {
	key := r.sessionKey(language, sessionID)
	if _, err := releaseSessionScript.Run(ctx, r.client, []string{key}, r.instanceID).Result(); err != nil {
		return fmt.Errorf("release session failed: %w", err)
	}
	return nil
}

func (r *RedisRegistry) Close() error {
	r.stopOnce.Do(func() {
		close(r.stopCh)
	})
	return r.client.Close()
}

func (r *RedisRegistry) heartbeatLoop() {
	ticker := time.NewTicker(r.heartbeatInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			_ = r.refreshInstance(ctx)
			cancel()
		case <-r.stopCh:
			return
		}
	}
}

func (r *RedisRegistry) refreshInstance(ctx context.Context) error {
	key := r.instanceKey(r.instanceID)
	now := time.Now().UTC().Format(time.RFC3339Nano)
	if err := r.client.HSet(ctx, key, map[string]any{
		"endpoint":  r.instanceEndpoint,
		"updatedAt": now,
	}).Err(); err != nil {
		return fmt.Errorf("refresh instance metadata failed: %w", err)
	}
	if err := r.client.Expire(ctx, key, r.instanceTTL).Err(); err != nil {
		return fmt.Errorf("refresh instance ttl failed: %w", err)
	}
	return nil
}

func (r *RedisRegistry) resolveInstanceEndpoint(ctx context.Context, ownerID string) (string, error) {
	key := r.instanceKey(ownerID)
	endpoint, err := r.client.HGet(ctx, key, "endpoint").Result()
	if err == redis.Nil {
		return "", nil
	}
	if err != nil {
		return "", fmt.Errorf("resolve instance endpoint failed: %w", err)
	}
	return strings.TrimSpace(endpoint), nil
}

func (r *RedisRegistry) sessionKey(language, sessionID string) string {
	return fmt.Sprintf("%s:sessions:%s:%s", r.keyPrefix, normalizePart(language), normalizePart(sessionID))
}

func (r *RedisRegistry) instanceKey(instanceID string) string {
	return fmt.Sprintf("%s:instances:%s", r.keyPrefix, normalizePart(instanceID))
}

func normalizePart(value string) string {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return "default"
	}
	return trimmed
}