121 lines
2.9 KiB
Go
121 lines
2.9 KiB
Go
package manager
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
|
|
"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/worker"
|
|
)
|
|
|
|
var (
|
|
ErrWorkerAlreadyRegistered = errors.New("worker type already registered")
|
|
ErrWorkerTypeUnknown = errors.New("unknown worker type")
|
|
ErrNilFactory = errors.New("nil worker factory")
|
|
ErrNilNormalizer = errors.New("nil worker normalizer")
|
|
)
|
|
|
|
type registryEntry struct {
|
|
factory worker.Factory
|
|
normalizer worker.Normalizer
|
|
}
|
|
|
|
type WorkerRegistry struct {
|
|
mu sync.RWMutex
|
|
m map[string]registryEntry
|
|
}
|
|
|
|
func NewWorkerRegistry() *WorkerRegistry {
|
|
return &WorkerRegistry{m: make(map[string]registryEntry)}
|
|
}
|
|
|
|
// Register a worker type with its factory and keyer.
|
|
func (wr *WorkerRegistry) Register(workerType string, factory worker.Factory, normalizer worker.Normalizer) error {
|
|
if factory == nil {
|
|
return ErrNilFactory
|
|
}
|
|
if normalizer == nil {
|
|
return ErrNilNormalizer
|
|
}
|
|
wr.mu.Lock()
|
|
defer wr.mu.Unlock()
|
|
if _, ok := wr.m[workerType]; ok {
|
|
return ErrWorkerAlreadyRegistered
|
|
}
|
|
wr.m[workerType] = registryEntry{factory: factory, normalizer: normalizer}
|
|
return nil
|
|
}
|
|
|
|
// Deregister removes a worker type.
|
|
func (wr *WorkerRegistry) Deregister(workerType string) error {
|
|
wr.mu.Lock()
|
|
defer wr.mu.Unlock()
|
|
if _, ok := wr.m[workerType]; !ok {
|
|
return ErrWorkerTypeUnknown
|
|
}
|
|
delete(wr.m, workerType)
|
|
return nil
|
|
}
|
|
|
|
// Spawn constructs a new worker instance for the given type.
|
|
func (wr *WorkerRegistry) Spawn(workerType string) (worker.Worker, error) {
|
|
wr.mu.RLock()
|
|
entry, ok := wr.m[workerType]
|
|
wr.mu.RUnlock()
|
|
if !ok {
|
|
return nil, ErrWorkerTypeUnknown
|
|
}
|
|
return entry.factory(), nil
|
|
}
|
|
|
|
func (wr *WorkerRegistry) NormalizeSpecificationBytes(workerType string, spec []byte) ([]byte, error) {
|
|
wr.mu.RLock()
|
|
entry, ok := wr.m[workerType]
|
|
wr.mu.RUnlock()
|
|
if !ok {
|
|
return nil, ErrWorkerTypeUnknown
|
|
}
|
|
return entry.normalizer.NormalizeSpecification(spec)
|
|
}
|
|
|
|
func (wr *WorkerRegistry) NormalizeUnitBytes(workerType string, unit []byte) ([]byte, error) {
|
|
wr.mu.RLock()
|
|
entry, ok := wr.m[workerType]
|
|
wr.mu.RUnlock()
|
|
if !ok {
|
|
return nil, ErrWorkerTypeUnknown
|
|
}
|
|
return entry.normalizer.NormalizeUnit(unit)
|
|
}
|
|
|
|
// Factory returns the registered factory.
|
|
func (wr *WorkerRegistry) Factory(workerType string) (worker.Factory, error) {
|
|
wr.mu.RLock()
|
|
entry, ok := wr.m[workerType]
|
|
wr.mu.RUnlock()
|
|
if !ok {
|
|
return nil, ErrWorkerTypeUnknown
|
|
}
|
|
return entry.factory, nil
|
|
}
|
|
|
|
func (wr *WorkerRegistry) Normalizer(workerType string) (worker.Normalizer, error) {
|
|
wr.mu.RLock()
|
|
entry, ok := wr.m[workerType]
|
|
wr.mu.RUnlock()
|
|
if !ok {
|
|
return nil, ErrWorkerTypeUnknown
|
|
}
|
|
return entry.normalizer, nil
|
|
}
|
|
|
|
// RegisteredTypes lists all worker types.
|
|
func (wr *WorkerRegistry) RegisteredTypes() []string {
|
|
wr.mu.RLock()
|
|
defer wr.mu.RUnlock()
|
|
out := make([]string, 0, len(wr.m))
|
|
for t := range wr.m {
|
|
out = append(out, t)
|
|
}
|
|
return out
|
|
}
|