org.apache.samza.storage.kv.Entry Java Examples

/**
 * Processes the incoming side input message envelope and updates the last processed offset for its SSP.
 * Synchronized inorder to be exclusive with flush().
 *
 * @param envelope incoming envelope to be processed
 */
public synchronized void process(IncomingMessageEnvelope envelope) {
  SystemStreamPartition envelopeSSP = envelope.getSystemStreamPartition();
  String envelopeOffset = envelope.getOffset();

  for (String store: this.sspToStores.get(envelopeSSP)) {
    SideInputsProcessor storeProcessor = this.storeToProcessor.get(store);
    KeyValueStore keyValueStore = (KeyValueStore) this.taskSideInputStorageManager.getStore(store);
    Collection<Entry<?, ?>> entriesToBeWritten = storeProcessor.process(envelope, keyValueStore);

    // TODO: SAMZA-2255: optimize writes to side input stores
    for (Entry entry : entriesToBeWritten) {
      // If the key is null we ignore, if the value is null, we issue a delete, else we issue a put
      if (entry.getKey() != null) {
        if (entry.getValue() != null) {
          keyValueStore.put(entry.getKey(), entry.getValue());
        } else {
          keyValueStore.delete(entry.getKey());
        }
      }
    }
  }

  this.lastProcessedOffsets.put(envelopeSSP, envelopeOffset);
}
 

@Override
public void window(MessageCollector collector,
  TaskCoordinator coordinator) {

  KeyValueIterator<String, String> entries = store.all();
  while (entries.hasNext()) {                                        // c
    Entry<String, String> entry = entries.next();
    String key = entry.getKey();
    String value = entry.getValue();
    if (isTimestampKey(key) && Cart.isAbandoned(value)) {            // d
      String shopper = extractShopper(key);
      String cart = store.get(asCartKey(shopper));
      
      AbandonedCartEvent event =
        new AbandonedCartEvent(shopper, cart);
      collector.send(new OutgoingMessageEnvelope(
        new SystemStream("kafka", "derived-events-ch04"), event));    // e
      
      resetShopper(shopper);
    }
  }
}
 

@SuppressWarnings("unchecked")
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("mystore");
  System.out.println("Contents of store: ");
  KeyValueIterator<String, String> iter = store.all();
  while (iter.hasNext()) {
    Entry<String, String> entry = iter.next();
    System.out.println(entry.getKey() + " => " + entry.getValue());
  }
  iter.close();
}
 

/**
 * For each standby task, we remove its changeLogSSPs from changelogSSP map and add it to the task's taskSideInputSSPs.
 * The task's sideInputManager will consume and restore these as well.
 *
 * @param containerModel the container's model
 * @param changelogSystemStreams the passed in set of changelogSystemStreams
 * @return A map of changeLogSSP to storeName across all tasks, assuming no two stores have the same changelogSSP
 */
private Map<String, SystemStream> getChangelogSystemStreams(ContainerModel containerModel, Map<String, SystemStream> changelogSystemStreams) {

  if (MapUtils.invertMap(changelogSystemStreams).size() != changelogSystemStreams.size()) {
    throw new SamzaException("Two stores cannot have the same changelog system-stream");
  }

  Map<SystemStreamPartition, String> changelogSSPToStore = new HashMap<>();
  changelogSystemStreams.forEach((storeName, systemStream) ->
      containerModel.getTasks().forEach((taskName, taskModel) -> { changelogSSPToStore.put(new SystemStreamPartition(systemStream, taskModel.getChangelogPartition()), storeName); })
  );

  getTasks(containerModel, TaskMode.Standby).forEach((taskName, taskModel) -> {
    this.taskSideInputStoreSSPs.putIfAbsent(taskName, new HashMap<>());
    changelogSystemStreams.forEach((storeName, systemStream) -> {
      SystemStreamPartition ssp = new SystemStreamPartition(systemStream, taskModel.getChangelogPartition());
      changelogSSPToStore.remove(ssp);
      this.taskSideInputStoreSSPs.get(taskName).put(storeName, Collections.singleton(ssp));
    });
  });

  // changelogSystemStreams correspond only to active tasks (since those of standby-tasks moved to sideInputs above)
  return MapUtils.invertMap(changelogSSPToStore).entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, x -> x.getValue().getSystemStream()));
}
 

