27925388821

Committed 22 Jun 2026 02:14AM UTC coverage: 92.403% (-1.2%) from 93.615%

Build # 27925388821

Build Type

push

github

Committed by

Hyshmily

Commit Message

fix:fix known bugs and promote performance

Signed-off-by: Hyshmily <cxm8607@outlook.com>

Coverage Stats

1143 of 1289 branches covered (88.67%)

Branch coverage included in aggregate %.

184 of 236 new or added lines in 19 files covered. (77.97%)

6 existing lines in 3 files now uncovered.

3321 of 3542 relevant lines covered (93.76%)

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

Hyshmily / hotkey / 27925388821

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