Extend with more methods. Use Optional in return type.

2 files changed, 22 insertions, 14 deletions

diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
index 9a6465cce3b..a734ba96c66 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
@@ -33,7 +33,6 @@ import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Date;
 import java.util.List;
-import java.util.Objects;
 import java.util.UUID;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -88,11 +87,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void onOpened(Connection connection) {
         handleListenerInvocation("Connection.Listener", "onOpened", "%h", List.of(connection), () -> {
             SocketChannelEndPoint endpoint = findUnderlyingSocketEndpoint(connection.getEndPoint());
-            ConnectionInfo info = connectionInfo.get(endpoint);
-            if (info == null) {
-                info = ConnectionInfo.from(endpoint);
-                connectionInfo.put(endpoint, info);
-            }
+            ConnectionInfo info = connectionInfo.computeIfAbsent(endpoint, ConnectionInfo::from);
             String connectionClassName = connection.getClass().getSimpleName(); // For hidden implementations of Connection
             if (connection instanceof SslConnection) {
                 SSLEngine sslEngine = ((SslConnection) connection).getSSLEngine();
@@ -117,7 +112,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void onClosed(Connection connection) {
         handleListenerInvocation("Connection.Listener", "onClosed", "%h", List.of(connection), () -> {
             SocketChannelEndPoint endpoint = findUnderlyingSocketEndpoint(connection.getEndPoint());
-            ConnectionInfo info = connectionInfo.get(endpoint);
+            ConnectionInfo info = connectionInfo.get(endpoint).orElse(null);
             if (info == null) return; // Closed connection already handled
             if (connection instanceof HttpConnection) {
                 info.setHttpBytes(connection.getBytesIn(), connection.getBytesOut());
@@ -148,7 +143,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void onRequestBegin(Request request) {
         handleListenerInvocation("HttpChannel.Listener", "onRequestBegin", "%h", List.of(request), () -> {
             SocketChannelEndPoint endpoint = findUnderlyingSocketEndpoint(request.getHttpChannel().getEndPoint());
-            ConnectionInfo info = Objects.requireNonNull(connectionInfo.get(endpoint));
+            ConnectionInfo info = connectionInfo.get(endpoint).get();
             info.incrementRequests();
             request.setAttribute(CONNECTION_ID_REQUEST_ATTRIBUTE, info.uuid());
         });
@@ -158,7 +153,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void onResponseBegin(Request request) {
         handleListenerInvocation("HttpChannel.Listener", "onResponseBegin", "%h", List.of(request), () -> {
             SocketChannelEndPoint endpoint = findUnderlyingSocketEndpoint(request.getHttpChannel().getEndPoint());
-            ConnectionInfo info = connectionInfo.get(endpoint);
+            ConnectionInfo info = connectionInfo.get(endpoint).orElse(null);
             if (info == null) return; // Connection closed before response started - observed during Jetty server shutdown
             info.incrementResponses();
         });
@@ -174,7 +169,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void handshakeSucceeded(Event event) {
         SSLEngine sslEngine = event.getSSLEngine();
         handleListenerInvocation("SslHandshakeListener", "handshakeSucceeded", "sslEngine=%h", List.of(sslEngine), () -> {
-            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine);
+            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine).orElse(null);
             if (info == null) return;
             info.setSslSessionDetails(sslEngine.getSession());
         });
@@ -185,7 +180,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
         SSLEngine sslEngine = event.getSSLEngine();
         handleListenerInvocation("SslHandshakeListener", "handshakeFailed", "sslEngine=%h,failure=%s", List.of(sslEngine, failure), () -> {
             log.log(Level.FINE, failure, failure::toString);
-            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine);
+            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine).orElse(null);
             if (info == null) return;
             info.setSslHandshakeFailure((SSLHandshakeException)failure);
         });
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
index 52142e534ba..b2bfa2a5dda 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
@@ -1,8 +1,11 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.http.server.jetty;
 
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 /**
  * A simplified {@link ConcurrentMap} using reference-equality to compare keys (similarly to {@link java.util.IdentityHashMap})
@@ -13,11 +16,21 @@ class SimpleConcurrentIdentityHashMap<K, V> {
 
     private final ConcurrentMap<IdentityKey<K>, V> wrappedMap = new ConcurrentHashMap<>();
 
-    V get(K key) { return wrappedMap.get(IdentityKey.of(key)); }
+    Optional<V> get(K key) { return Optional.ofNullable(wrappedMap.get(identityKey(key))); }
 
-    V remove(K key) { return wrappedMap.remove(IdentityKey.of(key)); }
+    Optional<V> remove(K key) { return Optional.ofNullable(wrappedMap.remove(identityKey(key))); }
 
-    void put(K key, V value) { wrappedMap.put(IdentityKey.of(key), value); }
+    void put(K key, V value) { wrappedMap.put(identityKey(key), value); }
+
+    V computeIfAbsent(K key, Supplier<V> supplier) {
+        return wrappedMap.computeIfAbsent(identityKey(key), ignored -> supplier.get());
+    }
+
+    V computeIfAbsent(K key, Function<K, V> factory) {
+        return wrappedMap.computeIfAbsent(identityKey(key), k -> factory.apply(k.instance));
+    }
+
+    private static <K> IdentityKey<K> identityKey(K key) { return IdentityKey.of(key); }
 
     private static class IdentityKey<K> {
         final K instance;