@Override
public ReadableState<Iterator<T>> readIterator() {
  final Iterator<Map.Entry<T, Boolean>> iter = mapState.readIterator().read();
  return new ReadableState<Iterator<T>>() {
    @Nullable
    @Override
    public Iterator<T> read() {
      return new Iterator<T>() {
        @Override
        public boolean hasNext() {
          return iter.hasNext();
        }

        @Override
        public T next() {
          return iter.next().getKey();
        }
      };
    }

    @Override
    public ReadableState<Iterator<T>> readLater() {
      return this;
    }
  };
}
 

@Override
public CompletableFuture<Void> putAllAsync(List<Entry<K, V>> records, Object ... args) {
  incCounter(metrics.numPutAlls);
  long startNs = clock.nanoTime();
  Preconditions.checkNotNull(table, "Cannot write to a read-only table: " + table);
  return table.putAllAsync(records, args).handle((result, e) -> {
    if (e != null) {
      throw new SamzaException("Failed to put records " + records, e);
    } else if (!isWriteAround) {
      cache.putAll(records, args);
    }

    updateTimer(metrics.putAllNs, clock.nanoTime() - startNs);
    return result;
  });
}
 

@Override
public ReadableState<Iterable<Map.Entry<KeyT, ValueT>>> entries() {
  return new ReadableState<Iterable<Map.Entry<KeyT, ValueT>>>() {
    @Override
    public Iterable<Map.Entry<KeyT, ValueT>> read() {
      return createIterable(
          entry ->
              new AbstractMap.SimpleEntry<>(
                  decodeKey(entry.getKey()), decodeValue(entry.getValue())));
    }

    @Override
    public ReadableState<Iterable<Map.Entry<KeyT, ValueT>>> readLater() {
      return this;
    }
  };
}
 

/**
 * Create taskStores for all stores in storageEngineFactories.
 * The store mode is chosen as bulk-load if its a non-sideinput store, and readWrite if its a sideInput store
 */
private Map<TaskName, Map<String, StorageEngine>> createTaskStores(ContainerModel containerModel, JobContext jobContext, ContainerContext containerContext,
    Map<String, StorageEngineFactory<Object, Object>> storageEngineFactories, Map<String, Serde<Object>> serdes,
    Map<TaskName, TaskInstanceMetrics> taskInstanceMetrics,
    Map<TaskName, TaskInstanceCollector> taskInstanceCollectors) {

  Map<TaskName, Map<String, StorageEngine>> taskStores = new HashMap<>();

  // iterate over each task in the containerModel, and each store in storageEngineFactories
  for (Map.Entry<TaskName, TaskModel> task : containerModel.getTasks().entrySet()) {
    TaskName taskName = task.getKey();
    TaskModel taskModel = task.getValue();

    if (!taskStores.containsKey(taskName)) {
      taskStores.put(taskName, new HashMap<>());
    }

    for (String storeName : storageEngineFactories.keySet()) {

      StorageEngineFactory.StoreMode storeMode = this.taskSideInputStoreSSPs.get(taskName).containsKey(storeName) ?
          StorageEngineFactory.StoreMode.ReadWrite : StorageEngineFactory.StoreMode.BulkLoad;

      StorageEngine storageEngine =
          createStore(storeName, taskName, taskModel, jobContext, containerContext, storageEngineFactories, serdes, taskInstanceMetrics, taskInstanceCollectors, storeMode);

      // add created store to map
      taskStores.get(taskName).put(storeName, storageEngine);

      LOG.info("Created store {} for task {} in mode {}", storeName, taskName, storeMode);
    }
  }

  return taskStores;
}
 

