/*
 * 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.
 */

/*
 * Protocol Buffers describing the Fn API and boostrapping.
 *
 * TODO: Usage of plural names in lists looks awkward in Java
 * e.g. getOutputsMap, addCodersBuilder
 *
 * TODO: gRPC / proto field names conflict with generated code
 * e.g. "class" in java, "output" in python
 */

syntax = "proto3";

/* TODO: Consider consolidating common components in another package
 * and lanaguage namespaces for re-use with Runner Api.
 */

package org.apache.beam.model.fn_execution.v1;

option go_package = "fnexecution_v1";
option java_package = "org.apache.beam.model.fnexecution.v1";
option java_outer_classname = "BeamFnApi";

import "beam_runner_api.proto";
import "endpoints.proto";
import "google/protobuf/timestamp.proto";
import "google/protobuf/wrappers.proto";

/*
 * Constructs that define the pipeline shape.
 *
 * These are mostly unstable due to the missing pieces to be shared with
 * the Runner Api like windowing strategy, display data, .... There are still
 * some modelling questions related to whether a side input is modelled
 * as another field on a PrimitiveTransform or as part of inputs and we
 * still are missing things like the CompositeTransform.
 */

// A representation of an input or output definition on a primitive transform.
// Stable
message Target {
  // A repeated list of target definitions.
  message List {
    repeated Target target = 1;
  }

  // (Required) The id of the PrimitiveTransform which is the target.
  string primitive_transform_reference = 1;

  // (Required) The local name of an input or output defined on the primitive
  // transform.
  string name = 2;
}

// A descriptor for connecting to a remote port using the Beam Fn Data API.
// Allows for communication between two environments (for example between the
// runner and the SDK).
// Stable
message RemoteGrpcPort {
  // (Required) An API descriptor which describes where to
  // connect to including any authentication that is required.
  org.apache.beam.model.pipeline.v1.ApiServiceDescriptor api_service_descriptor = 1;

  // (Required) The ID of the Coder that will be used to encode and decode data sent over this port.
  string coder_id = 2;
}

/*
 * Control Plane API
 *
 * Progress reporting and splitting still need further vetting. Also, this may change
 * with the addition of new types of instructions/responses related to metrics.
 */

// An API that describes the work that a SDK harness is meant to do.
// Stable
service BeamFnControl {
  // Instructions sent by the runner to the SDK requesting different types
  // of work.
  rpc Control(
    // A stream of responses to instructions the SDK was asked to be performed.
    stream InstructionResponse
  ) returns (
    // A stream of instructions requested of the SDK to be performed.
    stream InstructionRequest
  ) {}
}

// A request sent by a runner which the SDK is asked to fulfill.
// For any unsupported request type, an error should be returned with a
// matching instruction id.
// Stable
message InstructionRequest {
  // (Required) An unique identifier provided by the runner which represents
  // this requests execution. The InstructionResponse MUST have the matching id.
  string instruction_id = 1;

  // (Required) A request that the SDK Harness needs to interpret.
  oneof request {
    RegisterRequest register = 1000;
    ProcessBundleRequest process_bundle = 1001;
    ProcessBundleProgressRequest process_bundle_progress = 1002;
    ProcessBundleSplitRequest process_bundle_split = 1003;
  }
}

// The response for an associated request the SDK had been asked to fulfill.
// Stable
message InstructionResponse {
  // (Required) A reference provided by the runner which represents a requests
  // execution. The InstructionResponse MUST have the matching id when
  // responding to the runner.
  string instruction_id = 1;

  // If this is specified, then this instruction has failed.
  // A human readable string representing the reason as to why processing has
  // failed.
  string error = 2;

  // If the instruction did not fail, it is required to return an equivalent
  // response type depending on the request this matches.
  oneof response {
    RegisterResponse register = 1000;
    ProcessBundleResponse process_bundle = 1001;
    ProcessBundleProgressResponse process_bundle_progress = 1002;
    ProcessBundleSplitResponse process_bundle_split = 1003;
  }
}

// A list of objects which can be referred to by the runner in
// future requests.
// Stable
message RegisterRequest {
  // (Optional) The set of descriptors used to process bundles.
  repeated ProcessBundleDescriptor process_bundle_descriptor = 1;
}

