28159619024

Committed 25 Jun 2026 09:10AM UTC coverage: 92.667% (-0.02%) from 92.685%

Build # 28159619024

Build Type

push

github

Committed by

Hyshmily

Commit Message

release: v1.1.52

Coverage Stats

1199 of 1351 branches covered (88.75%)

Branch coverage included in aggregate %.

3426 of 3640 relevant lines covered (94.12%)

4.25 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.55

/worker/src/main/java/io/github/hyshmily/hotkey/worker/dispatch/WorkerHeartbeatProducer.java

/*
 * Copyright 2026 Hyshmily. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.github.hyshmily.hotkey.worker.dispatch;

import static io.github.hyshmily.hotkey.constants.HotKeyConstants.ROUTING_KEY_HEARTBEAT;

import io.github.hyshmily.hotkey.detection.HotKeyStateMachine;
import io.github.hyshmily.hotkey.sync.worker.WorkerHeartbeatMessage;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import java.lang.management.ManagementFactory;
import java.lang.management.OperatingSystemMXBean;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import lombok.AccessLevel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.context.annotation.DependsOn;
import org.springframework.data.redis.core.StringRedisTemplate;

/**
 * Enhanced Worker heartbeat sender.
 *
 * <p>Replaces the old ping-only heartbeat approach with
 * structured heartbeats containing epoch, decisionVersionHwm, loadFactor,
 * and readyToServe flags. Broadcast on the dedicated heartbeat exchange.
 *
 * <p>Epoch is persisted to Redis (fallback to local file), incremented
 * atomically on each process start for restart detection by Apps.
 */
@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
@Slf4j
@DependsOn("rabbitAdmin")
public class WorkerHeartbeatProducer {

  /** RabbitMQ template for publishing heartbeat messages. */
  private final RabbitTemplate rabbitTemplate;
  /** Target topic exchange for heartbeat messages. */
  private final String heartbeatExchange;
  /** Unique identity of this Worker node. */
  private final String workerId;
  /** State machine providing config-gossip fields (confirm/cool/grace counts). */
  private final HotKeyStateMachine stateMachine;
  /** Broadcaster for reading the current decision version watermark. */
  private final WorkerBroadcaster broadcaster;
  /** Monotonically increasing epoch, persisted across restarts. */
  private final long epoch;
  /** JVM start timestamp used for the ready-to-serve grace period. */
  private final long startTime;
  /** Scheduler for periodic heartbeat sends. */
  private final ScheduledExecutorService scheduler;
  /** Whether this instance owns the scheduler (self-created). */
  private final boolean ownsScheduler;
  /** Shared monotonic counter for config-change timestamps embedded in heartbeats. */
  private final AtomicLong configTimestampCounter;
  /** Interval between consecutive heartbeat sends (milliseconds). */
  private final long pingIntervalMs;
  /** Handle for the scheduled heartbeat task. */
  private ScheduledFuture<?> heartbeatTask;

  private static final String EPOCH_REDIS_KEY_PREFIX = "hotkey:worker:epoch:";

  /**
   * Creates a new heartbeat producer with a shared external scheduler.
   */
  public WorkerHeartbeatProducer(
    RabbitTemplate rabbitTemplate,
    String heartbeatExchange,
    String workerId,
    HotKeyStateMachine stateMachine,
    WorkerBroadcaster broadcaster,
    AtomicLong configTimestampCounter,
    RedisConnectionFactory redisConnectionFactory,
    long pingIntervalMs,
    ScheduledExecutorService scheduler
  ) {
    this(
      rabbitTemplate,
      heartbeatExchange,
      workerId,
      stateMachine,
      broadcaster,
      initEpoch(workerId, redisConnectionFactory),
      System.currentTimeMillis(),
      scheduler,
      false,
      configTimestampCounter,
      pingIntervalMs
    );
  }