@Test
public void testSnapshot() throws Exception {
  InMemoryKeyValueStore store = new InMemoryKeyValueStore(
      new KeyValueStoreMetrics("testInMemory", new MetricsRegistryMap()));
  ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
  String prefix = "prefix";
  for (int i = 0; i < 100; i++) {
    store.put(genKey(outputStream, prefix, i), genValue());
  }

  byte[] firstKey = genKey(outputStream, prefix, 0);
  byte[] lastKey = genKey(outputStream, prefix, 100);
  KeyValueSnapshot<byte[], byte[]> snapshot = store.snapshot(firstKey, lastKey);
  // Make sure the cached Iterable won't change when new elements are added
  store.put(genKey(outputStream, prefix, 200), genValue());
  assertTrue(Iterators.size(snapshot.iterator()) == 100);

  List<Integer> keys = new ArrayList<>();
  KeyValueIterator<byte[], byte[]> iter = snapshot.iterator();
  while (iter.hasNext()) {
    Entry<byte[], byte[]> entry = iter.next();
    int key = Ints.fromByteArray(Arrays.copyOfRange(entry.getKey(), prefix.getBytes().length, entry.getKey().length));
    keys.add(key);
  }
  assertEquals(keys, IntStream.rangeClosed(0, 99).boxed().collect(Collectors.toList()));

  outputStream.close();
  store.close();
}
 

@Override
protected TableDescriptor<Integer, Profile, ?> getTableDescriptor() {
  return new RocksDbTableDescriptor(PROFILE_TABLE, KVSerde.of(new IntegerSerde(), new ProfileJsonSerde()))
      .withSideInputs(ImmutableList.of(PROFILE_STREAM))
      .withSideInputsProcessor((msg, store) -> {
        TestTableData.Profile profile = (TestTableData.Profile) msg.getMessage();
        int key = profile.getMemberId();
        return ImmutableList.of(new Entry<>(key, profile));
      });
}
 

protected TableDescriptor<Integer, Profile, ?> getTableDescriptor() {
  return new InMemoryTableDescriptor(PROFILE_TABLE, KVSerde.of(new IntegerSerde(), new ProfileJsonSerde()))
      .withSideInputs(ImmutableList.of(PROFILE_STREAM))
      .withSideInputsProcessor((msg, store) -> {
        Profile profile = (Profile) msg.getMessage();
        int key = profile.getMemberId();
        return ImmutableList.of(new Entry<>(key, profile));
      });
}
 

/**
 * Since we are not able to track the instances of the iterators created here and close them
 * properly, we need to load the content into memory.
 */
private <OutputT> Iterable<OutputT> createIterable(
    SerializableFunction<org.apache.samza.storage.kv.Entry<ByteArray, byte[]>, OutputT> fn) {
  final ByteArray maxKey = createMaxKey();
  final KeyValueIterator<ByteArray, byte[]> kvIter = store.range(getEncodedStoreKey(), maxKey);
  final List<Entry<ByteArray, byte[]>> iterable = ImmutableList.copyOf(kvIter);
  kvIter.close();

  return new Iterable<OutputT>() {
    @Override
    public Iterator<OutputT> iterator() {
      final Iterator<Entry<ByteArray, byte[]>> iter = iterable.iterator();

      return new Iterator<OutputT>() {
        @Override
        public boolean hasNext() {
          return iter.hasNext();
        }

        @Override
        public OutputT next() {
          return fn.apply(iter.next());
        }
      };
    }
  };
}
 

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  String currentEpoch = this.store.get(CURRENT_EPOCH);
  logger.info("Checking if epoch " + currentEpoch + " is complete.");
  int count = 0;
  KeyValueIterator<String, String> iter = this.store.all();

  while (iter.hasNext()) {
    Entry<String, String> entry = iter.next();
    String foundEpoch = entry.getValue();
    if (foundEpoch.equals(currentEpoch)) {
      count += 1;
    } else {
      logger.info("####### Found a different epoch! - " + foundEpoch + " Current epoch is " + currentEpoch);
    }
  }
  iter.close();
  if (count == expectedKeys + 1) {
    logger.info("Epoch " + currentEpoch + " is complete.");
    int nextEpoch = Integer.parseInt(currentEpoch) + 1;
    for (int i = 0; i < numPartitions; i++) {
      logger.info("Emitting next epoch - " + Integer.toString(i) + " -> " + Integer.toString(nextEpoch));
      collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "epoch"), Integer.toString(i), Integer.toString(nextEpoch)));
    }
    this.store.put(CURRENT_EPOCH, Integer.toString(nextEpoch));
  } else if (count > expectedKeys + 1) {
    throw new IllegalStateException("Got " + count + " keys, which is more than the expected " + (expectedKeys + 1));
  } else {
    logger.info("Only found " + count + " valid keys, try again later.");
  }
}
 

