128 lines
3.7 KiB
Go
128 lines
3.7 KiB
Go
package state
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"sync"
|
|
|
|
"github.com/hohn/mrvacommander/pkg/common"
|
|
"github.com/hohn/mrvacommander/pkg/queue"
|
|
)
|
|
|
|
type LocalState struct {
|
|
jobs map[int][]queue.AnalyzeJob
|
|
info map[common.JobSpec]common.JobInfo
|
|
status map[common.JobSpec]common.Status
|
|
result map[common.JobSpec]queue.AnalyzeResult
|
|
sessionToJobIdToSpec map[int]map[int]common.JobSpec
|
|
mutex sync.Mutex
|
|
currentID int
|
|
}
|
|
|
|
func NewLocalState(startingID int) *LocalState {
|
|
state := &LocalState{
|
|
jobs: make(map[int][]queue.AnalyzeJob),
|
|
info: make(map[common.JobSpec]common.JobInfo),
|
|
status: make(map[common.JobSpec]common.Status),
|
|
result: make(map[common.JobSpec]queue.AnalyzeResult),
|
|
sessionToJobIdToSpec: make(map[int]map[int]common.JobSpec),
|
|
currentID: startingID,
|
|
}
|
|
state.sessionToJobIdToSpec[startingID] = make(map[int]common.JobSpec)
|
|
state.jobs[startingID] = []queue.AnalyzeJob{}
|
|
return state
|
|
}
|
|
|
|
func (s *LocalState) NextID() int {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
s.currentID++
|
|
s.jobs[s.currentID] = []queue.AnalyzeJob{}
|
|
return s.currentID
|
|
}
|
|
|
|
func (s *LocalState) GetResult(js common.JobSpec) (queue.AnalyzeResult, error) {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
if _, ok := s.result[js]; !ok {
|
|
return queue.AnalyzeResult{}, fmt.Errorf("result not found for job spec %v", js)
|
|
}
|
|
return s.result[js], nil
|
|
}
|
|
|
|
func (s *LocalState) GetJobSpecByRepoId(sessionId, jobRepoId int) (common.JobSpec, error) {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
idToSpec, ok := s.sessionToJobIdToSpec[sessionId]
|
|
if !ok {
|
|
return common.JobSpec{}, fmt.Errorf("job ids not found for session %v", sessionId)
|
|
}
|
|
spec, ok := idToSpec[jobRepoId]
|
|
if !ok {
|
|
return common.JobSpec{}, fmt.Errorf("job spec not found for job repo id %v", jobRepoId)
|
|
}
|
|
return spec, nil
|
|
}
|
|
|
|
func (s *LocalState) SetResult(js common.JobSpec, analyzeResult queue.AnalyzeResult) {
|
|
s.mutex.Lock()
|
|
s.result[js] = analyzeResult
|
|
s.mutex.Unlock()
|
|
}
|
|
|
|
func (s *LocalState) GetJobList(sessionID int) ([]queue.AnalyzeJob, error) {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
if _, ok := s.jobs[sessionID]; !ok {
|
|
return nil, fmt.Errorf("job list not found for session %v", sessionID)
|
|
}
|
|
return s.jobs[sessionID], nil
|
|
}
|
|
|
|
func (s *LocalState) GetJobInfo(js common.JobSpec) (common.JobInfo, error) {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
if _, ok := s.info[js]; !ok {
|
|
return common.JobInfo{}, fmt.Errorf("job info not found for job spec %v", js)
|
|
}
|
|
return s.info[js], nil
|
|
}
|
|
|
|
func (s *LocalState) SetJobInfo(js common.JobSpec, ji common.JobInfo) {
|
|
s.mutex.Lock()
|
|
s.info[js] = ji
|
|
s.mutex.Unlock()
|
|
}
|
|
|
|
func (s *LocalState) GetStatus(js common.JobSpec) (common.Status, error) {
|
|
s.mutex.Lock()
|
|
defer s.mutex.Unlock()
|
|
if _, ok := s.status[js]; !ok {
|
|
return common.StatusFailed, fmt.Errorf("status not found for job spec %v", js)
|
|
}
|
|
return s.status[js], nil
|
|
}
|
|
|
|
func (s *LocalState) SetStatus(js common.JobSpec, status common.Status) {
|
|
s.mutex.Lock()
|
|
s.status[js] = status
|
|
s.mutex.Unlock()
|
|
}
|
|
|
|
func (s *LocalState) AddJob(job queue.AnalyzeJob) {
|
|
s.mutex.Lock()
|
|
sessionID := job.Spec.SessionID
|
|
s.jobs[sessionID] = append(s.jobs[sessionID], job)
|
|
// Map the job index to JobSpec for quick result lookup
|
|
if _, ok := s.sessionToJobIdToSpec[sessionID]; !ok {
|
|
s.sessionToJobIdToSpec[sessionID] = make(map[int]common.JobSpec)
|
|
}
|
|
s.sessionToJobIdToSpec[sessionID][len(s.sessionToJobIdToSpec[sessionID])] = job.Spec
|
|
if len(s.jobs[sessionID]) != len(s.sessionToJobIdToSpec[sessionID]) {
|
|
msg := fmt.Sprintf("Unequal job list and job id map length. Session ID: %v", sessionID)
|
|
slog.Error(msg)
|
|
panic(msg)
|
|
}
|
|
s.mutex.Unlock()
|
|
}
|