sink.h

/* Copyright 2017 Carnegie Mellon University
 *
 * Licensed 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.
 */

#pragma once

#include "scanner/api/kernel.h"
#include "scanner/util/common.h"
#include "scanner/util/profiler.h"
#include "scanner/util/storehouse.h"

#include <vector>

namespace scanner {

struct SinkConfig {
  std::vector<std::string> input_columns;
  std::vector<proto::ColumnType> input_column_types;
  std::vector<u8> args;  
  i32 node_id;
  storehouse::StorageConfig* storage_config;
};

class Sink {
 public:
  Sink(const SinkConfig& config) {}

  virtual ~Sink(){};

  virtual void validate(proto::Result* result) { result->set_success(true); }

  virtual void new_stream(const std::vector<u8>& args) {};

  virtual void write(const BatchedElements& input_columns) = 0;

  virtual void finished() {};

  virtual void set_profiler(Profiler* profiler) { profiler_ = profiler; }

  Profiler* profiler_ = nullptr;
};

namespace internal {

class SinkBuilder;

using SinkConstructor =
    std::function<Sink*(const SinkConfig& config)>;

class SinkRegistration {
 public:
  SinkRegistration(const SinkBuilder& builder);
};

class SinkBuilder {
 public:
  friend class SinkRegistration;

  SinkBuilder(const std::string& name, SinkConstructor constructor)
    : name_(name),
      constructor_(constructor),
      variadic_inputs_(false),
      per_element_output_(false),
      entire_stream_output_(false) {}

  SinkBuilder& variadic_inputs() {
    if (input_columns_.size() > 0) {
      LOG(FATAL) << "Sink " << name_ << " cannot have both fixed and variadic "
                 << "inputs";
    }
    variadic_inputs_ = true;
    return *this;
  }

  SinkBuilder& input(const std::string& name,
                   ColumnType type = ColumnType::Bytes) {
    if (variadic_inputs_) {
      LOG(FATAL) << "Sink " << name_ << " cannot have both fixed and variadic "
                 << "inputs";
    }
    input_columns_.push_back(std::make_tuple(name, type));
    return *this;
  }

  SinkBuilder& frame_input(const std::string& name) {
    return input(name, ColumnType::Video);
  }

  SinkBuilder& per_element_output() {
    if (entire_stream_output_) {
      LOG(FATAL) << "Sink " << name_
                 << " cannot specify both per element and entire stream output";
    }
    per_element_output_ = true;
    return *this;
  }

  SinkBuilder& entire_stream_output() {
    LOG(FATAL) << "Entire stream output is not implemented yet.";

    if (per_element_output_) {
      LOG(FATAL) << "Sink " << name_
                 << " cannot specify both per element and entire stream output";
    }
    entire_stream_output_ = true;
    return *this;
  }

  SinkBuilder& protobuf_name(const std::string& name) {
    protobuf_name_ = name;
    return *this;
  }

  SinkBuilder& stream_protobuf_name(const std::string& name) {
    stream_protobuf_name_ = name;
    return *this;
  }

 private:
  std::string name_;
  SinkConstructor constructor_;
  bool variadic_inputs_;
  std::vector<std::tuple<std::string, ColumnType>> input_columns_;
  bool per_element_output_;
  bool entire_stream_output_;
  std::string protobuf_name_;
  std::string stream_protobuf_name_;
};
}

#define REGISTER_SINK(name__, sink__) \
  REGISTER_SINK_HELPER(__COUNTER__, name__, sink__)

#define REGISTER_SINK_HELPER(uid__, name__, sink__)    \
  REGISTER_SINK_UID(uid__, name__, sink__)

#define REGISTER_SINK_UID(uid__, name__, sink__)                              \
  static ::scanner::internal::SinkRegistration sink_registration_##uid__ \
      __attribute__((unused)) = ::scanner::internal::SinkBuilder(          \
          #name__, [](const ::scanner::SinkConfig& config) {               \
            return new sink__(config);                                     \
          })
}