// Stable
message RegisterResponse {
}

// Definitions that should be used to construct the bundle processing graph.
message ProcessBundleDescriptor {
  // (Required) A pipeline level unique id which can be used as a reference to
  // refer to this.
  string id = 1;

  // (Required) A map from pipeline-scoped id to PTransform.
  map<string, org.apache.beam.model.pipeline.v1.PTransform> transforms = 2;

  // (Required) A map from pipeline-scoped id to PCollection.
  map<string, org.apache.beam.model.pipeline.v1.PCollection> pcollections = 3;

  // (Required) A map from pipeline-scoped id to WindowingStrategy.
  map<string, org.apache.beam.model.pipeline.v1.WindowingStrategy> windowing_strategies = 4;

  // (Required) A map from pipeline-scoped id to Coder.
  map<string, org.apache.beam.model.pipeline.v1.Coder> coders = 5;

  // (Required) A map from pipeline-scoped id to Environment.
  map<string, org.apache.beam.model.pipeline.v1.Environment> environments = 6;

  // A descriptor describing the end point to use for State API
  // calls. Required if the Runner intends to send remote references over the
  // data plane or if any of the transforms rely on user state or side inputs.
  org.apache.beam.model.pipeline.v1.ApiServiceDescriptor state_api_service_descriptor = 7;
}

// Represents a partition of the bundle into two bundles: a "primary" and
// a "residual", with the following properties:
// - The work in primary and residual doesn't overlap, and combined, adds up
//   to the work in the current bundle if the split hadn't happened.
// - The current bundle, if it keeps executing, will have done none of the
//   work under residual roots.
// - The current bundle, if no further splits happen, will have done exactly
//   the work under primary_roots.
// For more rigorous definitions see https://s.apache.org/beam-breaking-fusion
message BundleSplit {
  // One of the root applications specifying the scope of work for a bundle.
  message Application {
    // (Required) The primitive transform to which to pass the element
    string ptransform_id = 1;

    // (Required) Name of the transform's input to which to pass the element.
    string input_id = 2;

    // (Required) The encoded element to pass to the transform.
    bytes element = 3;

    // Approximate lower bounds on timestamps of elements that this PTransform
    // will produce into each of its output PCollections, when invoked on this
    // element. Keyed by the transform's local output name.
    map<string, int64> output_watermarks = 4;

    // Approximate fraction of all work of the current bundle (before split)
    // represented by invoking this Application and its downstream applications.
    // The sum of fraction_of_work between all primary_roots and residual_roots
    // must add up to approximately 1.0.
    google.protobuf.DoubleValue fraction_of_work = 5;
  }

  // Root applications that should replace the current bundle.
  repeated Application primary_roots = 1;

  // Root applications that have been removed from the current bundle and
  // have to be executed in a separate bundle (e.g. in parallel on a different
  // worker, or after the current bundle completes, etc.)
  repeated Application residual_roots = 2;
}

// A request to process a given bundle.
// Stable
message ProcessBundleRequest {
  // (Required) A reference to the process bundle descriptor that must be
  // instantiated and executed by the SDK harness.
  string process_bundle_descriptor_reference = 1;

  // (Optional) A list of cache tokens that can be used by an SDK to reuse
  // cached data returned by the State API across multiple bundles.
  repeated bytes cache_tokens = 2;
}

// Stable
message ProcessBundleResponse {
  // (Optional) If metrics reporting is supported by the SDK, this represents
  // the final metrics to record for this bundle.
  Metrics metrics = 1;

  // (Optional) Specifies that the bundle has been split since the last
  // ProcessBundleProgressResponse was sent.
  BundleSplit split = 2;
}

// A request to report progress information for a given bundle.
// This is an optional request to be handled and is used to support advanced
// SDK features such as SplittableDoFn, user level metrics etc.
message ProcessBundleProgressRequest {
  // (Required) A reference to an active process bundle request with the given
  // instruction id.
  string instruction_reference = 1;
}

