[Feat] [history server] Add node endpoint

historyserver/pkg/eventserver/eventserver.go

@@ -22,6 +22,7 @@ type EventHandler struct {
 
 	ClusterTaskMap  *types.ClusterTaskMap
 	ClusterActorMap *types.ClusterActorMap
+	ClusterNodeMap  *types.ClusterNodeMap
 }
 
 var eventFilePattern = regexp.MustCompile(`-\d{4}-\d{2}-\d{2}-\d{2}$`)
@@ -44,6 +45,9 @@ func NewEventHandler(reader storage.StorageReader) *EventHandler {
 		ClusterActorMap: &types.ClusterActorMap{
 			ClusterActorMap: make(map[string]*types.ActorMap),
 		},
+		ClusterNodeMap: &types.ClusterNodeMap{
+			ClusterNodeMap: make(map[string]*types.NodeMap),
+		},
 	}
 }
 
@@ -202,6 +206,16 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 		return fmt.Errorf("clusterName is not a string, got %T", clusterNameVal)
 	}
 
+	sessionNameVal, ok := eventMap["sessionName"]
+	if !ok {
+		return fmt.Errorf("event missing 'sessionName' field")
+	}
+	sessionName, ok := sessionNameVal.(string)
+	if !ok {
+		return fmt.Errorf("sessionName is not a string, got %T", sessionNameVal)
+	}
+	clusterSessionID := currentClusterName + "_" + sessionName
+
 	logrus.Infof("current eventType: %v", eventType)
 	switch eventType {
 	case types.TASK_DEFINITION_EVENT:
@@ -229,7 +243,6 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 				t.State = existingEvents[len(existingEvents)-1].State
 			}
 		})
-
 	case types.TASK_LIFECYCLE_EVENT:
 		lifecycleEvent, ok := eventMap["taskLifecycleEvent"].(map[string]any)
 		if !ok {
@@ -324,7 +337,6 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 				t.EndTime = lastEvent.Timestamp
 			}
 		})
-
 	case types.ACTOR_DEFINITION_EVENT:
 		actorDef, ok := eventMap["actorDefinitionEvent"]
 		if !ok {
@@ -520,11 +532,14 @@ func (h *EventHandler) storeEvent(eventMap map[string]any) error {
 			}
 			a.NumRestarts = restartCount
 		})
-
 	case types.ACTOR_TASK_DEFINITION_EVENT:
 		// TODO: Handle actor task definition event
 		// This is related to GET /api/v0/tasks (type=ACTOR_TASK)
 		logrus.Debugf("ACTOR_TASK_DEFINITION_EVENT received, not yet implemented")
+	case types.NODE_DEFINITION_EVENT:
+		return h.handleNodeDefinitionEvent(eventMap, clusterSessionID)
+	case types.NODE_LIFECYCLE_EVENT:
+		return h.handleNodeLifecycleEvent(eventMap, clusterSessionID)
 	default:
 		logrus.Infof("Event not supported, skipping: %v", eventMap)
 	}
@@ -708,3 +723,93 @@ func (h *EventHandler) GetActorsMap(clusterName string) map[string]types.Actor {
 	}
 	return actors
 }
