From 30a2eb78fba59bf6e4c419bd0e3d9f1e799fbe24 Mon Sep 17 00:00:00 2001
From: Peter Lawrey <peter.lawrey@higherfrequencytrading.com>
Date: Sat, 15 Nov 2014 13:43:54 +0000
Subject: [PATCH] HCOLL-81 Update tests for lambdas in Java 8.

---
 .../openhft/chronicle/map/ChronicleMap.java   |  4 +++-
 .../chronicle/map/StatelessChronicleMap.java  | 24 +++++++++++++++++--
 .../chronicle/map/VanillaChronicleMap.java    |  4 ++--
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/main/java/net/openhft/chronicle/map/ChronicleMap.java b/src/main/java/net/openhft/chronicle/map/ChronicleMap.java
index 74cc5200e..9e998312a 100644
--- a/src/main/java/net/openhft/chronicle/map/ChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/ChronicleMap.java
@@ -189,13 +189,15 @@ public interface ChronicleMap<K, V> extends ConcurrentMap<K, V>, ChronicleHash {
      * @param key      to apply the mapping to
      * @param function to calculate a result
      * @param <R>      return type.
-     * @return the result of the function.
+     * @return the result of the function, or null if there is no entry for the key.
      */
     <R> R mapForKey(K key, @NotNull Function<? super V, R> function);
 
     /**
      * Apply a mutator to the value for a key and return a result.
      * A write lock is assumed.
+     * <p>
+     *     If there is no entry for this key it will be created and the value empty.
      *
      * @param key     to apply the mapping to
      * @param mutator to alter the value and calculate a result
diff --git a/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java b/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
index 2724725f4..0fd25b71e 100644
--- a/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java
@@ -22,6 +22,7 @@ import com.sun.jdi.connect.spi.ClosedConnectionException;
 import net.openhft.chronicle.hash.RemoteCallTimeoutException;
 import net.openhft.lang.io.ByteBufferBytes;
 import net.openhft.lang.io.Bytes;
+import net.openhft.lang.io.NativeBytes;
 import net.openhft.lang.threadlocal.ThreadLocalCopies;
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
@@ -30,6 +31,7 @@ import org.slf4j.LoggerFactory;
 import java.io.Closeable;
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -972,8 +974,26 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable {
             throw new IllegalStateException("the received transaction-id=" + inTransactionId +
                     ", does not match the expected transaction-id=" + transactionId);
 
-        if (isException)
-            throw (RuntimeException) bytes.readObject();
+        if (isException) {
+            Throwable throwable = (Throwable) bytes.readObject();
+            try {
+                Field stackTrace = Throwable.class.getDeclaredField("stackTrace");
+                stackTrace.setAccessible(true);
+                List<StackTraceElement> stes = new ArrayList<>(Arrays.asList((StackTraceElement[]) stackTrace.get(throwable)));
+                // prune the end of the stack.
+                for (int i = stes.size() - 1; i > 0 && stes.get(i).getClassName().startsWith("Thread"); i--) {
+                    stes.remove(i);
+                }
+                InetSocketAddress address = builder.remoteAddress();
+                stes.add(new StackTraceElement("~ remote", "tcp ~", address.getHostName(), address.getPort()));
+                StackTraceElement[] stackTrace2 = Thread.currentThread().getStackTrace();
+                for (int i = 4; i < stackTrace2.length; i++)
+                    stes.add(stackTrace2[i]);
+                stackTrace.set(throwable, stes.toArray(new StackTraceElement[stes.size()]));
+            } catch (Exception ignore) {
+            }
+            NativeBytes.UNSAFE.throwException(throwable);
+        }
 
         return bytes;
     }
diff --git a/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java b/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
index f1cb4eb73..f19a7e31c 100755
--- a/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java
@@ -367,9 +367,9 @@ class VanillaChronicleMap<K, KI, MKI extends MetaBytesInterop<K, KI>,
 
     @Override
     public <R> R mapForKey(K key, @NotNull Function<? super V, R> function) {
-        try (Context<K, V> entry = lookupUsing(key, null,
+        try (ReadContext<K, V> entry = lookupUsing(key, null,
                 LockType.READ_LOCK, false)) {
-            return function.apply(entry.value());
+            return entry.present() ? function.apply(entry.value()) : null;
         }
     }