mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
perf(coderd/x/chatd): add process-wide config cache for hot DB queries (#23272)
## Summary Adds a process-wide cache for three hot database queries in `chatd` that were hitting Postgres on **every chat turn** despite returning rarely-changing configuration data: | Query | Before (50k turns) | After | Reduction | |---|---|---|---| | `GetEnabledChatProviders` | ~98.6k calls | ~500-1000 | ~99% | | `GetChatModelConfigByID` | ~49.2k calls | ~500-1000 | ~98% | | `GetUserChatCustomPrompt` | ~46.7k calls | ~1000-2000 | ~97% | These were identified via `coder exp scaletest chat` (5000 concurrent chats × 10 turns) as the dominant source of Postgres load during chat processing. ## Design Follows the established **webpush subscription cache pattern** (`coderd/webpush/webpush.go`): - `sync.RWMutex` + `tailscale.com/util/singleflight` (generic) + generation-based stale prevention + TTL - 10s TTL for provider/model config, 5s TTL for user prompts - Negative caching for `sql.ErrNoRows` on user prompts (the common case — most users don't set custom prompts) - Deep-clones `ChatModelConfig.Options` (`json.RawMessage` = `[]byte`) on both store and read paths ### Invalidation Single pubsub channel (`chat:config_change`) with kind discriminator for cross-replica cache invalidation. Seven publish points in `coderd/chats.go` cover all admin mutation endpoints (create/update/delete for providers and model configs, put for user prompts). _This PR was generated with mux and was reviewed by a human_
This commit is contained in:
@@ -33,6 +33,7 @@ import (
|
|||||||
"github.com/coder/coder/v2/coderd/database"
|
"github.com/coder/coder/v2/coderd/database"
|
||||||
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
||||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||||
|
dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
|
||||||
"github.com/coder/coder/v2/coderd/externalauth"
|
"github.com/coder/coder/v2/coderd/externalauth"
|
||||||
"github.com/coder/coder/v2/coderd/externalauth/gitprovider"
|
"github.com/coder/coder/v2/coderd/externalauth/gitprovider"
|
||||||
"github.com/coder/coder/v2/coderd/httpapi"
|
"github.com/coder/coder/v2/coderd/httpapi"
|
||||||
@@ -110,6 +111,28 @@ func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error)
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func publishChatConfigEvent(logger slog.Logger, ps dbpubsub.Pubsub, kind pubsub.ChatConfigEventKind, entityID uuid.UUID) {
|
||||||
|
payload, err := json.Marshal(pubsub.ChatConfigEvent{
|
||||||
|
Kind: kind,
|
||||||
|
EntityID: entityID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
logger.Error(context.Background(), "failed to marshal chat config event",
|
||||||
|
slog.F("kind", kind),
|
||||||
|
slog.F("entity_id", entityID),
|
||||||
|
slog.Error(err),
|
||||||
|
)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := ps.Publish(pubsub.ChatConfigEventChannel, payload); err != nil {
|
||||||
|
logger.Error(context.Background(), "failed to publish chat config event",
|
||||||
|
slog.F("kind", kind),
|
||||||
|
slog.F("entity_id", entityID),
|
||||||
|
slog.Error(err),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
|
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
|
||||||
func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) {
|
func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) {
|
||||||
ctx := r.Context()
|
ctx := r.Context()
|
||||||
@@ -3052,6 +3075,8 @@ func (api *API) putUserChatCustomPrompt(rw http.ResponseWriter, r *http.Request)
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventUserPrompt, apiKey.UserID)
|
||||||
|
|
||||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.UserChatCustomPrompt{
|
httpapi.Write(ctx, rw, http.StatusOK, codersdk.UserChatCustomPrompt{
|
||||||
CustomPrompt: updatedConfig.Value,
|
CustomPrompt: updatedConfig.Value,
|
||||||
})
|
})
|
||||||
@@ -3886,6 +3911,8 @@ func (api *API) createChatProvider(rw http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil)
|
||||||
|
|
||||||
httpapi.Write(
|
httpapi.Write(
|
||||||
ctx,
|
ctx,
|
||||||
rw,
|
rw,
|
||||||
@@ -3972,6 +3999,8 @@ func (api *API) updateChatProvider(rw http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil)
|
||||||
|
|
||||||
httpapi.Write(
|
httpapi.Write(
|
||||||
ctx,
|
ctx,
|
||||||
rw,
|
rw,
|
||||||
@@ -4026,6 +4055,8 @@ func (api *API) deleteChatProvider(rw http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil)
|
||||||
|
|
||||||
rw.WriteHeader(http.StatusNoContent)
|
rw.WriteHeader(http.StatusNoContent)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4205,6 +4236,8 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, inserted.ID)
|
||||||
|
|
||||||
httpapi.Write(ctx, rw, http.StatusCreated, convertChatModelConfig(inserted))
|
httpapi.Write(ctx, rw, http.StatusCreated, convertChatModelConfig(inserted))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4376,6 +4409,8 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, updated.ID)
|
||||||
|
|
||||||
httpapi.Write(ctx, rw, http.StatusOK, convertChatModelConfig(updated))
|
httpapi.Write(ctx, rw, http.StatusOK, convertChatModelConfig(updated))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4416,6 +4451,8 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, modelConfigID)
|
||||||
|
|
||||||
rw.WriteHeader(http.StatusNoContent)
|
rw.WriteHeader(http.StatusNoContent)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,52 @@
|
|||||||
|
package pubsub
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"golang.org/x/xerrors"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ChatConfigEventChannel is the pubsub channel for chat config
|
||||||
|
// changes (providers, model configs, user prompts). All replicas
|
||||||
|
// subscribe to this channel to invalidate their local caches.
|
||||||
|
const ChatConfigEventChannel = "chat:config_change"
|
||||||
|
|
||||||
|
// HandleChatConfigEvent wraps a typed callback for ChatConfigEvent
|
||||||
|
// messages, following the same pattern as HandleChatEvent.
|
||||||
|
func HandleChatConfigEvent(cb func(ctx context.Context, payload ChatConfigEvent, err error)) func(ctx context.Context, message []byte, err error) {
|
||||||
|
return func(ctx context.Context, message []byte, err error) {
|
||||||
|
if err != nil {
|
||||||
|
cb(ctx, ChatConfigEvent{}, xerrors.Errorf("chat config event pubsub: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var payload ChatConfigEvent
|
||||||
|
if err := json.Unmarshal(message, &payload); err != nil {
|
||||||
|
cb(ctx, ChatConfigEvent{}, xerrors.Errorf("unmarshal chat config event: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
cb(ctx, payload, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ChatConfigEvent is published when chat configuration changes
|
||||||
|
// (provider CRUD, model config CRUD, or user prompt updates).
|
||||||
|
// Subscribers use this to invalidate their local caches.
|
||||||
|
type ChatConfigEvent struct {
|
||||||
|
Kind ChatConfigEventKind `json:"kind"`
|
||||||
|
// EntityID carries context for the invalidation:
|
||||||
|
// - For providers: uuid.Nil (all providers are invalidated).
|
||||||
|
// - For model configs: the specific config ID.
|
||||||
|
// - For user prompts: the user ID.
|
||||||
|
EntityID uuid.UUID `json:"entity_id"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ChatConfigEventKind string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ChatConfigEventProviders ChatConfigEventKind = "providers"
|
||||||
|
ChatConfigEventModelConfig ChatConfigEventKind = "model_config"
|
||||||
|
ChatConfigEventUserPrompt ChatConfigEventKind = "user_prompt"
|
||||||
|
)
|
||||||
+35
-4
@@ -108,6 +108,8 @@ type Server struct {
|
|||||||
pubsub pubsub.Pubsub
|
pubsub pubsub.Pubsub
|
||||||
webpushDispatcher webpush.Dispatcher
|
webpushDispatcher webpush.Dispatcher
|
||||||
providerAPIKeys chatprovider.ProviderAPIKeys
|
providerAPIKeys chatprovider.ProviderAPIKeys
|
||||||
|
configCache *chatConfigCache
|
||||||
|
configCacheUnsubscribe func()
|
||||||
|
|
||||||
// chatStreams stores per-chat stream state. Using sync.Map
|
// chatStreams stores per-chat stream state. Using sync.Map
|
||||||
// gives each chat independent locking — concurrent chats
|
// gives each chat independent locking — concurrent chats
|
||||||
@@ -1740,6 +1742,31 @@ func New(cfg Config) *Server {
|
|||||||
|
|
||||||
//nolint:gocritic // The chat processor uses a scoped chatd context.
|
//nolint:gocritic // The chat processor uses a scoped chatd context.
|
||||||
ctx = dbauthz.AsChatd(ctx)
|
ctx = dbauthz.AsChatd(ctx)
|
||||||
|
|
||||||
|
p.configCache = newChatConfigCache(ctx, cfg.Database, clk)
|
||||||
|
if p.pubsub != nil {
|
||||||
|
cancelConfigSub, err := p.pubsub.SubscribeWithErr(
|
||||||
|
coderdpubsub.ChatConfigEventChannel,
|
||||||
|
coderdpubsub.HandleChatConfigEvent(func(ctx context.Context, ev coderdpubsub.ChatConfigEvent, err error) {
|
||||||
|
if err != nil {
|
||||||
|
p.logger.Warn(ctx, "chat config event error", slog.Error(err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
switch ev.Kind {
|
||||||
|
case coderdpubsub.ChatConfigEventProviders:
|
||||||
|
p.configCache.InvalidateProviders()
|
||||||
|
case coderdpubsub.ChatConfigEventModelConfig:
|
||||||
|
p.configCache.InvalidateModelConfig(ev.EntityID)
|
||||||
|
case coderdpubsub.ChatConfigEventUserPrompt:
|
||||||
|
p.configCache.InvalidateUserPrompt(ev.EntityID)
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
p.logger.Error(ctx, "subscribe to chat config events", slog.Error(err))
|
||||||
|
}
|
||||||
|
p.configCacheUnsubscribe = cancelConfigSub
|
||||||
|
}
|
||||||
go p.start(ctx)
|
go p.start(ctx)
|
||||||
|
|
||||||
return p
|
return p
|
||||||
@@ -4017,7 +4044,7 @@ func (p *Server) resolveChatModel(
|
|||||||
})
|
})
|
||||||
g.Go(func() error {
|
g.Go(func() error {
|
||||||
var err error
|
var err error
|
||||||
providers, err = p.db.GetEnabledChatProviders(ctx)
|
providers, err = p.configCache.EnabledProviders(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return xerrors.Errorf("get enabled chat providers: %w", err)
|
return xerrors.Errorf("get enabled chat providers: %w", err)
|
||||||
}
|
}
|
||||||
@@ -4061,7 +4088,7 @@ func (p *Server) resolveModelConfig(
|
|||||||
chat database.Chat,
|
chat database.Chat,
|
||||||
) (database.ChatModelConfig, error) {
|
) (database.ChatModelConfig, error) {
|
||||||
if chat.LastModelConfigID != uuid.Nil {
|
if chat.LastModelConfigID != uuid.Nil {
|
||||||
modelConfig, err := p.db.GetChatModelConfigByID(
|
modelConfig, err := p.configCache.ModelConfigByID(
|
||||||
ctx, chat.LastModelConfigID,
|
ctx, chat.LastModelConfigID,
|
||||||
)
|
)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
@@ -4076,7 +4103,7 @@ func (p *Server) resolveModelConfig(
|
|||||||
// Model config was deleted, fall through to default.
|
// Model config was deleted, fall through to default.
|
||||||
}
|
}
|
||||||
|
|
||||||
defaultConfig, err := p.db.GetDefaultChatModelConfig(ctx)
|
defaultConfig, err := p.configCache.DefaultModelConfig(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if xerrors.Is(err, sql.ErrNoRows) {
|
if xerrors.Is(err, sql.ErrNoRows) {
|
||||||
return database.ChatModelConfig{}, xerrors.New(
|
return database.ChatModelConfig{}, xerrors.New(
|
||||||
@@ -4286,7 +4313,7 @@ func (p *Server) resolveUserCompactionThreshold(ctx context.Context, userID uuid
|
|||||||
// database and wraps it in <user-instructions> tags. Returns empty
|
// database and wraps it in <user-instructions> tags. Returns empty
|
||||||
// string if no prompt is set.
|
// string if no prompt is set.
|
||||||
func (p *Server) resolveUserPrompt(ctx context.Context, userID uuid.UUID) string {
|
func (p *Server) resolveUserPrompt(ctx context.Context, userID uuid.UUID) string {
|
||||||
raw, err := p.db.GetUserChatCustomPrompt(ctx, userID)
|
raw, err := p.configCache.UserPrompt(ctx, userID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// sql.ErrNoRows is the normal "not set" case.
|
// sql.ErrNoRows is the normal "not set" case.
|
||||||
return ""
|
return ""
|
||||||
@@ -4447,6 +4474,10 @@ func (p *Server) dispatchPush(
|
|||||||
|
|
||||||
// Close stops the processor and waits for it to finish.
|
// Close stops the processor and waits for it to finish.
|
||||||
func (p *Server) Close() error {
|
func (p *Server) Close() error {
|
||||||
|
if unsub := p.configCacheUnsubscribe; unsub != nil {
|
||||||
|
p.configCacheUnsubscribe = nil
|
||||||
|
unsub()
|
||||||
|
}
|
||||||
p.cancel()
|
p.cancel()
|
||||||
<-p.closed
|
<-p.closed
|
||||||
p.inflight.Wait()
|
p.inflight.Wait()
|
||||||
|
|||||||
@@ -0,0 +1,412 @@
|
|||||||
|
package chatd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"slices"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/ammario/tlru"
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"tailscale.com/util/singleflight"
|
||||||
|
|
||||||
|
"github.com/coder/coder/v2/coderd/database"
|
||||||
|
"github.com/coder/quartz"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
chatConfigProvidersTTL = 10 * time.Second
|
||||||
|
chatConfigModelConfigTTL = 10 * time.Second
|
||||||
|
chatConfigUserPromptTTL = 5 * time.Second
|
||||||
|
// Bound user-prompt cache cardinality so one-shot users do not
|
||||||
|
// accumulate forever in long-lived chatd processes.
|
||||||
|
chatConfigUserPromptEntryLimit = 64 * 1024
|
||||||
|
)
|
||||||
|
|
||||||
|
type cachedProviders struct {
|
||||||
|
providers []database.ChatProvider
|
||||||
|
expiresAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type cachedModelConfig struct {
|
||||||
|
config database.ChatModelConfig
|
||||||
|
expiresAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type modelConfigSnapshot struct {
|
||||||
|
epoch uint64
|
||||||
|
generation uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
// cloneModelConfig returns a shallow copy of cfg with Options
|
||||||
|
// deep-cloned so the cache owns its own backing array.
|
||||||
|
func cloneModelConfig(cfg database.ChatModelConfig) database.ChatModelConfig {
|
||||||
|
cfg.Options = slices.Clone(cfg.Options)
|
||||||
|
return cfg
|
||||||
|
}
|
||||||
|
|
||||||
|
type chatConfigCache struct {
|
||||||
|
db database.Store
|
||||||
|
clock quartz.Clock
|
||||||
|
// ctx is the server-scoped context used for all DB fills.
|
||||||
|
// Cache fills run inside singleflight.Do where one caller
|
||||||
|
// becomes the leader for all coalesced waiters. Using a
|
||||||
|
// per-request context would mean the leader's cancellation
|
||||||
|
// (timeout, user disconnect) fans the error to every waiter.
|
||||||
|
// Storing the server context here makes that impossible by
|
||||||
|
// construction — callers cannot pass a request context into
|
||||||
|
// the shared fill path.
|
||||||
|
ctx context.Context
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
|
||||||
|
// Providers (singleton).
|
||||||
|
providers *cachedProviders
|
||||||
|
providerGeneration uint64
|
||||||
|
providerFetches singleflight.Group[string, []database.ChatProvider]
|
||||||
|
|
||||||
|
// Model configs (keyed by ID).
|
||||||
|
modelTopologyEpoch uint64
|
||||||
|
modelConfigs map[uuid.UUID]cachedModelConfig
|
||||||
|
modelConfigFetches singleflight.Group[string, database.ChatModelConfig]
|
||||||
|
|
||||||
|
// Default model config (singleton).
|
||||||
|
defaultModelConfig *cachedModelConfig
|
||||||
|
defaultModelConfigGeneration uint64
|
||||||
|
defaultModelConfigFetches singleflight.Group[string, database.ChatModelConfig]
|
||||||
|
|
||||||
|
// User custom prompts (keyed by user ID).
|
||||||
|
userPromptEpoch uint64
|
||||||
|
userPrompts *tlru.Cache[uuid.UUID, string]
|
||||||
|
userPromptFetches singleflight.Group[string, string]
|
||||||
|
}
|
||||||
|
|
||||||
|
func newChatConfigCache(ctx context.Context, db database.Store, clock quartz.Clock) *chatConfigCache {
|
||||||
|
return &chatConfigCache{
|
||||||
|
db: db,
|
||||||
|
clock: clock,
|
||||||
|
ctx: ctx,
|
||||||
|
modelConfigs: make(map[uuid.UUID]cachedModelConfig),
|
||||||
|
userPrompts: tlru.New[uuid.UUID](
|
||||||
|
tlru.ConstantCost[string],
|
||||||
|
chatConfigUserPromptEntryLimit,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// singleflightDoChan wraps a singleflight group's DoChan method,
|
||||||
|
// allowing the caller to abandon the wait if their context is
|
||||||
|
// canceled while the shared fill continues running to completion.
|
||||||
|
// This separates two lifetimes: the fill runs under the server-scoped
|
||||||
|
// context, while each caller waits under its own request-scoped context.
|
||||||
|
func singleflightDoChan[K comparable, V any](
|
||||||
|
ctx context.Context,
|
||||||
|
group *singleflight.Group[K, V],
|
||||||
|
key K,
|
||||||
|
fn func() (V, error),
|
||||||
|
) (V, error) {
|
||||||
|
ch := group.DoChan(key, fn)
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
var zero V
|
||||||
|
return zero, ctx.Err()
|
||||||
|
case res := <-ch:
|
||||||
|
return res.Val, res.Err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) EnabledProviders(ctx context.Context) ([]database.ChatProvider, error) {
|
||||||
|
if providers, ok := c.cachedProviders(); ok {
|
||||||
|
return providers, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
generation := c.providersGeneration()
|
||||||
|
providers, err := singleflightDoChan(
|
||||||
|
ctx,
|
||||||
|
&c.providerFetches,
|
||||||
|
fmt.Sprintf("%d:providers", generation),
|
||||||
|
func() ([]database.ChatProvider, error) {
|
||||||
|
if cached, ok := c.cachedProviders(); ok {
|
||||||
|
return cached, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fetched, err := c.db.GetEnabledChatProviders(c.ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
c.storeProviders(generation, fetched)
|
||||||
|
return slices.Clone(fetched), nil
|
||||||
|
},
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return slices.Clone(providers), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) cachedProviders() ([]database.ChatProvider, bool) {
|
||||||
|
c.mu.RLock()
|
||||||
|
entry := c.providers
|
||||||
|
c.mu.RUnlock()
|
||||||
|
if entry == nil {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
if c.clock.Now().Before(entry.expiresAt) {
|
||||||
|
return slices.Clone(entry.providers), true
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
if current := c.providers; current != nil && !c.clock.Now().Before(current.expiresAt) {
|
||||||
|
c.providers = nil
|
||||||
|
}
|
||||||
|
c.mu.Unlock()
|
||||||
|
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) providersGeneration() uint64 {
|
||||||
|
c.mu.RLock()
|
||||||
|
generation := c.providerGeneration
|
||||||
|
c.mu.RUnlock()
|
||||||
|
return generation
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) storeProviders(generation uint64, providers []database.ChatProvider) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
|
||||||
|
if c.providerGeneration != generation {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
c.providers = &cachedProviders{
|
||||||
|
providers: slices.Clone(providers),
|
||||||
|
expiresAt: c.clock.Now().Add(chatConfigProvidersTTL),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) InvalidateProviders() {
|
||||||
|
c.mu.Lock()
|
||||||
|
c.providers = nil
|
||||||
|
c.providerGeneration++
|
||||||
|
// Provider topology changed — model selections depend on
|
||||||
|
// provider existence, so flush all model-config state.
|
||||||
|
clear(c.modelConfigs)
|
||||||
|
c.modelTopologyEpoch++
|
||||||
|
c.defaultModelConfig = nil
|
||||||
|
c.defaultModelConfigGeneration++
|
||||||
|
c.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) ModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
if config, ok := c.cachedModelConfig(id); ok {
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
snap := c.modelConfigSnapshot()
|
||||||
|
config, err := singleflightDoChan(ctx, &c.modelConfigFetches, fmt.Sprintf("%d:%s", snap.epoch, id), func() (database.ChatModelConfig, error) {
|
||||||
|
if cached, ok := c.cachedModelConfig(id); ok {
|
||||||
|
return cached, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fetched, err := c.db.GetChatModelConfigByID(c.ctx, id)
|
||||||
|
if err != nil {
|
||||||
|
return database.ChatModelConfig{}, err
|
||||||
|
}
|
||||||
|
c.storeModelConfig(snap, fetched)
|
||||||
|
return cloneModelConfig(fetched), nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return database.ChatModelConfig{}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) cachedModelConfig(id uuid.UUID) (database.ChatModelConfig, bool) {
|
||||||
|
c.mu.RLock()
|
||||||
|
entry, ok := c.modelConfigs[id]
|
||||||
|
c.mu.RUnlock()
|
||||||
|
if !ok {
|
||||||
|
return database.ChatModelConfig{}, false
|
||||||
|
}
|
||||||
|
if c.clock.Now().Before(entry.expiresAt) {
|
||||||
|
return cloneModelConfig(entry.config), true
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
if current, ok := c.modelConfigs[id]; ok && !c.clock.Now().Before(current.expiresAt) {
|
||||||
|
delete(c.modelConfigs, id)
|
||||||
|
}
|
||||||
|
c.mu.Unlock()
|
||||||
|
|
||||||
|
return database.ChatModelConfig{}, false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) modelConfigSnapshot() modelConfigSnapshot {
|
||||||
|
c.mu.RLock()
|
||||||
|
snap := modelConfigSnapshot{epoch: c.modelTopologyEpoch}
|
||||||
|
c.mu.RUnlock()
|
||||||
|
return snap
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) storeModelConfig(snap modelConfigSnapshot, config database.ChatModelConfig) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
|
||||||
|
if c.modelTopologyEpoch != snap.epoch {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
c.modelConfigs[config.ID] = cachedModelConfig{
|
||||||
|
config: cloneModelConfig(config),
|
||||||
|
expiresAt: c.clock.Now().Add(chatConfigModelConfigTTL),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) DefaultModelConfig(ctx context.Context) (database.ChatModelConfig, error) {
|
||||||
|
if config, ok := c.cachedDefaultModelConfig(); ok {
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
snap := c.defaultModelConfigSnapshot()
|
||||||
|
config, err := singleflightDoChan(ctx, &c.defaultModelConfigFetches, fmt.Sprintf("%d:default", snap.epoch), func() (database.ChatModelConfig, error) {
|
||||||
|
if cached, ok := c.cachedDefaultModelConfig(); ok {
|
||||||
|
return cached, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fetched, err := c.db.GetDefaultChatModelConfig(c.ctx)
|
||||||
|
if err != nil {
|
||||||
|
return database.ChatModelConfig{}, err
|
||||||
|
}
|
||||||
|
c.storeDefaultModelConfig(snap, fetched)
|
||||||
|
return cloneModelConfig(fetched), nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return database.ChatModelConfig{}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) cachedDefaultModelConfig() (database.ChatModelConfig, bool) {
|
||||||
|
c.mu.RLock()
|
||||||
|
entry := c.defaultModelConfig
|
||||||
|
c.mu.RUnlock()
|
||||||
|
if entry == nil {
|
||||||
|
return database.ChatModelConfig{}, false
|
||||||
|
}
|
||||||
|
if c.clock.Now().Before(entry.expiresAt) {
|
||||||
|
return cloneModelConfig(entry.config), true
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
if current := c.defaultModelConfig; current != nil && !c.clock.Now().Before(current.expiresAt) {
|
||||||
|
c.defaultModelConfig = nil
|
||||||
|
}
|
||||||
|
c.mu.Unlock()
|
||||||
|
|
||||||
|
return database.ChatModelConfig{}, false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) defaultModelConfigSnapshot() modelConfigSnapshot {
|
||||||
|
c.mu.RLock()
|
||||||
|
snap := modelConfigSnapshot{
|
||||||
|
epoch: c.modelTopologyEpoch,
|
||||||
|
generation: c.defaultModelConfigGeneration,
|
||||||
|
}
|
||||||
|
c.mu.RUnlock()
|
||||||
|
return snap
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) storeDefaultModelConfig(snap modelConfigSnapshot, config database.ChatModelConfig) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
|
||||||
|
if c.modelTopologyEpoch != snap.epoch {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if c.defaultModelConfigGeneration != snap.generation {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
c.defaultModelConfig = &cachedModelConfig{
|
||||||
|
config: cloneModelConfig(config),
|
||||||
|
expiresAt: c.clock.Now().Add(chatConfigModelConfigTTL),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) UserPrompt(ctx context.Context, userID uuid.UUID) (string, error) {
|
||||||
|
if prompt, ok := c.cachedUserPrompt(userID); ok {
|
||||||
|
return prompt, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
epoch := c.currentUserPromptEpoch()
|
||||||
|
prompt, err := singleflightDoChan(ctx, &c.userPromptFetches, fmt.Sprintf("%d:%s", epoch, userID), func() (string, error) {
|
||||||
|
if cached, ok := c.cachedUserPrompt(userID); ok {
|
||||||
|
return cached, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
fetched, err := c.db.GetUserChatCustomPrompt(c.ctx, userID)
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, sql.ErrNoRows) {
|
||||||
|
c.storeUserPrompt(epoch, userID, "")
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
c.storeUserPrompt(epoch, userID, fetched)
|
||||||
|
return fetched, nil
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
return prompt, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) cachedUserPrompt(userID uuid.UUID) (string, bool) {
|
||||||
|
prompt, _, ok := c.userPrompts.Get(userID)
|
||||||
|
if !ok {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
return prompt, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) currentUserPromptEpoch() uint64 {
|
||||||
|
c.mu.RLock()
|
||||||
|
epoch := c.userPromptEpoch
|
||||||
|
c.mu.RUnlock()
|
||||||
|
return epoch
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) storeUserPrompt(epoch uint64, userID uuid.UUID, prompt string) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
|
||||||
|
if c.userPromptEpoch != epoch {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
c.userPrompts.Set(userID, prompt, chatConfigUserPromptTTL)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) InvalidateModelConfig(id uuid.UUID) {
|
||||||
|
c.mu.Lock()
|
||||||
|
delete(c.modelConfigs, id)
|
||||||
|
c.modelTopologyEpoch++
|
||||||
|
c.defaultModelConfig = nil
|
||||||
|
c.defaultModelConfigGeneration++
|
||||||
|
c.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *chatConfigCache) InvalidateUserPrompt(userID uuid.UUID) {
|
||||||
|
c.mu.Lock()
|
||||||
|
c.userPrompts.Delete(userID)
|
||||||
|
c.userPromptEpoch++
|
||||||
|
c.mu.Unlock()
|
||||||
|
}
|
||||||
@@ -0,0 +1,978 @@
|
|||||||
|
package chatd //nolint:testpackage // Uses internal cache state.
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
"golang.org/x/xerrors"
|
||||||
|
|
||||||
|
"github.com/coder/coder/v2/coderd/database"
|
||||||
|
"github.com/coder/coder/v2/testutil"
|
||||||
|
"github.com/coder/quartz"
|
||||||
|
)
|
||||||
|
|
||||||
|
type stubChatConfigStore struct {
|
||||||
|
database.Store
|
||||||
|
|
||||||
|
getEnabledChatProviders func(context.Context) ([]database.ChatProvider, error)
|
||||||
|
getChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error)
|
||||||
|
getDefaultChatModelConfig func(context.Context) (database.ChatModelConfig, error)
|
||||||
|
getUserChatCustomPrompt func(context.Context, uuid.UUID) (string, error)
|
||||||
|
|
||||||
|
enabledProvidersCalls atomic.Int32
|
||||||
|
modelConfigByIDCalls atomic.Int32
|
||||||
|
defaultModelConfigCall atomic.Int32
|
||||||
|
userPromptCalls atomic.Int32
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubChatConfigStore) GetEnabledChatProviders(ctx context.Context) ([]database.ChatProvider, error) {
|
||||||
|
s.enabledProvidersCalls.Add(1)
|
||||||
|
if s.getEnabledChatProviders == nil {
|
||||||
|
panic("unexpected GetEnabledChatProviders call")
|
||||||
|
}
|
||||||
|
return s.getEnabledChatProviders(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubChatConfigStore) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
s.modelConfigByIDCalls.Add(1)
|
||||||
|
if s.getChatModelConfigByID == nil {
|
||||||
|
panic("unexpected GetChatModelConfigByID call")
|
||||||
|
}
|
||||||
|
return s.getChatModelConfigByID(ctx, id)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubChatConfigStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) {
|
||||||
|
s.defaultModelConfigCall.Add(1)
|
||||||
|
if s.getDefaultChatModelConfig == nil {
|
||||||
|
panic("unexpected GetDefaultChatModelConfig call")
|
||||||
|
}
|
||||||
|
return s.getDefaultChatModelConfig(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stubChatConfigStore) GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error) {
|
||||||
|
s.userPromptCalls.Add(1)
|
||||||
|
if s.getUserChatCustomPrompt == nil {
|
||||||
|
panic("unexpected GetUserChatCustomPrompt call")
|
||||||
|
}
|
||||||
|
return s.getUserChatCustomPrompt(ctx, userID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_EnabledProviders_CacheHit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
providers := []database.ChatProvider{testChatProvider("provider-a")}
|
||||||
|
store := &stubChatConfigStore{
|
||||||
|
getEnabledChatProviders: func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
return providers, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
second, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.Equal(t, providers, first)
|
||||||
|
require.Equal(t, providers, second)
|
||||||
|
require.Equal(t, int32(1), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_EnabledProviders_TTLExpiry(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
call := store.enabledProvidersCalls.Load()
|
||||||
|
return []database.ChatProvider{testChatProvider(fmt.Sprintf("provider-%d", call))}, nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
clock.Advance(chatConfigProvidersTTL).MustWait(ctx)
|
||||||
|
second, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, first, second)
|
||||||
|
require.Equal(t, int32(2), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_EnabledProviders_Invalidation(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
call := store.enabledProvidersCalls.Load()
|
||||||
|
return []database.ChatProvider{testChatProvider(fmt.Sprintf("provider-%d", call))}, nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
second, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, first, second)
|
||||||
|
require.Equal(t, int32(2), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_ModelConfigByID_CacheHit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
config := testChatModelConfig(configID, "model-a")
|
||||||
|
store := &stubChatConfigStore{
|
||||||
|
getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
return config, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
second, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.Equal(t, config, first)
|
||||||
|
require.Equal(t, config, second)
|
||||||
|
require.Equal(t, int32(1), store.modelConfigByIDCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_ModelConfigByID_ClonesOptionsForCache(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
const options = `{"temperature":0.1}`
|
||||||
|
config := testChatModelConfig(configID, "model-a")
|
||||||
|
config.Options = []byte(options)
|
||||||
|
store := &stubChatConfigStore{
|
||||||
|
getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
return config, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
// First call populates cache via singleflight.
|
||||||
|
first, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
first.Options[0] = 'x' // mutate singleflight return
|
||||||
|
|
||||||
|
// Second call is a cache hit.
|
||||||
|
second, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, options, string(second.Options))
|
||||||
|
second.Options[0] = 'y' // mutate cache-hit return
|
||||||
|
|
||||||
|
// Third call is another cache hit — must be unaffected.
|
||||||
|
third, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, options, string(third.Options))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_ModelConfigByID_NotFound(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
store := &stubChatConfigStore{
|
||||||
|
getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
return database.ChatModelConfig{}, sql.ErrNoRows
|
||||||
|
},
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
_, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.ErrorIs(t, err, sql.ErrNoRows)
|
||||||
|
_, err = cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.ErrorIs(t, err, sql.ErrNoRows)
|
||||||
|
|
||||||
|
require.Equal(t, int32(2), store.modelConfigByIDCalls.Load())
|
||||||
|
_, ok := cache.modelConfigs[configID]
|
||||||
|
require.False(t, ok)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateModelConfig_CascadesToDefault(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
config := testChatModelConfig(configID, "model-a")
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) {
|
||||||
|
call := store.defaultModelConfigCall.Load()
|
||||||
|
return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
_, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
firstDefault, err := cache.DefaultModelConfig(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
cache.InvalidateModelConfig(configID)
|
||||||
|
require.Nil(t, cache.defaultModelConfig)
|
||||||
|
|
||||||
|
secondDefault, err := cache.DefaultModelConfig(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, firstDefault, secondDefault)
|
||||||
|
require.Equal(t, int32(2), store.defaultModelConfigCall.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_UserPrompt_NegativeCaching(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
userID := uuid.New()
|
||||||
|
store := &stubChatConfigStore{
|
||||||
|
getUserChatCustomPrompt: func(context.Context, uuid.UUID) (string, error) {
|
||||||
|
return "", sql.ErrNoRows
|
||||||
|
},
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
second, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.Empty(t, first)
|
||||||
|
require.Empty(t, second)
|
||||||
|
require.Equal(t, int32(1), store.userPromptCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_UserPrompt_ExpiredEntryRefetches(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
userID := uuid.New()
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) {
|
||||||
|
call := store.userPromptCalls.Load()
|
||||||
|
return fmt.Sprintf("prompt-%d", call), nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
cache.userPrompts.Set(userID, "stale", 0)
|
||||||
|
|
||||||
|
first, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
second, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.Equal(t, "prompt-1", first)
|
||||||
|
require.Equal(t, first, second)
|
||||||
|
require.Equal(t, int32(1), store.userPromptCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateUserPrompt(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
userID := uuid.New()
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) {
|
||||||
|
call := store.userPromptCalls.Load()
|
||||||
|
return fmt.Sprintf("prompt-%d", call), nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
cache.InvalidateUserPrompt(userID)
|
||||||
|
second, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, first, second)
|
||||||
|
require.Equal(t, int32(2), store.userPromptCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateUserPrompt_BlocksStaleInFlightPrompt(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
userID := uuid.New()
|
||||||
|
const stalePrompt = "stale prompt"
|
||||||
|
const freshPrompt = "fresh prompt"
|
||||||
|
firstStarted := make(chan struct{})
|
||||||
|
secondStarted := make(chan struct{})
|
||||||
|
releaseFirst := make(chan struct{})
|
||||||
|
releaseSecond := make(chan struct{})
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getUserChatCustomPrompt = func(context.Context, uuid.UUID) (string, error) {
|
||||||
|
switch call := store.userPromptCalls.Load(); call {
|
||||||
|
case 1:
|
||||||
|
close(firstStarted)
|
||||||
|
<-releaseFirst
|
||||||
|
return stalePrompt, nil
|
||||||
|
case 2:
|
||||||
|
close(secondStarted)
|
||||||
|
<-releaseSecond
|
||||||
|
return freshPrompt, nil
|
||||||
|
default:
|
||||||
|
return "", xerrors.Errorf("unexpected user prompt call %d", call)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
type result struct {
|
||||||
|
prompt string
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
firstResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
prompt, err := cache.UserPrompt(ctx, userID)
|
||||||
|
firstResult <- result{prompt: prompt, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, firstStarted)
|
||||||
|
cache.InvalidateUserPrompt(userID)
|
||||||
|
|
||||||
|
secondResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
prompt, err := cache.UserPrompt(ctx, userID)
|
||||||
|
secondResult <- result{prompt: prompt, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, secondStarted)
|
||||||
|
close(releaseFirst)
|
||||||
|
first := <-firstResult
|
||||||
|
require.NoError(t, first.err)
|
||||||
|
require.Equal(t, stalePrompt, first.prompt)
|
||||||
|
_, _, ok := cache.userPrompts.Get(userID)
|
||||||
|
require.False(t, ok)
|
||||||
|
|
||||||
|
close(releaseSecond)
|
||||||
|
second := <-secondResult
|
||||||
|
require.NoError(t, second.err)
|
||||||
|
require.Equal(t, freshPrompt, second.prompt)
|
||||||
|
require.Equal(t, int32(2), store.userPromptCalls.Load())
|
||||||
|
|
||||||
|
third, err := cache.UserPrompt(ctx, userID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, freshPrompt, third)
|
||||||
|
require.Equal(t, int32(2), store.userPromptCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_Singleflight(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
providers := []database.ChatProvider{testChatProvider("provider-a")}
|
||||||
|
fetchStarted := make(chan struct{})
|
||||||
|
releaseFetch := make(chan struct{})
|
||||||
|
var startedOnce sync.Once
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
startedOnce.Do(func() { close(fetchStarted) })
|
||||||
|
<-releaseFetch
|
||||||
|
return providers, nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
const callers = 8
|
||||||
|
results := make([][]database.ChatProvider, callers)
|
||||||
|
errs := make([]error, callers)
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
start := make(chan struct{})
|
||||||
|
for i := 0; i < callers; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(i int) {
|
||||||
|
defer wg.Done()
|
||||||
|
<-start
|
||||||
|
results[i], errs[i] = cache.EnabledProviders(ctx)
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
|
||||||
|
close(start)
|
||||||
|
waitForSignal(t, fetchStarted)
|
||||||
|
close(releaseFetch)
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
for i := 0; i < callers; i++ {
|
||||||
|
require.NoError(t, errs[i])
|
||||||
|
require.Equal(t, providers, results[i])
|
||||||
|
}
|
||||||
|
require.Equal(t, int32(1), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_GenerationPreventsStaleWrite(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
firstProviders := []database.ChatProvider{testChatProvider("provider-a")}
|
||||||
|
secondProviders := []database.ChatProvider{testChatProvider("provider-b")}
|
||||||
|
fetchStarted := make(chan struct{})
|
||||||
|
releaseFetch := make(chan struct{})
|
||||||
|
var startedOnce sync.Once
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
call := store.enabledProvidersCalls.Load()
|
||||||
|
if call == 1 {
|
||||||
|
startedOnce.Do(func() { close(fetchStarted) })
|
||||||
|
<-releaseFetch
|
||||||
|
return firstProviders, nil
|
||||||
|
}
|
||||||
|
return secondProviders, nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
resultCh := make(chan []database.ChatProvider, 1)
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
providers, err := cache.EnabledProviders(ctx)
|
||||||
|
if err != nil {
|
||||||
|
errCh <- err
|
||||||
|
return
|
||||||
|
}
|
||||||
|
resultCh <- providers
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, fetchStarted)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
close(releaseFetch)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-errCh:
|
||||||
|
require.NoError(t, err)
|
||||||
|
case providers := <-resultCh:
|
||||||
|
require.Equal(t, firstProviders, providers)
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("timed out waiting for in-flight fetch")
|
||||||
|
}
|
||||||
|
|
||||||
|
require.Nil(t, cache.providers)
|
||||||
|
second, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, secondProviders, second)
|
||||||
|
require.Equal(t, int32(2), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateProviders_BlocksStaleInFlightProviders(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
staleProviders := []database.ChatProvider{testChatProvider("provider-stale")}
|
||||||
|
freshProviders := []database.ChatProvider{testChatProvider("provider-fresh")}
|
||||||
|
firstStarted := make(chan struct{})
|
||||||
|
secondStarted := make(chan struct{})
|
||||||
|
releaseFirst := make(chan struct{})
|
||||||
|
releaseSecond := make(chan struct{})
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getEnabledChatProviders = func(context.Context) ([]database.ChatProvider, error) {
|
||||||
|
switch call := store.enabledProvidersCalls.Load(); call {
|
||||||
|
case 1:
|
||||||
|
close(firstStarted)
|
||||||
|
<-releaseFirst
|
||||||
|
return staleProviders, nil
|
||||||
|
case 2:
|
||||||
|
close(secondStarted)
|
||||||
|
<-releaseSecond
|
||||||
|
return freshProviders, nil
|
||||||
|
default:
|
||||||
|
return nil, xerrors.Errorf("unexpected provider call %d", call)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
type result struct {
|
||||||
|
providers []database.ChatProvider
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
firstResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
providers, err := cache.EnabledProviders(ctx)
|
||||||
|
firstResult <- result{providers: providers, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, firstStarted)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
|
||||||
|
secondResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
providers, err := cache.EnabledProviders(ctx)
|
||||||
|
secondResult <- result{providers: providers, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, secondStarted)
|
||||||
|
close(releaseFirst)
|
||||||
|
first := <-firstResult
|
||||||
|
require.NoError(t, first.err)
|
||||||
|
require.Equal(t, staleProviders, first.providers)
|
||||||
|
require.Nil(t, cache.providers)
|
||||||
|
|
||||||
|
close(releaseSecond)
|
||||||
|
second := <-secondResult
|
||||||
|
require.NoError(t, second.err)
|
||||||
|
require.Equal(t, freshProviders, second.providers)
|
||||||
|
require.Equal(t, int32(2), store.enabledProvidersCalls.Load())
|
||||||
|
|
||||||
|
third, err := cache.EnabledProviders(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, freshProviders, third)
|
||||||
|
require.Equal(t, int32(2), store.enabledProvidersCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateProviders_CascadesToModelConfigs(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
call := store.modelConfigByIDCalls.Load()
|
||||||
|
return testChatModelConfig(configID, fmt.Sprintf("model-%d", call)), nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
second, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, first, second)
|
||||||
|
require.Equal(t, int32(2), store.modelConfigByIDCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateProviders_CascadesToDefaultModelConfig(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitShort)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) {
|
||||||
|
call := store.defaultModelConfigCall.Load()
|
||||||
|
return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
first, err := cache.DefaultModelConfig(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
second, err := cache.DefaultModelConfig(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
require.NotEqual(t, first, second)
|
||||||
|
require.Equal(t, int32(2), store.defaultModelConfigCall.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfigCache_InvalidateProviders_BlocksStaleInFlightModelConfig(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
configID := uuid.New()
|
||||||
|
staleConfig := testChatModelConfig(configID, "stale-model")
|
||||||
|
freshConfig := testChatModelConfig(configID, "fresh-model")
|
||||||
|
firstStarted := make(chan struct{})
|
||||||
|
secondStarted := make(chan struct{})
|
||||||
|
releaseFirst := make(chan struct{})
|
||||||
|
releaseSecond := make(chan struct{})
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
switch call := store.modelConfigByIDCalls.Load(); call {
|
||||||
|
case 1:
|
||||||
|
close(firstStarted)
|
||||||
|
<-releaseFirst
|
||||||
|
return staleConfig, nil
|
||||||
|
case 2:
|
||||||
|
close(secondStarted)
|
||||||
|
<-releaseSecond
|
||||||
|
return freshConfig, nil
|
||||||
|
default:
|
||||||
|
return database.ChatModelConfig{}, xerrors.Errorf("unexpected model config call %d", call)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
type result struct {
|
||||||
|
config database.ChatModelConfig
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
firstResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
config, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
firstResult <- result{config: config, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, firstStarted)
|
||||||
|
cache.InvalidateProviders()
|
||||||
|
|
||||||
|
secondResult := make(chan result, 1)
|
||||||
|
go func() {
|
||||||
|
config, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
secondResult <- result{config: config, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, secondStarted)
|
||||||
|
close(releaseFirst)
|
||||||
|
first := <-firstResult
|
||||||
|
require.NoError(t, first.err)
|
||||||
|
require.Equal(t, staleConfig, first.config)
|
||||||
|
_, ok := cache.modelConfigs[configID]
|
||||||
|
require.False(t, ok)
|
||||||
|
|
||||||
|
close(releaseSecond)
|
||||||
|
second := <-secondResult
|
||||||
|
require.NoError(t, second.err)
|
||||||
|
require.Equal(t, freshConfig, second.config)
|
||||||
|
require.Equal(t, int32(2), store.modelConfigByIDCalls.Load())
|
||||||
|
|
||||||
|
third, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, freshConfig, third)
|
||||||
|
require.Equal(t, int32(2), store.modelConfigByIDCalls.Load())
|
||||||
|
}
|
||||||
|
|
||||||
|
func testChatProvider(name string) database.ChatProvider {
|
||||||
|
return database.ChatProvider{
|
||||||
|
ID: uuid.New(),
|
||||||
|
Provider: name,
|
||||||
|
DisplayName: name,
|
||||||
|
Enabled: true,
|
||||||
|
CreatedAt: time.Unix(0, 0).UTC(),
|
||||||
|
UpdatedAt: time.Unix(0, 0).UTC(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func testChatModelConfig(id uuid.UUID, model string) database.ChatModelConfig {
|
||||||
|
return database.ChatModelConfig{
|
||||||
|
ID: id,
|
||||||
|
Provider: "openai",
|
||||||
|
Model: model,
|
||||||
|
DisplayName: model,
|
||||||
|
Enabled: true,
|
||||||
|
CreatedAt: time.Unix(0, 0).UTC(),
|
||||||
|
UpdatedAt: time.Unix(0, 0).UTC(),
|
||||||
|
ContextLimit: 128000,
|
||||||
|
CompressionThreshold: 64000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func waitForSignal(t *testing.T, ch <-chan struct{}) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-ch:
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("timed out waiting for signal")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestConfigCache_CallerCancellation verifies the DoChan-based
|
||||||
|
// cancellation semantics across all four cache methods:
|
||||||
|
// - A canceled caller returns immediately without waiting for the
|
||||||
|
// shared fill to complete.
|
||||||
|
// - One canceled waiter does not poison other coalesced waiters.
|
||||||
|
// - Server context cancellation propagates through the fill.
|
||||||
|
func TestConfigCache_CallerCancellation(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
type cacheMethod struct {
|
||||||
|
name string
|
||||||
|
// setupBlocked configures the store to block on release.
|
||||||
|
// The started channel is closed when the fill enters the
|
||||||
|
// store. The release channel unblocks the store.
|
||||||
|
setupBlocked func(store *stubChatConfigStore, started, release chan struct{})
|
||||||
|
// setupCtxSensitive configures the store to block until
|
||||||
|
// its context is canceled (for server-shutdown testing).
|
||||||
|
setupCtxSensitive func(store *stubChatConfigStore, started chan struct{})
|
||||||
|
// call invokes the cache method under test.
|
||||||
|
call func(ctx context.Context, cache *chatConfigCache) error
|
||||||
|
// storeCalls returns the number of underlying store calls.
|
||||||
|
storeCalls func(store *stubChatConfigStore) int32
|
||||||
|
}
|
||||||
|
|
||||||
|
configID := uuid.New()
|
||||||
|
userID := uuid.New()
|
||||||
|
|
||||||
|
methods := []cacheMethod{
|
||||||
|
{
|
||||||
|
name: "EnabledProviders",
|
||||||
|
setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getEnabledChatProviders = func(ctx context.Context) ([]database.ChatProvider, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return nil, ctx.Err()
|
||||||
|
case <-release:
|
||||||
|
return []database.ChatProvider{testChatProvider("p")}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getEnabledChatProviders = func(ctx context.Context) ([]database.ChatProvider, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
<-ctx.Done()
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
},
|
||||||
|
call: func(ctx context.Context, cache *chatConfigCache) error {
|
||||||
|
_, err := cache.EnabledProviders(ctx)
|
||||||
|
return err
|
||||||
|
},
|
||||||
|
storeCalls: func(store *stubChatConfigStore) int32 {
|
||||||
|
return store.enabledProvidersCalls.Load()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "ModelConfigByID",
|
||||||
|
setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getChatModelConfigByID = func(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return database.ChatModelConfig{}, ctx.Err()
|
||||||
|
case <-release:
|
||||||
|
return testChatModelConfig(id, "model"), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getChatModelConfigByID = func(ctx context.Context, _ uuid.UUID) (database.ChatModelConfig, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
<-ctx.Done()
|
||||||
|
return database.ChatModelConfig{}, ctx.Err()
|
||||||
|
}
|
||||||
|
},
|
||||||
|
call: func(ctx context.Context, cache *chatConfigCache) error {
|
||||||
|
_, err := cache.ModelConfigByID(ctx, configID)
|
||||||
|
return err
|
||||||
|
},
|
||||||
|
storeCalls: func(store *stubChatConfigStore) int32 {
|
||||||
|
return store.modelConfigByIDCalls.Load()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "DefaultModelConfig",
|
||||||
|
setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return database.ChatModelConfig{}, ctx.Err()
|
||||||
|
case <-release:
|
||||||
|
return testChatModelConfig(uuid.New(), "default"), nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
<-ctx.Done()
|
||||||
|
return database.ChatModelConfig{}, ctx.Err()
|
||||||
|
}
|
||||||
|
},
|
||||||
|
call: func(ctx context.Context, cache *chatConfigCache) error {
|
||||||
|
_, err := cache.DefaultModelConfig(ctx)
|
||||||
|
return err
|
||||||
|
},
|
||||||
|
storeCalls: func(store *stubChatConfigStore) int32 {
|
||||||
|
return store.defaultModelConfigCall.Load()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "UserPrompt",
|
||||||
|
setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getUserChatCustomPrompt = func(ctx context.Context, _ uuid.UUID) (string, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
return "", ctx.Err()
|
||||||
|
case <-release:
|
||||||
|
return "custom prompt", nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) {
|
||||||
|
var once sync.Once
|
||||||
|
store.getUserChatCustomPrompt = func(ctx context.Context, _ uuid.UUID) (string, error) {
|
||||||
|
once.Do(func() { close(started) })
|
||||||
|
<-ctx.Done()
|
||||||
|
return "", ctx.Err()
|
||||||
|
}
|
||||||
|
},
|
||||||
|
call: func(ctx context.Context, cache *chatConfigCache) error {
|
||||||
|
_, err := cache.UserPrompt(ctx, userID)
|
||||||
|
return err
|
||||||
|
},
|
||||||
|
storeCalls: func(store *stubChatConfigStore) int32 {
|
||||||
|
return store.userPromptCalls.Load()
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test A: A canceled caller stops waiting immediately; the
|
||||||
|
// shared fill still completes and populates the cache.
|
||||||
|
t.Run("CanceledCallerStopsWaiting", func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
for _, m := range methods {
|
||||||
|
t.Run(m.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
started := make(chan struct{})
|
||||||
|
release := make(chan struct{})
|
||||||
|
m.setupBlocked(store, started, release)
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
callerCtx, callerCancel := context.WithCancel(ctx)
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
errCh <- m.call(callerCtx, cache)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Wait for the fill to enter the store, then
|
||||||
|
// cancel the caller's context.
|
||||||
|
waitForSignal(t, started)
|
||||||
|
callerCancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-errCh:
|
||||||
|
require.ErrorIs(t, err, context.Canceled)
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("canceled caller did not return promptly")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Release the store so the fill can complete.
|
||||||
|
close(release)
|
||||||
|
|
||||||
|
// A fresh call must succeed — either a cache
|
||||||
|
// hit or by joining the still-in-flight fill.
|
||||||
|
// Only one store call should have occurred.
|
||||||
|
require.NoError(t, m.call(ctx, cache))
|
||||||
|
require.Equal(t, int32(1), m.storeCalls(store))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
// Test B: One canceled waiter does not poison other coalesced
|
||||||
|
// waiters sharing the same singleflight entry.
|
||||||
|
t.Run("CanceledWaiterDoesNotPoisonOthers", func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
for _, m := range methods {
|
||||||
|
t.Run(m.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
started := make(chan struct{})
|
||||||
|
release := make(chan struct{})
|
||||||
|
m.setupBlocked(store, started, release)
|
||||||
|
cache := newChatConfigCache(ctx, store, clock)
|
||||||
|
|
||||||
|
cancelCtx, cancel := context.WithCancel(ctx)
|
||||||
|
cancelErrCh := make(chan error, 1)
|
||||||
|
survivorErrCh := make(chan error, 1)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
cancelErrCh <- m.call(cancelCtx, cache)
|
||||||
|
}()
|
||||||
|
go func() {
|
||||||
|
survivorErrCh <- m.call(ctx, cache)
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, started)
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-cancelErrCh:
|
||||||
|
require.ErrorIs(t, err, context.Canceled)
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("canceled caller did not return promptly")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Release the store; the surviving waiter
|
||||||
|
// must receive the successful result.
|
||||||
|
close(release)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-survivorErrCh:
|
||||||
|
require.NoError(t, err)
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("survivor caller did not return")
|
||||||
|
}
|
||||||
|
|
||||||
|
require.Equal(t, int32(1), m.storeCalls(store))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
// Test C: Server context cancellation propagates through the
|
||||||
|
// fill, ensuring graceful shutdown behavior is preserved.
|
||||||
|
t.Run("ServerCancellation", func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
for _, m := range methods {
|
||||||
|
t.Run(m.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
clock := quartz.NewMock(t)
|
||||||
|
store := &stubChatConfigStore{}
|
||||||
|
started := make(chan struct{})
|
||||||
|
m.setupCtxSensitive(store, started)
|
||||||
|
|
||||||
|
serverCtx, serverCancel := context.WithCancel(context.Background())
|
||||||
|
defer serverCancel()
|
||||||
|
cache := newChatConfigCache(serverCtx, store, clock)
|
||||||
|
|
||||||
|
callerCtx := testutil.Context(t, testutil.WaitMedium)
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
errCh <- m.call(callerCtx, cache)
|
||||||
|
}()
|
||||||
|
|
||||||
|
waitForSignal(t, started)
|
||||||
|
serverCancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-errCh:
|
||||||
|
require.ErrorIs(t, err, context.Canceled)
|
||||||
|
case <-time.After(testutil.WaitShort):
|
||||||
|
t.Fatal("caller did not return after server cancel")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -72,7 +72,7 @@ func (p *Server) isAnthropicConfigured(ctx context.Context) bool {
|
|||||||
if p.providerAPIKeys.APIKey("anthropic") != "" {
|
if p.providerAPIKeys.APIKey("anthropic") != "" {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
dbProviders, err := p.db.GetEnabledChatProviders(ctx)
|
dbProviders, err := p.configCache.EnabledProviders(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user