NTaskStateManager.h

#ifndef N_TASK_STATE_MANAGER_H
#define N_TASK_STATE_MANAGER_H
 
#include <queue>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <cstddef>
#include <string>
#include <utility>
 
namespace Ndmspc {

class NTaskStateManager {
  public:
  using TaskId = size_t;
  using WorkerId = std::string;
  using TaskPayload = std::vector<int>; // coordinates
  
  NTaskStateManager() = default;
  ~NTaskStateManager() = default;
  
  void AddPending(TaskId id, const TaskPayload & payload);
  
  bool AssignToWorker(const WorkerId & worker, TaskId id);

  bool ClaimNextPendingForWorker(const WorkerId & worker, TaskId & id, TaskPayload & payload);
  
  bool MarkDone(TaskId id);
  
  bool MarkFailed(TaskId id);

  bool RequeueTask(TaskId id);
  
  std::set<TaskId> GetWorkerTasks(const WorkerId & worker) const;
  
  std::pair<TaskId, TaskPayload> GetNextPending();
  
  bool HasPending() const;
  
  std::vector<std::pair<TaskId, TaskPayload>> RecoverWorkerTasks(const WorkerId & worker);
  
  bool RemoveTaskFromWorker(const WorkerId & worker, TaskId id);
  
  WorkerId GetTaskWorker(TaskId id) const;
  
  bool IsDone(TaskId id) const;
  
  size_t PendingCount() const { return fPending.size(); }
  
  size_t RunningCount() const { return fRunning.size(); }
  
  size_t DoneCount() const { return fDone.size(); }
  
  size_t TotalCount() const { return PendingCount() + RunningCount() + DoneCount(); }
  
  void Clear();
  
  private:
  // State buckets: tasks flow pending → running → done
  std::queue<std::pair<TaskId, TaskPayload>> fPending;  // Not yet dispatched
  std::unordered_set<TaskId>                fPendingIds;
  std::unordered_map<TaskId, TaskPayload>   fRunning;   // Assigned to workers
  std::set<TaskId>                          fDone;      // Successfully completed
  
  // Mappings for efficient lookup
  std::unordered_map<WorkerId, std::set<TaskId>> fWorkerToTasks;  // Current assignments
  std::unordered_map<TaskId, WorkerId>           fTaskToWorker;   // Reverse mapping
  std::unordered_map<TaskId, TaskPayload>        fTaskPayloads;   // All payloads (for recovery)
  
  bool TaskExists(TaskId id) const;

  void EnqueuePending(TaskId id, const TaskPayload & payload);
};
 
} // namespace Ndmspc
 
#endif // N_TASK_STATE_MANAGER_H