content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.runners.samza.runtime;
  
  import java.io.ByteArrayInputStream;
  import java.io.ByteArrayOutputStream;
  import java.io.DataOutputStream;
  import java.io.IOException;
  import java.io.Serializable;
  import java.lang.ref.SoftReference;
  import java.util.AbstractMap;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import javax.annotation.Nonnull;
  import org.apache.beam.runners.core.StateInternals;
  import org.apache.beam.runners.core.StateInternalsFactory;
  import org.apache.beam.runners.core.StateNamespace;
  import org.apache.beam.runners.core.StateTag;
  import org.apache.beam.runners.samza.SamzaPipelineOptions;
  import org.apache.beam.runners.samza.state.SamzaMapState;
  import org.apache.beam.runners.samza.state.SamzaSetState;
  import org.apache.beam.runners.samza.transforms.UpdatingCombineFn;
  import org.apache.beam.sdk.coders.BooleanCoder;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.InstantCoder;
  import org.apache.beam.sdk.coders.VoidCoder;
  import org.apache.beam.sdk.state.BagState;
  import org.apache.beam.sdk.state.CombiningState;
  import org.apache.beam.sdk.state.MapState;
  import org.apache.beam.sdk.state.MultimapState;
  import org.apache.beam.sdk.state.OrderedListState;
  import org.apache.beam.sdk.state.ReadableState;
  import org.apache.beam.sdk.state.ReadableStates;
  import org.apache.beam.sdk.state.SetState;
  import org.apache.beam.sdk.state.State;
  import org.apache.beam.sdk.state.StateContext;
  import org.apache.beam.sdk.state.StateContexts;
  import org.apache.beam.sdk.state.ValueState;
  import org.apache.beam.sdk.state.WatermarkHoldState;
  import org.apache.beam.sdk.transforms.Combine;
  import org.apache.beam.sdk.transforms.CombineWithContext;
  import org.apache.beam.sdk.transforms.SerializableFunction;
  import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
  import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
  import org.apache.beam.sdk.util.construction.graph.UserStateReference;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Ints;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedBytes;
  import org.apache.samza.config.Config;
  import org.apache.samza.context.TaskContext;
  import org.apache.samza.serializers.Serde;
  import org.apache.samza.serializers.SerdeFactory;
  import org.apache.samza.storage.kv.Entry;
  import org.apache.samza.storage.kv.KeyValueIterator;
  import org.apache.samza.storage.kv.KeyValueStore;
  import org.checkerframework.checker.initialization.qual.Initialized;
  import org.checkerframework.checker.nullness.qual.NonNull;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
  import org.joda.time.Instant;
  
  /** {@link StateInternals} that uses Samza local {@link KeyValueStore} to manage state. */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "keyfor",
    "nullness"
  }) // TODO(https://github.com/apache/beam/issues/20497)
  public class SamzaStoreStateInternals<K> implements StateInternals {
    static final String BEAM_STORE = "beamStore";
  
    private static final ThreadLocal<SoftReference<ByteArrayOutputStream>> threadLocalBaos =
        new ThreadLocal<>();
  
    // the stores include both beamStore for system states as well as stores for user state
    private final Map<String, KeyValueStore<ByteArray, StateValue<?>>> stores;
    private final K key;
    private final byte[] keyBytes;
    private final int batchGetSize;
    private final String stageId;
  
    private SamzaStoreStateInternals(
        Map<String, KeyValueStore<ByteArray, StateValue<?>>> stores,
        @Nullable K key,
        byte @Nullable [] keyBytes,
        String stageId,
        int batchGetSize) {
      this.stores = stores;
      this.key = key;
      this.keyBytes = keyBytes;
      this.batchGetSize = batchGetSize;
      this.stageId = stageId;
    }
  
    @SuppressWarnings("unchecked")
    static KeyValueStore<ByteArray, StateValue<?>> getBeamStore(TaskContext context) {
      return (KeyValueStore<ByteArray, StateValue<?>>)
          context.getStore(SamzaStoreStateInternals.BEAM_STORE);
    }
  
    /**
     * Creates non keyed state internal factory to persist states in {@link
     * SamzaStoreStateInternals#BEAM_STORE}.
     */
    static <K> Factory<K> createNonKeyedStateInternalsFactory(
        String id, TaskContext context, SamzaPipelineOptions pipelineOptions) {
      return createStateInternalsFactory(id, null, context, pipelineOptions, Collections.emptyMap());
    }
  
    static <K> Factory<K> createStateInternalsFactory(
        String id,
        Coder<K> keyCoder,
        TaskContext context,
        SamzaPipelineOptions pipelineOptions,
        ExecutableStage executableStage) {
  
      Map<String, String> stateIdToStoreMap =
          executableStage.getUserStates().stream()
              .collect(
                  Collectors.toMap(UserStateReference::localName, UserStateReference::localName));
  
      return createStateInternalsFactory(id, keyCoder, context, pipelineOptions, stateIdToStoreMap);
    }
  
    @SuppressWarnings("unchecked")
    static <K> Factory<K> createStateInternalsFactory(
        String id,
        @Nullable Coder<K> keyCoder,
        TaskContext context,
        SamzaPipelineOptions pipelineOptions,
        Map<String, String> stateIdToStoreMap) {
      final int batchGetSize = pipelineOptions.getStoreBatchGetSize();
      final Map<String, KeyValueStore<ByteArray, StateValue<?>>> stores = new HashMap<>();
      stores.put(BEAM_STORE, getBeamStore(context));
  
      final Coder<K> stateKeyCoder;
      if (keyCoder != null) {
        stateIdToStoreMap
            .keySet()
            .forEach(
                stateId ->
                    stores.put(
                        stateId,
                        (KeyValueStore<ByteArray, StateValue<?>>)
                            context.getStore(stateIdToStoreMap.get(stateId))));
        stateKeyCoder = keyCoder;
      } else {
        stateKeyCoder = (Coder<K>) VoidCoder.of();
      }
      return new Factory<>(Objects.toString(id), stores, stateKeyCoder, batchGetSize);
    }
  
    @Override
    public K getKey() {
      return key;
    }
  
    @Override
    public <T extends State> T state(StateNamespace stateNamespace, StateTag<T> stateTag) {
      return state(stateNamespace, stateTag, StateContexts.nullContext());
    }
  
    @Override
    public <V extends State> V state(
        StateNamespace namespace, StateTag<V> address, StateContext<?> stateContext) {
      return address.bind(
          new StateTag.StateBinder() {
            @Override
            public <T> ValueState<T> bindValue(StateTag<ValueState<T>> spec, Coder<T> coder) {
              return new SamzaValueState<>(namespace, address, coder);
            }
  
            @Override
            public <T> BagState<T> bindBag(StateTag<BagState<T>> spec, Coder<T> elemCoder) {
              return new SamzaBagState<>(namespace, address, elemCoder);
            }
  
            @Override
            public <T> SetState<T> bindSet(StateTag<SetState<T>> spec, Coder<T> elemCoder) {
              return new SamzaSetStateImpl<>(namespace, address, elemCoder);
            }
  
            @Override
            public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
                StateTag<MapState<KeyT, ValueT>> spec,
                Coder<KeyT> mapKeyCoder,
                Coder<ValueT> mapValueCoder) {
              return new SamzaMapStateImpl<>(namespace, address, mapKeyCoder, mapValueCoder);
            }
  
            @Override
            public <KeyT, ValueT> MultimapState<KeyT, ValueT> bindMultimap(
                StateTag<MultimapState<KeyT, ValueT>> spec,
                Coder<KeyT> keyCoder,
                Coder<ValueT> valueCoder) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", MultimapState.class.getSimpleName()));
            }
  
            @Override
            public <T> OrderedListState<T> bindOrderedList(
                StateTag<OrderedListState<T>> spec, Coder<T> elemCoder) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", OrderedListState.class.getSimpleName()));
            }
  
            @Override
            public <InputT, AccumT, OutputT>
                CombiningState<InputT, AccumT, OutputT> bindCombiningValue(
                    StateTag<CombiningState<InputT, AccumT, OutputT>> spec,
                    Coder<AccumT> accumCoder,
                    Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {
              return new SamzaAccumulatorCombiningState<>(namespace, address, accumCoder, combineFn);
            }
  
            @Override
            public <InputT, AccumT, OutputT>
                CombiningState<InputT, AccumT, OutputT> bindCombiningValueWithContext(
                    StateTag<CombiningState<InputT, AccumT, OutputT>> spec,
                    Coder<AccumT> accumCoder,
                    CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", CombiningState.class.getSimpleName()));
            }
  
            @Override
            public WatermarkHoldState bindWatermark(
                StateTag<WatermarkHoldState> spec, TimestampCombiner timestampCombiner) {
              return new SamzaWatermarkHoldState(namespace, address, timestampCombiner);
            }
          });
    }
  
    /** Reuse the ByteArrayOutputStream buffer. */
    private static ByteArrayOutputStream getThreadLocalBaos() {
      final SoftReference<ByteArrayOutputStream> refBaos = threadLocalBaos.get();
      ByteArrayOutputStream baos = refBaos == null ? null : refBaos.get();
      if (baos == null) {
        baos = new ByteArrayOutputStream();
        threadLocalBaos.set(new SoftReference<>(baos));
      }
  
      baos.reset();
      return baos;
    }
  
    /** Factory class to create {@link SamzaStoreStateInternals}. */
    public static class Factory<K> implements StateInternalsFactory<K> {
      private final String stageId;
      private final Map<String, KeyValueStore<ByteArray, StateValue<?>>> stores;
      private final Coder<K> keyCoder;
      private final int batchGetSize;
  
      public Factory(
          String stageId,
          Map<String, KeyValueStore<ByteArray, StateValue<?>>> stores,
          Coder<K> keyCoder,
          int batchGetSize) {
        this.stageId = stageId;
        this.stores = stores;
        this.keyCoder = keyCoder;
        this.batchGetSize = batchGetSize;
      }
  
      @Override
      public StateInternals stateInternalsForKey(@Nullable K key) {
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        final DataOutputStream dos = new DataOutputStream(baos);
  
        try {
          if (key != null) {
            keyCoder.encode(key, baos);
          }
          final byte[] keyBytes = baos.toByteArray();
          baos.reset();
  
          dos.write(keyBytes.length);
          dos.write(keyBytes);
        } catch (IOException e) {
          throw new RuntimeException("Cannot encode key for state store", e);
        }
  
        return new SamzaStoreStateInternals<>(stores, key, baos.toByteArray(), stageId, batchGetSize);
      }
    }
  
    private abstract class AbstractSamzaState<T> {
      private final StateNamespace namespace;
      private final String addressId;
      private final boolean isBeamStore;
      private final String stageId;
      private final byte[] keyBytes;
      private byte[] encodedStoreKey;
      protected final Coder<T> coder;
      protected final KeyValueStore<ByteArray, StateValue<T>> store;
  
      @SuppressWarnings({"unchecked", "rawtypes"})
      protected AbstractSamzaState(
          StateNamespace namespace, StateTag<? extends State> address, Coder<T> coder) {
        this.coder = coder;
        this.namespace = namespace;
        this.addressId = address.getId();
        this.isBeamStore = !stores.containsKey(address.getId());
        this.store =
            isBeamStore
                ? (KeyValueStore) stores.get(BEAM_STORE)
                : (KeyValueStore) stores.get(address.getId());
        this.stageId = SamzaStoreStateInternals.this.stageId;
        this.keyBytes = SamzaStoreStateInternals.this.keyBytes;
      }
  
      protected void clearInternal() {
        store.delete(getEncodedStoreKey());
      }
  
      protected void writeInternal(T value) {
        store.put(getEncodedStoreKey(), StateValue.of(value, coder));
      }
  
      protected T readInternal() {
        final StateValue<T> stateValue = store.get(getEncodedStoreKey());
        return decodeValue(stateValue);
      }
  
      protected ReadableState<Boolean> isEmptyInternal() {
        return new ReadableState<Boolean>() {
          @Override
          public Boolean read() {
            return store.get(getEncodedStoreKey()) == null;
          }
  
          @Override
          public ReadableState<Boolean> readLater() {
            return this;
          }
        };
      }
  
      protected ByteArray getEncodedStoreKey() {
        return ByteArray.of(getEncodedStoreKeyBytes());
      }
  
      protected byte[] getEncodedStoreKeyBytes() {
        if (encodedStoreKey == null) {
          final ByteArrayOutputStream baos = getThreadLocalBaos();
          try (DataOutputStream dos = new DataOutputStream(baos)) {
            dos.write(keyBytes);
            dos.writeUTF(namespace.stringKey());
  
            if (isBeamStore) {
              // for system state, we need to differentiate based on the following:
              dos.writeUTF(stageId);
              dos.writeUTF(addressId);
            }
          } catch (IOException e) {
            throw new RuntimeException("Could not encode full address for state: " + addressId, e);
          }
          this.encodedStoreKey = baos.toByteArray();
        }
        return encodedStoreKey;
      }
  
      protected T decodeValue(StateValue<T> stateValue) {
        return stateValue == null ? null : stateValue.getValue(coder);
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (this == o) {
          return true;
        }
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
  
        @SuppressWarnings("unchecked")
        final AbstractSamzaState<?> that = (AbstractSamzaState<?>) o;
        if (isBeamStore || that.isBeamStore) {
          if (!isBeamStore || !that.isBeamStore || !stageId.equals(that.stageId)) {
            return false;
          }
        }
        return Arrays.equals(keyBytes, that.keyBytes)
            && addressId.equals(that.addressId)
            && this.namespace.equals(that.namespace);
      }
  
      @Override
      public int hashCode() {
        int result = namespace.hashCode();
        result = 31 * result + Arrays.hashCode(getEncodedStoreKeyBytes());
        return result;
      }
    }
  
    private class SamzaValueState<T> extends AbstractSamzaState<T> implements ValueState<T> {
      private SamzaValueState(
          StateNamespace namespace, StateTag<? extends State> address, Coder<T> coder) {
        super(namespace, address, coder);
      }
  
      @Override
      public void write(T input) {
        writeInternal(input);
      }
  
      @Override
      public T read() {
        return readInternal();
      }
  
      @Override
      public ValueState<T> readLater() {
        return this;
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
    }
  
    private class SamzaBagState<T> extends AbstractSamzaState<T> implements BagState<T> {
  
      private SamzaBagState(
          StateNamespace namespace, StateTag<? extends State> address, Coder<T> coder) {
        super(namespace, address, coder);
      }
  
      @Override
      public void add(T value) {
        synchronized (store) {
          final int size = getSize();
          final ByteArray encodedKey = encodeKey(size);
          store.put(encodedKey, StateValue.of(value, coder));
          store.put(getEncodedStoreKey(), StateValue.of(Ints.toByteArray(size + 1)));
        }
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        synchronized (store) {
          return isEmptyInternal();
        }
      }
  
      @Override
      @Nonnull
      public List<T> read() {
        synchronized (store) {
          final int size = getSize();
          if (size == 0) {
            return Collections.emptyList();
          }
  
          final List<T> values = new ArrayList<>(size);
          final List<ByteArray> keys = new ArrayList<>(size);
          int start = 0;
          while (start < size) {
            final int end = Math.min(size, start + batchGetSize);
            for (int i = start; i < end; i++) {
              keys.add(encodeKey(i));
            }
            store.getAll(keys).values().forEach(value -> values.add(decodeValue(value)));
  
            start += batchGetSize;
            keys.clear();
          }
          return values;
        }
      }
  
      @Override
      public BagState<T> readLater() {
        return this;
      }
  
      @Override
      public void clear() {
        synchronized (store) {
          final int size = getSize();
          if (size != 0) {
            final List<ByteArray> keys = new ArrayList<>(size);
            for (int i = 0; i < size; i++) {
              keys.add(encodeKey(i));
            }
            store.deleteAll(keys);
            store.delete(getEncodedStoreKey());
          }
        }
      }
  
      private int getSize() {
        final StateValue stateSize = store.get(getEncodedStoreKey());
        return (stateSize == null || stateSize.valueBytes == null)
            ? 0
            : Ints.fromByteArray(stateSize.valueBytes);
      }
  
      private ByteArray encodeKey(int size) {
        final ByteArrayOutputStream baos = getThreadLocalBaos();
        try (DataOutputStream dos = new DataOutputStream(baos)) {
          dos.write(getEncodedStoreKeyBytes());
          dos.writeInt(size);
          return ByteArray.of(baos.toByteArray());
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }
    }
  
    private class SamzaSetStateImpl<T> implements SamzaSetState<T> {
      private final SamzaMapStateImpl<T, Boolean> mapState;
  
      private SamzaSetStateImpl(
          StateNamespace namespace, StateTag<? extends State> address, Coder<T> coder) {
        mapState = new SamzaMapStateImpl<>(namespace, address, coder, BooleanCoder.of());
      }
  
      @Override
      public ReadableState<Boolean> contains(T t) {
        return mapState.get(t);
      }
  
      @Override
      public @Nullable ReadableState<Boolean> addIfAbsent(T t) {
        return mapState.putIfAbsent(t, true);
      }
  
      @Override
      public void remove(T t) {
        mapState.remove(t);
      }
  
      @Override
      public void add(T value) {
        mapState.put(value, true);
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return new ReadableState<Boolean>() {
  
          @Override
          public Boolean read() {
            return Iterables.isEmpty(mapState.entries().read());
          }
  
          @Override
          public ReadableState<Boolean> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public Iterable<T> read() {
        return mapState.keys().read();
      }
  
      @Override
      public SetState<T> readLater() {
        return this;
      }
  
      @Override
      public void clear() {
        mapState.clear();
      }
  
      @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 void closeIterators() {
        mapState.closeIterators();
      }
    }
  
    private class SamzaMapStateImpl<KeyT, ValueT> extends AbstractSamzaState<ValueT>
        implements SamzaMapState<KeyT, ValueT> {
  
      private final Coder<KeyT> keyCoder;
      private final int storeKeySize;
      private final List<KeyValueIterator<ByteArray, StateValue<ValueT>>> openIterators =
          Collections.synchronizedList(new ArrayList<>());
  
      private int maxKeySize;
  
      protected SamzaMapStateImpl(
          StateNamespace namespace,
          StateTag<? extends State> address,
          Coder<KeyT> keyCoder,
          Coder<ValueT> valueCoder) {
        super(namespace, address, valueCoder);
  
        this.keyCoder = keyCoder;
        this.storeKeySize = getEncodedStoreKeyBytes().length;
        // initial max key size is around 100k, so we can restore timer keys
        this.maxKeySize = this.storeKeySize + 100_000;
      }
  
      @Override
      public void put(KeyT key, ValueT value) {
        final ByteArray encodedKey = encodeKey(key);
        maxKeySize = Math.max(maxKeySize, encodedKey.getValue().length);
        store.put(encodedKey, StateValue.of(value, coder));
      }
  
      @Override
      public @Nullable ReadableState<ValueT> computeIfAbsent(
          KeyT key, Function<? super KeyT, ? extends ValueT> mappingFunction) {
        final ByteArray encodedKey = encodeKey(key);
        final ValueT current = decodeValue(store.get(encodedKey));
        if (current == null) {
          put(key, mappingFunction.apply(key));
        }
  
        return current == null ? null : ReadableStates.immediate(current);
      }
  
      @Override
      public void remove(KeyT key) {
        store.delete(encodeKey(key));
      }
  
      @Override
      public ReadableState<ValueT> get(KeyT key) {
        return getOrDefault(key, null);
      }
  
      @Override
      public @UnknownKeyFor @NonNull @Initialized ReadableState<ValueT> getOrDefault(
          KeyT key, @Nullable ValueT defaultValue) {
        return new ReadableState<ValueT>() {
          @Override
          public @Nullable ValueT read() {
            ValueT value = decodeValue(store.get(encodeKey(key)));
            return value != null ? value : defaultValue;
          }
  
          @Override
          public @UnknownKeyFor @NonNull @Initialized ReadableState<ValueT> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public ReadableState<Iterable<KeyT>> keys() {
        return new ReadableState<Iterable<KeyT>>() {
          @Override
          public Iterable<KeyT> read() {
            return createIterable(entry -> decodeKey(entry.getKey()));
          }
  
          @Override
          public ReadableState<Iterable<KeyT>> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public ReadableState<Iterable<ValueT>> values() {
        return new ReadableState<Iterable<ValueT>>() {
          @Override
          public Iterable<ValueT> read() {
            return createIterable(entry -> decodeValue(entry.getValue()));
          }
  
          @Override
          public ReadableState<Iterable<ValueT>> readLater() {
            return this;
          }
        };
      }
  
      @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;
          }
        };
      }
  
      @Override
      public @UnknownKeyFor @NonNull @Initialized ReadableState<
              @UnknownKeyFor @NonNull @Initialized Boolean>
          isEmpty() {
        ReadableState<Iterable<KeyT>> keys = this.keys();
        return new ReadableState<Boolean>() {
          @Override
          public @Nullable Boolean read() {
            return Iterables.isEmpty(keys.read());
          }
  
          @Override
          public @UnknownKeyFor @NonNull @Initialized ReadableState<Boolean> readLater() {
            keys.readLater();
            return this;
          }
        };
      }
  
      @Override
      public ReadableState<Iterator<Map.Entry<KeyT, ValueT>>> readIterator() {
        final ByteArray maxKey = createMaxKey();
        final KeyValueIterator<ByteArray, StateValue<ValueT>> kvIter =
            store.range(getEncodedStoreKey(), maxKey);
        openIterators.add(kvIter);
  
        return new ReadableState<Iterator<Map.Entry<KeyT, ValueT>>>() {
          @Nullable
          @Override
          public Iterator<Map.Entry<KeyT, ValueT>> read() {
            return new Iterator<Map.Entry<KeyT, ValueT>>() {
              @Override
              public boolean hasNext() {
                boolean hasNext = kvIter.hasNext();
                if (!hasNext) {
                  kvIter.close();
                  openIterators.remove(kvIter);
                }
                return hasNext;
              }
  
              @Override
              public Map.Entry<KeyT, ValueT> next() {
                Entry<ByteArray, StateValue<ValueT>> entry = kvIter.next();
                return new AbstractMap.SimpleEntry<>(
                    decodeKey(entry.getKey()), decodeValue(entry.getValue()));
              }
            };
          }
  
          @Override
          public ReadableState<Iterator<Map.Entry<KeyT, ValueT>>> readLater() {
            return this;
          }
        };
      }
  
      /**
       * 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, StateValue<ValueT>>, OutputT>
              fn) {
        final ByteArray maxKey = createMaxKey();
        final KeyValueIterator<ByteArray, StateValue<ValueT>> kvIter =
            store.range(getEncodedStoreKey(), maxKey);
        final List<Entry<ByteArray, StateValue<ValueT>>> iterable = ImmutableList.copyOf(kvIter);
        kvIter.close();
  
        return new Iterable<OutputT>() {
          @Override
          public Iterator<OutputT> iterator() {
            final Iterator<Entry<ByteArray, StateValue<ValueT>>> 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 clear() {
        final ByteArray maxKey = createMaxKey();
        final KeyValueIterator<ByteArray, StateValue<ValueT>> kvIter =
            store.range(getEncodedStoreKey(), maxKey);
        while (kvIter.hasNext()) {
          store.delete(kvIter.next().getKey());
        }
        kvIter.close();
      }
  
      private ByteArray encodeKey(KeyT key) {
        try {
          final ByteArrayOutputStream baos = getThreadLocalBaos();
          baos.write(getEncodedStoreKeyBytes());
          keyCoder.encode(key, baos);
          return ByteArray.of(baos.toByteArray());
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }
  
      private KeyT decodeKey(ByteArray keyBytes) {
        try {
          final byte[] realKey =
              Arrays.copyOfRange(keyBytes.value, storeKeySize, keyBytes.value.length);
          return keyCoder.decode(new ByteArrayInputStream(realKey));
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }
  
      private ByteArray createMaxKey() {
        byte[] maxKey = new byte[maxKeySize];
        Arrays.fill(maxKey, (byte) 0xff);
  
        final byte[] encodedKey = getEncodedStoreKeyBytes();
        System.arraycopy(encodedKey, 0, maxKey, 0, encodedKey.length);
        return ByteArray.of(maxKey);
      }
  
      @Override
      public void closeIterators() {
        openIterators.forEach(KeyValueIterator::close);
        openIterators.clear();
      }
    }
  
    private class SamzaAccumulatorCombiningState<InT, AccumT, OutT> extends AbstractSamzaState<AccumT>
        implements CombiningState<InT, AccumT, OutT> {
  
      private final Combine.CombineFn<InT, AccumT, OutT> combineFn;
  
      protected SamzaAccumulatorCombiningState(
          StateNamespace namespace,
          StateTag<? extends State> address,
          Coder<AccumT> coder,
          Combine.CombineFn<InT, AccumT, OutT> combineFn) {
        super(namespace, address, coder);
  
        this.combineFn = combineFn;
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
  
      @Override
      public void add(InT value) {
        final AccumT accum = getAccum();
        final AccumT current = combineFn.addInput(accum, value);
        writeInternal(current);
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return isEmptyInternal();
      }
  
      @Override
      public AccumT getAccum() {
        final AccumT accum = readInternal();
        return accum != null ? accum : combineFn.createAccumulator();
      }
  
      @Override
      public void addAccum(AccumT accum) {
        final AccumT currentAccum = getAccum();
        final AccumT mergedAccum = mergeAccumulators(Arrays.asList(currentAccum, accum));
        writeInternal(mergedAccum);
      }
  
      @Override
      public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
        return combineFn.mergeAccumulators(accumulators);
      }
  
      @Override
      public CombiningState<InT, AccumT, OutT> readLater() {
        return this;
      }
  
      @Override
      @Nonnull
      public OutT read() {
        AccumT accum = getAccum();
        OutT output = combineFn.extractOutput(accum);
        if (combineFn instanceof UpdatingCombineFn) {
          AccumT updatedAccum =
              ((UpdatingCombineFn<InT, AccumT, OutT>) combineFn).updateAfterFiring(accum);
          writeInternal(updatedAccum);
        }
        return output;
      }
    }
  
    private class SamzaWatermarkHoldState extends AbstractSamzaState<Instant>
        implements WatermarkHoldState {
  
      private final TimestampCombiner timestampCombiner;
  
      public <V extends State> SamzaWatermarkHoldState(
          StateNamespace namespace, StateTag<V> address, TimestampCombiner timestampCombiner) {
        super(namespace, address, InstantCoder.of());
        this.timestampCombiner = timestampCombiner;
      }
  
      @Override
      public void add(Instant value) {
        final Instant currentValue = readInternal();
        final Instant combinedValue =
            currentValue == null ? value : timestampCombiner.combine(currentValue, value);
  
        if (!combinedValue.equals(currentValue)) {
          writeInternal(combinedValue);
        }
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return isEmptyInternal();
      }
  
      @Override
      public Instant read() {
        return readInternal();
      }
  
      @Override
      public TimestampCombiner getTimestampCombiner() {
        return this.timestampCombiner;
      }
  
      @Override
      public WatermarkHoldState readLater() {
        return this;
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
    }
  
    /** Wrapper of byte[] so it can used as key in the KeyValueStore for caching. */
    public static class ByteArray implements Serializable, Comparable<ByteArray> {
  
      private final byte[] value;
  
      public static ByteArray of(byte[] value) {
        return new ByteArray(value);
      }
  
      private ByteArray(byte[] value) {
        this.value = value;
      }
  
      public byte[] getValue() {
        return value;
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
        ByteArray byteArray = (ByteArray) o;
        return Arrays.equals(value, byteArray.value);
      }
  
      @Override
      public int hashCode() {
        return value != null ? Arrays.hashCode(value) : 0;
      }
  
      @Override
      public int compareTo(ByteArray other) {
        return UnsignedBytes.lexicographicalComparator().compare(value, other.value);
      }
    }
  
    /** Factory class to provide {@link ByteArraySerde}. */
    public static class ByteArraySerdeFactory implements SerdeFactory<ByteArray> {
  
      @Override
      public Serde<ByteArray> getSerde(String name, Config config) {
        return new ByteArraySerde();
      }
  
      /** Serde for {@link ByteArray}. */
      public static class ByteArraySerde implements Serde<ByteArray> {
  
        @Override
        public byte[] toBytes(ByteArray byteArray) {
          return byteArray.value;
        }
  
        @Override
        public ByteArray fromBytes(byte[] bytes) {
          return ByteArray.of(bytes);
        }
      }
    }
  
    /**
     * Wrapper for state value so that unencoded value can be read directly from the cache of
     * KeyValueStore.
     */
    public static class StateValue<T> implements Serializable {
      private T value;
      private Coder<T> valueCoder;
      private byte[] valueBytes;
  
      private StateValue(T value, Coder<T> valueCoder, byte[] valueBytes) {
        this.value = value;
        this.valueCoder = valueCoder;
        this.valueBytes = valueBytes;
      }
  
      public static <T> StateValue<T> of(T value, Coder<T> valueCoder) {
        return new StateValue<>(value, valueCoder, null);
      }
  
      public static <T> StateValue<T> of(byte[] valueBytes) {
        return new StateValue<>(null, null, valueBytes);
      }
  
      public T getValue(Coder<T> coder) {
        if (value == null && valueBytes != null) {
          if (valueCoder == null) {
            valueCoder = coder;
          }
          try {
            value = valueCoder.decode(new ByteArrayInputStream(valueBytes));
          } catch (IOException e) {
            throw new RuntimeException("Could not decode state", e);
          }
        }
        return value;
      }
  
      public byte[] getValueBytes() {
        if (valueBytes == null && value != null) {
          final ByteArrayOutputStream baos = new ByteArrayOutputStream();
          try {
            valueCoder.encode(value, baos);
          } catch (IOException e) {
            throw new RuntimeException("Could not encode state value: " + value, e);
          }
          valueBytes = baos.toByteArray();
        }
        return valueBytes;
      }
    }
  
    /** Factory class to provide {@link StateValueSerdeFactory.StateValueSerde}. */
    public static class StateValueSerdeFactory implements SerdeFactory<StateValue<?>> {
      @Override
      public Serde<StateValue<?>> getSerde(String name, Config config) {
        return new StateValueSerde();
      }
  
      /** Serde for {@link StateValue}. */
      public static class StateValueSerde implements Serde<StateValue<?>> {
        @Override
        public StateValue<?> fromBytes(byte[] bytes) {
          return StateValue.of(bytes);
        }
  
        @Override
        public byte[] toBytes(StateValue<?> stateValue) {
          return stateValue == null ? null : stateValue.getValueBytes();
        }
      }
    }
  }
  