From f6862071b9364949152986df9223b2f44a5261b0 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Wed, 17 Sep 2025 16:15:20 -0700 Subject: [PATCH] Update registry.go --- go/internal/registry/registry.go | 88 +++++++++++++++++++------------- 1 file changed, 52 insertions(+), 36 deletions(-) diff --git a/go/internal/registry/registry.go b/go/internal/registry/registry.go index 6851d2ef6c..5545383dc4 100644 --- a/go/internal/registry/registry.go +++ b/go/internal/registry/registry.go @@ -31,12 +31,12 @@ import ( // Registry holds all registered actions and associated types, // and provides methods to register, query, and look up actions. type Registry struct { - mu sync.Mutex - parent api.Registry // parent registry for hierarchical lookups - actions map[string]api.Action - plugins map[string]api.Plugin - values map[string]any // Values can truly be anything. - + mu sync.RWMutex + resolveMu sync.RWMutex + parent api.Registry + actions map[string]api.Action + plugins map[string]api.Plugin + values map[string]any // Values can truly be anything. dotprompt *dotprompt.Dotprompt } @@ -102,8 +102,8 @@ func (r *Registry) RegisterAction(key string, action api.Action) { // It first checks the current registry, then falls back to the parent if not found. // Returns nil if the plugin is not found in the registry hierarchy. func (r *Registry) LookupPlugin(name string) api.Plugin { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() if plugin, ok := r.plugins[name]; ok { return plugin @@ -132,8 +132,8 @@ func (r *Registry) RegisterValue(name string, value any) { // It first checks the current registry, then falls back to the parent if not found. // Returns nil if the value is not found in the registry hierarchy. func (r *Registry) LookupValue(name string) any { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() if value, ok := r.values[name]; ok { return value @@ -149,8 +149,8 @@ func (r *Registry) LookupValue(name string) any { // LookupAction returns the action for the given key. // It first checks the current registry, then falls back to the parent if not found. func (r *Registry) LookupAction(key string) api.Action { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() if action, ok := r.actions[key]; ok { return action @@ -165,35 +165,47 @@ func (r *Registry) LookupAction(key string) api.Action { // ResolveAction looks up an action by key. If the action is not found, it attempts dynamic resolution. // Returns the action if found, or nil if not found. +// This method is safe to call concurrently and uses a single mutex to serialize all resolution operations. func (r *Registry) ResolveAction(key string) api.Action { action := r.LookupAction(key) - if action == nil { - typ, provider, name := api.ParseKey(key) - if typ == "" || name == "" { - slog.Debug("ResolveAction: failed to parse action key", "key", key) - return nil - } - plugins := r.ListPlugins() - for _, plugin := range plugins { - if dp, ok := plugin.(api.DynamicPlugin); ok && dp.Name() == provider { - resolvedAction := dp.ResolveAction(typ, name) - if resolvedAction != nil { - resolvedAction.Register(r) - } - break + if action != nil { + return action + } + + r.resolveMu.Lock() + defer r.resolveMu.Unlock() + + action = r.LookupAction(key) + if action != nil { + return action + } + + typ, provider, name := api.ParseKey(key) + if typ == "" || name == "" { + slog.Debug("ResolveAction: failed to parse action key", "key", key) + return nil + } + + plugins := r.ListPlugins() + for _, plugin := range plugins { + if dp, ok := plugin.(api.DynamicPlugin); ok && dp.Name() == provider { + resolvedAction := dp.ResolveAction(typ, name) + if resolvedAction != nil { + resolvedAction.Register(r) } + break } - action = r.LookupAction(key) } - return action + + return r.LookupAction(key) } // ListActions returns a list of all registered actions. // This includes actions from both the current registry and its parent hierarchy. // Child registry actions take precedence over parent actions with the same key. func (r *Registry) ListActions() []api.Action { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() var actions []api.Action for _, v := range r.actions { actions = append(actions, v) @@ -203,8 +215,8 @@ func (r *Registry) ListActions() []api.Action { // ListPlugins returns a list of all registered plugins. func (r *Registry) ListPlugins() []api.Plugin { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() var plugins []api.Plugin for _, p := range r.plugins { plugins = append(plugins, p) @@ -216,8 +228,8 @@ func (r *Registry) ListPlugins() []api.Plugin { // This includes values from both the current registry and its parent hierarchy. // Child registry values take precedence over parent values with the same key. func (r *Registry) ListValues() map[string]any { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() allValues := make(map[string]any) @@ -233,6 +245,8 @@ func (r *Registry) ListValues() map[string]any { // RegisterPartial adds the partial to the list of partials to the dotprompt instance func (r *Registry) RegisterPartial(name string, source string) { + r.mu.Lock() + defer r.mu.Unlock() if r.dotprompt == nil { r.dotprompt = dotprompt.NewDotprompt(nil) } @@ -247,6 +261,8 @@ func (r *Registry) RegisterPartial(name string, source string) { // RegisterHelper adds a helper function to the dotprompt instance func (r *Registry) RegisterHelper(name string, fn any) { + r.mu.Lock() + defer r.mu.Unlock() if r.dotprompt == nil { r.dotprompt = dotprompt.NewDotprompt(nil) } @@ -261,7 +277,7 @@ func (r *Registry) RegisterHelper(name string, fn any) { // Dotprompt returns the dotprompt instance. func (r *Registry) Dotprompt() *dotprompt.Dotprompt { - r.mu.Lock() - defer r.mu.Unlock() + r.mu.RLock() + defer r.mu.RUnlock() return r.dotprompt }