+
+// handleNodeDefinitionEvent processes NODE_DEFINITION_EVENT and merges it with the existing node map for a given cluster session.
+func (h *EventHandler) handleNodeDefinitionEvent(eventMap map[string]any, clusterSessionID string) error {
+	nodeDef, exists := eventMap["nodeDefinitionEvent"]
+	if !exists {
+		return fmt.Errorf("event does not have 'nodeDefinitionEvent' field")
+	}
+
+	jsonNodeDefinition, err := json.Marshal(nodeDef)
+	if err != nil {
+		return fmt.Errorf("failed to marshal node definition event: %w", err)
+	}
+
+	var currNode types.Node
+	if err := json.Unmarshal(jsonNodeDefinition, &currNode); err != nil {
+		return fmt.Errorf("failed to unmarshal node definition event: %w", err)
+	}
+
+	currNode.NodeID, err = utils.ConvertBase64ToHex(currNode.NodeID)
+	if err != nil {
+		return fmt.Errorf("failed to convert node ID from base64 to hex: %w", err)
+	}
+
+	nodeMap := h.ClusterNodeMap.GetOrCreateNodeMap(clusterSessionID)
+	nodeMap.CreateOrMergeNode(currNode.NodeID, func(n *types.Node) {
+		// TODO(jwj): Handle merging of node definition event to prevent overwriting lifecycle-derived fields.
+		existingStateTransitions := n.StateTransitions
+
+		*n = currNode
+
+		if len(existingStateTransitions) > 0 {
+			n.StateTransitions = existingStateTransitions
+		}
+	})
+
+	return nil
+}
+
+// handleNodeLifecycleEvent processes NODE_LIFECYCLE_EVENT and merges state transitions.
+func (h *EventHandler) handleNodeLifecycleEvent(eventMap map[string]any, clusterSessionID string) error {
+	nodeLifecycle, exists := eventMap["nodeLifecycleEvent"]
+	if !exists {
+		return fmt.Errorf("event does not have 'nodeLifecycleEvent' field")
+	}
+
+	jsonNodeLifecycle, err := json.Marshal(nodeLifecycle)
+	if err != nil {
+		return fmt.Errorf("failed to marshal node lifecycle event: %w", err)
+	}
+
+	var currNode types.Node
+	if err := json.Unmarshal(jsonNodeLifecycle, &currNode); err != nil {
+		return fmt.Errorf("failed to unmarshal node lifecycle event: %w", err)
+	}
+
+	currNode.NodeID, err = utils.ConvertBase64ToHex(currNode.NodeID)
+	if err != nil {
+		return fmt.Errorf("failed to convert node ID from base64 to hex: %w", err)
+	}
+
+	nodeMap := h.ClusterNodeMap.GetOrCreateNodeMap(clusterSessionID)
+	nodeMap.CreateOrMergeNode(currNode.NodeID, func(n *types.Node) {
+		// Merge state transitions.
+		n.StateTransitions = MergeStateTransitions[types.NodeStateTransition](
+			n.StateTransitions,
+			currNode.StateTransitions,
+		)
+	})
+
+	return nil
+}
+
+func (h *EventHandler) GetNodeMap(clusterSessionID string) map[string]types.Node {
+	h.ClusterNodeMap.RLock()
+	defer h.ClusterNodeMap.RUnlock()
+
+	nodeMap, ok := h.ClusterNodeMap.ClusterNodeMap[clusterSessionID]
+	if !ok {
+		return map[string]types.Node{}
+	}
+
+	nodeMap.Lock()
+	defer nodeMap.Unlock()
+
+	nodes := make(map[string]types.Node, len(nodeMap.NodeMap))
+	for id, node := range nodeMap.NodeMap {
+		nodes[id] = node.DeepCopy()
+	}
+	return nodes
+}

historyserver/pkg/eventserver/state_transition.go

@@ -0,0 +1,52 @@
+package eventserver
+
+import (
+	"sort"
+	"time"
+)
+
+// StateTransition defines a common interface for state transitions.
+type StateTransition interface {
+	GetState() string
+	GetTimestamp() time.Time
+}
+
+// StateTransitionKey is a unique key for deduplication.
+type StateTransitionKey struct {
+	State     string
+	Timestamp int64
+}
+
+// MergeStateTransitions merges new state transitions into existing ones, avoiding duplicates and maintaining chronological order.
+func MergeStateTransitions[T StateTransition](existingStateTransitions []T, newStateTransitions []T) []T {
+	// Build a map of existing transition keys for deduplication.
+	existingKeys := make(map[StateTransitionKey]bool, len(existingStateTransitions))
+	for _, tr := range existingStateTransitions {
+		key := StateTransitionKey{
+			State:     tr.GetState(),
+			Timestamp: tr.GetTimestamp().UnixNano(),
+		}
+		existingKeys[key] = true
+	}
+
+	// Append new state transitions that haven't seen before.
+	mergedStateTransitions := make([]T, len(existingStateTransitions))
+	copy(mergedStateTransitions, existingStateTransitions)
+	for _, tr := range newStateTransitions {
+		key := StateTransitionKey{
+			State:     tr.GetState(),
+			Timestamp: tr.GetTimestamp().UnixNano(),
+		}
+		if !existingKeys[key] {
+			mergedStateTransitions = append(mergedStateTransitions, tr)
+			existingKeys[key] = true
+		}
+	}
+
+	// Sort state transitions by timestamp to maintain chronological order.
+	sort.Slice(mergedStateTransitions, func(i, j int) bool {
+		return mergedStateTransitions[i].GetTimestamp().Before(mergedStateTransitions[j].GetTimestamp())
+	})
+
+	return mergedStateTransitions
+}

historyserver/pkg/eventserver/types/node.go

