From d536e98cca716d0b59e68692e7cee5a61d482c7c Mon Sep 17 00:00:00 2001
From: ableegoldman <sophie@confluent.io>
Date: Tue, 26 Feb 2019 23:37:26 -0800
Subject: [PATCH] Swapped TreeMap for ConcurrentSkipListMap

---
 .../internals/InMemoryKeyValueStore.java      | 29 +++++++++----------
 .../internals/AbstractKeyValueStoreTest.java  | 11 +++++++
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
index cc28d6487bb39..b37c39e3b2076 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
+++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
@@ -17,6 +17,8 @@
 package org.apache.kafka.streams.state.internals;
 
 import java.util.List;
+import java.util.concurrent.ConcurrentNavigableMap;
+import java.util.concurrent.ConcurrentSkipListMap;
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.processor.ProcessorContext;
@@ -26,18 +28,16 @@
 
 import java.util.Iterator;
 import java.util.Map;
-import java.util.NavigableMap;
-import java.util.TreeMap;
 
 public class InMemoryKeyValueStore implements KeyValueStore<Bytes, byte[]> {
     private final String name;
-    private final NavigableMap<Bytes, byte[]> map;
+    private final ConcurrentNavigableMap<Bytes, byte[]> map;
     private volatile boolean open = false;
 
     public InMemoryKeyValueStore(final String name) {
         this.name = name;
 
-        this.map = new TreeMap<>();
+        this.map = new ConcurrentSkipListMap<>();
     }
 
     @Override
@@ -46,7 +46,6 @@ public String name() {
     }
 
     @Override
-    @SuppressWarnings("unchecked")
     public void init(final ProcessorContext context,
                      final StateStore root) {
 
@@ -76,12 +75,12 @@ public boolean isOpen() {
     }
 
     @Override
-    public synchronized byte[] get(final Bytes key) {
+    public byte[] get(final Bytes key) {
         return this.map.get(key);
     }
 
     @Override
-    public synchronized void put(final Bytes key, final byte[] value) {
+    public void put(final Bytes key, final byte[] value) {
         if (value == null) {
             this.map.remove(key);
         } else {
@@ -90,7 +89,7 @@ public synchronized void put(final Bytes key, final byte[] value) {
     }
 
     @Override
-    public synchronized byte[] putIfAbsent(final Bytes key, final byte[] value) {
+    public byte[] putIfAbsent(final Bytes key, final byte[] value) {
         final byte[] originalValue = get(key);
         if (originalValue == null) {
             put(key, value);
@@ -99,29 +98,29 @@ public synchronized byte[] putIfAbsent(final Bytes key, final byte[] value) {
     }
 
     @Override
-    public synchronized void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
+    public void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
         for (final KeyValue<Bytes, byte[]> entry : entries) {
             put(entry.key, entry.value);
         }
     }
 
     @Override
-    public synchronized byte[] delete(final Bytes key) {
+    public byte[] delete(final Bytes key) {
         return this.map.remove(key);
     }
 
     @Override
-    public synchronized KeyValueIterator<Bytes, byte[]> range(final Bytes from,
-                                                              final Bytes to) {
+    public KeyValueIterator<Bytes, byte[]> range(final Bytes from, final Bytes to) {
         return new DelegatingPeekingKeyValueIterator<>(
             name,
             new InMemoryKeyValueIterator(this.map.subMap(from, true, to, true).entrySet().iterator()));
     }
 
     @Override
-    public synchronized KeyValueIterator<Bytes, byte[]> all() {
-        final TreeMap<Bytes, byte[]> copy = new TreeMap<>(this.map);
-        return new DelegatingPeekingKeyValueIterator<>(name, new InMemoryKeyValueIterator(copy.entrySet().iterator()));
+    public KeyValueIterator<Bytes, byte[]> all() {
+        return new DelegatingPeekingKeyValueIterator<>(
+            name,
+            new InMemoryKeyValueIterator(this.map.entrySet().iterator()));
     }
 
     @Override
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java
index ad87b608aa803..32a699ccc2b00 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractKeyValueStoreTest.java
@@ -378,4 +378,15 @@ public void shouldDeleteFromStore() {
         store.delete(2);
         assertNull(store.get(2));
     }
+
+    @Test
+    public void shouldNotThrowConcurrentModificationException() {
+        store.put(0, "zero");
+
+        final KeyValueIterator<Integer, String> results = store.range(0, 2);
+
+        store.put(1, "one");
+
+        assertEquals(new KeyValue<>(0, "zero"), results.next());
+    }
 }