github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/compute/store/inmemory/store.go (about)

     1  package inmemory
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	sync "github.com/bacalhau-project/golang-mutex-tracer"
     9  	"github.com/filecoin-project/bacalhau/pkg/compute/store"
    10  )
    11  
    12  const newExecutionComment = "Execution created"
    13  
    14  type Store struct {
    15  	executionMap map[string]store.Execution
    16  	shardMap     map[string][]string
    17  	history      map[string][]store.ExecutionHistory
    18  	mu           sync.RWMutex
    19  }
    20  
    21  func NewStore() *Store {
    22  	res := &Store{
    23  		executionMap: make(map[string]store.Execution),
    24  		shardMap:     make(map[string][]string),
    25  		history:      make(map[string][]store.ExecutionHistory),
    26  	}
    27  	res.mu.EnableTracerWithOpts(sync.Opts{
    28  		Threshold: 10 * time.Millisecond,
    29  		Id:        "InMemoryExecutionStore.mu",
    30  	})
    31  	return res
    32  }
    33  
    34  func (s *Store) GetExecution(ctx context.Context, id string) (store.Execution, error) {
    35  	s.mu.RLock()
    36  	defer s.mu.RUnlock()
    37  	execution, ok := s.executionMap[id]
    38  	if !ok {
    39  		return execution, store.NewErrExecutionNotFound(id)
    40  	}
    41  	return execution, nil
    42  }
    43  
    44  func (s *Store) GetExecutions(ctx context.Context, shardID string) ([]store.Execution, error) {
    45  	s.mu.RLock()
    46  	defer s.mu.RUnlock()
    47  	executionIDs, ok := s.shardMap[shardID]
    48  	if !ok {
    49  		return []store.Execution{}, store.NewErrExecutionsNotFound(shardID)
    50  	}
    51  	executions := make([]store.Execution, len(executionIDs))
    52  	for i, id := range executionIDs {
    53  		executions[i] = s.executionMap[id]
    54  	}
    55  	return executions, nil
    56  }
    57  
    58  func (s *Store) GetExecutionHistory(ctx context.Context, id string) ([]store.ExecutionHistory, error) {
    59  	s.mu.RLock()
    60  	defer s.mu.RUnlock()
    61  	history, ok := s.history[id]
    62  	if !ok {
    63  		return history, store.NewErrExecutionHistoryNotFound(id)
    64  	}
    65  	return history, nil
    66  }
    67  
    68  func (s *Store) CreateExecution(ctx context.Context, execution store.Execution) error {
    69  	s.mu.Lock()
    70  	defer s.mu.Unlock()
    71  	if _, ok := s.executionMap[execution.ID]; ok {
    72  		return store.NewErrExecutionAlreadyExists(execution.ID)
    73  	}
    74  	if err := store.ValidateNewExecution(ctx, execution); err != nil {
    75  		return fmt.Errorf("CreateExecution failure: %w", err)
    76  	}
    77  
    78  	s.executionMap[execution.ID] = execution
    79  	s.shardMap[execution.Shard.ID()] = append(s.shardMap[execution.Shard.ID()], execution.ID)
    80  	s.appendHistory(execution, store.ExecutionStateUndefined, newExecutionComment)
    81  	return nil
    82  }
    83  
    84  func (s *Store) UpdateExecutionState(ctx context.Context, request store.UpdateExecutionStateRequest) error {
    85  	s.mu.Lock()
    86  	defer s.mu.Unlock()
    87  	execution, ok := s.executionMap[request.ExecutionID]
    88  	if !ok {
    89  		return store.NewErrExecutionNotFound(request.ExecutionID)
    90  	}
    91  	if request.ExpectedState != store.ExecutionStateUndefined && execution.State != request.ExpectedState {
    92  		return store.NewErrInvalidExecutionState(request.ExecutionID, execution.State, request.ExpectedState)
    93  	}
    94  	if request.ExpectedVersion != 0 && execution.Version != request.ExpectedVersion {
    95  		return store.NewErrInvalidExecutionVersion(request.ExecutionID, execution.Version, request.ExpectedVersion)
    96  	}
    97  	if execution.State.IsTerminal() {
    98  		return store.NewErrExecutionAlreadyTerminal(request.ExecutionID, execution.State, request.NewState)
    99  	}
   100  	previousState := execution.State
   101  	execution.State = request.NewState
   102  	execution.Version += 1
   103  	execution.UpdateTime = time.Now()
   104  	s.executionMap[execution.ID] = execution
   105  	s.appendHistory(execution, previousState, request.Comment)
   106  	return nil
   107  }
   108  
   109  func (s *Store) appendHistory(updatedExecution store.Execution, previousState store.ExecutionState, comment string) {
   110  	historyEntry := store.ExecutionHistory{
   111  		ExecutionID:   updatedExecution.ID,
   112  		PreviousState: previousState,
   113  		NewState:      updatedExecution.State,
   114  		NewVersion:    updatedExecution.Version,
   115  		Comment:       comment,
   116  		Time:          updatedExecution.UpdateTime,
   117  	}
   118  	s.history[updatedExecution.ID] = append(s.history[updatedExecution.ID], historyEntry)
   119  }
   120  
   121  func (s *Store) DeleteExecution(ctx context.Context, id string) error {
   122  	s.mu.Lock()
   123  	defer s.mu.Unlock()
   124  	execution, ok := s.executionMap[id]
   125  	if ok {
   126  		delete(s.executionMap, id)
   127  		delete(s.history, id)
   128  		shardID := execution.Shard.ID()
   129  		shardExecutions := s.shardMap[shardID]
   130  		if len(shardExecutions) == 1 {
   131  			delete(s.shardMap, shardID)
   132  		} else {
   133  			for i, executionID := range shardExecutions {
   134  				if executionID == id {
   135  					s.shardMap[shardID] = append(shardExecutions[:i], shardExecutions[i+1:]...)
   136  					break
   137  				}
   138  			}
   139  		}
   140  	}
   141  	return nil
   142  }
   143  
   144  // compile-time check that we implement the interface ExecutionStore
   145  var _ store.ExecutionStore = (*Store)(nil)