@@ -0,0 +1,155 @@
+package types
+
+import (
+	"sync"
+	"time"
+)
+
+// NodeState is the state of a node.
+type NodeState string
+
+// TODO(jwj): Handle redeclaration of states with the actor type.
+// actor.go also defines ALIVE and DEAD.
+const (
+	NODE_ALIVE NodeState = "ALIVE"
+	NODE_DEAD  NodeState = "DEAD"
+)
+
+// NodeAliveSubState provides more granular state information for nodes in the ALIVE state.
+type NodeAliveSubState string
+
+const (
+	UNSPECIFIED NodeAliveSubState = "UNSPECIFIED"
+	DRAINING    NodeAliveSubState = "DRAINING"
+)
+
+// NodeDeathInfoReason specifies the reason why a node died.
+type NodeDeathInfoReason string
+
+const (
+	DEATH_REASON_UNSPECIFIED   NodeDeathInfoReason = "UNSPECIFIED"
+	EXPECTED_TERMINATION       NodeDeathInfoReason = "EXPECTED_TERMINATION"
+	UNEXPECTED_TERMINATION     NodeDeathInfoReason = "UNEXPECTED_TERMINATION"
+	AUTOSCALER_DRAIN_PREEMPTED NodeDeathInfoReason = "AUTOSCALER_DRAIN_PREEMPTED"
+	AUTOSCALER_DRAIN_IDLE      NodeDeathInfoReason = "AUTOSCALER_DRAIN_IDLE"
+)
+
+type NodeDeathInfo struct {
+	Reason        NodeDeathInfoReason `json:"reason"`
+	ReasonMessage string              `json:"reasonMessage"`
+}
+
+// NodeStateTransition represents a change in a node's state at a specific timestamp.
+type NodeStateTransition struct {
+	// State of a node (ALIVE, DEAD).
+	State NodeState `json:"state"`
+
+	// Timestamp when the state transition occurred.
+	Timestamp time.Time `json:"timestamp"`
+
+	// Resources available on a node (cpu, gpu, etc.), available only in the ALIVE state.
+	Resources map[string]float64 `json:"resources,omitempty"`
+
+	// Reason why a node died (UNSPECIFIED, EXPECTED_TERMINATION, UNEXPECTED_TERMINATION, AUTOSCALER_DRAIN_PREEMPTED, AUTOSCALER_DRAIN_IDLE),
+	// available only in the DEAD state
+	DeathInfo *NodeDeathInfo `json:"deathInfo,omitempty"`
+
+	// Sub-state of a node in the ALIVE state (UNSPECIFIED, DRAINING), available only in the ALIVE state.
+	AliveSubState NodeAliveSubState `json:"aliveSubState,omitempty"`
+}
+
+func (n NodeStateTransition) GetState() string {
+	return string(n.State)
+}
+
+func (n NodeStateTransition) GetTimestamp() time.Time {
+	return n.Timestamp
+}
+
+type Node struct {
+	// NodeID is the hexadecimal representation of the node ID.
+	NodeID         string            `json:"nodeId"`
+	NodeIPAddress  string            `json:"nodeIpAddress,omitempty"`
+	StartTimestamp time.Time         `json:"startTimestamp,omitempty"`
+	Labels         map[string]string `json:"labels,omitempty"`
+
+	// TODO(jwj): Make it clearer.
+	// Available only when there's at least one NODE_LIFECYCLE_EVENT.
+	StateTransitions []NodeStateTransition `json:"stateTransitions,omitempty"`
+}
+
+type NodeMap struct {
+	NodeMap map[string]Node
+	Mu      sync.Mutex
+}
+
+func (n *NodeMap) Lock() {
+	n.Mu.Lock()
+}
+
+func (n *NodeMap) Unlock() {
+	n.Mu.Unlock()
+}
+
+func NewNodeMap() *NodeMap {
+	return &NodeMap{
+		NodeMap: make(map[string]Node),
+	}
+}
+
+type ClusterNodeMap struct {
+	ClusterNodeMap map[string]*NodeMap
+	Mu             sync.RWMutex
+}
+
+func (c *ClusterNodeMap) RLock() {
+	c.Mu.RLock()
+}
+
+func (c *ClusterNodeMap) RUnlock() {
+	c.Mu.RUnlock()
+}
+
+func (c *ClusterNodeMap) Lock() {
+	c.Mu.Lock()
+}
+
+func (c *ClusterNodeMap) Unlock() {
+	c.Mu.Unlock()
+}
+
+// GetOrCreateNodeMap retrieves the NodeMap for the given cluster session, creating it if it doesn't exist.
+func (c *ClusterNodeMap) GetOrCreateNodeMap(clusterSessionID string) *NodeMap {
+	c.Lock()
+	defer c.Unlock()
+
+	nodeMap, exists := c.ClusterNodeMap[clusterSessionID]
+	if !exists {
+		nodeMap = NewNodeMap()
+		c.ClusterNodeMap[clusterSessionID] = nodeMap
+	}
+	return nodeMap
+}
+
+// CreateOrMergeNode retrieves or creates a Node and applies the merge function.
+func (n *NodeMap) CreateOrMergeNode(nodeId string, mergeFn func(*Node)) {
+	n.Lock()
+	defer n.Unlock()
+
+	node, exists := n.NodeMap[nodeId]
+	if !exists {
+		node = Node{NodeID: nodeId}
+	}
+
+	mergeFn(&node)
+	n.NodeMap[nodeId] = node
+}
+
+func (n Node) DeepCopy() Node {
+	cp := n
+	if len(n.StateTransitions) > 0 {
+		cp.StateTransitions = make([]NodeStateTransition, len(n.StateTransitions))
+		copy(cp.StateTransitions, n.StateTransitions)
+	}
+	return cp
+}