fix: address review comments on admin_script seeding

- Replace TOCTOU race (separate Load+Save) with atomic
  SaveJobTypeConfigIfNotExists on ConfigStore and Plugin
- Replace ineffective polling loop with single GetMaster call using
  30s context timeout, since GetMaster respects context cancellation
- Add unit tests for SaveJobTypeConfigIfNotExists (in-memory + on-disk)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Chris LuandClaude Opus 4.6 committed 2026-03-04 21:27:52 -08:00
1 parent 5432c03ba5
commit 4b2a231913
4 files changed
+146 -18

No files matched your search

+47
View File
@@ -256,6 +256,53 @@ func (s *ConfigStore) SaveJobTypeConfig(config *plugin_pb.PersistedJobTypeConfig
return nil
}
// SaveJobTypeConfigIfNotExists atomically checks whether a config for the
// given job type already exists and only persists config when none is found.
// Returns true if the config was saved, false if a config already existed.
func (s *ConfigStore) SaveJobTypeConfigIfNotExists(config *plugin_pb.PersistedJobTypeConfig) (bool, error) {
if config == nil {
return false, fmt.Errorf("job type config is nil")
}
if config.JobType == "" {
return false, fmt.Errorf("job type config has empty job_type")
}
sanitizedJobType, err := sanitizeJobType(config.JobType)
if err != nil {
return false, err
}
config.JobType = sanitizedJobType
clone := proto.Clone(config).(*plugin_pb.PersistedJobTypeConfig)
s.mu.Lock()
defer s.mu.Unlock()
if !s.configured {
if _, exists := s.memConfigs[config.JobType]; exists {
return false, nil
}
s.memConfigs[config.JobType] = clone
return true, nil
}
pbPath := filepath.Join(s.baseDir, jobTypesDirName, config.JobType, configPBFileName)
if _, statErr := os.Stat(pbPath); statErr == nil {
return false, nil
}
jobTypeDir, err := s.ensureJobTypeDir(config.JobType)
if err != nil {
return false, err
}
jsonPath := filepath.Join(jobTypeDir, configJSONFileName)
if err := writeProtoFiles(clone, filepath.Join(jobTypeDir, configPBFileName), jsonPath); err != nil {
return false, fmt.Errorf("save job type config for %s: %w", config.JobType, err)
}
return true, nil
}
func (s *ConfigStore) LoadJobTypeConfig(jobType string) (*plugin_pb.PersistedJobTypeConfig, error) {
if _, err := sanitizeJobType(jobType); err != nil {
return nil, err
+75
View File
@@ -208,6 +208,81 @@ func TestConfigStoreMonitorStateRoundTrip(t *testing.T) {
}
}
func TestConfigStoreSaveJobTypeConfigIfNotExists(t *testing.T) {
t.Parallel()
t.Run("in-memory", func(t *testing.T) {
t.Parallel()
store, err := NewConfigStore("")
if err != nil {
t.Fatalf("NewConfigStore: %v", err)
}
testSaveJobTypeConfigIfNotExists(t, store)
})
t.Run("on-disk", func(t *testing.T) {
t.Parallel()
store, err := NewConfigStore(t.TempDir())
if err != nil {
t.Fatalf("NewConfigStore: %v", err)
}
testSaveJobTypeConfigIfNotExists(t, store)
})
}
func testSaveJobTypeConfigIfNotExists(t *testing.T, store *ConfigStore) {
t.Helper()
cfg := &plugin_pb.PersistedJobTypeConfig{
JobType: "admin_script",
AdminRuntime: &plugin_pb.AdminRuntimeConfig{Enabled: true},
}
// First call should save.
saved, err := store.SaveJobTypeConfigIfNotExists(cfg)
if err != nil {
t.Fatalf("first SaveJobTypeConfigIfNotExists: %v", err)
}
if !saved {
t.Fatal("expected first call to save the config")
}
// Second call with same job type should not save.
saved, err = store.SaveJobTypeConfigIfNotExists(&plugin_pb.PersistedJobTypeConfig{
JobType: "admin_script",
AdminRuntime: &plugin_pb.AdminRuntimeConfig{Enabled: false},
})
if err != nil {
t.Fatalf("second SaveJobTypeConfigIfNotExists: %v", err)
}
if saved {
t.Fatal("expected second call to be a no-op")
}
// Verify the original config was preserved.
loaded, err := store.LoadJobTypeConfig("admin_script")
if err != nil {
t.Fatalf("LoadJobTypeConfig: %v", err)
}
if loaded == nil {
t.Fatal("expected config to exist")
}
if !loaded.AdminRuntime.Enabled {
t.Fatal("expected original config (Enabled=true) to be preserved")
}
// Different job type should still save.
saved, err = store.SaveJobTypeConfigIfNotExists(&plugin_pb.PersistedJobTypeConfig{
JobType: "vacuum",
})
if err != nil {
t.Fatalf("SaveJobTypeConfigIfNotExists for different type: %v", err)
}
if !saved {
t.Fatal("expected save for a different job type")
}
}
func TestConfigStoreJobDetailRoundTrip(t *testing.T) {
t.Parallel()
+13
View File
@@ -402,6 +402,19 @@ func (r *Plugin) SaveJobTypeConfig(config *plugin_pb.PersistedJobTypeConfig) err
return nil
}
// SaveJobTypeConfigIfNotExists atomically saves config only if no config
// exists yet for the job type. Returns true if saved, false if already exists.
func (r *Plugin) SaveJobTypeConfigIfNotExists(config *plugin_pb.PersistedJobTypeConfig) (bool, error) {
saved, err := r.store.SaveJobTypeConfigIfNotExists(config)
if err != nil {
return false, err
}
if saved {
r.wakeScheduler()
}
return saved, nil
}
func (r *Plugin) LoadDescriptor(jobType string) (*plugin_pb.JobTypeDescriptor, error) {
return r.store.LoadDescriptor(jobType)
}