Polishing

Author: akup

decentralized-api/cmd/devshardd/engine.go

@@ -249,7 +249,7 @@ func (e *devshardEngine) doWithFallbackNodes(
 		if !ok {
 			if limit && len(capacityExcluded) > 0 {
 				// Every known node is at its local bound — wait and retry.
-				clearMap(capacityExcluded)
+				clear(capacityExcluded)
 				select {
 				case <-ctx.Done():
 					lastReason = observability.ReasonTimeout
@@ -332,12 +332,6 @@ func mergeExcluded(a, b map[string]struct{}) map[string]struct{} {
 	return out
 }
 
-func clearMap(m map[string]struct{}) {
-	for k := range m {
-		delete(m, k)
-	}
-}
-
 // shouldFallback reports whether an Acquire error means dapi is unreachable
 // and the passive cache should be used. ResourceExhausted is not a fallback
 // trigger — dapi is up and remains authoritative for load balancing.

decentralized-api/internal/event_listener/escrow_events.go

@@ -41,7 +41,12 @@ func (e *DevshardEscrowCreatedEventHandler) Handle(event *chainevents.JSONRPCRes
 	if err != nil {
 		return fmt.Errorf("parse escrow_id: %w", err)
 	}
-	epochIndex, _ := strconv.ParseUint(firstAttr(ev, escrowCreatedEvent+".epoch_index"), 10, 64)
+	epochIndex, err := strconv.ParseUint(firstAttr(ev, escrowCreatedEvent+".epoch_index"), 10, 64)
+	if err != nil {
+		logging.Warn("host_events: malformed epoch_index on escrow_created; defaulting to 0", types.EventProcessing,
+			"escrow_id", escrowID, "raw", firstAttr(ev, escrowCreatedEvent+".epoch_index"), "error", err)
+		epochIndex = 0
+	}
 	payload := &apiconfig.EscrowPayload{
 		EscrowID:   escrowID,
 		EpochIndex: epochIndex,
@@ -50,7 +55,7 @@ func (e *DevshardEscrowCreatedEventHandler) Handle(event *chainevents.JSONRPCRes
 		Amount:     firstAttr(ev, escrowCreatedEvent+".amount"),
 	}
 
-	hold, ok := el.localHoldsEscrowSlot(fmt.Sprint(escrowID))
+	hold, ok := el.localHoldsEscrowSlot(strconv.FormatUint(escrowID, 10))
 	if ok && !hold {
 		logging.Debug("host_events: skip escrow_created; local node not in slots", types.EventProcessing,
 			"escrow_id", escrowID)
@@ -93,7 +98,7 @@ func (e *DevshardEscrowSettledEventHandler) Handle(event *chainevents.JSONRPCRes
 		Remainder:   firstAttr(ev, escrowSettledEvent+".remainder"),
 	}
 
-	hold, ok := el.localHoldsEscrowSlot(fmt.Sprint(escrowID))
+	hold, ok := el.localHoldsEscrowSlot(strconv.FormatUint(escrowID, 10))
 	if ok && !hold {
 		logging.Debug("host_events: skip escrow_settled; local node not in slots", types.EventProcessing,
 			"escrow_id", escrowID)

decentralized-api/nodemanager/host_events_rpc.go

@@ -30,70 +30,79 @@ func (s *Server) GetHostEvents(ctx context.Context, req *gen.GetHostEventsReques
 	cursor := req.GetCursor()
 
 	for {
-		var wake <-chan struct{}
-		var release func()
-		if maxWait > 0 {
-			// Subscribe (with the client's kind filter) before Since to avoid
-			// lost wake-ups. Only subscribed kinds wake this waiter, so an
-			// unsubscribed kind (e.g. maintenance) cannot reset the deadline.
-			wake, release = s.hostEvents.Subscribe(subscribe)
+		resp, err, again := s.getHostEventsOnce(ctx, cursor, clientGen, subscribe, maxWait)
+		if !again {
+			return resp, err
 		}
+		// Notified by a subscribed kind: loop and re-check Since.
+	}
+}
 
-		got := s.hostEvents.Since(cursor, clientGen, subscribe)
-		if got.Reset {
-			if release != nil {
-				release()
-			}
-			logging.Info("host_events: GetHostEvents needs_reset", types.Config,
-				"cursor", cursor,
-				"clientGeneration", clientGen,
-				"serverGeneration", got.Generation,
-				"nextCursor", got.NextCursor,
-			)
-			return s.hostEventsResponse(got, nil, false, true), nil
-		}
-		if len(got.Events) > 0 {
-			if release != nil {
-				release()
-			}
-			logging.Debug("host_events: GetHostEvents returning events", types.Config,
-				"cursor", cursor,
-				"count", len(got.Events),
-				"nextCursor", got.NextCursor,
-				"generation", got.Generation,
-			)
-			return s.hostEventsResponse(got, hostEventsToProto(got.Events), false, false), nil
-		}
+// getHostEventsOnce runs one subscribe → Since → wait cycle. release is deferred
+// so every exit path (return or loop-continue via again=true) deregisters the waiter.
+func (s *Server) getHostEventsOnce(
+	ctx context.Context,
+	cursor, clientGen uint64,
+	subscribe []apiconfig.HostEventKind,
+	maxWait time.Duration,
+) (resp *gen.GetHostEventsResponse, err error, again bool) {
+	var wake <-chan struct{}
+	var release func()
+	if maxWait > 0 {
+		// Subscribe (with the client's kind filter) before Since to avoid
+		// lost wake-ups. Only subscribed kinds wake this waiter, so an
+		// unsubscribed kind (e.g. maintenance) cannot reset the deadline.
+		wake, release = s.hostEvents.Subscribe(subscribe)
+	}
+	if release != nil {
+		defer release()
+	}
 
-		if maxWait <= 0 {
-			return s.hostEventsResponse(got, nil, true, false), nil
-		}
+	got := s.hostEvents.Since(cursor, clientGen, subscribe)
+	if got.Reset {
+		logging.Info("host_events: GetHostEvents needs_reset", types.Config,
+			"cursor", cursor,
+			"clientGeneration", clientGen,
+			"serverGeneration", got.Generation,
+			"nextCursor", got.NextCursor,
+		)
+		return s.hostEventsResponse(got, nil, false, true), nil, false
+	}
+	if len(got.Events) > 0 {
+		logging.Debug("host_events: GetHostEvents returning events", types.Config,
+			"cursor", cursor,
+			"count", len(got.Events),
+			"nextCursor", got.NextCursor,
+			"generation", got.Generation,
+		)
+		return s.hostEventsResponse(got, hostEventsToProto(got.Events), false, false), nil, false
+	}
+
+	if maxWait <= 0 {
+		return s.hostEventsResponse(got, nil, true, false), nil, false
+	}
 
-		logging.Debug("host_events: GetHostEvents long-poll waiting", types.Config,
+	logging.Debug("host_events: GetHostEvents long-poll waiting", types.Config,
+		"cursor", cursor,
+		"nextCursor", got.NextCursor,
+		"generation", got.Generation,
+		"maxWait", maxWait,
+	)
+	outcome, waitErr := longpoll.Wait(ctx, wake, maxWait)
+	if waitErr != nil {
+		return nil, status.FromContextError(waitErr).Err(), false
+	}
+	if outcome == longpoll.TimedOut {
+		got = s.hostEvents.Since(cursor, clientGen, subscribe)
+		logging.Debug("host_events: GetHostEvents long-poll timed out", types.Config,
 			"cursor", cursor,
 			"nextCursor", got.NextCursor,
 			"generation", got.Generation,
 			"maxWait", maxWait,
 		)
-		outcome, err := longpoll.Wait(ctx, wake, maxWait)
-		if release != nil {
-			release()
-		}
-		if err != nil {
-			return nil, status.FromContextError(err).Err()
-		}
-		if outcome == longpoll.TimedOut {
-			got = s.hostEvents.Since(cursor, clientGen, subscribe)
-			logging.Debug("host_events: GetHostEvents long-poll timed out", types.Config,
-				"cursor", cursor,
-				"nextCursor", got.NextCursor,
-				"generation", got.Generation,
-				"maxWait", maxWait,
-			)
-			return s.hostEventsResponse(got, nil, true, got.Reset), nil
-		}
-		// Notified by a subscribed kind: loop and re-check Since.
+		return s.hostEventsResponse(got, nil, true, got.Reset), nil, false
 	}
+	return nil, nil, true
 }
 
 func (s *Server) hostEventsResponse(got apiconfig.HostEventSince, events []*gen.HostEvent, unchanged, needsReset bool) *gen.GetHostEventsResponse {

devshard/mlnode/capacity.go

@@ -88,12 +88,12 @@ type Cache struct {
 	mu    sync.Mutex
 	nodes map[string]*nodeCap // key: nodeID
 
-	// localInFlight is keyed by nodeID\x00model (physical vLLM occupancy).
-	inFlight map[string]int
+	// inFlight is keyed by (nodeID, model) — physical vLLM occupancy.
+	inFlight map[nodeModelKey]int
 	// unknownInFlight tracks in-flight slots for capacity-unknown fallback nodes
 	// (not in c.nodes). Kept separate so applyPoll's prune never wipes live
 	// counts for a node dapi has not (yet) reported.
-	unknownInFlight map[string]int
+	unknownInFlight map[nodeModelKey]int
 
 	activeLoad ActiveLoadFunc
 	now        func() time.Time
@@ -135,8 +135,8 @@ func NewCache(client CapacityClient, opts CacheOptions) *Cache {
 	return &Cache{
 		client:                   client,
 		nodes:                    make(map[string]*nodeCap),
-		inFlight:                 make(map[string]int),
-		unknownInFlight:          make(map[string]int),
+		inFlight:                 make(map[nodeModelKey]int),
+		unknownInFlight:          make(map[nodeModelKey]int),
 		activeLoad:               active,
 		now:                      opts.Now,
 		log:                      opts.Log,
@@ -372,8 +372,7 @@ func (c *Cache) TryAcquire(nodeID, model string) bool {
 	if c.availableSlotsLocked(nodeID, div) <= 0 {
 		return false
 	}
-	key := inFlightKey(nodeID, model)
-	c.inFlight[key]++
+	c.inFlight[nodeModelKey{nodeID: nodeID, model: model}]++
 	return true
 }
 
@@ -422,7 +421,7 @@ func (c *Cache) Release(nodeID, model string) {
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	key := inFlightKey(nodeID, model)
+	key := nodeModelKey{nodeID: nodeID, model: model}
 	if c.inFlight[key] <= 1 {
 		delete(c.inFlight, key)
 		return
@@ -458,10 +457,10 @@ func (c *Cache) TryAcquireUnknown(nodeID, model string) bool {
 	eff := c.unknownEffectiveMax() // reads load map; must not hold c.mu
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if unknownInFlightSum(c.unknownInFlight, nodeID) >= eff {
+	if inFlightSum(c.unknownInFlight, nodeID) >= eff {
 		return false
 	}
-	c.unknownInFlight[inFlightKey(nodeID, model)]++
+	c.unknownInFlight[nodeModelKey{nodeID: nodeID, model: model}]++
 	return true
 }
 
@@ -472,60 +471,40 @@ func (c *Cache) ReleaseUnknown(nodeID, model string) {
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	key := inFlightKey(nodeID, model)
+	key := nodeModelKey{nodeID: nodeID, model: model}
 	if c.unknownInFlight[key] <= 1 {
 		delete(c.unknownInFlight, key)
 		return
 	}
 	c.unknownInFlight[key]--
 }
 
-func unknownInFlightSum(m map[string]int, nodeID string) int {
-	prefix := nodeID + "\x00"
+// nodeModelKey identifies one (node, model) in-flight counter.
+type nodeModelKey struct {
+	nodeID string
+	model  string
+}
+
+func inFlightSum(m map[nodeModelKey]int, nodeID string) int {
 	sum := 0
 	for k, v := range m {
-		if len(k) >= len(prefix) && k[:len(prefix)] == prefix {
+		if k.nodeID == nodeID {
 			sum += v
 		}
 	}
 	return sum
 }
 
 func (c *Cache) inFlightSumLocked(nodeID string) int {
-	prefix := nodeID + "\x00"
-	sum := 0
-	for k, v := range c.inFlight {
-		if len(k) >= len(prefix) && k[:len(prefix)] == prefix {
-			sum += v
-		}
-	}
-	return sum
+	return inFlightSum(c.inFlight, nodeID)
 }
 
 func (c *Cache) pruneInFlightLocked() {
 	for k := range c.inFlight {
-		nodeID, _, ok := splitInFlightKey(k)
-		if !ok {
+		if _, exists := c.nodes[k.nodeID]; !exists {
 			delete(c.inFlight, k)
-			continue
-		}
-		if _, exists := c.nodes[nodeID]; !exists {
-			delete(c.inFlight, k)
-		}
-	}
-}
-
-func inFlightKey(nodeID, model string) string {
-	return nodeID + "\x00" + model
-}
-
-func splitInFlightKey(key string) (nodeID, model string, ok bool) {
-	for i := 0; i < len(key); i++ {
-		if key[i] == 0 {
-			return key[:i], key[i+1:], true
 		}
 	}
-	return "", "", false
 }
 
 // ApplyPollForTest applies a successful poll snapshot (tests only).