geo/server/internal/tenant/app/browser_fetch_queue.go

package app

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/geo-platform/tenant-api/internal/shared/config"
	"go.uber.org/zap"
)

var errBrowserFetchQueueFull = errors.New("browser fetch queue is full")

type browserFetchResult struct {
	URL      string `json:"url,omitempty"`
	Title    string `json:"title,omitempty"`
	Text     string `json:"text,omitempty"`
	Markdown string `json:"markdown,omitempty"`
}

type browserFetchRequest struct {
	URL       string                         `json:"url"`
	TimeoutMS int64                          `json:"timeout_ms,omitempty"`
	Options   map[string]interface{}         `json:"options,omitempty"`
	ctx       context.Context                `json:"-"`
	cfg       config.BrowserFetchConfig      `json:"-"`
	resp      chan browserFetchRequestResult `json:"-"`
}

type browserFetchRequestResult struct {
	result browserFetchResult
	err    error
}

type browserFetchQueue struct {
	logger *zap.Logger
	client *http.Client
	jobs   chan browserFetchRequest
	cache  *browserFetchMemoryCache
}

type browserFetchCacheEntry struct {
	result    browserFetchResult
	expiresAt time.Time
}

type browserFetchMemoryCache struct {
	mu      sync.Mutex
	entries map[string]browserFetchCacheEntry
}

func (s *KnowledgeService) browserFetchQueue(cfg config.BrowserFetchConfig) *browserFetchQueue {
	s.browserFetchMu.Lock()
	defer s.browserFetchMu.Unlock()
	if s.browserFetchQ != nil {
		return s.browserFetchQ
	}

	config.NormalizeBrowserFetchConfig(&cfg)
	queue := &browserFetchQueue{
		logger: s.logger,
		client: &http.Client{},
		jobs:   make(chan browserFetchRequest, cfg.QueueSize),
		cache: &browserFetchMemoryCache{
			entries: make(map[string]browserFetchCacheEntry),
		},
	}

	for i := 0; i < cfg.WorkerConcurrency; i++ {
		go queue.runWorker(i)
	}
	s.browserFetchQ = queue
	return queue
}

func (q *browserFetchQueue) Fetch(ctx context.Context, rawURL string, cfg config.BrowserFetchConfig) (browserFetchResult, error) {
	if q == nil {
		return browserFetchResult{}, errBrowserFetchNotAvailable
	}
	config.NormalizeBrowserFetchConfig(&cfg)
	cacheKey := browserFetchCacheKey(rawURL)
	if cfg.CacheTTL > 0 {
		if result, ok := q.cache.Get(cacheKey); ok {
			return result, nil
		}
	}

	request := browserFetchRequest{
		URL:       rawURL,
		TimeoutMS: int64(cfg.RequestTimeout / time.Millisecond),
		Options: map[string]interface{}{
			"provider": cfg.Provider,
		},
		ctx:  ctx,
		cfg:  cfg,
		resp: make(chan browserFetchRequestResult, 1),
	}

	select {
	case q.jobs <- request:
	case <-ctx.Done():
		return browserFetchResult{}, ctx.Err()
	default:
		return browserFetchResult{}, errBrowserFetchQueueFull
	}

	select {
	case result := <-request.resp:
		if result.err != nil {
			return browserFetchResult{}, result.err
		}
		if cfg.CacheTTL > 0 {
			q.cache.Set(cacheKey, result.result, cfg.CacheTTL)
		}
		return result.result, nil
	case <-ctx.Done():
		return browserFetchResult{}, ctx.Err()
	}
}

func (q *browserFetchQueue) runWorker(workerID int) {
	for request := range q.jobs {
		result, err := q.execute(context.Background(), request)
		if err != nil && q.logger != nil {
			q.logger.Warn("browser fetch request failed",
				zap.Int("worker_id", workerID),
				zap.String("url", request.URL),
				zap.Error(err),
			)
		}
		request.resp <- browserFetchRequestResult{result: result, err: err}
	}
}

func (q *browserFetchQueue) execute(parent context.Context, request browserFetchRequest) (browserFetchResult, error) {
	if request.ctx != nil {
		parent = request.ctx
	}
	timeout := time.Duration(request.TimeoutMS) * time.Millisecond
	if timeout <= 0 {
		timeout = 45 * time.Second
	}
	ctx, cancel := context.WithTimeout(parent, timeout)
	defer cancel()

	return q.callService(ctx, request.URL, request.cfg)
}

func (q *browserFetchQueue) callService(ctx context.Context, rawURL string, cfg config.BrowserFetchConfig) (browserFetchResult, error) {
	config.NormalizeBrowserFetchConfig(&cfg)
	endpoint := strings.TrimRight(cfg.ServiceURL, "/") + "/fetch"
	payload := browserFetchRequest{
		URL:       rawURL,
		TimeoutMS: int64(cfg.RequestTimeout / time.Millisecond),
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return browserFetchResult{}, fmt.Errorf("marshal browser fetch request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return browserFetchResult{}, fmt.Errorf("create browser fetch request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	if cfg.Token != "" {
		req.Header.Set("Authorization", "Bearer "+cfg.Token)
	}

	resp, err := q.client.Do(req)
	if err != nil {
		return browserFetchResult{}, fmt.Errorf("execute browser fetch request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= http.StatusBadRequest {
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return browserFetchResult{}, fmt.Errorf("browser fetch service failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
	}

	var result browserFetchResult
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return browserFetchResult{}, fmt.Errorf("decode browser fetch response: %w", err)
	}
	return result, nil
}

func (c *browserFetchMemoryCache) Get(key string) (browserFetchResult, bool) {
	if c == nil || key == "" {
		return browserFetchResult{}, false
	}
	now := time.Now()
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, ok := c.entries[key]
	if !ok {
		return browserFetchResult{}, false
	}
	if now.After(entry.expiresAt) {
		delete(c.entries, key)
		return browserFetchResult{}, false
	}
	return entry.result, true
}

func (c *browserFetchMemoryCache) Set(key string, result browserFetchResult, ttl time.Duration) {
	if c == nil || key == "" || ttl <= 0 {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries[key] = browserFetchCacheEntry{
		result:    result,
		expiresAt: time.Now().Add(ttl),
	}
}

func browserFetchCacheKey(rawURL string) string {
	normalized := strings.TrimSpace(rawURL)
	if parsed, err := url.Parse(normalized); err == nil && parsed.Host != "" {
		parsed.Fragment = ""
		normalized = parsed.String()
	}
	sum := sha256.Sum256([]byte(normalized))
	return hex.EncodeToString(sum[:])
}