feat: enhance completion service with session management and language support
- Introduced session management using Redis for tracking active sessions. - Added session claiming and releasing functionality in the completion manager. - Enhanced HTTP and WebSocket completion endpoints to support multiple languages. - Implemented request timeout and maximum body size configurations for API routes. - Updated client-side code to handle session IDs and language parameters in completion requests. - Improved error handling for unsupported languages and session conflicts. - Added tests for the completion manager to ensure proper session handling and cleanup.
This commit is contained in:
222
backend/internal/cluster/redis_registry.go
Normal file
222
backend/internal/cluster/redis_registry.go
Normal file
@@ -0,0 +1,222 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user