@SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/21068)
  })
  /*
   * Returns an iterables containing all distinct keys in this multimap.
   */
  public PrefetchableIterable<K> keys() {
    checkState(
        !isClosed,
        "Multimap user state is no longer usable because it is closed for %s",
        keysStateRequest.getStateKey());
    if (isCleared) {
      List<K> keys = new ArrayList<>(pendingAdds.size());
      for (Map.Entry<?, KV<K, List<V>>> entry : pendingAdds.entrySet()) {
        keys.add(entry.getValue().getKey());
      }
      return PrefetchableIterables.concat(keys);
    }

    Set<Object> pendingRemovesNow = new HashSet<>(pendingRemoves.keySet());
    Map<Object, K> pendingAddsNow = new HashMap<>();
    for (Map.Entry<Object, KV<K, List<V>>> entry : pendingAdds.entrySet()) {
      pendingAddsNow.put(entry.getKey(), entry.getValue().getKey());
    }
    return new PrefetchableIterables.Default<K>() {
      @Override
      public PrefetchableIterator<K> createIterator() {
        return new PrefetchableIterator<K>() {
          PrefetchableIterator<K> persistedKeysIterator = persistedKeys.iterator();
          Iterator<K> pendingAddsNowIterator;
          boolean hasNext;
          K nextKey;

          @Override
          public boolean isReady() {
            return persistedKeysIterator.isReady();
          }

          @Override
          public void prefetch() {
            if (!isReady()) {
              persistedKeysIterator.prefetch();
            }
          }

          @Override
          public boolean hasNext() {
            if (hasNext) {
              return true;
            }

            while (persistedKeysIterator.hasNext()) {
              nextKey = persistedKeysIterator.next();
              Object nextKeyStructuralValue = mapKeyCoder.structuralValue(nextKey);
              if (!pendingRemovesNow.contains(nextKeyStructuralValue)) {
                // Remove all keys that we will visit when passing over the persistedKeysIterator
                // so we do not revisit them when passing over the pendingAddsNowIterator
                if (pendingAddsNow.containsKey(nextKeyStructuralValue)) {
                  pendingAddsNow.remove(nextKeyStructuralValue);
                }
                hasNext = true;
                return true;
              }
            }

            if (pendingAddsNowIterator == null) {
              pendingAddsNowIterator = pendingAddsNow.values().iterator();
            }
            while (pendingAddsNowIterator.hasNext()) {
              nextKey = pendingAddsNowIterator.next();
              hasNext = true;
              return true;
            }

            return false;
          }

          @Override
          public K next() {
            if (!hasNext()) {
              throw new NoSuchElementException();
            }
            hasNext = false;
            return nextKey;
          }
        };
      }
    };
  }
@Test
  public void testGetKeysPrefetch() throws Exception {
    FakeBeamFnStateClient fakeClient =
        new FakeBeamFnStateClient(
            ImmutableMap.of(
                createMultimapKeyStateKey(),
                KV.of(ByteArrayCoder.of(), singletonList(A1)),
                createMultimapValueStateKey(A1),
                KV.of(StringUtf8Coder.of(), asList("V1", "V2"))));
    MultimapUserState<byte[], String> userState =
        new MultimapUserState<>(
            Caches.eternal(),
            fakeClient,
            "instructionId",
            createMultimapKeyStateKey(),
            ByteArrayCoder.of(),
            StringUtf8Coder.of());

    PrefetchableIterable<byte[]> keys = userState.keys();
    assertEquals(0, fakeClient.getCallCount());
    keys.prefetch();
    assertEquals(1, fakeClient.getCallCount());
    assertArrayEquals(new byte[][] {A1}, Iterables.toArray(keys, byte[].class));
    assertEquals(1, fakeClient.getCallCount());
  }