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.translation;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import org.apache.beam.runners.samza.runtime.OpMessage;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.values.WindowedValues;
  import org.apache.samza.Partition;
  import org.apache.samza.config.Config;
  import org.apache.samza.metrics.MetricsRegistry;
  import org.apache.samza.system.IncomingMessageEnvelope;
  import org.apache.samza.system.SystemAdmin;
  import org.apache.samza.system.SystemConsumer;
  import org.apache.samza.system.SystemFactory;
  import org.apache.samza.system.SystemProducer;
  import org.apache.samza.system.SystemStreamMetadata;
  import org.apache.samza.system.SystemStreamPartition;
  
  /**
   * This is a trivial system for generating impulse event in Samza when translating IMPULSE transform
   * in portable api.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class SamzaImpulseSystemFactory implements SystemFactory {
    @Override
    public SystemConsumer getConsumer(
        String systemName, Config config, MetricsRegistry metricsRegistry) {
      return new SamzaImpulseSystemConsumer();
    }
  
    @Override
    public SystemProducer getProducer(
        String systemName, Config config, MetricsRegistry metricsRegistry) {
      throw new UnsupportedOperationException("SamzaImpulseSystem doesn't support producing");
    }
  
    @Override
    public SystemAdmin getAdmin(String systemName, Config config) {
      return new SamzaImpulseSystemAdmin();
    }
  
    private static final String DUMMY_OFFSET = "0";
  
    /** System admin for ImpulseSystem. */
    public static class SamzaImpulseSystemAdmin implements SystemAdmin {
      @Override
      public Map<SystemStreamPartition, String> getOffsetsAfter(
          Map<SystemStreamPartition, String> offset) {
        return offset.keySet().stream()
            .collect(Collectors.toMap(Function.identity(), k -> DUMMY_OFFSET));
      }
  
      @Override
      public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
        return streamNames.stream()
            .collect(
                Collectors.toMap(
                    Function.identity(),
                    stream -> {
                      // Impulse system will always be single partition
                      Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata>
                          partitionMetadata =
                              Collections.singletonMap(
                                  new Partition(0),
                                  new SystemStreamMetadata.SystemStreamPartitionMetadata(
                                      DUMMY_OFFSET, DUMMY_OFFSET, DUMMY_OFFSET));
                      return new SystemStreamMetadata(stream, partitionMetadata);
                    }));
      }
  
      @Override
      public Integer offsetComparator(String offset1, String offset2) {
        return 0;
      }
    }
  
    /** System consumer for ImpulseSystem. */
    public static class SamzaImpulseSystemConsumer implements SystemConsumer {
      private AtomicBoolean isEnd = new AtomicBoolean(false);
  
      @Override
      public void start() {}
  
      @Override
      public void stop() {}
  
      @Override
      public void register(SystemStreamPartition ssp, String offset) {}
  
      private static List<IncomingMessageEnvelope> constructMessages(SystemStreamPartition ssp) {
        final IncomingMessageEnvelope impulseMessage =
            new IncomingMessageEnvelope(
                ssp,
                DUMMY_OFFSET,
                /* key */ null,
                OpMessage.ofElement(WindowedValues.valueInGlobalWindow(new byte[0])));
  
        final IncomingMessageEnvelope watermarkMessage =
            IncomingMessageEnvelope.buildWatermarkEnvelope(
                ssp, BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
  
        final IncomingMessageEnvelope endOfStreamMessage =
            IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp);
  
        return Arrays.asList(impulseMessage, watermarkMessage, endOfStreamMessage);
      }
  
      @Override
      public Map<SystemStreamPartition, List<IncomingMessageEnvelope>> poll(
          Set<SystemStreamPartition> ssps, long timeout) throws InterruptedException {
        if (isEnd.compareAndSet(false, true)) {
          return ssps.stream()
              .collect(
                  Collectors.toMap(
                      Function.identity(), SamzaImpulseSystemConsumer::constructMessages));
        } else {
          return Collections.emptyMap();
        }
      }
    }
  }
  