Fix session affinity management (stickiness) in wildfly-http-client

common/src/main/java/org/wildfly/httpclient/common/HttpClientMessages.java

@@ -78,4 +78,6 @@ interface HttpClientMessages extends BasicLogger {
     @Message(id = 14, value = "JavaEE to JakartaEE backward compatibility layer have been installed")
     void javaeeToJakartaeeBackwardCompatibilityLayerInstalled();
 
+    @Message(id = 15, value = "Failed to acquire backend server")
+    RuntimeException failedToAcquireBackendServer(@Cause Throwable e);
 }

common/src/main/java/org/wildfly/httpclient/common/HttpTargetContext.java

@@ -27,6 +27,7 @@
 import io.undertow.util.Cookies;
 import io.undertow.util.HeaderValues;
 import io.undertow.util.Headers;
+import io.undertow.util.HttpString;
 import io.undertow.util.Methods;
 import io.undertow.util.StatusCodes;
 import org.jboss.marshalling.InputStreamByteInput;
@@ -46,12 +47,14 @@
 import java.io.ObjectInput;
 import java.io.OutputStream;
 import java.net.URI;
+import java.net.URISyntaxException;
 import java.security.AccessController;
 import java.security.GeneralSecurityException;
 import java.security.PrivilegedAction;
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.zip.GZIPInputStream;
@@ -71,7 +74,7 @@ public class HttpTargetContext extends AbstractAttachable {
     }
 
     private static final String EXCEPTION_TYPE = "application/x-wf-jbmar-exception";
-
+    private static final HttpString BACKEND_HEADER = new HttpString("Backend");
     private static final String JSESSIONID = "JSESSIONID";
 
     private final HttpConnectionPool connectionPool;
@@ -144,6 +147,49 @@ private void acquireSessionAffinity(CountDownLatch latch, AuthenticationConfigur
         }, null, latch::countDown);
     }
 
+    public URI acquireBackendServer() throws Exception {
+        // HttpClientMessages.MESSAGES.info("HtpTargetContext: acquireBackendServer()");
+        return acquireBackendServer(AUTH_CONTEXT_CLIENT.getAuthenticationConfiguration(uri, AuthenticationContext.captureCurrent()));
+    }
+
+    private URI acquireBackendServer(AuthenticationConfiguration authenticationConfiguration) throws Exception {
+        ClientRequest clientRequest = new ClientRequest();
+        clientRequest.setMethod(Methods.GET);
+        clientRequest.setPath(uri.getPath() + "/common/v1/backend");
+        AuthenticationContext context = AuthenticationContext.captureCurrent();
+        SSLContext sslContext;
+        try {
+            sslContext = AUTH_CONTEXT_CLIENT.getSSLContext(uri, context);
+        } catch(GeneralSecurityException e) {
+            HttpClientMessages.MESSAGES.failedToAcquireBackendServer(e);
+            return null;
+        }
+
+        // returns a URI of the form <scheme>://<host>:<port>?name=<jboss.node.name>
+        // this permits having access to *both* the IP:port and the hostname identifiers for the server
+        CompletableFuture<URI> result = new CompletableFuture<>();
+        sendRequest(clientRequest, sslContext, authenticationConfiguration,
+                null,
+                null,
+                ((resultStream, response, closeable) -> {
+                    HeaderValues backends = response.getResponseHeaders().get(BACKEND_HEADER);
+                    if (backends == null) {
+                        result.completeExceptionally(HttpClientMessages.MESSAGES.failedToAcquireBackendServer(new Exception("Missing backend header on response")));
+                    }
+                    try {
+                        String backendString = backends.getFirst();
+                        URI backendURI = new URI(backendString);
+                        result.complete(backendURI);
+                    } catch(URISyntaxException use) {
+                        result.completeExceptionally(HttpClientMessages.MESSAGES.failedToAcquireBackendServer(use));
+                    } finally {
+                        IoUtils.safeClose(closeable);
+                    }
+                }),
+                result::completeExceptionally, null, null);
+        return result.get();
+    }
+
     public void sendRequest(ClientRequest request, SSLContext sslContext, AuthenticationConfiguration authenticationConfiguration, HttpMarshaller httpMarshaller, HttpStickinessHandler httpStickinessHandler, HttpResultHandler httpResultHandler, HttpFailureHandler failureHandler, ContentType expectedResponse, Runnable completedTask) {
         sendRequest(request, sslContext, authenticationConfiguration, httpMarshaller, httpStickinessHandler, httpResultHandler, failureHandler, expectedResponse, completedTask, false);
     }

ejb/src/main/java/org/wildfly/httpclient/ejb/HttpEJBReceiver.java

@@ -284,6 +284,10 @@ protected SessionID createSession(final EJBReceiverSessionCreationContext receiv
         targetContext.awaitSessionId(true, authenticationConfiguration);
         CompletableFuture<SessionID> result = new CompletableFuture<>();
 
+        // debugging
+        URI backendURI = targetContext.acquireBackendServer();
+        EjbHttpClientMessages.MESSAGES.infof("HttpEJBReceiver: Getting backend server URI: %s", backendURI);
+
         HttpEJBInvocationBuilder builder = new HttpEJBInvocationBuilder()
                 .setInvocationType(HttpEJBInvocationBuilder.InvocationType.STATEFUL_CREATE)
                 .setAppName(locator.getAppName())

ejb/src/test/java/org/wildfly/httpclient/ejb/SimpleInvocationTestCase.java

@@ -19,6 +19,7 @@
 package org.wildfly.httpclient.ejb;
 
 import io.undertow.util.Headers;
+import io.undertow.util.HttpString;
 import org.jboss.ejb.client.EJBClient;
 import org.jboss.ejb.client.EJBClientContext;
 import org.jboss.ejb.client.EJBClientInvocationContext;
@@ -57,6 +58,7 @@ public class SimpleInvocationTestCase {
     @Before
     public void before() {
         EJBTestServer.registerServicesHandler("common/v1/affinity", httpServerExchange -> httpServerExchange.getResponseHeaders().put(Headers.SET_COOKIE, "JSESSIONID=" + EJBTestServer.INITIAL_SESSION_AFFINITY));
+        EJBTestServer.registerServicesHandler("common/v1/backend", httpServerExchange -> httpServerExchange.getResponseHeaders().put(new HttpString("Backend"), EJBTestServer.getDefaultServerURL()+"?node=localhost" ));
         StringBuilder sb = new StringBuilder();
         for (int i = 0; i < 10000; ++i) {
             sb.append("Hello World ");