ragflow/internal/server/variable.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package server

import (
	"context"
	"fmt"
	"ragflow/internal/utility"
	"sync"
	"time"

	"go.uber.org/zap"

	"ragflow/internal/logger"
)

// Variables holds all runtime variables that can be changed during system operation
// Unlike Config, these can be modified at runtime
type Variables struct {
	SecretKey string `json:"secret_key"`
}

// VariableStore interface for persistent storage (e.g., Redis)
type VariableStore interface {
	Get(key string) (string, error)
	Set(key string, value string, exp time.Duration) bool
	SetNX(key string, value string, exp time.Duration) bool
}

var (
	globalVariables *Variables
	variablesOnce   sync.Once
	variablesMu     sync.RWMutex
)

const (
	// DefaultSecretKey is used when no secret key is found in storage
	DefaultSecretKey = "infiniflow-token"
	// SecretKeyRedisKey is the Redis key for storing secret key
	SecretKeyRedisKey = "ragflow:system:secret_key"
	// SecretKeyTTL is the TTL for secret key in Redis (0 = no expiration)
	SecretKeyTTL = 0
)

// InitVariables initializes all runtime variables from persistent storage
// This should be called after Config and Cache are initialized
func InitVariables(store VariableStore) error {
	var initErr error
	variablesOnce.Do(func() {
		globalVariables = &Variables{}

		generatedKey, err := utility.GenerateSecretKey()
		if err != nil {
			initErr = fmt.Errorf("failed to generate secret key: %w", err)
		}

		// Initialize SecretKey
		secretKey, err := GetOrCreateKey(store, SecretKeyRedisKey, generatedKey)
		if err != nil {
			initErr = fmt.Errorf("failed to initialize secret key: %w", err)
		} else {
			globalVariables.SecretKey = secretKey
			logger.Info("Secret key initialized from store")
		}

		logger.Info("Server variables initialized successfully")
	})
	return initErr
}

// GetVariables returns the global variables instance
func GetVariables() *Variables {
	variablesMu.RLock()
	defer variablesMu.RUnlock()
	return globalVariables
}

// GetSecretKey returns the current secret key
func GetSecretKey() string {
	variablesMu.RLock()
	defer variablesMu.RUnlock()
	if globalVariables == nil {
		return DefaultSecretKey
	}
	return globalVariables.SecretKey
}

// SetSecretKey updates the secret key at runtime
func SetSecretKey(key string) {
	variablesMu.Lock()
	defer variablesMu.Unlock()
	if globalVariables != nil {
		globalVariables.SecretKey = key
		logger.Info("Secret key updated at runtime")
	}
}

// GetOrCreateKey gets a key from store, or creates it if not exists
// - If key exists in store, returns the stored value
// - If key doesn't exist, calls createFn to generate value, stores it, and returns it
// - Uses SetNX to ensure atomic creation (only one caller succeeds when key doesn't exist)
func GetOrCreateKey(store VariableStore, key string, newValue string) (string, error) {
	if store == nil {
		err := fmt.Errorf("store is nil")
		logger.Warn("VariableStore is nil, cannot get or create key", zap.String("key", key))
		return "store is nil", err
	}

	// Try to get existing value
	value, err := store.Get(key)
	if err != nil {
		logger.Warn("Failed to get key from store", zap.String("key", key), zap.Error(err))
		return "", err
	}

	// Key exists, return the value
	if value != "" {
		logger.Debug("Key found in store", zap.String("key", key))
		return value, nil
	}

	// Key doesn't exist, generate new value
	logger.Info("Generating new value for key", zap.String("key", key))

	// Try to set with NX (only if not exists) - ensures atomicity
	if store.SetNX(key, newValue, SecretKeyTTL) {
		logger.Info("New value stored successfully", zap.String("key", key))
		return newValue, nil
	}

	// Another process might have set it, try to get again
	value, err = store.Get(key)
	if err != nil {
		logger.Warn("Failed to get key after SetNX", zap.String("key", key), zap.Error(err))
		return newValue, nil // Return our generated value as fallback
	}

	if value != "" {
		logger.Info("Using value set by another process", zap.String("key", key))
		return value, nil
	}

	// If still empty, use our generated value
	return newValue, nil
}

// RefreshVariables refreshes all variables from storage
// Call this when you want to reload variables from persistent storage
func RefreshVariables(store VariableStore) error {
	if store == nil {
		return fmt.Errorf("store is nil")
	}

	variablesMu.Lock()
	defer variablesMu.Unlock()

	if globalVariables == nil {
		globalVariables = &Variables{}
	}

	// Refresh SecretKey
	secretKey, err := store.Get(SecretKeyRedisKey)
	if err != nil {
		logger.Warn("Failed to refresh secret key from store", zap.Error(err))
		return err
	}
	if secretKey != "" {
		globalVariables.SecretKey = secretKey
		logger.Info("Secret key refreshed from store")
	}

	return nil
}

// VariableWatcher watches for variable changes in storage
// This can be used to detect changes made by other instances
type VariableWatcher struct {
	store    VariableStore
	stopChan chan struct{}
	wg       sync.WaitGroup
}

// NewVariableWatcher creates a new variable watcher
func NewVariableWatcher(store VariableStore) *VariableWatcher {
	return &VariableWatcher{
		store:    store,
		stopChan: make(chan struct{}),
	}
}

// Start starts watching for variable changes
func (w *VariableWatcher) Start(interval time.Duration) {
	w.wg.Add(1)
	go func() {
		defer w.wg.Done()
		ticker := time.NewTicker(interval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				if err := RefreshVariables(w.store); err != nil {
					logger.Debug("Failed to refresh variables", zap.Error(err))
				}
			case <-w.stopChan:
				return
			}
		}
	}()
	logger.Info("Variable watcher started", zap.Duration("interval", interval))
}

// Stop stops the variable watcher
func (w *VariableWatcher) Stop() {
	close(w.stopChan)
	w.wg.Wait()
	logger.Info("Variable watcher stopped")
}

// SaveToStorage saves current variables to persistent storage
func SaveToStorage(store VariableStore) error {
	if store == nil {
		return fmt.Errorf("store is nil")
	}

	variablesMu.RLock()
	defer variablesMu.RUnlock()

	if globalVariables == nil {
		return fmt.Errorf("variables not initialized")
	}

	// Save SecretKey
	if !store.Set(SecretKeyRedisKey, globalVariables.SecretKey, SecretKeyTTL) {
		return fmt.Errorf("failed to save secret key to store")
	}

	logger.Info("Variables saved to storage")
	return nil
}

// WithTimeout creates a context with timeout for variable operations
func WithTimeout(timeout time.Duration) (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), timeout)
}