message MonitoringInfo {
  // The name defining the metric or monitored state.
  string urn = 1;

  // This is specified as a URN that implies:
  // A message class: (Distribution, Counter, Extrema, MonitoringDataTable).
  // Sub types like field formats - int64, double, string.
  // Aggregation methods - SUM, LATEST, TOP-N, BOTTOM-N, DISTRIBUTION
  // valid values are:
  // beam:metrics:[SumInt64|LatestInt64|Top-NInt64|Bottom-NInt64|
  //     SumDouble|LatestDouble|Top-NDouble|Bottom-NDouble|DistributionInt64|
  //     DistributionDouble|MonitoringDataTable]
  string type = 2;

  // The Metric or monitored state.
  oneof data {
    MonitoringTableData monitoring_table_data = 3;
    Metric metric = 4;
  }

  enum MonitoringInfoLabels {
    TRANSFORM = 0;
    PCOLLECTION = 1;
    WINDOWING_STRATEGY = 2;
    CODER = 3;
    ENVIRONMENT = 4;
  }
  // A set of key+value labels which define the scope of the metric.
  // Either a well defined entity id for matching the enum names in
  // the MonitoringInfoLabels enum or any arbitrary label
  // set by a custom metric or user metric.
  // A monitoring system is expected to be able to aggregate the metrics together
  // for all updates having the same URN and labels.
  // Some systems such as Stackdriver will be able to aggregate the metrics
  // using a subset of the provided labels
  map<string, string> labels = 5;
}

message Metric {
  // (Required) The data for this metric.
  oneof data {
    CounterData counter_data = 1;
    DistributionData distribution_data = 2;
    ExtremaData extrema_data = 3;
  }
}

// Data associated with a Counter or Gauge metric.
// This is designed to be compatible with metric collection
// systems such as DropWizard.
message CounterData {
   oneof value {
     int64 int64_value = 1;
     double double_value = 2;
     string string_value = 3;
   }
}

// Extrema messages are used for calculating
// Top-N/Bottom-N metrics.
message ExtremaData {
  oneof extrema {
    IntExtremaData int_extrema_data = 1;
    DoubleExtremaData double_extrema_data = 2;
  }
}

message IntExtremaData {
  repeated int64 int_values = 1;
}

message DoubleExtremaData {
  repeated double double_values = 2;
}

// Data associated with a distribution metric.
// This is based off of the current DistributionData metric.
// This is not a stackdriver or dropwizard compatible
// style of distribution metric.
message DistributionData {
  oneof distribution {
    IntDistributionData int_distribution_data = 1;
    DoubleDistributionData double_distribution_data = 2;
  }
}

message IntDistributionData {
  int64 count = 1;
  int64 sum = 2;
  int64 min = 3;
  int64 max = 4;
}

message DoubleDistributionData {
  int64 count = 1;
  double sum = 2;
  double min = 3;
  double max = 4;
}

// General MonitoredState information which contains
// structured information which does not fit into a typical
// metric format. For example, a table of important files
// and metadata which an I/O source is reading.
// Note: Since MonitoredState is designed to be
// customizable, and allow engines to aggregate these
// metrics in custom ways.
// Engines can use custom aggregation functions for specific URNs
// by inspecting the column values.
// An SDK should always report its current state, that is all
// relevant MonitoredState for its PTransform at the current moment
// and this should be kept small.
// For example, an SDK can emit the oldest three files which
// have been waiting for data for over 1 hour.
// If an engine supports the URN with a custom aggregation then
// it can filter these and keep only the Top-3 rows based on
// how long the files have been waiting for data.
// Otherwise an engine can ignore the MonitoringTableData
// or union all the rows together into one large table and display
// them in a UI.
message MonitoringTableData {
  message MonitoringColumnValue {
    oneof value {
      int64 int64_value = 1;
      double double_value = 2;
      string string_value = 3;
      google.protobuf.Timestamp timestamp = 4;
    }
  }

  message MonitoringRow {
    repeated MonitoringColumnValue values = 1;
  }

 // The number of column names must match the
 // number of values in each MonitoringRow.
 repeated string column_names = 1;
 repeated MonitoringRow row_data = 2;
}

