embulk-output-fluentd 0.1.0

data/gradlew.bat

@@ -0,0 +1,84 @@
+@if "%DEBUG%" == "" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+set DIRNAME=%~dp0
+if "%DIRNAME%" == "" set DIRNAME=.
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS=
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if "%ERRORLEVEL%" == "0" goto init
+
+echo.
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto init
+
+echo.
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:init
+@rem Get command-line arguments, handling Windows variants
+
+if not "%OS%" == "Windows_NT" goto win9xME_args
+
+:win9xME_args
+@rem Slurp the command line arguments.
+set CMD_LINE_ARGS=
+set _SKIP=2
+
+:win9xME_args_slurp
+if "x%~1" == "x" goto execute
+
+set CMD_LINE_ARGS=%*
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+
+:end
+@rem End local scope for the variables with windows NT shell
+if "%ERRORLEVEL%"=="0" goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+if  not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
+exit /b 1
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega

data/lib/embulk/output/fluentd.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_output(
+  "fluentd", "org.embulk.output.fluentd.FluentdOutputPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/project/build.properties

@@ -0,0 +1 @@
+sbt.version=0.13.15

data/project/plugins.sbt

@@ -0,0 +1,2 @@
+addSbtPlugin("com.lucidchart" % "sbt-scalafmt" % "1.7")
+addSbtPlugin("io.get-coursier" % "sbt-coursier" % "1.0.0-RC6")

data/settings.gradle

@@ -0,0 +1 @@
+rootProject.name = 'embulk-output-fluentd'

data/src/main/resources/application.conf

@@ -0,0 +1,8 @@
+blocking-dispatcher {
+  type = Dispatcher
+  executor = "thread-pool-executor"
+  thread-pool-executor {
+    fixed-pool-size = 4
+  }
+  throughput = 100
+}

data/src/main/scala/org/embulk/output/fluentd/ColumnVisitor.scala

@@ -0,0 +1,43 @@
+package org.embulk.output.fluentd
+
+import org.embulk.spi.time.TimestampFormatter
+import org.embulk.spi.{Column, PageReader, ColumnVisitor => EmbulkColumnVisitor}
+
+case class ColumnVisitor(reader: PageReader, timestampFormatter: TimestampFormatter) extends EmbulkColumnVisitor {
+  import scala.collection.mutable
+
+  private val record = mutable.Map[String, AnyRef]()
+
+  override def timestampColumn(column: Column): Unit =
+    value(column, reader.getTimestamp).foreach(v => put(column, timestampFormatter.format(v)))
+
+  override def stringColumn(column: Column): Unit =
+    value(column, reader.getString).foreach(v => put(column, v))
+
+  override def longColumn(column: Column): Unit =
+    value(column, reader.getLong).foreach(v => put(column, Long.box(v)))
+
+  override def doubleColumn(column: Column): Unit =
+    value(column, reader.getDouble).foreach(v => put(column, Double.box(v)))
+
+  override def booleanColumn(column: Column): Unit =
+    value(column, reader.getBoolean).foreach(v => put(column, Boolean.box(v)))
+
+  override def jsonColumn(column: Column): Unit =
+    value(column, reader.getJson).foreach(v => put(column, v.toJson))
+
+  def value[A](column: Column, method: => (Column => A)): Option[A] =
+    if (reader.isNull(column)) {
+      None
+    } else {
+      Some(method(column))
+    }
+
+  def put[A <: AnyRef](column: Column, value: A): Unit = {
+    record.put(column.getName, value)
+    ()
+  }
+
+  def getRecord: Map[String, AnyRef] = record.toMap
+
+}

data/src/main/scala/org/embulk/output/fluentd/FluentdOutputPlugin.scala

@@ -0,0 +1,55 @@
+package org.embulk.output.fluentd
+
+import java.util
+
+import org.embulk.config._
+import org.embulk.output.fluentd.sender._
+import org.embulk.spi._
+import wvlet.log._
+
+class FluentdOutputPlugin extends OutputPlugin {
+
+  override def transaction(config: ConfigSource,
+                           schema: Schema,
+                           taskCount: Int,
+                           control: OutputPlugin.Control): ConfigDiff = {
+    Logger.setDefaultLogLevel(LogLevel.OFF)
+    val task = config.loadConfig(classOf[PluginTask])
+    FluentdOutputPlugin.taskCountOpt = Some(taskCount)
+    control.run(task.dump())
+    FluentdOutputPlugin.sender.foreach(_.close())
+    Exec.newConfigDiff
+  }
+
+  override def resume(taskSource: TaskSource,
+                      schema: Schema,
+                      taskCount: Int,
+                      control: OutputPlugin.Control): ConfigDiff =
+    throw new UnsupportedOperationException("fluentd output plugin does not support resuming")
+
+  override def cleanup(taskSource: TaskSource,
+                       schema: Schema,
+                       taskCount: Int,
+                       successTaskReports: util.List[TaskReport]): Unit = {}
+
+  override def open(taskSource: TaskSource, schema: Schema, taskIndex: Int): TransactionalPageOutput = {
+    FluentdOutputPlugin.sender.synchronized {
+      FluentdOutputPlugin.sender match {
+        case Some(sender) =>
+          FluentdTransactionalPageOutput(taskSource, schema, taskIndex, FluentdOutputPlugin.taskCountOpt, sender)
+        case None =>
+          val task = taskSource.loadTask(classOf[PluginTask])
+          SenderBuilder(task).withSession { session =>
+            val sender = session.build[Sender]
+            FluentdOutputPlugin.sender = Option(sender)
+            FluentdTransactionalPageOutput(taskSource, schema, taskIndex, FluentdOutputPlugin.taskCountOpt, sender)
+          }
+      }
+    }
+  }
+}
+
+object FluentdOutputPlugin {
+  var sender: Option[Sender]    = None
+  var taskCountOpt: Option[Int] = None
+}

data/src/main/scala/org/embulk/output/fluentd/FluentdTransactionalPageOutput.scala

@@ -0,0 +1,52 @@
+package org.embulk.output.fluentd
+
+import com.google.common.base.Optional
+import org.embulk.config.{TaskReport, TaskSource}
+import org.embulk.output.fluentd.sender.Sender
+import org.embulk.spi._
+import org.embulk.spi.time.TimestampFormatter
+
+case class FluentdTransactionalPageOutput(taskSource: TaskSource,
+                                          schema: Schema,
+                                          taskIndex: Int,
+                                          taskCountOpt: Option[Int],
+                                          sender: Sender)
+    extends TransactionalPageOutput {
+
+  val task: PluginTask = taskSource.loadTask(classOf[PluginTask])
+  val logger           = Exec.getLogger(classOf[FluentdTransactionalPageOutput])
+
+  def timestampFormatter(): TimestampFormatter =
+    new TimestampFormatter(task, Optional.absent())
+
+  override def add(page: Page): Unit = {
+    sender(asIterator(page).toSeq)
+  }
+
+  def asIterator(page: Page): Iterator[Map[String, AnyRef]] = {
+    val reader: PageReader = new PageReader(schema)
+    reader.setPage(page)
+    Iterator.continually {
+      if (reader.nextRecord()) {
+        val visitor = ColumnVisitor(reader, timestampFormatter())
+        schema.visitColumns(visitor)
+        visitor.getRecord
+      } else {
+        reader.close()
+        Map.empty[String, AnyRef]
+      }
+    } takeWhile (_ != Map.empty[String, AnyRef])
+  }
+
+  override def commit(): TaskReport = Exec.newTaskReport
+  override def abort(): Unit        = ()
+  override def finish(): Unit = {
+    logger.debug(s"finished at " + this)
+    // for map/reduce executor.
+    if (taskCountOpt.isEmpty) {
+      // close immediately.
+      sender.close()
+    }
+  }
+  override def close(): Unit = ()
+}

data/src/main/scala/org/embulk/output/fluentd/PluginTask.scala

@@ -0,0 +1,36 @@
+package org.embulk.output.fluentd
+
+import com.google.common.base.Optional
+import org.embulk.config.{Config, ConfigDefault, Task}
+import org.embulk.spi.time.TimestampFormatter
+
+trait PluginTask extends Task with TimestampFormatter.Task {
+
+  @Config("host")
+  @ConfigDefault("\"127.0.0.1\"")
+  def getHost: String
+
+  @Config("port")
+  @ConfigDefault("24224")
+  def getPort: Int
+
+  @Config("async_size")
+  @ConfigDefault("1")
+  def getAsyncSize: Int
+
+  @Config("request_per_seconds")
+  @ConfigDefault("0")
+  def getRequestPerSeconds: Int
+
+  @Config("request_grouping_size")
+  @ConfigDefault("100")
+  def getRequestGroupingSize: Int
+
+  @Config("tag")
+  def getTag: String
+
+  @Config("time_key")
+  @ConfigDefault("null")
+  def getTimeKey: Optional[String]
+
+}

data/src/main/scala/org/embulk/output/fluentd/sender/ActorManager.scala

@@ -0,0 +1,29 @@
+package org.embulk.output.fluentd.sender
+
+import akka.actor.{ActorRef, ActorSystem, Props, Terminated}
+import akka.stream.{ActorMaterializer, ActorMaterializerSettings, Supervision}
+
+import scala.concurrent.{ExecutionContext, Future}
+
+case class ActorManagerImpl(implicit val system: ActorSystem) extends ActorManager {
+  val supervisor: ActorRef = system.actorOf(Props[SuperVisor])
+  val decider: Supervision.Decider = {
+    case _: Exception => Supervision.Resume
+    case _            => Supervision.Stop
+  }
+  implicit val materializer = ActorMaterializer(
+    ActorMaterializerSettings(system)
+      .withSupervisionStrategy(decider)
+      .withDispatcher("blocking-dispatcher"))
+
+  implicit val dispatcher: ExecutionContext =
+    system.dispatchers.lookup("blocking-dispatcher")
+}
+
+trait ActorManager {
+  implicit val system: ActorSystem
+  val supervisor: ActorRef
+  implicit val materializer: ActorMaterializer
+  def terminate(): Future[Terminated] = system.terminate()
+  implicit val dispatcher: ExecutionContext
+}

data/src/main/scala/org/embulk/output/fluentd/sender/Sender.scala

@@ -0,0 +1,131 @@
+package org.embulk.output.fluentd.sender
+
+import akka._
+import akka.pattern.ask
+import akka.stream._
+import akka.stream.scaladsl._
+import akka.util.{ByteString, Timeout}
+import org.slf4j.Logger
+
+import scala.concurrent._
+import scala.concurrent.duration._
+import scala.util._
+
+trait Sender {
+  def close(): Unit
+  val instance: SourceQueueWithComplete[Seq[Map[String, AnyRef]]]
+  def apply(value: Seq[Map[String, AnyRef]]): Future[QueueOfferResult]
+  def tcpHandling(size: Int, byteString: ByteString): Future[Done]
+  def waitForComplete(): Result
+}
+
+case class SenderImpl private[sender] (host: String,
+                                       port: Int,
+                                       groupedSize: Int,
+                                       asyncSize: Int,
+                                       senderFlow: SenderFlow,
+                                       actorManager: ActorManager,
+                                       asyncSizeRequestPerSecond: Int = 0,
+                                       retryCount: Int = 0,
+                                       retryDelayIntervalSecond: Int = 10)(implicit logger: Logger)
+    extends Sender {
+  import actorManager._
+  system.scheduler.schedule(0.seconds, 30.seconds, supervisor, LogStatus(logger))
+
+  val retryDelayIntervalSecondDuration: FiniteDuration = retryDelayIntervalSecond.seconds
+
+  def apply(value: Seq[Map[String, AnyRef]]): Future[QueueOfferResult] = {
+    actorManager.supervisor ! Record(value.size)
+    instance.offer(value)
+  }
+
+  def close(): Unit = {
+    implicit val timeout        = Timeout(5.seconds)
+    val f: Future[ClosedStatus] = (actorManager.supervisor ? Close).mapTo[ClosedStatus]
+    val result                  = Await.result(f, Duration.Inf)
+    if (!result.alreadyClosed) {
+      logger.debug("wait for closing.")
+      // wait for akka-stream termination.
+      instance.complete()
+      val result = waitForComplete()
+      Await.result(actorManager.terminate(), Duration.Inf)
+      actorManager.system.terminate()
+      logger.info(
+        s"Completed RecordCount:${result.record} completedCount:${result.complete} retriedRecordCount:${result.retried}")
+    }
+  }
+
+  def waitForComplete(): Result = {
+    logger.debug("wait for complete.")
+    var result: Option[Result] = None
+    implicit val timeout       = Timeout(5.seconds)
+    while (result.isEmpty) {
+      (actorManager.supervisor ? GetStatus).onComplete {
+        case Success(Result(recordCount, complete, failed, retried)) =>
+          logger.debug(s"current status ${Result(recordCount, complete, failed, retried)}")
+          if (recordCount == (complete + failed)) {
+            result = Some(Result(recordCount, complete, failed, retried))
+          }
+        case Success(Stop(recordCount, complete, failed, retried)) =>
+          result = Some(Result(recordCount, complete, failed, retried))
+        case _ =>
+          sys.error("fail of wait complete.")
+      }
+      Thread.sleep(1000)
+    }
+    result.get
+  }
+
+  val instance: SourceQueueWithComplete[Seq[Map[String, AnyRef]]] = {
+    val base = Source
+      .queue(Int.MaxValue, OverflowStrategy.backpressure)
+      .grouped(groupedSize)
+      .via(senderFlow.msgPackFlow)
+    val withThrottle = if (asyncSizeRequestPerSecond > 0) {
+      base.throttle(asyncSize, asyncSizeRequestPerSecond.seconds, 0, ThrottleMode.Shaping)
+    } else base
+    withThrottle
+      .mapAsync(asyncSize) {
+        case (size, byteString) =>
+          tcpHandling(size, byteString)
+      }
+      .to(Sink.ignore)
+      .run()
+  }
+
+  def sendCommand(byteString: ByteString): Future[Done] =
+    Source
+      .single(byteString)
+      .via(senderFlow.tcpConnectionFlow(host, port))
+      .runWith(Sink.ignore)
+
+  def tcpHandling(size: Int, byteString: ByteString): Future[Done] = {
+    def _tcpHandling(size: Int, byteString: ByteString, c: Int)(retried: Boolean): Future[Done] = {
+      val futureCommand = sendCommand(byteString)
+      futureCommand.onComplete {
+        case Success(_) =>
+          actorManager.supervisor ! Complete(size)
+        case Failure(e) if c > 0 =>
+          logger.info(
+            s"Sending fluentd ${size.toString} records was failed. - will retry ${c - 1} more times ${retryDelayIntervalSecondDuration.toSeconds} seconds later.",
+            e)
+          actorManager.supervisor ! Retried(size)
+          akka.pattern.after(retryDelayIntervalSecondDuration, actorManager.system.scheduler)(
+            _tcpHandling(size, byteString, c - 1)(retried = true))
+        case Failure(e) =>
+          actorManager.supervisor ! Failed(size)
+          logger.error(
+            s"Sending fluentd retry count is over and will be terminate soon. Please check your fluentd environment.",
+            e)
+          sys.error("Sending fluentd was terminated cause of retry count over.")
+          instance.complete()
+      }
+      futureCommand
+    }
+    _tcpHandling(size, byteString, retryCount)(retried = false).recoverWith {
+      case _: Exception =>
+        Future.successful(Done)
+    }
+  }
+
+}