96 lines
2.1 KiB
Go
96 lines
2.1 KiB
Go
package state
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"sync"
|
|
)
|
|
|
|
type WorkerRuntimeStore struct {
|
|
mu sync.Mutex
|
|
file AtomicJSON
|
|
modes map[string]string
|
|
}
|
|
|
|
type workerRuntimeFile struct {
|
|
Workers map[string]string `json:"workers"`
|
|
}
|
|
|
|
func NewWorkerRuntimeStore(path string) (*WorkerRuntimeStore, error) {
|
|
s := &WorkerRuntimeStore{file: AtomicJSON{Path: path, Mode: 0600}, modes: map[string]string{}}
|
|
var f workerRuntimeFile
|
|
if err := s.file.Load(&f); err != nil && !errors.Is(err, os.ErrNotExist) {
|
|
return nil, err
|
|
}
|
|
for k, v := range f.Workers {
|
|
if k != "" {
|
|
s.modes[k] = v
|
|
}
|
|
}
|
|
return s, nil
|
|
}
|
|
func (s *WorkerRuntimeStore) List(context.Context) (map[string]string, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
out := map[string]string{}
|
|
for k, v := range s.modes {
|
|
out[k] = v
|
|
}
|
|
return out, nil
|
|
}
|
|
func (s *WorkerRuntimeStore) Put(_ context.Context, name, mode string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
old, had := s.modes[name]
|
|
s.modes[name] = mode
|
|
if err := s.saveLocked(); err != nil {
|
|
if had {
|
|
s.modes[name] = old
|
|
} else {
|
|
delete(s.modes, name)
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
func (s *WorkerRuntimeStore) Delete(_ context.Context, name string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
old, had := s.modes[name]
|
|
delete(s.modes, name)
|
|
if err := s.saveLocked(); err != nil {
|
|
if had {
|
|
s.modes[name] = old
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
func (s *WorkerRuntimeStore) saveLocked() error {
|
|
keys := make([]string, 0, len(s.modes))
|
|
for k := range s.modes {
|
|
keys = append(keys, k)
|
|
}
|
|
sort.Strings(keys)
|
|
m := map[string]string{}
|
|
for _, k := range keys {
|
|
m[k] = s.modes[k]
|
|
}
|
|
return s.file.Save(workerRuntimeFile{Workers: m})
|
|
}
|
|
func (s *WorkerRuntimeStore) Health(context.Context) error {
|
|
if err := os.MkdirAll(filepath.Dir(s.file.Path), 0750); err != nil {
|
|
return err
|
|
}
|
|
f, err := os.OpenFile(s.file.Path+".health", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_ = f.Close()
|
|
return os.Remove(s.file.Path + ".health")
|
|
}
|
|
func (s *WorkerRuntimeStore) Path() string { return s.file.Path }
|