@Test
public void testPutAllAsync() {
  final List<Entry<Integer, Integer>> entries = new LinkedList<>();

  for (int i = 0; i < BATCH_SIZE; i++) {
    entries.add(new Entry<>(i, i));
  }

  CompletableFuture<Void> future = asyncBatchingTable.putAllAsync(entries);
  final BatchProcessor<Integer, Integer> batchProcessor = asyncBatchingTable.getBatchProcessor();

  sleep();

  // Verify that putAllAsync is finished.
  Assert.assertTrue(future.isDone());

  // There should be no pending operations.
  Assert.assertEquals(0, batchProcessor.size());

  // The addBatchUpdates batch operations propagates to the table.
  verify(table, times(1)).putAllAsync(anyList());

  // This new addBatchUpdates will make the batch size to be 1.
  asyncBatchingTable.putAsync(BATCH_SIZE, BATCH_SIZE);

  Assert.assertEquals(1, batchProcessor.size());
}
 

@Test
public void testPutWithPermFailureOnMaxCount() {
  TableRetryPolicy policy = new TableRetryPolicy();
  policy.withFixedBackoff(Duration.ofMillis(5));
  policy.withStopAfterAttempts(10);
  TableReadFunction<String, String> readFn = mock(TableReadFunction.class);
  TableWriteFunction<String, String> writeFn = mock(TableWriteFunction.class);
  doReturn(true).when(writeFn).isRetriable(any());
  CompletableFuture<String> future = new CompletableFuture();
  future.completeExceptionally(new RuntimeException("test exception"));
  doReturn(future).when(writeFn).putAllAsync(any());
  AsyncReadWriteTable delegate = new AsyncRemoteTable(readFn, writeFn);
  AsyncRetriableTable table = new AsyncRetriableTable("t1", delegate, null, policy, schedExec, readFn, writeFn);
  table.init(TestRemoteTable.getMockContext());

  try {
    table.putAllAsync(Arrays.asList(new Entry(1, 2))).join();
    fail();
  } catch (Throwable t) {
  }

  verify(writeFn, atLeast(11)).putAllAsync(any());
  assertEquals(10, table.writeRetryMetrics.retryCount.getCount());
  assertEquals(0, table.writeRetryMetrics.successCount.getCount());
  assertEquals(1, table.writeRetryMetrics.permFailureCount.getCount());
  assertTrue(table.writeRetryMetrics.retryTimer.getSnapshot().getMax() > 0);
}
 

