mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +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/db2sdk"
|
||||
"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/gitprovider"
|
||||
"github.com/coder/coder/v2/coderd/httpapi"
|
||||
@@ -110,6 +111,28 @@ func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error)
|
||||
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.
|
||||
func (api *API) watchChats(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
@@ -3052,6 +3075,8 @@ func (api *API) putUserChatCustomPrompt(rw http.ResponseWriter, r *http.Request)
|
||||
return
|
||||
}
|
||||
|
||||
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventUserPrompt, apiKey.UserID)
|
||||
|
||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.UserChatCustomPrompt{
|
||||
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(
|
||||
ctx,
|
||||
rw,
|
||||
@@ -3972,6 +3999,8 @@ func (api *API) updateChatProvider(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil)
|
||||
|
||||
httpapi.Write(
|
||||
ctx,
|
||||
rw,
|
||||
@@ -4026,6 +4055,8 @@ func (api *API) deleteChatProvider(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventProviders, uuid.Nil)
|
||||
|
||||
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))
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -4416,6 +4451,8 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventModelConfig, modelConfigID)
|
||||
|
||||
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
|
||||
webpushDispatcher webpush.Dispatcher
|
||||
providerAPIKeys chatprovider.ProviderAPIKeys
|
||||
configCache *chatConfigCache
|
||||
configCacheUnsubscribe func()
|
||||
|
||||
// chatStreams stores per-chat stream state. Using sync.Map
|
||||
// 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.
|
||||
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)
|
||||
|
||||
return p
|
||||
@@ -4017,7 +4044,7 @@ func (p *Server) resolveChatModel(
|
||||
})
|
||||
g.Go(func() error {
|
||||
var err error
|
||||
providers, err = p.db.GetEnabledChatProviders(ctx)
|
||||
providers, err = p.configCache.EnabledProviders(ctx)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get enabled chat providers: %w", err)
|
||||
}
|
||||
@@ -4061,7 +4088,7 @@ func (p *Server) resolveModelConfig(
|
||||
chat database.Chat,
|
||||
) (database.ChatModelConfig, error) {
|
||||
if chat.LastModelConfigID != uuid.Nil {
|
||||
modelConfig, err := p.db.GetChatModelConfigByID(
|
||||
modelConfig, err := p.configCache.ModelConfigByID(
|
||||
ctx, chat.LastModelConfigID,
|
||||
)
|
||||
if err == nil {
|
||||
@@ -4076,7 +4103,7 @@ func (p *Server) resolveModelConfig(
|
||||
// Model config was deleted, fall through to default.
|
||||
}
|
||||
|
||||
defaultConfig, err := p.db.GetDefaultChatModelConfig(ctx)
|
||||
defaultConfig, err := p.configCache.DefaultModelConfig(ctx)
|
||||
if err != nil {
|
||||
if xerrors.Is(err, sql.ErrNoRows) {
|
||||
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
|
||||
// string if no prompt is set.
|
||||
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 {
|
||||
// sql.ErrNoRows is the normal "not set" case.
|
||||
return ""
|
||||
@@ -4447,6 +4474,10 @@ func (p *Server) dispatchPush(
|
||||
|
||||
// Close stops the processor and waits for it to finish.
|
||||
func (p *Server) Close() error {
|
||||
if unsub := p.configCacheUnsubscribe; unsub != nil {
|
||||
p.configCacheUnsubscribe = nil
|
||||
unsub()
|
||||
}
|
||||
p.cancel()
|
||||
<-p.closed
|
||||
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") != "" {
|
||||
return true
|
||||
}
|
||||
dbProviders, err := p.db.GetEnabledChatProviders(ctx)
|
||||
dbProviders, err := p.configCache.EnabledProviders(ctx)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user