  /**
   * Creates a new heartbeat producer with its own internal scheduler
   * (primarily for testing, with mocked Redis).
   */
  public WorkerHeartbeatProducer(
    RabbitTemplate rabbitTemplate,
    String heartbeatExchange,
    String workerId,
    HotKeyStateMachine stateMachine,
    WorkerBroadcaster broadcaster,
    AtomicLong configTimestampCounter,
    long pingIntervalMs
  ) {
    this(
      rabbitTemplate,
      heartbeatExchange,
      workerId,
      stateMachine,
      broadcaster,
      initEpoch(workerId, null),
      System.currentTimeMillis(),
      Executors.newSingleThreadScheduledExecutor(r -> {
        Thread t = new Thread(r, "hb-producer");
        t.setDaemon(true);
        return t;
      }),
      true,
      configTimestampCounter,
      pingIntervalMs
    );
  }

  private static long initEpoch(String workerId, RedisConnectionFactory redisConnectionFactory) {
    StringRedisTemplate template = new StringRedisTemplate(redisConnectionFactory);
    template.afterPropertiesSet();
    return doInitEpoch(template, workerId);
  }

  /**
   * Initializes the epoch by atomically incrementing a Redis counter.
   * Falls back to a local file if Redis is unavailable.
   *
   * @param redis the Redis template (created ad-hoc for this one-time operation)
   * @return the new epoch value for this Worker incarnation
   */
  private static long doInitEpoch(StringRedisTemplate redis, String workerId) {
    try {
      String key = EPOCH_REDIS_KEY_PREFIX + workerId;
      String val = redis.opsForValue().get(key);
      long next = (val != null) ? Long.parseLong(val) + 1 : 1;
      redis.opsForValue().set(key, String.valueOf(next));
      log.info("Epoch initialized via Redis: workerId={}, epoch={}", workerId, next);
      return next;
    } catch (Exception e) {
      log.warn("Redis unavailable for epoch, using local file fallback", e);
      return initEpochFromLocalFile(workerId);
    }
  }

  /**
   * Fallback epoch initialization using a local temp file when Redis is
   * unavailable.  Reads the previous value, increments, and writes back.
   *
   * <p>If both Redis and local file fallback fail, falls back to
   * {@code System.currentTimeMillis()} as a last resort epoch value
   * (non-monotonic risk across restarts is logged at error level).
   *
   * @return the new epoch value, or {@code System.currentTimeMillis()} on complete failure
   */
  private static long initEpochFromLocalFile(String workerId) {
    try {
      Path path = Path.of(System.getProperty("java.io.tmpdir"), "hotkey-epoch-" + workerId);
      long next = 1;
      if (Files.exists(path)) {
        String content = Files.readString(path);
        next = Long.parseLong(content.trim()) + 1;
      }
      Files.writeString(path, String.valueOf(next));
      log.info("Epoch initialized via local file: workerId={}, epoch={}", workerId, next);
      return next;
    } catch (Exception e) {
      log.error("Epoch local file fallback failed, using System.currentTimeMillis()", e);
      return System.currentTimeMillis();
    }
  }

  /**
   * Computes the current CPU load factor (0.0 – 1.0) via the platform MXBean.
   *
   * @return CPU load clamped to [0, 1], or 0.0 if unavailable
   */
  private double computeLoadFactor() {
    double cpuLoad = 0.0;
    try {
      OperatingSystemMXBean osBean = ManagementFactory.getOperatingSystemMXBean();
      if (osBean instanceof com.sun.management.OperatingSystemMXBean sunOsBean) {
        cpuLoad = sunOsBean.getCpuLoad();
      }
    } catch (Exception ignored) {}
    return Math.min(1.0, cpuLoad);
  }

  /**
   * Whether this Worker is ready to serve hot-key detection.
   * Returns {@code true} after a 3-second startup grace period or as soon as
   * the state machine has tracked at least one key.
   *
   * @return {@code true} if the Worker considers itself ready
   */
  private boolean isReadyToServe() {
    long uptime = System.currentTimeMillis() - startTime;
    return uptime > 3000 || stateMachine.getTrackedKeys() > 0;
  }