// DEPRECATED
message Metrics {
  // PTransform level metrics.
  // These metrics are split into processed and active element groups for
  // progress reporting purposes. This allows a Runner to see what is measured,
  // what is estimated and what can be extrapolated to be able to accurately
  // estimate the backlog of remaining work.
  message PTransform {
    // Metrics that are measured for processed and active element groups.
    message Measured {
      // (Optional) Map from local input name to number of elements processed
      // from this input.
      // If unset, assumed to be the sum of the outputs of all producers to
      // this transform (for ProcessedElements) and 0 (for ActiveElements).
      map<string, int64> input_element_counts = 1;

      // (Required) Map from local output name to number of elements produced
      // for this output.
      map<string, int64> output_element_counts = 2;

      // (Optional) The total time spent so far in processing the elements in
      // this group, in seconds.
      double total_time_spent = 3;

      // TODO: Add other element group level metrics.
    }

    // Metrics for fully processed elements.
    message ProcessedElements {
      // (Required)
      Measured measured = 1;
    }

    // Metrics for active elements.
    // An element is considered active if the SDK has started but not finished
    // processing it yet.
    message ActiveElements {
      // (Required)
      Measured measured = 1;

      // Estimated metrics.

      // (Optional) Sum of estimated fraction of known work remaining for all
      // active elements, as reported by this transform.
      // If not reported, a Runner could extrapolate this from the processed
      // elements.
      // TODO: Handle the case when known work is infinite.
      double fraction_remaining = 2;

      // (Optional) Map from local output name to sum of estimated number
      // of elements remaining for this output from all active elements,
      // as reported by this transform.
      // If not reported, a Runner could extrapolate this from the processed
      // elements.
      map<string, int64> output_elements_remaining = 3;
    }

    // (Required): Metrics for processed elements.
    ProcessedElements processed_elements = 1;
    // (Required): Metrics for active elements.
    ActiveElements active_elements = 2;

    // (Optional): Map from local output name to its watermark.
    // The watermarks reported are tentative, to get a better sense of progress
    // while processing a bundle but before it is committed. At bundle commit
    // time, a Runner needs to also take into account the timers set to compute
    // the actual watermarks.
    map<string, int64> watermarks = 3;

    repeated User user = 4;
    // TODO: Define other transform level system metrics.
  }

  // User defined metrics
  message User {

    // A key for identifying a metric at the most granular level.
    message MetricName {
      // (Required): The namespace of this metric.
      string namespace = 2;

      // (Required): The name of this metric.
      string name = 3;
    }

    // Data associated with a counter metric.
    message CounterData {
      int64 value = 1;
    }

    // Data associated with a distribution metric.
    message DistributionData {
      int64 count = 1;
      int64 sum = 2;
      int64 min = 3;
      int64 max = 4;
    }

    // Data associated with a Gauge metric.
    message GaugeData {
      int64 value = 1;
      google.protobuf.Timestamp timestamp = 2;
    }

    // (Required) The identifier for this metric.
    MetricName metric_name = 1;

    // (Required) The data for this metric.
    oneof data {
      CounterData counter_data = 1001;
      DistributionData distribution_data = 1002;
      GaugeData gauge_data = 1003;
    }
  }

  map<string, PTransform> ptransforms = 1;
}

message ProcessBundleProgressResponse {
  // (Required)
  Metrics metrics = 1;

  // (Optional) Specifies that the bundle has been split since the last
  // ProcessBundleProgressResponse was sent.
  BundleSplit split = 2;
}

message ProcessBundleSplitRequest {
  // (Required) A reference to an active process bundle request with the given
  // instruction id.
  string instruction_reference = 1;

  // Specifies that the runner would like the bundle to split itself using
  // BundleSplit, and give up some of the work that the bundle hasn't started
  // doing yet, so that it can be done in a separate bundle (perhaps in
  // parallel with the current bundle).
  //
  // The value is the fraction of unstarted work to keep. E.g. 0 means give up
  // as much as possible of unstarted work (e.g. checkpoint), 0.5 means give
  // up about half of the unstarted work, etc.
  // This is a hint and the value is approximate.
  //
  // The value is relative to the current scope of work of the bundle.
  google.protobuf.DoubleValue fraction_of_remainder = 2;
}

