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.flink.translation.wrappers.streaming.state;
  
  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 org.apache.beam.runners.core.StateInternals;
  import org.apache.beam.runners.core.StateNamespace;
  import org.apache.beam.runners.core.StateTag;
  import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
  import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.ListCoder;
  import org.apache.beam.sdk.coders.MapCoder;
  import org.apache.beam.sdk.coders.StringUtf8Coder;
  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.SetState;
  import org.apache.beam.sdk.state.State;
  import org.apache.beam.sdk.state.StateContext;
  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.windowing.TimestampCombiner;
  import org.apache.beam.sdk.util.CombineContextFactory;
  import org.apache.flink.api.common.ExecutionConfig;
  import org.apache.flink.api.common.state.ListState;
  import org.apache.flink.api.common.state.ListStateDescriptor;
  import org.apache.flink.api.common.state.OperatorStateStore;
  import org.apache.flink.runtime.state.OperatorStateBackend;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /**
   * {@link StateInternals} that uses a Flink {@link OperatorStateBackend} to manage the broadcast
   * state. The state is the same on all parallel instances of the operator. So we just need store
   * state of operator-0 in OperatorStateBackend.
   *
   * <p>Note: Ignore index of key. Mainly for SideInputs.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class FlinkBroadcastStateInternals<K> implements StateInternals {
  
    private int indexInSubtaskGroup;
    private final OperatorStateBackend stateBackend;
    // stateName -> <namespace, state>
    private Map<String, Map<String, ?>> stateForNonZeroOperator;
  
    private final SerializablePipelineOptions pipelineOptions;
  
    public FlinkBroadcastStateInternals(
        int indexInSubtaskGroup,
        OperatorStateBackend stateBackend,
        SerializablePipelineOptions pipelineOptions) {
      this.stateBackend = stateBackend;
      this.indexInSubtaskGroup = indexInSubtaskGroup;
      this.pipelineOptions = pipelineOptions;
      if (indexInSubtaskGroup != 0) {
        stateForNonZeroOperator = new HashMap<>();
      }
    }
  
    @Override
    public @Nullable K getKey() {
      return null;
    }
  
    @Override
    public <T extends State> T state(
        final StateNamespace namespace, StateTag<T> address, final StateContext<?> context) {
  
      return address.bind(
          new StateTag.StateBinder() {
  
            @Override
            public <T2> ValueState<T2> bindValue(StateTag<ValueState<T2>> address, Coder<T2> coder) {
  
              return new FlinkBroadcastValueState<>(
                  stateBackend, address, namespace, coder, pipelineOptions);
            }
  
            @Override
            public <T2> BagState<T2> bindBag(StateTag<BagState<T2>> address, Coder<T2> elemCoder) {
  
              return new FlinkBroadcastBagState<>(
                  stateBackend, address, namespace, elemCoder, pipelineOptions);
            }
  
            @Override
            public <T2> SetState<T2> bindSet(StateTag<SetState<T2>> address, Coder<T2> elemCoder) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", SetState.class.getSimpleName()));
            }
  
            @Override
            public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
                StateTag<MapState<KeyT, ValueT>> spec,
                Coder<KeyT> mapKeyCoder,
                Coder<ValueT> mapValueCoder) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", MapState.class.getSimpleName()));
            }
  
            @Override
            public <ElemT> OrderedListState<ElemT> bindOrderedList(
                StateTag<OrderedListState<ElemT>> spec, Coder<ElemT> elemCoder) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", OrderedListState.class.getSimpleName()));
            }
  
            @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 <InputT, AccumT, OutputT>
                CombiningState<InputT, AccumT, OutputT> bindCombiningValue(
                    StateTag<CombiningState<InputT, AccumT, OutputT>> address,
                    Coder<AccumT> accumCoder,
                    Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {
  
              return new FlinkCombiningState<>(
                  stateBackend, address, combineFn, namespace, accumCoder, pipelineOptions);
            }
  
            @Override
            public <InputT, AccumT, OutputT>
                CombiningState<InputT, AccumT, OutputT> bindCombiningValueWithContext(
                    StateTag<CombiningState<InputT, AccumT, OutputT>> address,
                    Coder<AccumT> accumCoder,
                    CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
              return new FlinkCombiningStateWithContext<>(
                  stateBackend,
                  address,
                  combineFn,
                  namespace,
                  accumCoder,
                  CombineContextFactory.createFromStateContext(context));
            }
  
            @Override
            public WatermarkHoldState bindWatermark(
                StateTag<WatermarkHoldState> address, TimestampCombiner timestampCombiner) {
              throw new UnsupportedOperationException(
                  String.format("%s is not supported", WatermarkHoldState.class.getSimpleName()));
            }
          });
    }
  
    /**
     * 1. The way we would use it is to only checkpoint anything from the operator with subtask index
     * 0 because we assume that the state is the same on all parallel instances of the operator.
     *
     * <p>2. Use map to support namespace.
     */
    private abstract class AbstractBroadcastState<T> {
  
      private String name;
      private final StateNamespace namespace;
      private final ListStateDescriptor<Map<String, T>> flinkStateDescriptor;
      private final OperatorStateStore flinkStateBackend;
  
      AbstractBroadcastState(
          OperatorStateBackend flinkStateBackend,
          String name,
          StateNamespace namespace,
          Coder<T> coder,
          SerializablePipelineOptions pipelineOptions) {
        this.name = name;
  
        this.namespace = namespace;
        this.flinkStateBackend = flinkStateBackend;
  
        CoderTypeInformation<Map<String, T>> typeInfo =
            new CoderTypeInformation<>(MapCoder.of(StringUtf8Coder.of(), coder), pipelineOptions);
  
        flinkStateDescriptor =
            new ListStateDescriptor<>(name, typeInfo.createSerializer(new ExecutionConfig()));
      }
  
      /** Get map(namespce->T) from index 0. */
      Map<String, T> getMap() throws Exception {
        if (indexInSubtaskGroup == 0) {
          return getMapFromBroadcastState();
        } else {
          Map<String, T> result = (Map<String, T>) stateForNonZeroOperator.get(name);
          // maybe restore from BroadcastState of Operator-0
          if (result == null) {
            result = getMapFromBroadcastState();
            if (result != null) {
              stateForNonZeroOperator.put(name, result);
              // we don't need it anymore, must clear it.
              flinkStateBackend.getUnionListState(flinkStateDescriptor).clear();
            }
          }
          return result;
        }
      }
  
      Map<String, T> getMapFromBroadcastState() throws Exception {
        ListState<Map<String, T>> state = flinkStateBackend.getUnionListState(flinkStateDescriptor);
        Iterable<Map<String, T>> iterable = state.get();
        Map<String, T> ret = null;
        if (iterable != null) {
          // just use index 0
          Iterator<Map<String, T>> iterator = iterable.iterator();
          if (iterator.hasNext()) {
            ret = iterator.next();
          }
        }
        return ret;
      }
  
      /** Update map(namespce->T) from index 0. */
      void updateMap(Map<String, T> map) throws Exception {
        if (indexInSubtaskGroup == 0) {
          ListState<Map<String, T>> state = flinkStateBackend.getUnionListState(flinkStateDescriptor);
          state.clear();
          if (map.size() > 0) {
            state.add(map);
          }
        } else {
          if (map.isEmpty()) {
            stateForNonZeroOperator.remove(name);
            // updateMap is always behind getMap,
            // getMap will clear map in BroadcastOperatorState,
            // we don't need clear here.
          } else {
            stateForNonZeroOperator.put(name, map);
          }
        }
      }
  
      void writeInternal(T input) {
        try {
          Map<String, T> map = getMap();
          if (map == null) {
            map = new HashMap<>();
          }
          map.put(namespace.stringKey(), input);
          updateMap(map);
        } catch (Exception e) {
          throw new RuntimeException("Error updating state.", e);
        }
      }
  
      T readInternal() {
        try {
          Map<String, T> map = getMap();
          if (map == null) {
            return null;
          } else {
            return map.get(namespace.stringKey());
          }
        } catch (Exception e) {
          throw new RuntimeException("Error reading state.", e);
        }
      }
  
      void clearInternal() {
        try {
          Map<String, T> map = getMap();
          if (map != null) {
            map.remove(namespace.stringKey());
            updateMap(map);
          }
        } catch (Exception e) {
          throw new RuntimeException("Error clearing state.", e);
        }
      }
    }
  
    private class FlinkBroadcastValueState<T> extends AbstractBroadcastState<T>
        implements ValueState<T> {
  
      private final StateNamespace namespace;
      private final StateTag<ValueState<T>> address;
  
      FlinkBroadcastValueState(
          OperatorStateBackend flinkStateBackend,
          StateTag<ValueState<T>> address,
          StateNamespace namespace,
          Coder<T> coder,
          SerializablePipelineOptions pipelineOptions) {
        super(flinkStateBackend, address.getId(), namespace, coder, pipelineOptions);
  
        this.namespace = namespace;
        this.address = address;
      }
  
      @Override
      public void write(T input) {
        writeInternal(input);
      }
  
      @Override
      public ValueState<T> readLater() {
        return this;
      }
  
      @Override
      public T read() {
        return readInternal();
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (this == o) {
          return true;
        }
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
  
        FlinkBroadcastValueState<?> that = (FlinkBroadcastValueState<?>) o;
  
        return namespace.equals(that.namespace) && address.equals(that.address);
      }
  
      @Override
      public int hashCode() {
        int result = namespace.hashCode();
        result = 31 * result + address.hashCode();
        return result;
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
    }
  
    private class FlinkBroadcastBagState<T> extends AbstractBroadcastState<List<T>>
        implements BagState<T> {
  
      private final StateNamespace namespace;
      private final StateTag<BagState<T>> address;
  
      FlinkBroadcastBagState(
          OperatorStateBackend flinkStateBackend,
          StateTag<BagState<T>> address,
          StateNamespace namespace,
          Coder<T> coder,
          SerializablePipelineOptions pipelineOptions) {
        super(flinkStateBackend, address.getId(), namespace, ListCoder.of(coder), pipelineOptions);
  
        this.namespace = namespace;
        this.address = address;
      }
  
      @Override
      public void add(T input) {
        List<T> list = readInternal();
        if (list == null) {
          list = new ArrayList<>();
        }
        list.add(input);
        writeInternal(list);
      }
  
      @Override
      public BagState<T> readLater() {
        return this;
      }
  
      @Override
      public Iterable<T> read() {
        List<T> result = readInternal();
        return result != null ? result : Collections.emptyList();
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return new ReadableState<Boolean>() {
          @Override
          public Boolean read() {
            try {
              List<T> result = readInternal();
              return result == null;
            } catch (Exception e) {
              throw new RuntimeException("Error reading state.", e);
            }
          }
  
          @Override
          public ReadableState<Boolean> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (this == o) {
          return true;
        }
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
  
        FlinkBroadcastBagState<?> that = (FlinkBroadcastBagState<?>) o;
  
        return namespace.equals(that.namespace) && address.equals(that.address);
      }
  
      @Override
      public int hashCode() {
        int result = namespace.hashCode();
        result = 31 * result + address.hashCode();
        return result;
      }
    }
  
    private class FlinkCombiningState<InputT, AccumT, OutputT> extends AbstractBroadcastState<AccumT>
        implements CombiningState<InputT, AccumT, OutputT> {
  
      private final StateNamespace namespace;
      private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;
      private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn;
  
      FlinkCombiningState(
          OperatorStateBackend flinkStateBackend,
          StateTag<CombiningState<InputT, AccumT, OutputT>> address,
          Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
          StateNamespace namespace,
          Coder<AccumT> accumCoder,
          SerializablePipelineOptions pipelineOptions) {
        super(flinkStateBackend, address.getId(), namespace, accumCoder, pipelineOptions);
  
        this.namespace = namespace;
        this.address = address;
        this.combineFn = combineFn;
      }
  
      @Override
      public CombiningState<InputT, AccumT, OutputT> readLater() {
        return this;
      }
  
      @Override
      public void add(InputT value) {
        AccumT current = readInternal();
        if (current == null) {
          current = combineFn.createAccumulator();
        }
        current = combineFn.addInput(current, value);
        writeInternal(current);
      }
  
      @Override
      public void addAccum(AccumT accum) {
        AccumT current = readInternal();
  
        if (current == null) {
          writeInternal(accum);
        } else {
          current = combineFn.mergeAccumulators(Arrays.asList(current, accum));
          writeInternal(current);
        }
      }
  
      @Override
      public AccumT getAccum() {
        AccumT accum = readInternal();
        return accum != null ? accum : combineFn.createAccumulator();
      }
  
      @Override
      public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
        return combineFn.mergeAccumulators(accumulators);
      }
  
      @Override
      public OutputT read() {
        AccumT accum = readInternal();
        if (accum != null) {
          return combineFn.extractOutput(accum);
        } else {
          return combineFn.extractOutput(combineFn.createAccumulator());
        }
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return new ReadableState<Boolean>() {
          @Override
          public Boolean read() {
            try {
              return readInternal() == null;
            } catch (Exception e) {
              throw new RuntimeException("Error reading state.", e);
            }
          }
  
          @Override
          public ReadableState<Boolean> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (this == o) {
          return true;
        }
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
  
        FlinkCombiningState<?, ?, ?> that = (FlinkCombiningState<?, ?, ?>) o;
  
        return namespace.equals(that.namespace) && address.equals(that.address);
      }
  
      @Override
      public int hashCode() {
        int result = namespace.hashCode();
        result = 31 * result + address.hashCode();
        return result;
      }
    }
  
    private class FlinkCombiningStateWithContext<K2, InputT, AccumT, OutputT>
        extends AbstractBroadcastState<AccumT> implements CombiningState<InputT, AccumT, OutputT> {
  
      private final StateNamespace namespace;
      private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;
      private final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn;
      private final CombineWithContext.Context context;
  
      FlinkCombiningStateWithContext(
          OperatorStateBackend flinkStateBackend,
          StateTag<CombiningState<InputT, AccumT, OutputT>> address,
          CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,
          StateNamespace namespace,
          Coder<AccumT> accumCoder,
          CombineWithContext.Context context) {
        super(flinkStateBackend, address.getId(), namespace, accumCoder, pipelineOptions);
  
        this.namespace = namespace;
        this.address = address;
        this.combineFn = combineFn;
        this.context = context;
      }
  
      @Override
      public CombiningState<InputT, AccumT, OutputT> readLater() {
        return this;
      }
  
      @Override
      public void add(InputT value) {
        try {
          AccumT current = readInternal();
          if (current == null) {
            current = combineFn.createAccumulator(context);
          }
          current = combineFn.addInput(current, value, context);
          writeInternal(current);
        } catch (Exception e) {
          throw new RuntimeException("Error adding to state.", e);
        }
      }
  
      @Override
      public void addAccum(AccumT accum) {
        try {
  
          AccumT current = readInternal();
          if (current == null) {
            writeInternal(accum);
          } else {
            current = combineFn.mergeAccumulators(Arrays.asList(current, accum), context);
            writeInternal(current);
          }
        } catch (Exception e) {
          throw new RuntimeException("Error adding to state.", e);
        }
      }
  
      @Override
      public AccumT getAccum() {
        try {
          AccumT accum = readInternal();
          return accum != null ? accum : combineFn.createAccumulator(context);
        } catch (Exception e) {
          throw new RuntimeException("Error reading state.", e);
        }
      }
  
      @Override
      public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
        return combineFn.mergeAccumulators(accumulators, context);
      }
  
      @Override
      public OutputT read() {
        try {
          AccumT accum = readInternal();
          if (accum == null) {
            accum = combineFn.createAccumulator(context);
          }
          return combineFn.extractOutput(accum, context);
        } catch (Exception e) {
          throw new RuntimeException("Error reading state.", e);
        }
      }
  
      @Override
      public ReadableState<Boolean> isEmpty() {
        return new ReadableState<Boolean>() {
          @Override
          public Boolean read() {
            try {
              return readInternal() == null;
            } catch (Exception e) {
              throw new RuntimeException("Error reading state.", e);
            }
          }
  
          @Override
          public ReadableState<Boolean> readLater() {
            return this;
          }
        };
      }
  
      @Override
      public void clear() {
        clearInternal();
      }
  
      @Override
      public boolean equals(@Nullable Object o) {
        if (this == o) {
          return true;
        }
        if (o == null || getClass() != o.getClass()) {
          return false;
        }
  
        FlinkCombiningStateWithContext<?, ?, ?, ?> that =
            (FlinkCombiningStateWithContext<?, ?, ?, ?>) o;
  
        return namespace.equals(that.namespace) && address.equals(that.address);
      }
  
      @Override
      public int hashCode() {
        int result = namespace.hashCode();
        result = 31 * result + address.hashCode();
        return result;
      }
    }
  }
  