@Test
public void testPutAllWithOneRetry() {
  TableRetryPolicy policy = new TableRetryPolicy();
  policy.withFixedBackoff(Duration.ofMillis(10));
  TableReadFunction<String, String> readFn = mock(TableReadFunction.class);
  TableWriteFunction<String, String> writeFn = mock(TableWriteFunction.class);
  doReturn(true).when(writeFn).isRetriable(any());

  AtomicInteger times = new AtomicInteger();
  doAnswer(invocation -> {
    CompletableFuture<Map<String, String>> future = new CompletableFuture();
    if (times.get() > 0) {
      future.complete(null);
    } else {
      times.incrementAndGet();
      future.completeExceptionally(new RuntimeException("test exception"));
    }
    return future;
  }).when(writeFn).putAllAsync(any());

  AsyncReadWriteTable delegate = new AsyncRemoteTable(readFn, writeFn);
  AsyncRetriableTable table = new AsyncRetriableTable("t1", delegate, null, policy, schedExec, readFn, writeFn);
  table.init(TestRemoteTable.getMockContext());

  table.putAllAsync(Arrays.asList(new Entry(1, 2))).join();
  verify(writeFn, times(2)).putAllAsync(any());
  assertEquals(1, table.writeRetryMetrics.retryCount.getCount());
  assertEquals(0, table.writeRetryMetrics.successCount.getCount());
  assertEquals(0, table.writeRetryMetrics.permFailureCount.getCount());
  assertTrue(table.writeRetryMetrics.retryTimer.getSnapshot().getMax() > 0);
}
 

@Test
public void testPutWithoutRetry() {
  TableRetryPolicy policy = new TableRetryPolicy();
  policy.withFixedBackoff(Duration.ofMillis(100));
  TableReadFunction readFn = mock(TableReadFunction.class);
  TableWriteFunction writeFn = mock(TableWriteFunction.class);
  doReturn(true).when(writeFn).isRetriable(any());
  doReturn(CompletableFuture.completedFuture(null)).when(writeFn).putAsync(any(), any());
  doReturn(CompletableFuture.completedFuture(null)).when(writeFn).putAllAsync(any());
  doReturn(CompletableFuture.completedFuture(null)).when(writeFn).deleteAsync(any());
  doReturn(CompletableFuture.completedFuture(null)).when(writeFn).deleteAllAsync(any());
  AsyncReadWriteTable delegate = new AsyncRemoteTable(readFn, writeFn);
  AsyncRetriableTable table = new AsyncRetriableTable("t1", delegate, null, policy, schedExec, readFn, writeFn);

  int times = 0;
  table.init(TestRemoteTable.getMockContext());
  verify(readFn, times(0)).init(any(), any());
  verify(writeFn, times(0)).init(any(), any());
  table.putAsync("foo", "bar").join();
  verify(writeFn, times(1)).putAsync(any(), any());
  assertEquals(++times, table.writeRetryMetrics.successCount.getCount());
  table.putAllAsync(Arrays.asList(new Entry("1", "2"))).join();
  verify(writeFn, times(1)).putAllAsync(any());
  assertEquals(++times, table.writeRetryMetrics.successCount.getCount());
  table.deleteAsync("1").join();
  verify(writeFn, times(1)).deleteAsync(any());
  assertEquals(++times, table.writeRetryMetrics.successCount.getCount());
  table.deleteAllAsync(Arrays.asList("1", "2")).join();
  verify(writeFn, times(1)).deleteAllAsync(any());
  assertEquals(++times, table.writeRetryMetrics.successCount.getCount());
  assertEquals(0, table.writeRetryMetrics.retryCount.getCount());
  assertEquals(0, table.writeRetryMetrics.retryTimer.getSnapshot().getMax());
  assertEquals(0, table.writeRetryMetrics.permFailureCount.getCount());
  assertNull(table.readRetryMetrics);
}
 

@Test
public void testPutAllAsyncWithArgs() {
  writeTable.putAllAsync(Arrays.asList(new Entry("1", "2")), Arrays.asList(1)).join();
  verify(writeFn, times(0)).putAllAsync(anyCollection());
  verify(writeFn, times(1)).putAllAsync(anyCollection(), any());
  verify(writeRateLimiter, times(0)).throttle(anyString());
  verify(writeRateLimiter, times(0)).throttle(anyString(), anyString());
  verify(writeRateLimiter, times(0)).throttle(anyCollection());
  verify(writeRateLimiter, times(1)).throttleRecords(anyCollection());
  verify(writeRateLimiter, times(0)).throttle(anyInt(), any());
  verifyReadPartNotCalled();
}
 