message ProcessBundleSplitResponse {
  // Empty.
}

/*
 * Data Plane API
 */

// Messages used to represent logical byte streams.
// Stable
message Elements {
  // Represents multiple encoded elements in nested context for a given named
  // instruction and target.
  message Data {
    // (Required) A reference to an active instruction request with the given
    // instruction id.
    string instruction_reference = 1;

    // (Required) A definition representing a consumer or producer of this data.
    // If received by a harness, this represents the consumer within that
    // harness that should consume these bytes. If sent by a harness, this
    // represents the producer of these bytes.
    //
    // Note that a single element may span multiple Data messages.
    //
    // Note that a sending/receiving pair should share the same target
    // identifier.
    Target target = 2;

    // (Optional) Represents a part of a logical byte stream. Elements within
    // the logical byte stream are encoded in the nested context and
    // concatenated together.
    //
    // An empty data block represents the end of stream for the given
    // instruction and target.
    bytes data = 3;
  }

  // (Required) A list containing parts of logical byte streams.
  repeated Data data = 1;
}

// Stable
service BeamFnData {
  // Used to send data between harnesses.
  rpc Data(
    // A stream of data representing input.
    stream Elements
  ) returns (
    // A stream of data representing output.
    stream Elements
  ) {}
}

/*
 * State API
 */

message StateRequest {
  // (Required) An unique identifier provided by the SDK which represents this
  // requests execution. The StateResponse corresponding with this request
  // will have the matching id.
  string id = 1;

  // (Required) The associated instruction id of the work that is currently
  // being processed. This allows for the runner to associate any modifications
  // to state to be committed with the appropriate work execution.
  string instruction_reference = 2;

  // (Required) The state key this request is for.
  StateKey state_key = 3;

  // (Required) The action to take on this request.
  oneof request {
    // A request to get state.
    StateGetRequest get = 1000;

    // A request to append to state.
    StateAppendRequest append = 1001;

    // A request to clear state.
    StateClearRequest clear = 1002;
  }
}

message StateResponse {
  // (Required) A reference provided by the SDK which represents a requests
  // execution. The StateResponse must have the matching id when responding
  // to the SDK.
  string id = 1;

  // (Optional) If this is specified, then the state request has failed.
  // A human readable string representing the reason as to why the request
  // failed.
  string error = 2;

  // (Optional) If this is specified, then the result of this state request
  // can be cached using the supplied token.
  bytes cache_token = 3;

  // A corresponding response matching the request will be populated.
  oneof response {
    // A response to getting state.
    StateGetResponse get = 1000;

    // A response to appending to state.
    StateAppendResponse append = 1001;

    // A response to clearing state.
    StateClearResponse clear = 1002;
  }
}

service BeamFnState {
  // Used to get/append/clear state stored by the runner on behalf of the SDK.
  rpc State(
    // A stream of state instructions requested of the runner.
    stream StateRequest
  ) returns (
    // A stream of responses to state instructions the runner was asked to be
    // performed.
    stream StateResponse
  ) {}
}

message StateKey {
  message Runner {
    // (Required) Opaque information supplied by the runner. Used to support
    // remote references.
    bytes key = 1;
  }

  message MultimapSideInput {
    // (Required) The id of the PTransform containing a side input.
    string ptransform_id = 1;
    // (Required) The id of the side input.
    string side_input_id = 2;
    // (Required) The window (after mapping the currently executing elements
    // window into the side input windows domain) encoded in a nested context.
    bytes window = 3;
    // (Required) The key encoded in a nested context.
    bytes key = 4;
  }

  message BagUserState {
    // (Required) The id of the PTransform containing user state.
    string ptransform_id = 1;
    // (Required) The id of the user state.
    string user_state_id = 2;
    // (Required) The window encoded in a nested context.
    bytes window = 3;
    // (Required) The key of the currently executing element encoded in a
    // nested context.
    bytes key = 4;
  }

  // (Required) One of the following state keys must be set.
  oneof type {
    Runner runner = 1;
    MultimapSideInput multimap_side_input = 2;
    BagUserState bag_user_state = 3;
    // TODO: represent a state key for user map state
  }
}

