diff --git a/core_service.go b/core_service.go index 1f42c56..fd7ce11 100644 --- a/core_service.go +++ b/core_service.go @@ -18,6 +18,7 @@ type CoreService struct { missingKeys []MissingKey missingKeysMu sync.Mutex + hookInstalled bool } // ServiceOptions configures the i18n Core service. @@ -81,11 +82,19 @@ func NewCoreService(opts ServiceOptions) func(*core.Core) (any, error) { // OnStartup initialises the i18n service. func (s *CoreService) OnStartup(_ context.Context) error { if s.svc.Mode() == ModeCollect { - OnMissingKey(s.handleMissingKey) + s.ensureMissingKeyCollector() } return nil } +func (s *CoreService) ensureMissingKeyCollector() { + if s.hookInstalled { + return + } + appendMissingKeyHandler(s.handleMissingKey) + s.hookInstalled = true +} + func (s *CoreService) handleMissingKey(mk MissingKey) { s.missingKeysMu.Lock() defer s.missingKeysMu.Unlock() @@ -112,9 +121,7 @@ func (s *CoreService) ClearMissingKeys() { func (s *CoreService) SetMode(mode Mode) { s.svc.SetMode(mode) if mode == ModeCollect { - OnMissingKey(s.handleMissingKey) - } else { - OnMissingKey(nil) + s.ensureMissingKeyCollector() } } diff --git a/hooks.go b/hooks.go index 643838b..de49b5f 100644 --- a/hooks.go +++ b/hooks.go @@ -11,6 +11,10 @@ import ( var missingKeyHandler atomic.Value +type missingKeyHandlersState struct { + handlers []MissingKeyHandler +} + type localeRegistration struct { fsys fs.FS dir string @@ -55,20 +59,46 @@ func loadRegisteredLocales(svc *Service) { // OnMissingKey registers a handler for missing translation keys. func OnMissingKey(h MissingKeyHandler) { - missingKeyHandler.Store(h) + if h == nil { + missingKeyHandler.Store(missingKeyHandlersState{}) + return + } + missingKeyHandler.Store(missingKeyHandlersState{handlers: []MissingKeyHandler{h}}) +} + +func appendMissingKeyHandler(h MissingKeyHandler) { + if h == nil { + return + } + current := missingKeyHandlers() + current.handlers = append(current.handlers, h) + missingKeyHandler.Store(current) +} + +func missingKeyHandlers() missingKeyHandlersState { + v := missingKeyHandler.Load() + if v == nil { + return missingKeyHandlersState{} + } + state, ok := v.(missingKeyHandlersState) + if !ok { + return missingKeyHandlersState{} + } + return state } func dispatchMissingKey(key string, args map[string]any) { - v := missingKeyHandler.Load() - if v == nil { - return - } - h, ok := v.(MissingKeyHandler) - if !ok || h == nil { + state := missingKeyHandlers() + if len(state.handlers) == 0 { return } file, line := missingKeyCaller() - h(MissingKey{Key: key, Args: args, CallerFile: file, CallerLine: line}) + mk := MissingKey{Key: key, Args: args, CallerFile: file, CallerLine: line} + for _, h := range state.handlers { + if h != nil { + h(mk) + } + } } func missingKeyCaller() (string, int) { diff --git a/hooks_test.go b/hooks_test.go index 22bd234..7d05b23 100644 --- a/hooks_test.go +++ b/hooks_test.go @@ -324,9 +324,44 @@ func TestOnMissingKey_Good_TranslationContextArgs(t *testing.T) { } func TestDispatchMissingKey_Good_NoHandler(t *testing.T) { - // Store nil handler (using correct type) - missingKeyHandler.Store(MissingKeyHandler(nil)) + // Reset to the empty handler set. + OnMissingKey(nil) // Should not panic when dispatching with nil handler dispatchMissingKey("test.key", nil) } + +func TestCoreServiceSetMode_Good_PreservesMissingKeyHandlers(t *testing.T) { + svc, err := New() + require.NoError(t, err) + + prev := missingKeyHandlers() + t.Cleanup(func() { + missingKeyHandler.Store(prev) + }) + + var observed int + OnMissingKey(func(MissingKey) { + observed++ + }) + t.Cleanup(func() { + OnMissingKey(nil) + }) + + coreSvc := &CoreService{svc: svc} + coreSvc.SetMode(ModeCollect) + + _ = svc.T("missing.core.service.key") + + if observed != 1 { + t.Fatalf("custom missing key handler called %d times, want 1", observed) + } + + missing := coreSvc.MissingKeys() + if len(missing) != 1 { + t.Fatalf("CoreService captured %d missing keys, want 1", len(missing)) + } + if missing[0].Key != "missing.core.service.key" { + t.Fatalf("captured missing key = %q, want %q", missing[0].Key, "missing.core.service.key") + } +}