28425041198

Committed 30 Jun 2026 06:30AM UTC coverage: 90.88% (-0.3%) from 91.186%

Build # 28425041198

Build Type

push

github

Committed by

Hyshmily

Commit Message

refactor: simplify Worker cluster management and failure detection

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

Coverage Stats

1250 of 1437 branches covered (86.99%)

Branch coverage included in aggregate %.

10 of 21 new or added lines in 4 files covered. (47.62%)

1 existing line in 1 file now uncovered.

3543 of 3837 relevant lines covered (92.34%)

4.22 hits per line

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

96.15

/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 io.github.hyshmily.hotkey.util.HotKeyThreadFactory;
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.context.annotation.DependsOn;
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
@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(new HotKeyThreadFactory("hotkey-hb-producer")),
      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;
  }

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

Hyshmily / hotkey / 28425041198

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