@Test
public void testPutAllAsync() {
  writeTable.putAllAsync(Arrays.asList(new Entry("1", "2"))).join();
  verify(writeFn, times(1)).putAllAsync(anyCollection());
  verify(writeFn, times(0)).putAllAsync(anyCollection(), any());
  verify(writeRateLimiter, times(0)).throttle(anyString());
  verify(writeRateLimiter, times(0)).throttle(anyString(), anyString());
  verify(writeRateLimiter, times(0)).throttle(anyCollection());
  verify(writeRateLimiter, times(1)).throttleRecords(anyCollection());
  verify(writeRateLimiter, times(0)).throttle(anyInt(), any());
  verifyReadPartNotCalled();
}
 

/**
 * Recreate all non-sideInput persistent stores in ReadWrite mode.
 *
 */
private void recreatePersistentTaskStoresInReadWriteMode(ContainerModel containerModel, JobContext jobContext,
    ContainerContext containerContext, Map<String, StorageEngineFactory<Object, Object>> storageEngineFactories,
    Map<String, Serde<Object>> serdes, Map<TaskName, TaskInstanceMetrics> taskInstanceMetrics,
    Map<TaskName, TaskInstanceCollector> taskInstanceCollectors) {

  // iterate over each task and each storeName
  for (Map.Entry<TaskName, TaskModel> task : containerModel.getTasks().entrySet()) {
    TaskName taskName = task.getKey();
    TaskModel taskModel = task.getValue();
    Map<String, StorageEngine> nonSideInputStores = getNonSideInputStores(taskName);

    for (String storeName : nonSideInputStores.keySet()) {

      // if this store has been already created then re-create and overwrite it only if it is a
      // persistentStore and a non-sideInputStore, because sideInputStores are always created in RW mode
      if (nonSideInputStores.get(storeName).getStoreProperties().isPersistedToDisk()) {

        StorageEngine storageEngine =
            createStore(storeName, taskName, taskModel, jobContext, containerContext, storageEngineFactories, serdes, taskInstanceMetrics, taskInstanceCollectors,
                StorageEngineFactory.StoreMode.ReadWrite);

        // add created store to map
        this.taskStores.get(taskName).put(storeName, storageEngine);

        LOG.info("Re-created store {} in read-write mode for task {} because it a persistent store", storeName, taskName);
      } else {
        LOG.info("Skipping re-creation of store {} for task {}", storeName, taskName);
      }
    }
  }
}
 

/**
 * Store the table {@code records} with specified {@code keys}. This method must be thread-safe.
 * The default implementation calls putAllAsync and blocks on the completion afterwards.
 * @param records table records to be written
 */
default void putAll(List<Entry<K, V>> records) {
  try {
    putAllAsync(records).get();
  } catch (InterruptedException | ExecutionException e) {
    throw new SamzaException("PUT_ALL failed for " + records, e);
  }
}
 

/**
 * Define how batch put should be done.
 *
 * @param operations The batch get operations.
 * @return A CompletableFuture to represent the status of the batch operation.
 */
private CompletableFuture<?> handleBatchPut(Collection<Operation<K, V>> operations) {
  Preconditions.checkNotNull(operations);

  final List<Entry<K, V>> puts = operations.stream()
      .map(op -> new Entry<>(op.getKey(), op.getValue()))
      .collect(Collectors.toList());
  if (puts.isEmpty()) {
    return CompletableFuture.completedFuture(Collections.EMPTY_MAP);
  }

  final Object[] args = getOperationArgs(operations);
  return args == null ? table.putAllAsync(puts) : table.putAllAsync(puts, args);
}
 

