package com.spotify.sparkey.extra;

import com.spotify.sparkey.CompressionType;
import com.spotify.sparkey.Sparkey;
import com.spotify.sparkey.SparkeyReader;
import com.spotify.sparkey.SparkeyWriter;
import com.spotify.sparkey.system.BaseSystemTest;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.mockito.stubbing.Answer;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.IntStream;

import static junit.framework.Assert.assertEquals;
import static junit.framework.Assert.assertNotSame;
import static org.mockito.Mockito.*;

public class ThreadLocalSparkeyReaderTest extends BaseSystemTest {
  private ThreadLocalSparkeyReader reader;

  @Before
  public void setUp() throws Exception {
    super.setUp();

    SparkeyWriter writer = Sparkey.createNew(indexFile, CompressionType.NONE, 1024);
    for (int i = 0; i < 10000; i++) {
      writer.put("key_" + i, "value_"+i);
    }
    writer.writeHash();
    writer.close();

    reader = new ThreadLocalSparkeyReader(indexFile);
  }

  @Override
  @After
  public void tearDown() throws Exception {
    if (reader != null) {
      reader.close();
    }
    super.tearDown();
  }

  @Test
  public void testCorrectness() throws IOException {
    for (int i=0; i<10; i++) {
      assertEquals("value_"+i, reader.getAsString("key_"+i));
    }
  }

  @Test
  public void testConcurrentUsageCorrectness() {
    IntStream stream = IntStream.range(0, 10000);
    stream.parallel().forEach(i -> {
      try {
        assertEquals("value_"+i, reader.getAsString("key_"+i));
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    });
  }

  @Test
  public void testDifferentReaders() throws Exception {
    final SparkeyReader localReader = reader.getDelegateReader();

    final ExecutorService executorService = Executors.newSingleThreadExecutor();
    final SparkeyReader remoteReader = executorService.submit(() -> reader.getDelegateReader()).get();
    executorService.shutdown();

    assertNotSame(localReader, remoteReader);
  }

  @Test(expected = IllegalStateException.class)
  public void testAlreadyClosed() throws IOException {
    reader.close();
    reader.getAsString("key_1");
  }

}

@@ -29,16 +29,57 @@
 extra;%0A%0A
+import com.google.common.base.Stopwatch;%0A
 import c
@@ -361,44 +361,36 @@
 st;%0A
+%0A
 import 
-org.mockito.stubbing.Answer;%0A
+java.io.IOException;
 %0Aimp
@@ -394,34 +394,33 @@
 import java.
-io.IOExcep
+time.Dura
 tion;%0Aimport
@@ -434,40 +434,71 @@
 til.
-ArrayList;%0Aimport java.util.List
+concurrent.Callable;%0Aimport java.util.concurrent.CountDownLatch
 ;%0Aim
@@ -579,16 +579,108 @@
 cutors;%0A
+import java.util.concurrent.ForkJoinPool;%0Aimport java.util.concurrent.atomic.AtomicInteger;%0A
 import j
@@ -832,25 +832,31 @@
 org.
-mockito.Mockito.*
+junit.Assert.assertTrue
 ;%0A%0Ap
@@ -1926,32 +1926,1644 @@
  %7D%0A    %7D);%0A  %7D%0A%0A
+  @Test%0A  public void testConcurrentPerformance() throws Exception %7B%0A    final int numRuns = 1_000_000;%0A    final AtomicInteger failures = new AtomicInteger();%0A    final Duration elapsed = min(() -%3E measure(numRuns, failures));%0A    double nanosPerRun = (double) elapsed.toNanos() / numRuns;%0A    System.out.println(%22Nanos per lookup: %22 + nanosPerRun);%0A    assertEquals(0, failures.get());%0A    assertTrue(nanosPerRun %3C 400);%0A  %7D%0A%0A  private static Duration min(Callable%3CDuration%3E callable) throws Exception %7B%0A    return IntStream.range(0, 10).mapToObj(i -%3E %7B%0A      try %7B%0A        return callable.call();%0A      %7D catch (Exception e) %7B%0A        throw new RuntimeException(e);%0A      %7D%0A    %7D).min(Duration::compareTo).get();%0A  %7D%0A%0A  private Duration measure(int numRuns, AtomicInteger failures) throws InterruptedException %7B%0A    final CountDownLatch latch = new CountDownLatch(numRuns);%0A    final Stopwatch stopwatch = Stopwatch.createStarted();%0A    for (int i = 0; i %3C numRuns; i++) %7B%0A      ForkJoinPool.commonPool().execute(lookup(failures, latch, i %25 10000));%0A    %7D%0A    latch.await();%0A    final Duration elapsed = stopwatch.stop().elapsed();%0A    System.out.println(%22Elapsed: %22 + elapsed);%0A    return elapsed;%0A  %7D%0A%0A  private Runnable lookup(AtomicInteger failures, CountDownLatch latch, int index) %7B%0A    return () -%3E %7B%0A      try %7B%0A        final String ans = reader.getAsString(%22key_%22 + index);%0A        if (!ans.equals(%22value_%22 + index)) %7B%0A          failures.incrementAndGet();%0A        %7D%0A      %7D catch (IOException e) %7B%0A        failures.incrementAndGet();%0A      %7D finally %7B%0A        latch.countDown();%0A      %7D%0A    %7D;%0A  %7D%0A%0A
   @Test%0A  public

Add ThreadLocal performance test