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.sdk.nexmark.queries;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.io.Serializable;
  import java.nio.channels.Channels;
  import java.nio.channels.WritableByteChannel;
  import java.nio.charset.StandardCharsets;
  import java.util.concurrent.ThreadLocalRandom;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
  import org.apache.beam.sdk.metrics.Counter;
  import org.apache.beam.sdk.metrics.Metrics;
  import org.apache.beam.sdk.nexmark.NexmarkConfiguration;
  import org.apache.beam.sdk.nexmark.model.Done;
  import org.apache.beam.sdk.nexmark.model.Event;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.GroupByKey;
  import org.apache.beam.sdk.transforms.ParDo;
  import org.apache.beam.sdk.transforms.windowing.AfterEach;
  import org.apache.beam.sdk.transforms.windowing.AfterFirst;
  import org.apache.beam.sdk.transforms.windowing.AfterPane;
  import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
  import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
  import org.apache.beam.sdk.transforms.windowing.FixedWindows;
  import org.apache.beam.sdk.transforms.windowing.PaneInfo;
  import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
  import org.apache.beam.sdk.transforms.windowing.Repeatedly;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Query "10", 'Log to sharded files' (Not in original suite.)
   *
   * <p>Every windowSizeSec, save all events from the last period into 2*maxWorkers log files.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class Query10 extends NexmarkQueryTransform<Done> {
    private static final Logger LOG = LoggerFactory.getLogger(Query10.class);
    private static final int NUM_SHARDS_PER_WORKER = 5;
    private static final Duration LATE_BATCHING_PERIOD = Duration.standardSeconds(10);
    private final NexmarkConfiguration configuration;
  
    /** Capture everything we need to know about the records in a single output file. */
    private static class OutputFile implements Serializable {
      /** Maximum possible timestamp of records in file. */
      private final Instant maxTimestamp;
      /** Shard within window. */
      private final String shard;
      /** Index of file in all files in shard. */
      private final long index;
      /** Timing of records in this file. */
      private final PaneInfo.Timing timing;
      /** Path to file containing records, or {@literal null} if no output required. */
      private final @Nullable String filename;
  
      public OutputFile(
          Instant maxTimestamp,
          String shard,
          long index,
          PaneInfo.Timing timing,
          @Nullable String filename) {
        this.maxTimestamp = maxTimestamp;
        this.shard = shard;
        this.index = index;
        this.timing = timing;
        this.filename = filename;
      }
  
      @Override
      public String toString() {
        return String.format("%s %s %d %s %s%n", maxTimestamp, shard, index, timing, filename);
      }
    }
  
    /** GCS uri prefix for all log and 'finished' files. If null they won't be written. */
    private @Nullable String outputPath;
  
    /** Maximum number of workers, used to determine log sharding factor. */
    private int maxNumWorkers;
  
    public Query10(NexmarkConfiguration configuration) {
      super("Query10");
      this.configuration = configuration;
    }
  
    public void setOutputPath(@Nullable String outputPath) {
      this.outputPath = outputPath;
    }
  
    public void setMaxNumWorkers(int maxNumWorkers) {
      this.maxNumWorkers = maxNumWorkers;
    }
  
    /** Return channel for writing bytes to GCS. */
    @SuppressWarnings("unused")
    private WritableByteChannel openWritableGcsFile(GcsOptions options, String filename)
        throws IOException {
      // TODO
      // [https://github.com/apache/beam/issues/20670] Fix after PR: right now this is a specific
      // Google added use case
      // Discuss it on ML: shall we keep GCS or use HDFS or use a generic beam filesystem way.
      throw new UnsupportedOperationException("Disabled after removal of GcsIOChannelFactory");
    }
  
    /** Return a short string to describe {@code timing}. */
    private String timingToString(PaneInfo.Timing timing) {
      switch (timing) {
        case EARLY:
          return "E";
        case ON_TIME:
          return "O";
        case LATE:
          return "L";
        case UNKNOWN:
          return "U";
      }
      throw new RuntimeException(); // cases are exhaustive
    }
  
    /** Construct an {@link OutputFile} for {@code pane} in {@code window} for {@code shard}. */
    private OutputFile outputFileFor(BoundedWindow window, String shard, PaneInfo paneInfo) {
      @Nullable
      String filename =
          outputPath == null
              ? null
              : String.format(
                  "%s/LOG-%s-%s-%03d-%s-%x",
                  outputPath,
                  window.maxTimestamp(),
                  shard,
                  paneInfo.getIndex(),
                  timingToString(paneInfo.getTiming()),
                  ThreadLocalRandom.current().nextLong());
      return new OutputFile(
          window.maxTimestamp(), shard, paneInfo.getIndex(), paneInfo.getTiming(), filename);
    }
  
    /**
     * Return path to which we should write the index for {@code window}, or {@literal null} if no
     * output required.
     */
    private @Nullable String indexPathFor(BoundedWindow window) {
      if (outputPath == null) {
        return null;
      }
      return String.format("%s/INDEX-%s", outputPath, window.maxTimestamp());
    }
  
    @Override
    public PCollection<Done> expand(PCollection<Event> events) {
      final int numLogShards = maxNumWorkers * NUM_SHARDS_PER_WORKER;
  
      return events
          .apply(
              name + ".ShardEvents",
              ParDo.of(
                  new DoFn<Event, KV<String, Event>>() {
                    private final Counter lateCounter = Metrics.counter(name, "actuallyLateEvent");
                    private final Counter onTimeCounter = Metrics.counter(name, "onTimeCounter");
  
                    @ProcessElement
                    public void processElement(ProcessContext c) {
                      if (c.element().hasAnnotation("LATE")) {
                        lateCounter.inc();
                        LOG.debug("Observed late: {}", c.element());
                      } else {
                        onTimeCounter.inc();
                      }
                      int shardNum = (int) Math.abs((long) c.element().hashCode() % numLogShards);
                      String shard = String.format("shard-%05d-of-%05d", shardNum, numLogShards);
                      c.output(KV.of(shard, c.element()));
                    }
                  }))
          .apply(
              name + ".WindowEvents",
              Window.<KV<String, Event>>into(
                      FixedWindows.of(Duration.standardSeconds(configuration.windowSizeSec)))
                  .triggering(
                      AfterEach.inOrder(
                          Repeatedly.forever(
                                  AfterPane.elementCountAtLeast(configuration.maxLogEvents))
                              .orFinally(AfterWatermark.pastEndOfWindow()),
                          Repeatedly.forever(
                              AfterFirst.of(
                                  AfterPane.elementCountAtLeast(configuration.maxLogEvents),
                                  AfterProcessingTime.pastFirstElementInPane()
                                      .plusDelayOf(LATE_BATCHING_PERIOD)))))
                  .discardingFiredPanes()
                  // Use a 1 day allowed lateness so that any forgotten hold will stall the
                  // pipeline for that period and be very noticeable.
                  .withAllowedLateness(Duration.standardDays(1)))
          .apply(name + ".GroupByKey", GroupByKey.create())
          .apply(
              name + ".CheckForLateEvents",
              ParDo.of(
                  new DoFn<KV<String, Iterable<Event>>, KV<String, Iterable<Event>>>() {
                    private final Counter earlyCounter = Metrics.counter(name, "earlyShard");
                    private final Counter onTimeCounter = Metrics.counter(name, "onTimeShard");
                    private final Counter lateCounter = Metrics.counter(name, "lateShard");
                    private final Counter unexpectedLatePaneCounter =
                        Metrics.counter(name, "ERROR_unexpectedLatePane");
                    private final Counter unexpectedOnTimeElementCounter =
                        Metrics.counter(name, "ERROR_unexpectedOnTimeElement");
  
                    @ProcessElement
                    public void processElement(ProcessContext c, BoundedWindow window) {
                      int numLate = 0;
                      int numOnTime = 0;
                      for (Event event : c.element().getValue()) {
                        if (event.hasAnnotation("LATE")) {
                          numLate++;
                        } else {
                          numOnTime++;
                        }
                      }
                      String shard = c.element().getKey();
                      LOG.debug(
                          String.format(
                              "%s with timestamp %s has %d actually late and %d on-time "
                                  + "elements in pane %s for window %s",
                              shard,
                              c.timestamp(),
                              numLate,
                              numOnTime,
                              c.pane(),
                              window.maxTimestamp()));
                      if (c.pane().getTiming() == PaneInfo.Timing.LATE) {
                        if (numLate == 0) {
                          LOG.error("ERROR! No late events in late pane for {}", shard);
                          unexpectedLatePaneCounter.inc();
                        }
                        if (numOnTime > 0) {
                          LOG.error(
                              "ERROR! Have {} on-time events in late pane for {}", numOnTime, shard);
                          unexpectedOnTimeElementCounter.inc();
                        }
                        lateCounter.inc();
                      } else if (c.pane().getTiming() == PaneInfo.Timing.EARLY) {
                        if (numOnTime + numLate < configuration.maxLogEvents) {
                          LOG.error(
                              "ERROR! Only have {} events in early pane for {}",
                              numOnTime + numLate,
                              shard);
                        }
                        earlyCounter.inc();
                      } else {
                        onTimeCounter.inc();
                      }
                      c.output(c.element());
                    }
                  }))
          .apply(
              name + ".UploadEvents",
              ParDo.of(
                  new DoFn<KV<String, Iterable<Event>>, KV<Void, OutputFile>>() {
                    private final Counter savedFileCounter = Metrics.counter(name, "savedFile");
                    private final Counter writtenRecordsCounter =
                        Metrics.counter(name, "writtenRecords");
  
                    @ProcessElement
                    public void processElement(ProcessContext c, BoundedWindow window)
                        throws IOException {
                      String shard = c.element().getKey();
                      GcsOptions options = c.getPipelineOptions().as(GcsOptions.class);
                      OutputFile outputFile = outputFileFor(window, shard, c.pane());
                      LOG.debug(
                          String.format(
                              "Writing %s with record timestamp %s, window timestamp %s, pane %s",
                              shard, c.timestamp(), window.maxTimestamp(), c.pane()));
                      if (outputFile.filename != null) {
                        LOG.info("Beginning write to '{}'", outputFile.filename);
                        int n = 0;
                        try (OutputStream output =
                            Channels.newOutputStream(
                                openWritableGcsFile(options, outputFile.filename))) {
                          for (Event event : c.element().getValue()) {
                            Event.CODER.encode(event, output, Coder.Context.OUTER);
                            writtenRecordsCounter.inc();
                            if (++n % 10000 == 0) {
                              LOG.info("So far written {} records to '{}'", n, outputFile.filename);
                            }
                          }
                        }
                        LOG.info("Written all {} records to '{}'", n, outputFile.filename);
                      }
                      savedFileCounter.inc();
                      c.output(KV.of(null, outputFile));
                    }
                  }))
          // Clear fancy triggering from above.
          .apply(
              name + ".WindowLogFiles",
              Window.<KV<Void, OutputFile>>into(
                      FixedWindows.of(Duration.standardSeconds(configuration.windowSizeSec)))
                  .triggering(DefaultTrigger.of())
                  // We expect no late data here, but we'll assume the worst so we can detect any.
                  .withAllowedLateness(Duration.standardDays(1))
                  .discardingFiredPanes())
          // this GroupByKey allows to have one file per window
          .apply(name + ".GroupByKey2", GroupByKey.create())
          .apply(
              name + ".Index",
              ParDo.of(
                  new DoFn<KV<Void, Iterable<OutputFile>>, Done>() {
                    private final Counter unexpectedLateCounter =
                        Metrics.counter(name, "ERROR_unexpectedLate");
                    private final Counter unexpectedEarlyCounter =
                        Metrics.counter(name, "ERROR_unexpectedEarly");
                    private final Counter unexpectedIndexCounter =
                        Metrics.counter(name, "ERROR_unexpectedIndex");
                    private final Counter finalizedCounter = Metrics.counter(name, "indexed");
  
                    @ProcessElement
                    public void processElement(ProcessContext c, BoundedWindow window)
                        throws IOException {
                      if (c.pane().getTiming() == Timing.LATE) {
                        unexpectedLateCounter.inc();
                        LOG.error("ERROR! Unexpected LATE pane: {}", c.pane());
                      } else if (c.pane().getTiming() == Timing.EARLY) {
                        unexpectedEarlyCounter.inc();
                        LOG.error("ERROR! Unexpected EARLY pane: {}", c.pane());
                      } else if (c.pane().getTiming() == Timing.ON_TIME && c.pane().getIndex() != 0) {
                        unexpectedIndexCounter.inc();
                        LOG.error("ERROR! Unexpected ON_TIME pane index: {}", c.pane());
                      } else {
                        GcsOptions options = c.getPipelineOptions().as(GcsOptions.class);
                        LOG.debug(
                            "Index with record timestamp {}, window timestamp {}, pane {}",
                            c.timestamp(),
                            window.maxTimestamp(),
                            c.pane());
  
                        @Nullable String filename = indexPathFor(window);
                        if (filename != null) {
                          LOG.info("Beginning write to '{}'", filename);
                          int n = 0;
                          try (OutputStream output =
                              Channels.newOutputStream(openWritableGcsFile(options, filename))) {
                            for (OutputFile outputFile : c.element().getValue()) {
                              output.write(outputFile.toString().getBytes(StandardCharsets.UTF_8));
                              n++;
                            }
                          }
                          LOG.info("Written all {} lines to '{}'", n, filename);
                        }
                        c.output(new Done("written for timestamp " + window.maxTimestamp()));
                        finalizedCounter.inc();
                      }
                    }
                  }));
    }
  }
  