Repository: squashql/squashql (java)

Issue:
ConcurrentModificationException due to the cache 
Multiple threads can read/write into a table result coming from the cache. This could cause the following issue:
```
Caused by: java.util.ConcurrentModificationException
	at java.base/java.util.ArrayList$Itr.checkForComodification(ArrayList.java:1013)
	at java.base/java.util.ArrayList$Itr.next(ArrayList.java:967)
	at io.squashql.table.Table.columnIndex(Table.java:63)
	at io.squashql.table.Table.getHeader(Table.java:58)
	at io.squashql.table.TableUtils.selectAndOrderColumns(TableUtils.java:166)
	at io.squashql.table.TableUtils.selectAndOrderColumns(TableUtils.java:155)
	at io.squashql.query.QueryExecutor.executeQuery(QueryExecutor.java:204)
	at io.squashql.query.QueryMergeExecutor.lambda$executeQueryMerge$1(QueryMergeExecutor.java:21)
	at io.squashql.query.QueryMergeExecutor.lambda$execute$8(QueryMergeExecutor.java:89)
	at java.base/java.util.concurrent.CompletableFuture$AsyncSupply.run(CompletableFuture.java:1768)
```

A simple but cumbersome and inefficient solution in to make all `Table` methods synchronized. 

Pull request:
Fix query cache concurrency issue
Fixes #318 

