thinh/tunnel-client
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add retry strategy on timeout and connection lost

Browse Source
This commit is contained in:
Thinh Ha
2025-07-06 16:48:55 +00:00
parent eea345e93e
commit cdb0790b98
2 changed files with 90 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/main/java/dev/thinhha/tunnel_client/TunnelClientApplication.java
View File

@ -20,6 +20,7 @@ public class TunnelClientApplication {
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
System.out.println("Shutting down tunnel client...");
tunnelClient.stop();
}));
while (true) {

94
src/main/java/dev/thinhha/tunnel_client/service/TunnelClient.java
View File

@ -17,8 +17,9 @@ import java.io.IOException;
import java.net.URI;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
@Service
@Slf4j
@ -35,6 +36,19 @@ public class TunnelClient extends BinaryWebSocketHandler {
// Track WebSocket connections: wsConnectionId -> target WebSocket session
private final Map<String, WebSocketSession> webSocketConnections = new ConcurrentHashMap<>();
// Retry mechanism fields
private final AtomicInteger retryCount = new AtomicInteger(0);
private final AtomicBoolean isConnecting = new AtomicBoolean(false);
private final AtomicBoolean shouldStop = new AtomicBoolean(false);
private final ScheduledExecutorService retryScheduler = Executors.newSingleThreadScheduledExecutor();
// Retry configuration
private static final int MAX_RETRY_ATTEMPTS = 10;
private static final long INITIAL_RETRY_DELAY_MS = 1500;
private static final long MAX_RETRY_DELAY_MS = 60000;
private static final double BACKOFF_MULTIPLIER = 2.0;
private static final long CONNECTION_TIMEOUT_MS = 5000;
public TunnelClient(TunnelConfig tunnelConfig, RouteResolver routeResolver, HeaderManipulationService headerManipulationService, ObjectMapper objectMapper, RestTemplate restTemplate) {
this.tunnelConfig = tunnelConfig;
this.routeResolver = routeResolver;
@ -44,6 +58,24 @@ public class TunnelClient extends BinaryWebSocketHandler {
}
public void connect() {
shouldStop.set(false);
retryCount.set(0);
connectWithRetry();
}
private void connectWithRetry() {
if (shouldStop.get()) {
log.info("Connection stopped by user");
return;
}
if (isConnecting.get()) {
log.debug("Connection already in progress");
return;
}
isConnecting.set(true);
try {
StandardWebSocketClient client = new StandardWebSocketClient();
WebSocketHttpHeaders headers = new WebSocketHttpHeaders();
@ -53,21 +85,64 @@ public class TunnelClient extends BinaryWebSocketHandler {
}
URI serverUri = URI.create(tunnelConfig.getServer().getUrl() + "/client");
log.info("Connecting to tunnel server at: {}", serverUri);
log.info("Connecting to tunnel server at: {} (attempt {}/{})", serverUri, retryCount.get() + 1, MAX_RETRY_ATTEMPTS);
client.execute(this, headers, serverUri);
connectionLatch.await();
boolean connected = connectionLatch.await(CONNECTION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
if (!connected) {
throw new RuntimeException("Connection timeout after " + CONNECTION_TIMEOUT_MS + "ms");
}
log.info("Connected to tunnel server as client: {}", tunnelConfig.getClient().getName());
retryCount.set(0);
} catch (Exception e) {
log.error("Failed to connect to tunnel server", e);
log.error("Failed to connect to tunnel server (attempt {}/{}): {}", retryCount.get() + 1, MAX_RETRY_ATTEMPTS, e.getMessage());
scheduleRetry();
} finally {
isConnecting.set(false);
}
}
private void scheduleRetry() {
int currentRetryCount = retryCount.incrementAndGet();
if (currentRetryCount >= MAX_RETRY_ATTEMPTS) {
log.error("Max retry attempts ({}) reached. Stopping connection attempts.", MAX_RETRY_ATTEMPTS);
return;
}
long delayMs = calculateRetryDelay(currentRetryCount);
log.info("Scheduling retry {} in {} ms", currentRetryCount, delayMs);
retryScheduler.schedule(() -> {
log.info("Attempting to reconnect (retry {}/{})", currentRetryCount, MAX_RETRY_ATTEMPTS);
connectWithRetry();
}, delayMs, TimeUnit.MILLISECONDS);
}
private long calculateRetryDelay(int attemptCount) {
long delay = (long) (INITIAL_RETRY_DELAY_MS * Math.pow(BACKOFF_MULTIPLIER, attemptCount - 1));
return Math.min(delay, MAX_RETRY_DELAY_MS);
}
public void stop() {
shouldStop.set(true);
if (session != null && session.isOpen()) {
try {
session.close();
} catch (Exception e) {
log.warn("Error closing WebSocket session: {}", e.getMessage());
}
}
retryScheduler.shutdown();
log.info("Tunnel client stopped");
}
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
this.session = session;
retryCount.set(0);
log.info("WebSocket connection established with session: {}", session.getId());
connectionLatch.countDown();
}
@ -166,12 +241,21 @@ public class TunnelClient extends BinaryWebSocketHandler {
@Override
public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
log.error("Transport error: {}", exception.getMessage());
if (!shouldStop.get()) {
log.info("Attempting to reconnect after transport error");
scheduleRetry();
}
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
log.info("Connection closed with status: {}", closeStatus);
this.session = null;
if (!shouldStop.get() && !closeStatus.equals(CloseStatus.NORMAL)) {
log.info("Connection closed unexpectedly, attempting to reconnect");
scheduleRetry();
}
}
@Override