@Override
public CompletableFuture<Void> putAllAsync(List<Entry<K, V>> entries, Object ... args) {
  Preconditions.checkNotNull(writeFn, "null writeFn");
  return args.length > 0
      ? writeFn.putAllAsync(entries, args)
      : writeFn.putAllAsync(entries);
}
 

int getEntryCredits(Collection<Entry<K, V>> entries, Object ... args) {
  if (creditFn == null) {
    return entries.size();
  } else {
    return entries.stream().mapToInt(e -> creditFn.getCredits(e.getKey(), e.getValue(), args)).sum();
  }
}
 

@Override
public CompletableFuture<Void> putAllAsync(List<Entry<K, V>> records, Object ... args) {

  Preconditions.checkNotNull(writeFn, "null write function");
  Preconditions.checkNotNull(records, "null records");

  if (records.isEmpty()) {
    return CompletableFuture.completedFuture(null);
  }

  List<K> deleteKeys = records.stream()
      .filter(e -> e.getValue() == null).map(Entry::getKey).collect(Collectors.toList());
  List<Entry<K, V>> putRecords = records.stream()
      .filter(e -> e.getValue() != null).collect(Collectors.toList());

  CompletableFuture<Void> deleteFuture = deleteKeys.isEmpty()
      ? CompletableFuture.completedFuture(null)
      : deleteAllAsync(deleteKeys, args);

  // Return the combined future
  return CompletableFuture.allOf(
      deleteFuture,
      instrument(() -> asyncTable.putAllAsync(putRecords, args), metrics.numPutAlls, metrics.putAllNs))
      .exceptionally(e -> {
        String strKeys = records.stream().map(r -> r.getKey().toString()).collect(Collectors.joining(","));
        throw new SamzaException(String.format("Failed to put records with keys=" + strKeys), e);
      });
}
 

@Override
public void putAll(List<Entry<K, V>> entries, Object ... args) {
  try {
    putAllAsync(entries, args).get();
  } catch (Exception e) {
    throw new SamzaException(e);
  }
}
 

@Override
public KeyValueIterator<K, V> all() {
  final Iterator<Map.Entry<K, V>> iterator = map.entrySet().iterator();
  return new KeyValueIterator<K, V>() {
    @Override
    public void close() {
      //not applicable
    }

    @Override
    public boolean hasNext() {
      return iterator.hasNext();
    }

    @Override
    public Entry<K, V> next() {
      Map.Entry<K, V> kv = iterator.next();
      return new Entry<>(kv.getKey(), kv.getValue());
    }

    @Override
    public void remove() {
      iterator.remove();
    }

  };
}
 

@Override
public void putAll(List<Entry<K, V>> entries, Object ... args) {
  try {
    putAllAsync(entries).get();
  } catch (Exception e) {
    throw new SamzaException("PUT_ALL failed", e);
  }
}
 

@Override
public CompletableFuture<Map<K, V>> getAllAsync(List<K> keys, Object ... args) {
  incCounter(metrics.numGetAlls);
  // Make a copy of entries which might be immutable
  Map<K, V> getAllResult = new HashMap<>();
  List<K> missingKeys = lookupCache(keys, getAllResult);

  if (missingKeys.isEmpty()) {
    return CompletableFuture.completedFuture(getAllResult);
  }

  long startNs = clock.nanoTime();
  return table.getAllAsync(missingKeys, args).handle((records, e) -> {
    if (e != null) {
      throw new SamzaException("Failed to get records for " + keys, e);
    } else {
      if (records != null) {
        cache.putAll(records.entrySet().stream()
            .map(r -> new Entry<>(r.getKey(), r.getValue()))
            .collect(Collectors.toList()), args);
        getAllResult.putAll(records);
      }
      updateTimer(metrics.getAllNs, clock.nanoTime() - startNs);
      return getAllResult;
    }
  });
}
 

@Override
public void putAll(List<Entry<K, V>> records, Object ... args) {
  try {
    putAllAsync(records, args).get();
  } catch (Exception e) {
    throw new SamzaException(e);
  }
}