scray.service.qservice.thrifscala.ScrayStatefulTService$FinagleService.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scray-service Show documentation
scray remote query service interface
There is a newer version: 1.1.2
/**
 * Generated by Scrooge
 *   version: 4.16.0
 *   rev: 0201cac9fdd6188248d42da91fd14c87744cc4a5
 *   built at: 20170421-124523
 */
package scray.service.qservice.thrifscala

import com.twitter.finagle.Thrift
import com.twitter.finagle.stats.{NullStatsReceiver, StatsReceiver}
import com.twitter.scrooge.{ThriftStruct, TReusableMemoryTransport}
import com.twitter.util.{Future, Return, Throw, Throwables}
import java.nio.ByteBuffer
import java.util.Arrays
import org.apache.thrift.protocol._
import org.apache.thrift.TApplicationException
import org.apache.thrift.transport.TMemoryInputTransport
import scala.collection.mutable.{
  ArrayBuffer => mutable$ArrayBuffer, HashMap => mutable$HashMap}
import scala.collection.{Map, Set}

import scala.language.higherKinds

/**
 * Stateful query service with continuous result pages
 */
@javax.annotation.Generated(value = Array("com.twitter.scrooge.Compiler"))
class ScrayStatefulTService$FinagleService(
  iface: ScrayStatefulTService[Future],
  protocolFactory: TProtocolFactory,
  stats: StatsReceiver,
  maxThriftBufferSize: Int,
  serviceName: String
) extends com.twitter.finagle.Service[Array[Byte], Array[Byte]] {
  import ScrayStatefulTService._

  def this(
    iface: ScrayStatefulTService[Future],
    protocolFactory: TProtocolFactory,
    stats: StatsReceiver,
    maxThriftBufferSize: Int
  ) = this(iface, protocolFactory, stats, maxThriftBufferSize, "ScrayStatefulTService")

  def this(
    iface: ScrayStatefulTService[Future],
    protocolFactory: TProtocolFactory
  ) = this(iface, protocolFactory, NullStatsReceiver, Thrift.maxThriftBufferSize)

  private[this] val tlReusableBuffer = new ThreadLocal[TReusableMemoryTransport] {
    override def initialValue() = TReusableMemoryTransport(512)
  }

  private[this] def reusableBuffer: TReusableMemoryTransport = {
    val buf = tlReusableBuffer.get()
    buf.reset()
    buf
  }

  private[this] val resetCounter = stats.scope("buffer").counter("resetCount")

  private[this] def resetBuffer(trans: TReusableMemoryTransport): Unit = {
    if (trans.currentCapacity > maxThriftBufferSize) {
      resetCounter.incr()
      tlReusableBuffer.remove()
    }
  }

  protected val functionMap = new mutable$HashMap[String, (TProtocol, Int) => Future[Array[Byte]]]()

  protected def addFunction(name: String, f: (TProtocol, Int) => Future[Array[Byte]]): Unit = {
    functionMap(name) = f
  }

  protected def exception(name: String, seqid: Int, code: Int, message: String): Future[Array[Byte]] = {
    try {
      val x = new TApplicationException(code, message)
      val memoryBuffer = reusableBuffer
      try {
        val oprot = protocolFactory.getProtocol(memoryBuffer)

        oprot.writeMessageBegin(new TMessage(name, TMessageType.EXCEPTION, seqid))
        x.write(oprot)
        oprot.writeMessageEnd()
        oprot.getTransport().flush()
        Future.value(Arrays.copyOfRange(memoryBuffer.getArray(), 0, memoryBuffer.length()))
      } finally {
        resetBuffer(memoryBuffer)
      }
    } catch {
      case e: Exception => Future.exception(e)
    }
  }

  protected def reply(name: String, seqid: Int, result: ThriftStruct): Future[Array[Byte]] = {
    try {
      val memoryBuffer = reusableBuffer
      try {
        val oprot = protocolFactory.getProtocol(memoryBuffer)

        oprot.writeMessageBegin(new TMessage(name, TMessageType.REPLY, seqid))
        result.write(oprot)
        oprot.writeMessageEnd()

        Future.value(Arrays.copyOfRange(memoryBuffer.getArray(), 0, memoryBuffer.length()))
      } finally {
        resetBuffer(memoryBuffer)
      }
    } catch {
      case e: Exception => Future.exception(e)
    }
  }

  final def apply(request: Array[Byte]): Future[Array[Byte]] = {
    val inputTransport = new TMemoryInputTransport(request)
    val iprot = protocolFactory.getProtocol(inputTransport)

    try {
      val msg = iprot.readMessageBegin()
      val func = functionMap.get(msg.name)
      func match {
        case _root_.scala.Some(fn) =>
          fn(iprot, msg.seqid)
        case _ =>
          TProtocolUtil.skip(iprot, TType.STRUCT)
          exception(msg.name, msg.seqid, TApplicationException.UNKNOWN_METHOD,
            "Invalid method name: '" + msg.name + "'")
      }
    } catch {
      case e: Exception => Future.exception(e)
    }
  }

  // ---- end boilerplate.

  private[this] val scopedStats = if (serviceName != "") stats.scope(serviceName) else stats
  private[this] object __stats_query {
    val RequestsCounter = scopedStats.scope("query").counter("requests")
    val SuccessCounter = scopedStats.scope("query").counter("success")
    val FailuresCounter = scopedStats.scope("query").counter("failures")
    val FailuresScope = scopedStats.scope("query").scope("failures")
  }
  addFunction("query", { (iprot: TProtocol, seqid: Int) =>
    try {
      __stats_query.RequestsCounter.incr()
      val args = Query.Args.decode(iprot)
      iprot.readMessageEnd()
      (try {
        iface.query(args.query)
      } catch {
        case e: Exception => Future.exception(e)
      }).flatMap { value: scray.service.qmodel.thrifscala.ScrayUUID =>
        reply("query", seqid, Query.Result(success = Some(value)))
      }.rescue {
        case e: scray.service.qservice.thrifscala.ScrayTException => {
          reply("query", seqid, Query.Result(ex = Some(e)))
        }
        case e => Future.exception(e)
      }.respond {
        case Return(_) =>
          __stats_query.SuccessCounter.incr()
        case Throw(ex) =>
          __stats_query.FailuresCounter.incr()
          __stats_query.FailuresScope.counter(Throwables.mkString(ex): _*).incr()
      }
    } catch {
      case e: TProtocolException => {
        iprot.readMessageEnd()
        exception("query", seqid, TApplicationException.PROTOCOL_ERROR, e.getMessage)
      }
      case e: Exception => Future.exception(e)
    }
  })
  private[this] object __stats_getResults {
    val RequestsCounter = scopedStats.scope("getResults").counter("requests")
    val SuccessCounter = scopedStats.scope("getResults").counter("success")
    val FailuresCounter = scopedStats.scope("getResults").counter("failures")
    val FailuresScope = scopedStats.scope("getResults").scope("failures")
  }
  addFunction("getResults", { (iprot: TProtocol, seqid: Int) =>
    try {
      __stats_getResults.RequestsCounter.incr()
      val args = GetResults.Args.decode(iprot)
      iprot.readMessageEnd()
      (try {
        iface.getResults(args.queryId)
      } catch {
        case e: Exception => Future.exception(e)
      }).flatMap { value: scray.service.qservice.thrifscala.ScrayTResultFrame =>
        reply("getResults", seqid, GetResults.Result(success = Some(value)))
      }.rescue {
        case e: scray.service.qservice.thrifscala.ScrayTException => {
          reply("getResults", seqid, GetResults.Result(ex = Some(e)))
        }
        case e => Future.exception(e)
      }.respond {
        case Return(_) =>
          __stats_getResults.SuccessCounter.incr()
        case Throw(ex) =>
          __stats_getResults.FailuresCounter.incr()
          __stats_getResults.FailuresScope.counter(Throwables.mkString(ex): _*).incr()
      }
    } catch {
      case e: TProtocolException => {
        iprot.readMessageEnd()
        exception("getResults", seqid, TApplicationException.PROTOCOL_ERROR, e.getMessage)
      }
      case e: Exception => Future.exception(e)
    }
  })
}