feat: support refreshing Iceberg tables

engine/table/src/main/java/io/deephaven/engine/table/impl/locations/TableLocationProvider.java

@@ -40,24 +40,6 @@ interface Listener extends BasicTableDataListener {
  */
  void endTransaction(@NotNull Object token);
 
- /**
- * <p>
- * Notify the listener of a {@link TableLocationKey} encountered while initiating or maintaining the location
- * subscription. This should occur at most once per location, but the order of delivery is <i>not</i>
- * guaranteed.
- * </p>
- * <p>
- * If transactionToken is {@code null}, the key will be added to the pending additions immediately.
- * </p>
- *
- * @param tableLocationKey The new table location key.
- * @param transactionToken The token identifying the transaction, or {@code null} if this addition is not part
- * of a transaction.
- */
- void handleTableLocationKeyAdded(
- @NotNull ImmutableTableLocationKey tableLocationKey,
- @Nullable Object transactionToken);
-
  /**
  * Notify the listener of a {@link TableLocationKey} encountered while initiating or maintaining the location
  * subscription. This should occur at most once per location, but the order of delivery is <i>not</i>
@@ -66,50 +48,33 @@ void handleTableLocationKeyAdded(
  *
  * @param tableLocationKey The new table location key.
  */
- default void handleTableLocationKeyAdded(@NotNull ImmutableTableLocationKey tableLocationKey) {
- handleTableLocationKeyAdded(tableLocationKey, null);
- }
+ void handleTableLocationKeyAdded(@NotNull ImmutableTableLocationKey tableLocationKey);
 
  /**
- * <p>
- * Notify the listener of a {@link TableLocationKey} that has been removed.
- * </p>
- * <p>
- * If transactionToken is {@code null}, the key will be added to the pending removals immediately.
- * </p>
- *
- * @param tableLocationKey The table location key that was removed.
- * @param transactionToken The token identifying the transaction, or {@code null} if this addition is not part
- * of a transaction.
- */
- void handleTableLocationKeyRemoved(
- @NotNull ImmutableTableLocationKey tableLocationKey,
- @Nullable Object transactionToken);
-
- /**
- * Notify the listener of a {@link TableLocationKey} that has been removed. This addition is not part of any
- * transaction, and is equivalent to {@code handleTableLocationKeyAdded(tableLocationKey, null);} by default.
+ * Notify the listener of a {@link TableLocationKey} that has been removed. This removal is not part of any
+ * transaction, and is equivalent to {@code handleTableLocationKeyRemoved(tableLocationKey, null);} by default.
  *
  * @param tableLocationKey The table location key that was removed.
  */
  @SuppressWarnings("unused")
- default void handleTableLocationKeyRemoved(@NotNull ImmutableTableLocationKey tableLocationKey) {
- handleTableLocationKeyRemoved(tableLocationKey, null);
- }
+ void handleTableLocationKeyRemoved(@NotNull ImmutableTableLocationKey tableLocationKey);
 
  /**
  * <p>
  * Notify the listener of collections of {@link TableLocationKey TableLocationKeys} added or removed while
- * initiating or maintaining the location subscription. This should occur at most once per location, but the
- * order of delivery is <i>not</i> guaranteed.
+ * initiating or maintaining the location subscription. Addition or removal should occur at most once per
+ * location, but the order of delivery is <i>not</i> guaranteed.
  * </p>
  *
  * @param addedKeys Collection of table location keys that were added.
  * @param removedKeys Collection of table location keys that were removed.
  */
- void handleTableLocationKeysUpdate(
- @Nullable Collection<ImmutableTableLocationKey> addedKeys,
- @Nullable Collection<ImmutableTableLocationKey> removedKeys);
+ default void handleTableLocationKeysUpdate(
+ @NotNull Collection<ImmutableTableLocationKey> addedKeys,
+ @NotNull Collection<ImmutableTableLocationKey> removedKeys) {
+ removedKeys.forEach(this::handleTableLocationKeyRemoved);
+ addedKeys.forEach(this::handleTableLocationKeyAdded);
+ }
  }
 
  /**

engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocation.java

@@ -258,7 +258,7 @@ public final void incrementReferenceCount() {
  /**
  * Decrement the reference count by one, when the reference count reaches zero this location will be cleared.
  *
- * @throws IllegalStateException If the reference count was not successfully incremented
+ * @throws IllegalStateException If the reference count was not successfully decremented
  */
  public void decrementReferenceCount() {
  referenceCounted.decrementReferenceCount();

engine/table/src/main/java/io/deephaven/engine/table/impl/locations/impl/AbstractTableLocationProvider.java

@@ -25,46 +25,55 @@ public abstract class AbstractTableLocationProvider
  extends SubscriptionAggregator<TableLocationProvider.Listener>
  implements TableLocationProvider {
 
- private static final Set<TableLocationKey> EMPTY_TABLE_LOCATION_KEYS = Collections.emptySet();
+ private static final Set<ImmutableTableLocationKey> EMPTY_TABLE_LOCATION_KEYS = Collections.emptySet();
 
  /**
  * Helper class to manage a transaction of added and removed location keys.
  */
  private static class Transaction {
- Set<TableLocationKey> locationsAdded = EMPTY_TABLE_LOCATION_KEYS;
- Set<TableLocationKey> locationsRemoved = EMPTY_TABLE_LOCATION_KEYS;
+ Set<ImmutableTableLocationKey> locationsAdded = EMPTY_TABLE_LOCATION_KEYS;
+ Set<ImmutableTableLocationKey> locationsRemoved = EMPTY_TABLE_LOCATION_KEYS;
 
- void addlocationKey(TableLocationKey locationKey) {
+ void addlocationKey(ImmutableTableLocationKey locationKey) {
  if (locationsAdded == EMPTY_TABLE_LOCATION_KEYS) {
  locationsAdded = new HashSet<>();
- } else if (locationsAdded.contains(locationKey)) {
+ }
+ // When adding a location in a transaction, check for logical consistency.
+ // 1. If the location was already added in this transaction, we have a problem. A transaction should not
+ // add the same location twice.
+ // 2. If the location was already removed in this transaction, we have a `replace` operation which is not a
+ // logical error (although it may not be supported by all consumers).
+ if (!locationsAdded.add(locationKey)) {
  throw new IllegalStateException("TableLocationKey " + locationKey
  + " was already added in this transaction.");
  }
- locationsAdded.add(locationKey);
  }
 
- void removeLocationKey(TableLocationKey locationKey) {
+ void removeLocationKey(ImmutableTableLocationKey locationKey) {
  if (locationsRemoved == EMPTY_TABLE_LOCATION_KEYS) {
  locationsRemoved = new HashSet<>();
- } else if (locationsRemoved.contains(locationKey)) {
+ }
+ // When removing a location in a transaction, check for logical consistency.
+ // 1. If the location was already removed in this transaction, we have a problem. A transaction should not
+ // remove the same location twice.
+ if (!locationsRemoved.add(locationKey)) {
  throw new IllegalStateException("TableLocationKey " + locationKey
  + " was already removed and has not been replaced.");
- } else if (locationsAdded.contains(locationKey)) {
+ }
+ // 2. If the location was already added in this transaction, we have a problem. A transaction should not
+ // add then remove the same location.
+ if (locationsAdded.contains(locationKey)) {
  throw new IllegalStateException("TableLocationKey " + locationKey
  + " was removed after being added in the same transaction.");
  }
- locationsRemoved.add(locationKey);
  }
  }
 
  private final ImmutableTableKey tableKey;
 
- // These sets represent open transactions that are being accumulated.
+ // Open transactions that are being accumulated
  private final Map<Object, Transaction> transactions = new HashMap<>();
 
- private final Object transactionLock = new Object();
-
  /**
  * Map from {@link TableLocationKey} to itself, or to a {@link TableLocation}. The values are {@link TableLocation}s
  * if:
@@ -122,7 +131,7 @@ public final ImmutableTableKey getKey() {
 
  @Override
  protected final void deliverInitialSnapshot(@NotNull final TableLocationProvider.Listener listener) {
- listener.handleTableLocationKeysUpdate(unmodifiableTableLocationKeys, null);
+ listener.handleTableLocationKeysUpdate(unmodifiableTableLocationKeys, List.of());
  }
 
  /**
@@ -131,12 +140,14 @@ protected final void deliverInitialSnapshot(@NotNull final TableLocationProvider
  * @param token A token to identify the transaction
  */
  protected void beginTransaction(@NotNull final Object token) {
- synchronized (transactionLock) {
+ synchronized (transactions) {
  // Verify that we can start a new transaction with this token.
- if (transactions.containsKey(token)) {
- throw new IllegalStateException("A transaction with token " + token + " is currently open.");
- }
- transactions.put(token, new Transaction());
+ transactions.compute(token, (key, val) -> {
+ if (val != null) {
+ throw new IllegalStateException("A transaction with token " + token + " is already open.");
+ }
+ return new Transaction();
+ });
  }
  }
 
@@ -147,52 +158,50 @@ protected void beginTransaction(@NotNull final Object token) {
  */
  protected void endTransaction(@NotNull final Object token) {
  final Transaction transaction;
- synchronized (transactionLock) {
+ synchronized (transactions) {
  // Verify that this transaction is open.
  transaction = transactions.remove(token);
  if (transaction == null) {
  throw new IllegalStateException("No transaction with token " + token + " is currently open.");
  }
  }
 
- final Collection<ImmutableTableLocationKey> addedImmutableKeys =
+ final Collection<ImmutableTableLocationKey> addedKeys =
  new ArrayList<>(transaction.locationsAdded.size());
- final Collection<ImmutableTableLocationKey> removedImmutableKeys =
+ final Collection<ImmutableTableLocationKey> removedKeys =
  new ArrayList<>(transaction.locationsRemoved.size());
 
- // Process the accumulated adds and removes under a lock on `tableLocations` to keep modifications atomic to
- // other holders of this lock.
- synchronized (tableLocations) {
- if (transaction.locationsAdded != EMPTY_TABLE_LOCATION_KEYS
- || transaction.locationsRemoved != EMPTY_TABLE_LOCATION_KEYS) {
- for (TableLocationKey locationKey : transaction.locationsAdded) {
- locationCreatedRecorder = false;
- final Object result = tableLocations.putIfAbsent(locationKey, this::observeInsert);
- visitLocationKey(locationKey);
- if (locationCreatedRecorder) {
- verifyPartitionKeys(locationKey);
- addedImmutableKeys.add(toKeyImmutable(result));
- }
- }
+ // Return early if there are no changes to process.
+ if (transaction.locationsAdded == EMPTY_TABLE_LOCATION_KEYS
+ && transaction.locationsRemoved == EMPTY_TABLE_LOCATION_KEYS) {
+ return;
+ }
 
- for (TableLocationKey locationKey : transaction.locationsRemoved) {
- final Object removedLocation = tableLocations.remove(locationKey);
- if (removedLocation != null) {
- removedImmutableKeys.add(toKeyImmutable(locationKey));
- }
- }
+ for (ImmutableTableLocationKey locationKey : transaction.locationsRemoved) {
+ final Object removedLocation = tableLocations.remove(locationKey);
+ if (removedLocation != null) {
+ removedKeys.add(locationKey);
  }
  }
 
- if (subscriptions != null) {
+ for (ImmutableTableLocationKey locationKey : transaction.locationsAdded) {
+ locationCreatedRecorder = false;
+ tableLocations.putIfAbsent(locationKey, this::observeInsert);
+ visitLocationKey(locationKey);
+ if (locationCreatedRecorder) {
+ verifyPartitionKeys(locationKey);
+ addedKeys.add(locationKey);
+ }
+ }
+
+ if (subscriptions != null && (!addedKeys.isEmpty() || !removedKeys.isEmpty())) {
  synchronized (subscriptions) {
  // Push the notifications to the subscribers.
- if ((!addedImmutableKeys.isEmpty() || !removedImmutableKeys.isEmpty())
- && subscriptions.deliverNotification(
- Listener::handleTableLocationKeysUpdate,
- addedImmutableKeys,
- removedImmutableKeys,
- true)) {
+ if (subscriptions.deliverNotification(
+ Listener::handleTableLocationKeysUpdate,
+ addedKeys,
+ removedKeys,
+ true)) {
  onEmpty();
  }
  }
@@ -220,29 +229,27 @@ protected final void handleTableLocationKeyAdded(
  @NotNull final TableLocationKey locationKey,
  @Nullable final Object transactionToken) {
 
- if (!supportsSubscriptions()) {
- tableLocations.putIfAbsent(locationKey, TableLocationKey::makeImmutable);
- visitLocationKey(locationKey);
- return;
- }
-
  if (transactionToken != null) {
- // When adding a location in a transaction, check for logical consistency.
- // 1. If the location was already added in this transaction, we have a problem. A transaction should not
- // add the same location twice.
- // 2. If the location was already removed in this transaction, we have a `replace` operation which is not a
- // logical error (although it may not be supported by all consumers).
- synchronized (transactionLock) {
- final Transaction transaction = transactions.get(transactionToken);
+ final Transaction transaction;
+ synchronized (transactions) {
+ transaction = transactions.get(transactionToken);
  if (transaction == null) {
  throw new IllegalStateException(
  "No transaction with token " + transactionToken + " is currently open.");
  }
- transaction.addlocationKey(locationKey);
  }
+ // Store an immutable key
+ transaction.addlocationKey(toKeyImmutable(locationKey));
  return;
  }
 
+ if (!supportsSubscriptions()) {
+ tableLocations.putIfAbsent(locationKey, TableLocationKey::makeImmutable);
+ visitLocationKey(locationKey);
+ return;
+ }
+
+
  // If we're not in a transaction, we should push this key immediately.
  synchronized (subscriptions) {
  // Since we're holding the lock on subscriptions, the following code is overly complicated - we could
@@ -328,12 +335,9 @@ protected void doInitialization() {
  @Override
  @NotNull
  public final Collection<ImmutableTableLocationKey> getTableLocationKeys() {
- // This lock is held while `endTransaction()` updates `tableLocations` with the accumulated adds/removes.
- // Locking here ensures that this call won't return while `tableLocations` (and `unmodifiableTableLocationKeys`)
- // contain a partial transaction.
- synchronized (tableLocations) {
- return unmodifiableTableLocationKeys;
- }
+ final List<ImmutableTableLocationKey> result = new ArrayList<>(tableLocations.size());
+ tableLocations.keySet().forEach(key -> result.add((ImmutableTableLocationKey) key));
+ return result;
  }
 
  @Override
@@ -413,25 +417,22 @@ protected final void handleTableLocationKeyRemoved(@NotNull final TableLocationK
  protected void handleTableLocationKeyRemoved(
  @NotNull final TableLocationKey locationKey,
  @Nullable final Object transactionToken) {
- if (!supportsSubscriptions()) {
- return;
- }
-
- // When removing a location in a transaction, check for logical consistency.
- // 1. If the location was already removed in this transaction, we have a problem. A transaction should not
- // remove the same location twice.
- // 2. If the location was already added in this transaction, we have a problem. A transaction should not
- // add then remove the same location.
  if (transactionToken != null) {
- synchronized (transactionLock) {
- final Transaction transaction = transactions.get(transactionToken);
+ final Transaction transaction;
+ synchronized (transactions) {
+ transaction = transactions.get(transactionToken);
  if (transaction == null) {
  throw new IllegalStateException(
  "No transaction with token " + transactionToken + " is currently open.");
  }
- transaction.removeLocationKey(locationKey);
- return;
  }
+ transaction.removeLocationKey(toKeyImmutable(locationKey));
+ return;
+ }
+
+ if (!supportsSubscriptions()) {
+ tableLocations.remove(locationKey);
+ return;
  }
 
  // If we're not in a transaction, we should push this key immediately.