Patch:
diff --git a/core/src/main/java/io/squashql/query/AComparisonExecutor.java b/core/src/main/java/io/squashql/query/AComparisonExecutor.java
index 2958972..101e508 100644
--- a/core/src/main/java/io/squashql/query/AComparisonExecutor.java
+++ b/core/src/main/java/io/squashql/query/AComparisonExecutor.java
@@ -42,7 +42,7 @@ public abstract class AComparisonExecutor<T extends CompiledComparisonMeasure> {
     int readFromTableColumnsCount = (int) readFromTable.headers().stream().filter(header -> !header.isMeasure()).count();
     Object[] buffer = new Object[readFromTableColumnsCount];
     Header[] headers = new Header[readFromTableColumnsCount];
-    List<Object> result = new ArrayList<>((int) writeToTable.count());
+    List<Object> result = new ArrayList<>(writeToTable.count());
     List<Object> readAggregateValues = readFromTable.getAggregateValues(cm.measure());
     List<Object> writeAggregateValues = writeToTable.getAggregateValues(cm.measure());
     BiFunction<Object, Object, Object> comparisonBiFunction = cm.comparisonMethod() != null
diff --git a/core/src/main/java/io/squashql/query/CaffeineQueryCache.java b/core/src/main/java/io/squashql/query/CaffeineQueryCache.java
index c9ecf90..a283475 100644
--- a/core/src/main/java/io/squashql/query/CaffeineQueryCache.java
+++ b/core/src/main/java/io/squashql/query/CaffeineQueryCache.java
@@ -6,16 +6,23 @@ import com.github.benmanes.caffeine.cache.RemovalListener;
 import com.github.benmanes.caffeine.cache.stats.CacheStats;
 import com.github.benmanes.caffeine.cache.stats.ConcurrentStatsCounter;
 import com.github.benmanes.caffeine.cache.stats.StatsCounter;
-import io.squashql.query.compiled.CompiledAggregatedMeasure;
+import com.google.common.util.concurrent.Striped;
 import io.squashql.query.compiled.CompiledMeasure;
 import io.squashql.query.database.SqlUtils;
+import io.squashql.query.dictionary.ObjectArrayDictionary;
 import io.squashql.query.dto.CacheStatsDto;
 import io.squashql.table.ColumnarTable;
 import io.squashql.table.Table;
 import io.squashql.type.TypedField;
+import lombok.AllArgsConstructor;
 
 import java.time.Duration;
 import java.util.*;
+import java.util.concurrent.Callable;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+
+import static io.squashql.query.compiled.CompiledAggregatedMeasure.COMPILED_COUNT;
 
 public class CaffeineQueryCache implements QueryCache {
 
@@ -23,11 +30,12 @@ public class CaffeineQueryCache implements QueryCache {
 
   private volatile StatsCounter scopeCounter = new ConcurrentStatsCounter();
   private volatile StatsCounter measureCounter = new ConcurrentStatsCounter();
+  private final Striped<ReadWriteLock> lock;
 
   /**
    * The cached results.
    */
-  private final Cache<QueryCacheKey, Table> results;
+  private final Cache<QueryCacheKey, DelegateTable> results;
 
   public CaffeineQueryCache() {
     this(MAX_SIZE, (a, b, c) -> {
@@ -42,6 +50,7 @@ public class CaffeineQueryCache implements QueryCache {
             // Use removalListener and not evictionListener because evictionListener is called before updating the stats
             .removalListener(evictionListener)
             .build();
+    this.lock = Striped.readWriteLock(Runtime.getRuntime().availableProcessors() * 4);
   }
 
   @Override
@@ -52,21 +61,20 @@ public class CaffeineQueryCache implements QueryCache {
 
     List<List<Object>> values = new ArrayList<>();
     Table table = this.results.getIfPresent(key);
-    for (TypedField f : columns) {
-      values.add(table.getColumnValues(SqlUtils.squashqlExpression(f)));
-    }
-    values.add(table.getAggregateValues(CompiledAggregatedMeasure.COMPILED_COUNT));
-    return new ColumnarTable(
-            headers,
-            Collections.singleton(CompiledAggregatedMeasure.COMPILED_COUNT),
-            values);
+    return executeRead(table, () -> {
+      for (TypedField f : columns) {
+        values.add(table.getColumnValues(SqlUtils.squashqlExpression(f)));
+      }
+      values.add(table.getColumnValues(COMPILED_COUNT.alias()));
+      return new ColumnarTable(headers, Collections.singleton(COMPILED_COUNT), values);
+    });
   }
 
   @Override
   public boolean contains(CompiledMeasure measure, QueryCacheKey scope) {
     Table table = this.results.getIfPresent(scope);
     if (table != null) {
-      return table.measures().contains(measure);
+      return executeRead(table, () -> table.measures().contains(measure));
     }
     return false;
   }
@@ -75,16 +83,18 @@ public class CaffeineQueryCache implements QueryCache {
   public void contributeToCache(Table result, Set<CompiledMeasure> measures, QueryCacheKey scope) {
     Table cache = this.results.get(scope, s -> {
       this.measureCounter.recordMisses(measures.size());
-      return result;
+      return new DelegateTable(((ColumnarTable) result).copy());
     });
 
-    for (CompiledMeasure measure : measures) {
-      if (!cache.measures().contains(measure)) {
-        // Not in the previousResult, add it.
-        cache.transferAggregates(result, measure);
-        this.measureCounter.recordMisses(1);
+    executeWrite(cache, () -> {
+      for (CompiledMeasure measure : measures) {
+        if (!cache.measures().contains(measure)) {
+          // Not in the previousResult, add it.
+          cache.transferAggregates(result, measure);
+          this.measureCounter.recordMisses(1);
+        }
       }
-    }
+    });
   }
 
   @Override
@@ -92,10 +102,15 @@ public class CaffeineQueryCache implements QueryCache {
     if (measures.isEmpty()) {
       return;
     }
-    Table cacheResult = this.results.getIfPresent(scope);
-    for (CompiledMeasure measure : measures) {
-      result.transferAggregates(cacheResult, measure);
-      this.measureCounter.recordHits(1);
+    Table cache = this.results.getIfPresent(scope);
+    if (cache != null) {
+      executeRead(cache, () -> {
+        for (CompiledMeasure measure : measures) {
+          result.transferAggregates(cache, measure);
+          this.measureCounter.recordHits(1);
+        }
+        return null;
+      });
     }
   }
 
@@ -130,4 +145,81 @@ public class CaffeineQueryCache implements QueryCache {
     this.measureCounter = new ConcurrentStatsCounter();
     this.scopeCounter = new ConcurrentStatsCounter();
   }
+
+  private <V> V executeRead(Table t, Callable<V> callable) {
+    Lock l = this.lock.get(t).readLock();
+    try {
+      l.lock();
+      return callable.call();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    } finally {
+      l.unlock();
+    }
+  }
+
+  private void executeWrite(Table t, Runnable runnable) {
+    Lock l = this.lock.get(t).writeLock();
+    try {
+      l.lock();
+      runnable.run();
+    } finally {
+      l.unlock();
+    }
+  }
+
+  /**
+   * A wrapper around another {@link Table} to make sure this implementation does not override {@link Object#hashCode()}
+   * and {@link Object#equals(Object)} to work with the {@link Striped striped lock}.
+   */
+  @AllArgsConstructor
+  private final class DelegateTable implements Table {
+
+    private final Table underlying;
+
+    @Override
+    public boolean equals(Object o) {
+      return super.equals(o); // DO NOT CHANGE IT
+    }
+
+    @Override
+    public int hashCode() {
+      return super.hashCode(); // DO NOT CHANGE IT
+    }
+
+    @Override
+    public List<Header> headers() {
+      return this.underlying.headers();
+    }
+
+    @Override
+    public Set<CompiledMeasure> measures() {
+      return this.underlying.measures();
+    }
+
+    @Override
+    public void transferAggregates(Table from, CompiledMeasure measure) {
+      this.underlying.transferAggregates(from, measure);
+    }
+
+    @Override
+    public ObjectArrayDictionary pointDictionary() {
+      return this.underlying.pointDictionary();
+    }
+
+    @Override
+    public Iterator<List<Object>> iterator() {
+      return this.underlying.iterator();
+    }
+
+    @Override
+    public void addAggregates(Header header, CompiledMeasure measure, List<Object> values) {
+      this.underlying.addAggregates(header, measure, values);
+    }
+
+    @Override
+    public int count() {
+      return this.underlying.count();
+    }
+  }
 }
diff --git a/core/src/main/java/io/squashql/query/compiled/Evaluator.java b/core/src/main/java/io/squashql/query/compiled/Evaluator.java
index 872b453..8f00c14 100644
--- a/core/src/main/java/io/squashql/query/compiled/Evaluator.java
+++ b/core/src/main/java/io/squashql/query/compiled/Evaluator.java
@@ -126,7 +126,7 @@ public class Evaluator implements BiConsumer<QueryPlanNodeKey, ExecutionContext>
       throw new IllegalArgumentException("Unexpected type " + cm.getClass() + ". Only double and long are supported");
     }
     Header header = new Header(cm.alias(), type, true);
-    List<Object> r = Collections.nCopies((int) intermediateResult.count(), v);
+    List<Object> r = Collections.nCopies(intermediateResult.count(), v);
     intermediateResult.addAggregates(header, cm, r);
   }
 
@@ -175,7 +175,7 @@ public class Evaluator implements BiConsumer<QueryPlanNodeKey, ExecutionContext>
       columnValues.add(readTable.getColumnValues(alias));
     }
 
-    List<Object> vectorValues = ListUtils.createListWithNulls((int) writeToTable.count());
+    List<Object> vectorValues = ListUtils.createListWithNulls(writeToTable.count());
     writeToTable.pointDictionary().forEach((point, index) -> {
       int position = readTable.pointDictionary().getPosition(point);
       if (position >= 0) {
diff --git a/core/src/main/java/io/squashql/table/ColumnarTable.java b/core/src/main/java/io/squashql/table/ColumnarTable.java
index 78bdd93..b0862e9 100644
--- a/core/src/main/java/io/squashql/table/ColumnarTable.java
+++ b/core/src/main/java/io/squashql/table/ColumnarTable.java
@@ -66,22 +66,18 @@ public class ColumnarTable implements Table {
                       + " but does not match the headers of the destination table " + toHeaderNames);
     }
 
-    if (from instanceof ColumnarTable ct) {
-      List<Object> values = ListUtils.createListWithNulls((int) count());
-      List<Object> aggregateValues = from.getAggregateValues(measure);
-      this.pointDictionary.get().forEach((point, index) -> {
-        int position = ct.pointDictionary.get().getPosition(point);
-        if (position >= 0) {
-          values.set(index, aggregateValues.get(position));
-        }
-      });
-      Header header = from.getHeader(measure);
-      this.headers.add(new Header(header.name(), header.type(), true));
-      this.measures.add(measure);
-      this.values.add(values);
-    } else {
-      throw new IllegalArgumentException();
-    }
+    List<Object> values = ListUtils.createListWithNulls(count());
+    List<Object> aggregateValues = from.getAggregateValues(measure);
+    this.pointDictionary.get().forEach((point, index) -> {
+      int position = from.pointDictionary().getPosition(point);
+      if (position >= 0) {
+        values.set(index, aggregateValues.get(position));
+      }
+    });
+    Header header = from.getHeader(measure);
+    this.headers.add(new Header(header.name(), header.type(), true));
+    this.measures.add(measure);
+    this.values.add(values);
   }
 
   @Override
@@ -94,7 +90,7 @@ public class ColumnarTable implements Table {
   }
 
   @Override
-  public long count() {
+  public int count() {
     return this.values.get(0).size();
   }
 
@@ -123,10 +119,18 @@ public class ColumnarTable implements Table {
     return new ColumnarTableIterator();
   }
 
+  public ColumnarTable copy() {
+    List<List<Object>> newValues = new ArrayList<>();
+    for (int i = 0; i < this.headers.size(); i++) {
+      newValues.add(new ArrayList<>(getColumn(i)));
+    }
+    return new ColumnarTable(this.headers, this.measures, newValues);
+  }
+
   private class ColumnarTableIterator implements Iterator<List<Object>> {
 
     int current = 0;
-    long size = count();
+    int size = count();
 
     @Override
     public boolean hasNext() {
diff --git a/core/src/main/java/io/squashql/table/PivotTableUtils.java b/core/src/main/java/io/squashql/table/PivotTableUtils.java
index 33a5e2d..1a16284 100644
--- a/core/src/main/java/io/squashql/table/PivotTableUtils.java
+++ b/core/src/main/java/io/squashql/table/PivotTableUtils.java
@@ -210,7 +210,7 @@ public final class PivotTableUtils {
             ? PivotTableUtils.findNullValuesOnEntireColumn(pivotTable)
             : empty;
 
-    List<Map<String, Object>> cells = new ArrayList<>((int) pivotTable.table.count());
+    List<Map<String, Object>> cells = new ArrayList<>(pivotTable.table.count());
     List<String> headerNames = pivotTable.table.headers().stream().map(Header::name).toList();
     int[] line = new int[1];
     pivotTable.table.forEach(row -> {
diff --git a/core/src/main/java/io/squashql/table/RowTable.java b/core/src/main/java/io/squashql/table/RowTable.java
index fba3e3a..51cfab8 100644
--- a/core/src/main/java/io/squashql/table/RowTable.java
+++ b/core/src/main/java/io/squashql/table/RowTable.java
@@ -37,7 +37,7 @@ public class RowTable implements Table {
   }
 
   @Override
-  public long count() {
+  public int count() {
     return this.rows.size();
   }
 
diff --git a/core/src/main/java/io/squashql/table/Table.java b/core/src/main/java/io/squashql/table/Table.java
index 22cb5ab..faa53f9 100644
--- a/core/src/main/java/io/squashql/table/Table.java
+++ b/core/src/main/java/io/squashql/table/Table.java
@@ -86,7 +86,7 @@ public interface Table extends Iterable<List<Object>> {
    *
    * @return the number of rows
    */
-  long count();
+  int count();
 
   /**
    * Returns the total number of rows before applying the query limit as in {@link QueryDto#limit}.
@@ -108,12 +108,12 @@ public interface Table extends Iterable<List<Object>> {
   default String toString(int numRows) {
     return TableUtils.toString(headers(), () -> new Iterator<>() {
 
-      Iterator<List<Object>> underlying = iterator();
-      int[] c = new int[1];
+      final Iterator<List<Object>> underlying = iterator();
+      final int[] c = new int[1];
 
       @Override
       public boolean hasNext() {
-        return this.c[0] < numRows ? this.underlying.hasNext() : false;
+        return this.c[0] < numRows && this.underlying.hasNext();
       }
 
       @Override
diff --git a/core/src/main/java/io/squashql/table/TableUtils.java b/core/src/main/java/io/squashql/table/TableUtils.java
index 8b0d9fe..bfda287 100644
--- a/core/src/main/java/io/squashql/table/TableUtils.java
+++ b/core/src/main/java/io/squashql/table/TableUtils.java
@@ -243,12 +243,8 @@ public class TableUtils {
     // To lazily copy the table when needed.
     boolean[] lazilyCreated = new boolean[1];
     Supplier<Table> finalTable = Suppliers.memoize(() -> {
-      List<List<Object>> newValues = new ArrayList<>();
-      for (int i = 0; i < table.headers().size(); i++) {
-        newValues.add(new ArrayList<>(table.getColumn(i)));
-      }
       lazilyCreated[0] = true;
-      return new ColumnarTable(table.headers(), table.measures(), newValues);
+      return table.copy();
     });
 
     for (int rowIndex = 0; rowIndex < table.count(); rowIndex++) {
diff --git a/js/typescript-library/package-lock.json b/js/typescript-library/package-lock.json
index 854512c..5ae185e 100644
--- a/js/typescript-library/package-lock.json
+++ b/js/typescript-library/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@squashql/squashql-js",
-  "version": "1.17.2",
+  "version": "1.18.1",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "@squashql/squashql-js",
-      "version": "1.17.2",
+      "version": "1.18.1",
       "license": "Apache-2.0",
       "dependencies": {
         "axios": "^1.6.8"


Test patch:
diff --git a/core/src/test/java/io/squashql/query/ATestQueryCache.java b/core/src/test/java/io/squashql/query/ATestQueryCache.java
index fe5436d..c41d136 100644
--- a/core/src/test/java/io/squashql/query/ATestQueryCache.java
+++ b/core/src/test/java/io/squashql/query/ATestQueryCache.java
@@ -14,11 +14,11 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
@@ -598,6 +598,38 @@ public abstract class ATestQueryCache extends ABaseTestQuery {
             List.of("food", 3d));
   }
 
+  @Test
+  void testQueryCacheConcurrency() throws Exception {
+    int nThreads = Runtime.getRuntime().availableProcessors();
+    ExecutorService executor = Executors.newFixedThreadPool(nThreads);
+
+    for (int j = 0; j < 1000; j++) {
+      this.executor.queryCache.clear();
+      List<Future<Table>> futures = new ArrayList<>();
+      for (int i = 0; i < nThreads; i++) {
+        String alias = Integer.toString(i); // use a different alias read/write different measures from the cache
+        Future<Table> result = executor.submit(() -> {
+          QueryDto query = Query
+                  .from(this.storeName)
+                  .select(tableFields(List.of("category")), List.of(sum(alias, "price")))
+                  .build();
+          return this.executor.executeQuery(query);
+        });
+        futures.add(result);
+      }
+
+      for (Future<Table> future : futures) {
+        Table table = future.get();
+        Assertions.assertThat(table).containsExactly(
+                List.of("cloth", 10d),
+                List.of("drink", 2d),
+                List.of("food", 3d));
+      }
+    }
+
+    executor.shutdownNow();
+  }
+
   private void assertCacheStats(int hitCount, int missCount) {
     TestUtil.assertCacheStats(this.queryCache, hitCount, missCount);
   }
diff --git a/core/src/test/java/io/squashql/query/ATestSQLFunction.java b/core/src/test/java/io/squashql/query/ATestSQLFunction.java
index 60ef8f7..732c318 100644
--- a/core/src/test/java/io/squashql/query/ATestSQLFunction.java
+++ b/core/src/test/java/io/squashql/query/ATestSQLFunction.java
@@ -149,7 +149,7 @@ public abstract class ATestSQLFunction extends ABaseTestQuery {
     Assertions.assertThat(result.headers().stream().map(Header::name)).containsExactly("current_date()", CountMeasure.INSTANCE.alias);
     // We cannot assert the date sent back because the test and the db are not guarantee to run in the same timezone
     //    Assertions.assertThat(result).containsExactly(List.of(this.now, 1L));
-    Assertions.assertThat(result.count()).isEqualTo(1L);
+    Assertions.assertThat(result.count()).isEqualTo(1);
 
     // currentDate in select and where + alias. This query does not make any sense...
     Field cr = currentDate().as("CR");
@@ -162,6 +162,6 @@ public abstract class ATestSQLFunction extends ABaseTestQuery {
     Assertions.assertThat(result.headers().stream().map(Header::name)).containsExactly("CR", CountMeasure.INSTANCE.alias);
     // We cannot assert the date sent back because the test and the db are not guarantee to run in the same timezone
     //    Assertions.assertThat(result).containsExactly(List.of(this.now, 1L));
-    Assertions.assertThat(result.count()).isEqualTo(1L);
+    Assertions.assertThat(result.count()).isEqualTo(1);
   }
 }
diff --git a/core/src/test/java/io/squashql/table/TestMergeTablesEdgeCases.java b/core/src/test/java/io/squashql/table/TestMergeTablesEdgeCases.java
index 66c73cc..9b300f2 100644
--- a/core/src/test/java/io/squashql/table/TestMergeTablesEdgeCases.java
+++ b/core/src/test/java/io/squashql/table/TestMergeTablesEdgeCases.java
@@ -36,7 +36,7 @@ class TestMergeTablesEdgeCases {
                     new ArrayList<>(Arrays.asList(20, 25, 12, 5))));
 
     Table emptyTable = Mockito.mock(Table.class);
-    Mockito.when(emptyTable.count()).thenReturn(0L);
+    Mockito.when(emptyTable.count()).thenReturn(0);
     Table mergedTable = MergeTables.mergeTables(emptyTable, table, JoinType.LEFT);
     Assertions.assertThat(mergedTable).isEqualTo(table);
 
diff --git a/duckdb/src/test/java/io/squashql/TestPivotTable.java b/duckdb/src/test/java/io/squashql/TestPivotTable.java
index 4c49805..4b4988e 100644
--- a/duckdb/src/test/java/io/squashql/TestPivotTable.java
+++ b/duckdb/src/test/java/io/squashql/TestPivotTable.java
@@ -357,7 +357,7 @@ public class TestPivotTable {
     List<Field> columns = List.of(new AliasedField("category"));
     PivotTable pivotTable = this.executor.executePivotQuery(new PivotTableQueryDto(query, rows, columns));
     // We check it does not throw and it returns a result
-    Assertions.assertThat(pivotTable.table.count()).isEqualTo(18L);
+    Assertions.assertThat(pivotTable.table.count()).isEqualTo(18);
   }
 
   @Test