From 3de687a78651a3cd14406a49f30e5d3afbe0ae24 Mon Sep 17 00:00:00 2001
From: Igor Abdrakhimov <octosyllabic@gmail.com>
Date: Thu, 28 Mar 2024 14:05:14 -0700
Subject: [PATCH] Add methods for set keep alive (#751)

---
 .../amazon/awssdk/crt/io/SocketOptions.java   | 46 ++++++++++++++++++-
 src/native/socket_options.c                   |  6 ++-
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java b/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java
index f588f5e0d..0313071a5 100644
--- a/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java
+++ b/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java
@@ -6,6 +6,7 @@
 
 import software.amazon.awssdk.crt.CrtResource;
 import software.amazon.awssdk.crt.CrtRuntimeException;
+import software.amazon.awssdk.crt.Log;
 
 /**
  * This class wraps the aws_socket_options from aws-c-io to provide
@@ -87,6 +88,12 @@ int getValue() {
      * 0 disables keepalive
      */
     public int keepAliveTimeoutSecs = 0;
+    /**
+     * Sets the number of keep alive probes allowed to fail before the connection is considered lost.
+     * If zero, OS defaults are used.
+     * On Windows, this option is meaningless until Windows 10 1703.
+     */
+    public int keepAliveMaxFailedProbes = 0;
 
     /**
      * If true, enables periodic transmits of keepalive messages for detecting a disconnected peer.
@@ -97,7 +104,42 @@ int getValue() {
      * Creates a new set of socket options
      */
     public SocketOptions() {
+    }
+
+    /**
+     * Enables TCP keepalive.
+     *
+     * @param keepAliveTimeoutSecs Sets the number of seconds to wait for a keepalive response before considering the
+     * connection timed out. 0 disables keepalive.
+     * @param keepAliveIntervalSecs Sets the number of seconds between TCP keepalive packets being sent to the peer.
+     * 0 disables keepalive.
+     */
+    public void setTcpKeepAlive(int keepAliveTimeoutSecs, int keepAliveIntervalSecs)
+    {
+        if (keepAliveTimeoutSecs == 0 || keepAliveIntervalSecs == 0) {
+            Log.log(Log.LogLevel.Warn, Log.LogSubject.IoSocket,
+                    "Both keepAliveTimeoutSecs and keepAliveIntervalSecs must be non-zero in order to enable TCP keepalive");
+        }
+        this.keepAliveTimeoutSecs = keepAliveTimeoutSecs;
+        this.keepAliveIntervalSecs = keepAliveIntervalSecs;
+        this.keepAlive = true;
+    }
 
+    /**
+     * Enables TCP keepalive.
+     *
+     * @param keepAliveTimeoutSecs Sets the number of seconds to wait for a keepalive response before considering the
+     * connection timed out. 0 disables keepalive.
+     * @param keepAliveIntervalSecs Sets the number of seconds between TCP keepalive packets being sent to the peer.
+     * 0 disables keepalive.
+     * @param keepAliveMaxFailedProbes Sets the number of keep alive probes allowed to fail before the connection is considered lost.
+     * If zero, OS defaults are used.
+     * On Windows, this option is meaningless until Windows 10 1703.
+     */
+    public void setTcpKeepAlive(int keepAliveTimeoutSecs, int keepAliveIntervalSecs, int keepAliveMaxFailedProbes)
+    {
+        this.keepAliveMaxFailedProbes = keepAliveMaxFailedProbes;
+        setTcpKeepAlive(keepAliveTimeoutSecs, keepAliveIntervalSecs);
     }
 
     @Override
@@ -109,6 +151,7 @@ public long getNativeHandle() {
                 connectTimeoutMs,
                 keepAliveIntervalSecs,
                 keepAliveTimeoutSecs,
+                keepAliveMaxFailedProbes,
                 keepAlive
             ));
         }
@@ -135,7 +178,8 @@ protected void releaseNativeHandle() {
     /*******************************************************************************
      * native methods
      ******************************************************************************/
-    private static native long socketOptionsNew(int domain, int type, int connectTimeoutMs, int keepAliveIntervalSecs, int keepAliveTimeoutSecs, boolean keepAlive);
+    private static native long socketOptionsNew(
+            int domain, int type, int connectTimeoutMs, int keepAliveIntervalSecs, int keepAliveTimeoutSecs, int keepAliveMaxFailedProbes, boolean keepAlive);
 
     private static native void socketOptionsDestroy(long elg);
 };
diff --git a/src/native/socket_options.c b/src/native/socket_options.c
index 67a6031c3..c02fa19ba 100644
--- a/src/native/socket_options.c
+++ b/src/native/socket_options.c
@@ -31,6 +31,7 @@ jlong JNICALL Java_software_amazon_awssdk_crt_io_SocketOptions_socketOptionsNew(
     jint connect_timeout_ms,
     jint keep_alive_interval_secs,
     jint keep_alive_timeout_secs,
+    jint keep_alive_max_failed_probes,
     jboolean keep_alive) {
     (void)jni_class;
     aws_cache_jni_ids(env);
@@ -43,8 +44,9 @@ jlong JNICALL Java_software_amazon_awssdk_crt_io_SocketOptions_socketOptionsNew(
     options->domain = domain;
     options->type = type;
     options->connect_timeout_ms = connect_timeout_ms;
-    options->keep_alive_interval_sec = (short)keep_alive_interval_secs;
-    options->keep_alive_timeout_sec = (short)keep_alive_timeout_secs;
+    options->keep_alive_interval_sec = (uint16_t)keep_alive_interval_secs;
+    options->keep_alive_timeout_sec = (uint16_t)keep_alive_timeout_secs;
+    options->keep_alive_max_failed_probes = (uint16_t)keep_alive_max_failed_probes;
     options->keepalive = keep_alive;
 
     return (jlong)options;