reqqueue.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"

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

// RequestStatus represents the state of a queued chat request.
type RequestStatus string

const (
	StatusQueued     RequestStatus = "queued"
	StatusProcessing RequestStatus = "processing"
	StatusCompleted  RequestStatus = "completed"
	StatusFailed     RequestStatus = "failed"
)

// QueuedRequest is a chat request waiting for model capacity.
type QueuedRequest struct {
	ID        string          `json:"id"`
	Model     string          `json:"model"`
	Status    RequestStatus   `json:"status"`
	Body      json.RawMessage `json:"body"`
	SessionID string          `json:"session_id,omitempty"`
	CreatedAt time.Time       `json:"created_at"`
	Result    json.RawMessage `json:"result,omitempty"`
	Error     string          `json:"error,omitempty"`
	PollURL   string          `json:"poll_url,omitempty"`
}

// RequestQueue manages queued chat requests in Dragonfly (Redis DB 3).
// When a model is unavailable, requests are enqueued. A background processor
// polls for model availability and forwards queued requests when capacity appears.
type RequestQueue struct {
	client      *redis.Client
	rayServeURL string
	httpClient  *http.Client
	discovery   *ModelDiscovery
	ttl         time.Duration
}

// NewRequestQueue creates a request queue connected to Dragonfly DB 3.
func NewRequestQueue(addr string, rayServeURL string, discovery *ModelDiscovery, ttlSeconds int) (*RequestQueue, error) {
	client := redis.NewClient(&redis.Options{
		Addr:         addr,
		DB:           3, // DB 0 = Ray GCS, DB 1 = batch jobs, DB 2 = sessions, DB 3 = request queue
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
		DialTimeout:  5 * time.Second,
	})

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("connect to redis (db=3) at %s: %w", addr, err)
	}

	log.Printf("request queue connected to redis at %s (db=3, ttl=%ds)", addr, ttlSeconds)

	return &RequestQueue{
		client:      client,
		rayServeURL: rayServeURL,
		httpClient:  &http.Client{Timeout: 5 * time.Minute},
		discovery:   discovery,
		ttl:         time.Duration(ttlSeconds) * time.Second,
	}, nil
}

func reqKey(id string) string          { return "chatreq:" + id }
func modelQueueKey(model string) string { return "chatqueue:" + model }
func demandCounterKey(model string) string { return "demand:" + model }

// Enqueue stores a request and adds it to the model's queue.
func (q *RequestQueue) Enqueue(ctx context.Context, id, model string, body json.RawMessage, sessionID string) error {
	req := QueuedRequest{
		ID:        id,
		Model:     model,
		Status:    StatusQueued,
		Body:      body,
		SessionID: sessionID,
		CreatedAt: time.Now(),
		PollURL:   "/v1/requests/" + id,
	}

	data, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("marshal queued request: %w", err)
	}

	pipe := q.client.Pipeline()
	pipe.Set(ctx, reqKey(id), data, q.ttl)
	pipe.LPush(ctx, modelQueueKey(model), id)
	pipe.Expire(ctx, modelQueueKey(model), q.ttl)
	// Increment demand counter — GPUScale reads this to decide if workers are idle
	pipe.Incr(ctx, demandCounterKey(model))
	_, err = pipe.Exec(ctx)
	if err != nil {
		return fmt.Errorf("enqueue request %s: %w", id, err)
	}

	queuedRequestsTotal.WithLabelValues(model).Inc()
	queuedRequestsActive.WithLabelValues(model).Inc()
	log.Printf("reqqueue: enqueued request %s for model %s", id, model)
	return nil
}

// GetRequest loads a queued request by ID.
func (q *RequestQueue) GetRequest(ctx context.Context, id string) (*QueuedRequest, error) {
	data, err := q.client.Get(ctx, reqKey(id)).Bytes()
	if err == redis.Nil {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get request %s: %w", id, err)
	}

	var req QueuedRequest
	if err := json.Unmarshal(data, &req); err != nil {
		return nil, fmt.Errorf("unmarshal request %s: %w", id, err)
	}
	return &req, nil
}

// RunProcessor processes queued requests when models become available.
func (q *RequestQueue) RunProcessor(ctx context.Context, pollInterval time.Duration) {
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			q.processQueues(ctx)
		}
	}
}

func (q *RequestQueue) processQueues(ctx context.Context) {
	loadedModels := q.discovery.ListLoadedModels()
	for _, model := range loadedModels {
		q.processModelQueue(ctx, model)
	}
}

func (q *RequestQueue) processModelQueue(ctx context.Context, model string) {
	for {
		// RPOP from the queue (FIFO: LPUSH + RPOP)
		requestID, err := q.client.RPop(ctx, modelQueueKey(model)).Result()
		if err == redis.Nil {
			return // queue empty
		}
		if err != nil {
			log.Printf("reqqueue: rpop %s: %v", modelQueueKey(model), err)
			return
		}

		req, err := q.GetRequest(ctx, requestID)
		if req == nil || err != nil {
			log.Printf("reqqueue: request %s not found or error: %v", requestID, err)
			continue
		}

		if req.Status != StatusQueued {
			continue // already processed (dedup by request ID)
		}

		// Mark as processing
		req.Status = StatusProcessing
		q.saveRequest(ctx, req)

		// Forward to Ray Serve
		result, err := q.forwardToRay(ctx, req.Body)
		if err != nil {
			req.Status = StatusFailed
			req.Error = err.Error()
			q.saveRequest(ctx, req)
			q.decrDemand(ctx, model)
			queuedRequestsActive.WithLabelValues(model).Dec()
			activeRequestsByModel.WithLabelValues(model).Dec()
			log.Printf("reqqueue: request %s failed: %v", requestID, err)
			continue
		}

		req.Status = StatusCompleted
		req.Result = result
		q.saveRequest(ctx, req)
		q.decrDemand(ctx, model)
		queuedRequestsActive.WithLabelValues(model).Dec()
		activeRequestsByModel.WithLabelValues(model).Dec()
		log.Printf("reqqueue: request %s completed", requestID)
	}
}

func (q *RequestQueue) forwardToRay(ctx context.Context, body json.RawMessage) (json.RawMessage, error) {
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, q.rayServeURL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := q.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("ray serve request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("ray serve returned %d: %s", resp.StatusCode, string(respBody))
	}

	return json.RawMessage(respBody), nil
}

// decrDemand decrements the demand counter for a model.
// Clamps to 0 to avoid negative counts from race conditions.
func (q *RequestQueue) decrDemand(ctx context.Context, model string) {
	val, err := q.client.Decr(ctx, demandCounterKey(model)).Result()
	if err != nil {
		log.Printf("reqqueue: failed to decr demand for %s: %v", model, err)
		return
	}
	if val < 0 {
		q.client.Set(ctx, demandCounterKey(model), 0, 0)
	}
}

func (q *RequestQueue) saveRequest(ctx context.Context, req *QueuedRequest) {
	data, err := json.Marshal(req)
	if err != nil {
		log.Printf("reqqueue: failed to marshal request %s: %v", req.ID, err)
		return
	}
	q.client.Set(ctx, reqKey(req.ID), data, q.ttl)
}

// Close shuts down the Redis connection.
func (q *RequestQueue) Close() error {
	return q.client.Close()
}