  /**
   * Computes an integer fingerprint of the current state-machine config
   * (confirm count, cool count, grace count) so receivers can quickly
   * detect config changes without comparing individual fields.
   *
   * @return config fingerprint hash
   */
  private int computeConfigFingerprint() {
    int hash = stateMachine.getConfirmCount();
    hash = 31 * hash + stateMachine.getCoolCount();
    hash = 31 * hash + stateMachine.getPreCoolGraceCount();
    return hash;
  }

  /**
   * Starts the periodic heartbeat sender at the configured ping interval.
   *
   * <p>First heartbeat is delayed by {@code pingIntervalMs} to allow the
   * RabbitMQ connection and exchange declarations (RabbitAdmin) to complete,
   * preventing channel-level NOT_FOUND errors when the heartbeat exchange
   * has not yet been declared.
   */
  @PostConstruct
  public void start() {
    try {
      heartbeatTask = scheduler.scheduleAtFixedRate(this::sendHeartbeat, pingIntervalMs, pingIntervalMs, TimeUnit.MILLISECONDS);
    } catch (Exception e) {
      log.error("Failed to start heartbeat scheduler; Worker heartbeat will not be sent. " +
          "Application continues but App instances may mark this Worker as dead.", e);
    }
  }

  /**
   * Gracefully stops the heartbeat sender.
   *
   * <p>Cancels the scheduled task; shuts down the scheduler only if owned.
   */
  @PreDestroy
  public void stop() {
    if (heartbeatTask != null) {
      heartbeatTask.cancel(false);
    }
    if (ownsScheduler) {
      scheduler.shutdown();
    }
  }

  /**
   * Builds and sends a single heartbeat message containing epoch, decision
   * version watermark, load factor, ready-to-serve flag, config fingerprint,
   * and state-machine config gossip fields.
   *
   * <p>Published to the configured heartbeat exchange with routing key
   * {@code heartbeat.<workerId>}.  Silently drops if the channel or
   * connection is unavailable (fire-and-forget).
   */
  void sendHeartbeat() {
    try {
      WorkerHeartbeatMessage hb = new WorkerHeartbeatMessage(
        workerId,
        epoch,
        System.currentTimeMillis(),
        broadcaster.getCurrentDecisionVersion(),
        computeLoadFactor(),
        isReadyToServe(),
        computeConfigFingerprint(),
        stateMachine.getConfirmCount(),
        stateMachine.getCoolCount(),
        stateMachine.getPreCoolGraceCount(),
        configTimestampCounter.get()
      );
      rabbitTemplate.send(heartbeatExchange, ROUTING_KEY_HEARTBEAT + workerId, hb.toMessage());
    } catch (Exception e) {
      log.error("Scheduled sendHeartbeat failed", e);
    }
  }
}