// A request to get state.
message StateGetRequest {
  // (Optional) If specified, signals to the runner that the response
  // should resume from the following continuation token.
  //
  // If unspecified, signals to the runner that the response should start
  // from the beginning of the logical continuable stream.
  bytes continuation_token = 1;
}

// A response to get state representing a logical byte stream which can be
// continued using the state API.
message StateGetResponse {
  // (Optional) If specified, represents a token which can be used with the
  // state API to get the next chunk of this logical byte stream. The end of
  // the logical byte stream is signalled by this field being unset.
  bytes continuation_token = 1;

  // Represents a part of a logical byte stream. Elements within
  // the logical byte stream are encoded in the nested context and
  // concatenated together.
  bytes data = 2;
}

// A request to append state.
message StateAppendRequest {
  // Represents a part of a logical byte stream. Elements within
  // the logical byte stream are encoded in the nested context and
  // multiple append requests are concatenated together.
  bytes data = 1;
}

// A response to append state.
message StateAppendResponse {
}

// A request to clear state.
message StateClearRequest {
}

// A response to clear state.
message StateClearResponse {
}

/*
 * Logging API
 *
 * This is very stable. There can be some changes to how we define a LogEntry,
 * to increase/decrease the severity types, the way we format an exception/stack
 * trace, or the log site.
 */

// A log entry
message LogEntry {
  // A list of log entries, enables buffering and batching of multiple
  // log messages using the logging API.
  message List {
    // (Required) One or or more log messages.
    repeated LogEntry log_entries = 1;
  }

  // The severity of the event described in a log entry, expressed as one of the
  // severity levels listed below. For your reference, the levels are
  // assigned the listed numeric values. The effect of using numeric values
  // other than those listed is undefined.
  //
  // If you are writing log entries, you should map other severity encodings to
  // one of these standard levels. For example, you might map all of
  // Java's FINE, FINER, and FINEST levels to `Severity.DEBUG`.
  //
  // This list is intentionally not comprehensive; the intent is to provide a
  // common set of "good enough" severity levels so that logging front ends
  // can provide filtering and searching across log types. Users of the API are
  // free not to use all severity levels in their log messages.
  message Severity {
    enum Enum {
      // Unspecified level information. Will be logged at the TRACE level.
      UNSPECIFIED = 0;
      // Trace level information.
      TRACE = 1;
      // Debugging information.
      DEBUG = 2;
      // Normal events.
      INFO = 3;
      // Normal but significant events, such as start up, shut down, or
      // configuration.
      NOTICE = 4;
      // Warning events might cause problems.
      WARN = 5;
      // Error events are likely to cause problems.
      ERROR = 6;
      // Critical events cause severe problems or brief outages and may
      // indicate that a person must take action.
      CRITICAL = 7;
    }
  }

  // (Required) The severity of the log statement.
  Severity.Enum severity = 1;

  // (Required) The time at which this log statement occurred.
  google.protobuf.Timestamp timestamp = 2;

  // (Required) A human readable message.
  string message = 3;

  // (Optional) An optional trace of the functions involved. For example, in
  // Java this can include multiple causes and multiple suppressed exceptions.
  string trace = 4;

  // (Optional) A reference to the instruction this log statement is associated
  // with.
  string instruction_reference = 5;

  // (Optional) A reference to the primitive transform this log statement is
  // associated with.
  string primitive_transform_reference = 6;

  // (Optional) Human-readable name of the function or method being invoked,
  // with optional context such as the class or package name. The format can
  // vary by language. For example:
  //   qual.if.ied.Class.method (Java)
  //   dir/package.func (Go)
  //   module.function (Python)
  //   file.cc:382 (C++)
  string log_location = 7;

  // (Optional) The name of the thread this log statement is associated with.
  string thread = 8;
}

message LogControl {
}

// Stable
service BeamFnLogging {
  // Allows for the SDK to emit log entries which the runner can
  // associate with the active job.
  rpc Logging(
    // A stream of log entries batched into lists emitted by the SDK harness.
    stream LogEntry.List
  ) returns (
    // A stream of log control messages used to configure the SDK.
    stream LogControl
  ) {}
}