1	/*
2	* Copyright 2026 Hyshmily. All Rights Reserved.
3	*
4	* Licensed under the Apache License, Version 2.0 (the "License");
5	* you may not use this file except in compliance with the License.
6	* You may obtain a copy of the License at
7	*
8	* http://www.apache.org/licenses/LICENSE-2.0
9	*
10	* Unless required by applicable law or agreed to in writing, software
11	* distributed under the License is distributed on an "AS IS" BASIS,
12	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	* See the License for the specific language governing permissions and
14	* limitations under the License.
15	*/
16	package io.github.hyshmily.hotkey.worker.dispatch;
17
18	import static io.github.hyshmily.hotkey.constants.HotKeyConstants.ROUTING_KEY_HEARTBEAT;
19
20	import io.github.hyshmily.hotkey.detection.HotKeyStateMachine;
21	import io.github.hyshmily.hotkey.sync.worker.WorkerHeartbeatMessage;
22	import jakarta.annotation.PostConstruct;
23	import jakarta.annotation.PreDestroy;
24	import java.lang.management.ManagementFactory;
25	import java.lang.management.OperatingSystemMXBean;
26	import java.nio.file.Files;
27	import java.nio.file.Path;
28	import java.util.concurrent.Executors;
29	import java.util.concurrent.ScheduledExecutorService;
30	import java.util.concurrent.ScheduledFuture;
31	import java.util.concurrent.TimeUnit;
32	import java.util.concurrent.atomic.AtomicLong;
33	import lombok.AccessLevel;
34	import lombok.RequiredArgsConstructor;
35	import lombok.extern.slf4j.Slf4j;
36	import org.springframework.amqp.rabbit.core.RabbitTemplate;
37	import org.springframework.data.redis.connection.RedisConnectionFactory;
38	import org.springframework.context.annotation.DependsOn;
39	import org.springframework.data.redis.core.StringRedisTemplate;
40
41	/**
42	* Enhanced Worker heartbeat sender.
43	*
44	* <p>Replaces the old ping-only heartbeat approach with
45	* structured heartbeats containing epoch, decisionVersionHwm, loadFactor,
46	* and readyToServe flags. Broadcast on the dedicated heartbeat exchange.
47	*
48	* <p>Epoch is persisted to Redis (fallback to local file), incremented
49	* atomically on each process start for restart detection by Apps.
50	*/
51	@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
52	@Slf4j	4✔
53	@DependsOn("rabbitAdmin")
54	public class WorkerHeartbeatProducer {
55
56	/** RabbitMQ template for publishing heartbeat messages. */
57	private final RabbitTemplate rabbitTemplate;
58	/** Target topic exchange for heartbeat messages. */
59	private final String heartbeatExchange;
60	/** Unique identity of this Worker node. */
61	private final String workerId;
62	/** State machine providing config-gossip fields (confirm/cool/grace counts). */
63	private final HotKeyStateMachine stateMachine;
64	/** Broadcaster for reading the current decision version watermark. */
65	private final WorkerBroadcaster broadcaster;
66	/** Monotonically increasing epoch, persisted across restarts. */
67	private final long epoch;
68	/** JVM start timestamp used for the ready-to-serve grace period. */
69	private final long startTime;
70	/** Scheduler for periodic heartbeat sends. */
71	private final ScheduledExecutorService scheduler;
72	/** Whether this instance owns the scheduler (self-created). */
73	private final boolean ownsScheduler;
74	/** Shared monotonic counter for config-change timestamps embedded in heartbeats. */
75	private final AtomicLong configTimestampCounter;
76	/** Interval between consecutive heartbeat sends (milliseconds). */
77	private final long pingIntervalMs;
78	/** Handle for the scheduled heartbeat task. */
79	private ScheduledFuture<?> heartbeatTask;
80
81	private static final String EPOCH_REDIS_KEY_PREFIX = "hotkey:worker:epoch:";
82
83	/**
84	* Creates a new heartbeat producer with a shared external scheduler.
85	*/
86	public WorkerHeartbeatProducer(
87	RabbitTemplate rabbitTemplate,
88	String heartbeatExchange,
89	String workerId,
90	HotKeyStateMachine stateMachine,
91	WorkerBroadcaster broadcaster,
92	AtomicLong configTimestampCounter,
93	RedisConnectionFactory redisConnectionFactory,
94	long pingIntervalMs,
95	ScheduledExecutorService scheduler
96	) {
97	this(	9✔
98	rabbitTemplate,
99	heartbeatExchange,
100	workerId,
101	stateMachine,
102	broadcaster,
103	initEpoch(workerId, redisConnectionFactory),	1✔
104	System.currentTimeMillis(),	5✔
105	scheduler,
106	false,
107	configTimestampCounter,
108	pingIntervalMs
109	);
110	}	1✔
111
112	/**
113	* Creates a new heartbeat producer with its own internal scheduler
114	* (primarily for testing, with mocked Redis).
115	*/
116	public WorkerHeartbeatProducer(
117	RabbitTemplate rabbitTemplate,
118	String heartbeatExchange,
119	String workerId,
120	HotKeyStateMachine stateMachine,
121	WorkerBroadcaster broadcaster,
122	AtomicLong configTimestampCounter,
123	long pingIntervalMs
124	) {
125	this(	9✔
126	rabbitTemplate,
127	heartbeatExchange,
128	workerId,
129	stateMachine,
130	broadcaster,
131	initEpoch(workerId, null),	1✔
132	System.currentTimeMillis(),	2✔
133	Executors.newSingleThreadScheduledExecutor(r -> {	4✔
134	Thread t = new Thread(r, "hb-producer");	6✔
135	t.setDaemon(true);	3✔
136	return t;	2✔
137	}),
138	true,
139	configTimestampCounter,
140	pingIntervalMs
141	);
142	}	1✔
143
144	private static long initEpoch(String workerId, RedisConnectionFactory redisConnectionFactory) {
145	StringRedisTemplate template = new StringRedisTemplate(redisConnectionFactory);	5✔
146	template.afterPropertiesSet();	2✔
147	return doInitEpoch(template, workerId);	4✔
148	}
149
150	/**
151	* Initializes the epoch by atomically incrementing a Redis counter.
152	* Falls back to a local file if Redis is unavailable.
153	*
154	* @param redis the Redis template (created ad-hoc for this one-time operation)
155	* @return the new epoch value for this Worker incarnation
156	*/
157	private static long doInitEpoch(StringRedisTemplate redis, String workerId) {
158	try {
159	String key = EPOCH_REDIS_KEY_PREFIX + workerId;	3✔
160	String val = redis.opsForValue().get(key);	6✔
161	long next = (val != null) ? Long.parseLong(val) + 1 : 1;	9✔
162	redis.opsForValue().set(key, String.valueOf(next));	6✔
163	log.info("Epoch initialized via Redis: workerId={}, epoch={}", workerId, next);	6✔
164	return next;	2✔
165	} catch (Exception e) {	1✔
166	log.warn("Redis unavailable for epoch, using local file fallback", e);	4✔
167	return initEpochFromLocalFile(workerId);	3✔
168	}
169	}
170
171	/**
172	* Fallback epoch initialization using a local temp file when Redis is
173	* unavailable. Reads the previous value, increments, and writes back.
174	*
175	* <p>If both Redis and local file fallback fail, falls back to
176	* {@code System.currentTimeMillis()} as a last resort epoch value
177	* (non-monotonic risk across restarts is logged at error level).
178	*
179	* @return the new epoch value, or {@code System.currentTimeMillis()} on complete failure
180	*/
181	private static long initEpochFromLocalFile(String workerId) {
182	try {
183	Path path = Path.of(System.getProperty("java.io.tmpdir"), "hotkey-epoch-" + workerId);	11✔
184	long next = 1;	2✔
185	if (Files.exists(path)) {	5✔
186	String content = Files.readString(path);	3✔
187	next = Long.parseLong(content.trim()) + 1;	6✔
188	}
189	Files.writeString(path, String.valueOf(next));	7✔
190	log.info("Epoch initialized via local file: workerId={}, epoch={}", workerId, next);	6✔
191	return next;	2✔
192	} catch (Exception e) {	1✔
193	log.error("Epoch local file fallback failed, using System.currentTimeMillis()", e);	4✔
194	return System.currentTimeMillis();	2✔
195	}
196	}
197
198	/**
199	* Computes the current CPU load factor (0.0 – 1.0) via the platform MXBean.
200	*
201	* @return CPU load clamped to [0, 1], or 0.0 if unavailable
202	*/
203	private double computeLoadFactor() {
204	double cpuLoad = 0.0;	2✔
205	try {
206	OperatingSystemMXBean osBean = ManagementFactory.getOperatingSystemMXBean();	2✔
207	if (osBean instanceof com.sun.management.OperatingSystemMXBean sunOsBean) {	6!
208	cpuLoad = sunOsBean.getCpuLoad();	3✔
209	}
210	} catch (Exception ignored) {}	1✔
211	return Math.min(1.0, cpuLoad);	4✔
212	}
213
214	/**
215	* Whether this Worker is ready to serve hot-key detection.
216	* Returns {@code true} after a 3-second startup grace period or as soon as
217	* the state machine has tracked at least one key.
218	*
219	* @return {@code true} if the Worker considers itself ready
220	*/
221	private boolean isReadyToServe() {
222	long uptime = System.currentTimeMillis() - startTime;	5✔
223	return uptime > 3000 \|\| stateMachine.getTrackedKeys() > 0;	12✔
224	}
225
226	/**
227	* Computes an integer fingerprint of the current state-machine config
228	* (confirm count, cool count, grace count) so receivers can quickly
229	* detect config changes without comparing individual fields.
230	*
231	* @return config fingerprint hash
232	*/
233	private int computeConfigFingerprint() {
234	int hash = stateMachine.getConfirmCount();	4✔
235	hash = 31 * hash + stateMachine.getCoolCount();	8✔
236	hash = 31 * hash + stateMachine.getPreCoolGraceCount();	8✔
237	return hash;	2✔
238	}
239
240	/**
241	* Starts the periodic heartbeat sender at the configured ping interval.
242	*
243	* <p>First heartbeat is delayed by {@code pingIntervalMs} to allow the
244	* RabbitMQ connection and exchange declarations (RabbitAdmin) to complete,
245	* preventing channel-level NOT_FOUND errors when the heartbeat exchange
246	* has not yet been declared.
247	*/
248	@PostConstruct
249	public void start() {
250	try {
251	heartbeatTask = scheduler.scheduleAtFixedRate(this::sendHeartbeat, pingIntervalMs, pingIntervalMs, TimeUnit.MILLISECONDS);	12✔
252	} catch (Exception e) {	×
253	log.error("Failed to start heartbeat scheduler; Worker heartbeat will not be sent. " +	×
254	"Application continues but App instances may mark this Worker as dead.", e);
255	}	1✔
256	}	1✔
257
258	/**
259	* Gracefully stops the heartbeat sender.
260	*
261	* <p>Cancels the scheduled task; shuts down the scheduler only if owned.
262	*/
263	@PreDestroy
264	public void stop() {
265	if (heartbeatTask != null) {	3✔
266	heartbeatTask.cancel(false);	5✔
267	}
268	if (ownsScheduler) {	3✔
269	scheduler.shutdown();	3✔
270	}
271	}	1✔
272
273	/**
274	* Builds and sends a single heartbeat message containing epoch, decision
275	* version watermark, load factor, ready-to-serve flag, config fingerprint,
276	* and state-machine config gossip fields.
277	*
278	* <p>Published to the configured heartbeat exchange with routing key
279	* {@code heartbeat.<workerId>}. Silently drops if the channel or
280	* connection is unavailable (fire-and-forget).
281	*/
282	void sendHeartbeat() {
283	try {
284	WorkerHeartbeatMessage hb = new WorkerHeartbeatMessage(	6✔
285	workerId,
286	epoch,
287	System.currentTimeMillis(),	3✔
288	broadcaster.getCurrentDecisionVersion(),	2✔
289	computeLoadFactor(),	2✔
290	isReadyToServe(),	2✔
291	computeConfigFingerprint(),	3✔
292	stateMachine.getConfirmCount(),	3✔
293	stateMachine.getCoolCount(),	3✔
294	stateMachine.getPreCoolGraceCount(),	3✔
295	configTimestampCounter.get()	3✔
296	);
297	rabbitTemplate.send(heartbeatExchange, ROUTING_KEY_HEARTBEAT + workerId, hb.toMessage());	10✔
298	} catch (Exception e) {	1✔
299	log.error("Scheduled sendHeartbeat failed", e);	4✔
300	}	1✔
301	}	1✔
302	}

Hyshmily / hotkey / 28159619024

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous