Scala GeneratedMessage类代码示例

//设置package包名称以及导入依赖的类
package walfie.gbf.raidfinder.server.persistence

import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion, Message}
import java.net.URI
import redis.clients.jedis.{BinaryJedis, Jedis}

trait ProtobufStorage {
  type CacheItem[T] = GeneratedMessage with Message[T]

  def set[T <: CacheItem[T]](key: String, value: T): Unit
  def get[T <: CacheItem[T]](
    key: String
  )(implicit companion: GeneratedMessageCompanion[T]): Option[T]

  def close(): Unit
}

object ProtobufStorage {
  def redis(uri: URI): RedisProtobufStorage = {
    new RedisProtobufStorage(new BinaryJedis(uri))
  }
}

// TODO: Write integration test
class RedisProtobufStorage(redis: BinaryJedis) extends ProtobufStorage {
  def set[T <: CacheItem[T]](key: String, value: T): Unit = {
    redis.set(key.getBytes, value.toByteArray)
  }

  def get[T <: CacheItem[T]](
    key: String
  )(implicit companion: GeneratedMessageCompanion[T]): Option[T] = {
    Option(redis.get(key.getBytes)).flatMap { bytes =>
      companion.validate(bytes).toOption
    }
  }

  def close(): Unit = redis.close()
}

object NoOpProtobufStorage extends ProtobufStorage {
  def set[T <: CacheItem[T]](key: String, value: T): Unit = ()

  def get[T <: CacheItem[T]](
    key: String
  )(implicit companion: GeneratedMessageCompanion[T]): Option[T] = None

  def close(): Unit = ()
} 

//设置package包名称以及导入依赖的类
package net.cakesolutions.strictify.scalapb.ops

import scala.util.Try

import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion, Message}
import net.cakesolutions.strictify.core.{Binding, StrictifyError}

trait StrictTypeDeserializer[S] {
  def parseFrom(s: Array[Byte]): Either[StrictifyError, S]
}
object StrictTypeDeserializer {
  implicit def instance[S, L <: GeneratedMessage with Message[L]](
    implicit
    binding: Binding.Aux[S, L],
    companion: GeneratedMessageCompanion[L]
  ) = new StrictTypeDeserializer[S] {
    override def parseFrom(s: Array[Byte]): Either[StrictifyError, S] =
      Try(companion.parseFrom(s)).fold(e => Left(StrictifyError(e)), binding.strictify.apply)
  }
}

trait StrictTypeSerializer[S] {
  def toByteArray(value: S): Array[Byte]
}
object StrictTypeSerializer {
  implicit def instance[S, L <: GeneratedMessage with Message[L]](
    implicit
    binding: Binding.Aux[S, L]
  ) = new StrictTypeSerializer[S] {
    override def toByteArray(value: S): Array[Byte] =
      binding.loosen(value).toByteArray
  }
}

object SerializationOps {
  def parseFrom[S: StrictTypeDeserializer](s: Array[Byte]): Either[StrictifyError, S] =
    implicitly[StrictTypeDeserializer[S]].parseFrom(s)

  def toByteArray[S: StrictTypeSerializer](value: S): Array[Byte] =
    implicitly[StrictTypeSerializer[S]].toByteArray(value)
} 

//设置package包名称以及导入依赖的类
package tmt.demo.connectors

import akka.cluster.pubsub.DistributedPubSub
import akka.cluster.pubsub.DistributedPubSubMediator.Subscribe
import akka.stream.scaladsl.Sink
import com.trueaccord.scalapb.GeneratedMessage
import tmt.app.library.Connector
import tmt.app.utils.ActorRuntime
import tmt.demo.zeromq_drivers.ZmqPublisherFactory

class AkkaToZmqFlow(actorRuntime: ActorRuntime, zmqPublisherFactory: ZmqPublisherFactory) {
  import actorRuntime._

  def connect[Msg <: GeneratedMessage](
    subscriberTopic: String,
    publishingPort: Int
  ) = {
    val (sourceLinkedRef, source) = Connector.coupling[Msg](Sink.asPublisher(fanout = false))
    DistributedPubSub(system).mediator ! Subscribe(subscriberTopic, sourceLinkedRef)
    val zmqPublisher = zmqPublisherFactory.make[Msg](publishingPort)

    zmqPublisher
      .publish(source)
      .onComplete { x =>
        zmqPublisher.shutdown()
      }
  }
} 

//设置package包名称以及导入依赖的类
package tmt.demo.zeromq_drivers

import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion}
import org.zeromq.ZMQ
import tmt.app.configs.AppSettings
import tmt.app.utils.{ActorRuntime, EC, PbMessage}

import scala.concurrent.ExecutionContext.global
import scala.concurrent.Future
import scala.util.control.NonFatal

class ZmqServer(settings: AppSettings, runtime: ActorRuntime) {

  import runtime._

  val address = s"tcp://*:${settings.mcsCommandsPort}"

  private val socket = zmqContext.socket(ZMQ.REP)
  socket.bind(address)
  private val ec = EC.singleThreadedEc()

  def start[Req <: PbMessage.Of[Req]](
    reqParser: GeneratedMessageCompanion[Req]
  )(f: Req => GeneratedMessage): Future[Unit] = Future {
    while (true) {
      try {
        val message = socket.recv(0)
        val request = reqParser.parseFrom(message)
        println(s"***** ZmqSever received : [$request]")
        Thread.sleep(100)
        socket.send(f(request).toByteArray, 0)
      } catch {
        case NonFatal(ex) =>
          ex.printStackTrace()
          socket.send("error".getBytes, 0)
      }

    }
  }(ec).map { x =>
    println(s"ZmqServer completed with value: $x")
  }(global)

  def shutdown(): Unit = {
    socket.close()
    ec.shutdown()
  }
} 

//设置package包名称以及导入依赖的类
package tmt.demo.zeromq_drivers

import akka.stream.scaladsl.{Sink, Source}
import com.trueaccord.scalapb.GeneratedMessage
import org.zeromq.ZMQ
import tmt.app.utils.{EC, ActorRuntime}

import scala.concurrent.Future

class ZmqPublisher[Msg <: GeneratedMessage](port: Int, runtime: ActorRuntime) {

  import runtime._

  val address = s"tcp://*:$port"

  private val socket = zmqContext.socket(ZMQ.PUB)
  socket.bind(address)

  private val ec = EC.singleThreadedEc()

  def publish(messages: Source[Msg, Any]): Future[Unit] = {
    messages
      .mapAsync(1)(publishSingle)
      .runWith(Sink.ignore)
      .map { x =>
        println(s"ZmqPublisher completed with value: $x")
      }(system.dispatcher)
  }

  def publishSingle(message: Msg) = Future {
    println(s"********** ZmqPublisher is publishing: [$message]")
    socket.send(message.toByteArray, 0)
  }(ec)

  def shutdown(): Unit = {
    socket.close()
    ec.shutdown()
  }
}

class ZmqPublisherFactory(runtime: ActorRuntime) {
  def make[Msg <: GeneratedMessage](port: Int) = new ZmqPublisher[Msg](port, runtime)
} 

//设置package包名称以及导入依赖的类
package tmt.demo.zeromq_drivers

import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion}
import org.zeromq.ZMQ
import tmt.app.configs.AppSettings
import tmt.app.utils.{PbMessage, EC, ActorRuntime}

import scala.concurrent.Future

class ZmqClient(settings: AppSettings, runtime: ActorRuntime) {

  import runtime._

  val address = s"tcp://${settings.mcsHostname}:${settings.mcsCommandsPort}"

  private val socket = zmqContext.socket(ZMQ.REQ)
  println(s"ZmqClient is connecting to server $address")
  socket.connect(address)
  private val ec = EC.singleThreadedEc()

  def query[Req <: GeneratedMessage, Res <: PbMessage.Of[Res]](
    request: Req,
    responseParser: GeneratedMessageCompanion[Res]
  ): Future[Res] = Future {
    println(s"***** ZmqClient sending $request to $address at ${System.currentTimeMillis()}")
    socket.send(request.toByteArray, 0)
    responseParser.parseFrom(socket.recv(0))
  }(ec)

  def shutdown(): Unit = {
    socket.close()
    ec.shutdown()
  }
} 

//设置package包名称以及导入依赖的类
package tmt.demo.hcd_drivers

import akka.cluster.pubsub.DistributedPubSub
import akka.cluster.pubsub.DistributedPubSubMediator.Publish
import akka.stream.scaladsl.Source
import com.trueaccord.scalapb.GeneratedMessage
import tmt.app.utils.ActorRuntime

class EventPublisher(actorRuntime: ActorRuntime) {

  import actorRuntime._

  def publish[Msg <: GeneratedMessage](
    messages: Source[Msg, Any],
    topic: String
  ) = {
    messages.runForeach { message =>
      println(s"********* PublisherClient is publishing: $message")
      DistributedPubSub(system).mediator ! Publish(topic, message)
    }
  }
} 

//设置package包名称以及导入依赖的类
package org.karps.structures

import scala.util.{Failure, Success, Try}

import com.trueaccord.scalapb.{GeneratedMessage, Message, GeneratedMessageCompanion}
import com.trueaccord.scalapb.json._

object ProtoUtils {
  def checkField[X](x: X, fieldName: String): Try[X] = {
    if (x == null) {
      missingField(fieldName)
    } else {
      Success(x)
    }
  }
  
  def checkField[X](x: Option[X], fieldName: String): Try[X] = {
    if (x == None) {
      missingField(fieldName)
    } else {
      Success(x.get)
    }
  }

  def missingField[X](fieldName: String): Try[X] = {
    Failure(new Exception(s"Missing field $fieldName"))
  }
  
  def unrecognized[X](fieldName: String, x: Int): Try[X] = {
    Failure(new Exception(s"unrecognized value for $fieldName: $x"))
  }
  
  def sequence[T](xs : Seq[Try[T]]) : Try[Seq[T]] = (Try(Seq[T]()) /: xs) {
    (a, b) => a flatMap (c => b map (d => c :+ d))
  }
  
  def fromExtra[A <: GeneratedMessage with Message[A]](extra: OpExtra)(
implicit cmp: GeneratedMessageCompanion[A]): Try[A] = {
    Try(JsonFormat.fromJsonString[A](extra.content))
  }
  
   def fromString[A <: GeneratedMessage with Message[A]](extra: String)(
implicit cmp: GeneratedMessageCompanion[A]): Try[A] = {
    Try(JsonFormat.fromJsonString[A](extra))
  }

  def toJsonString[A <: GeneratedMessage](m: A): String = printer.print(m)

  private val printer = new Printer(includingDefaultValueFields = false)
} 

//设置package包名称以及导入依赖的类
package com.example.utilities.serialization

import akka.http.scaladsl.marshalling.{Marshaller, ToEntityMarshaller}
import akka.http.scaladsl.model.MediaType.Compressible
import akka.http.scaladsl.model.{ContentType, ContentTypes, HttpEntity, MediaType}
import akka.http.scaladsl.unmarshalling.Unmarshaller.UnsupportedContentTypeException
import akka.http.scaladsl.unmarshalling.{FromEntityUnmarshaller, Unmarshaller}
import akka.http.scaladsl.util.FastFuture
import com.google.protobuf.CodedInputStream
import com.trueaccord.scalapb.json.JsonFormat
import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion, Message}

import scala.concurrent.Future

trait ScalaPBMarshalling {
  private val protobufContentType = ContentType(MediaType.applicationBinary("octet-stream", Compressible, "proto"))
  private val applicationJsonContentType = ContentTypes.`application/json`

  def scalaPBFromRequestUnmarshaller[O <: GeneratedMessage with Message[O]](companion: GeneratedMessageCompanion[O]): FromEntityUnmarshaller[O] = {
    Unmarshaller.withMaterializer[HttpEntity, O](_ => implicit mat => {
      case [email protected](`applicationJsonContentType`, data) =>
        val charBuffer = Unmarshaller.bestUnmarshallingCharsetFor(entity)
        FastFuture.successful(JsonFormat.fromJsonString(data.decodeString(charBuffer.nioCharset().name()))(companion))
      case [email protected](`protobufContentType`, data) =>
        FastFuture.successful(companion.parseFrom(CodedInputStream.newInstance(data.asByteBuffer)))
      case entity =>
        Future.failed(UnsupportedContentTypeException(applicationJsonContentType, protobufContentType))
    })
  }

  implicit def scalaPBToEntityMarshaller[U <: GeneratedMessage]: ToEntityMarshaller[U] = {
    def jsonMarshaller(): ToEntityMarshaller[U] = {
      val contentType = applicationJsonContentType
      Marshaller.withFixedContentType(contentType) { value =>
        HttpEntity(contentType, JsonFormat.toJsonString(value))
      }
    }

    def protobufMarshaller(): ToEntityMarshaller[U] = {
      Marshaller.withFixedContentType(protobufContentType) { value =>
        HttpEntity(protobufContentType, value.toByteArray)
      }
    }

    Marshaller.oneOf(jsonMarshaller(), protobufMarshaller())
  }

} 

//设置package包名称以及导入依赖的类
package com.trueaccord.scalapb.parquet

import java.util

import com.google.protobuf.Descriptors.Descriptor
import com.trueaccord.scalapb.{GeneratedMessage, Message}
import org.apache.hadoop.conf.Configuration
import org.apache.parquet.hadoop.BadConfigurationException
import org.apache.parquet.hadoop.api.WriteSupport
import org.apache.parquet.hadoop.api.WriteSupport.WriteContext
import org.apache.parquet.io.api.RecordConsumer
import org.apache.parquet.schema.MessageType

class ScalaPBWriteSupport[T <: GeneratedMessage with Message[T]] extends WriteSupport[T] {
  var pbClass: Class[T] = null
  var recordConsumer: RecordConsumer = null

  override def init(configuration: Configuration): WriteContext = {
    if (pbClass == null) {
      pbClass = configuration.getClass(ScalaPBWriteSupport.SCALAPB_CLASS_WRITE, null, classOf[GeneratedMessage]).asInstanceOf[Class[T]]
      if (pbClass == null) {
        throw new BadConfigurationException("ScalaPB class not specified. Please use ScalaPBOutputFormat.setMessageClass.")
      }
    }
    val descriptor: Descriptor = pbClass.getMethod("descriptor").invoke(null).asInstanceOf[Descriptor]
    val rootSchema: MessageType = SchemaConverter.convert(descriptor)
    val extraMetaDtata = new util.HashMap[String, String]
    extraMetaDtata.put(ScalaPBReadSupport.PB_CLASS, pbClass.getName)
    new WriteContext(rootSchema, extraMetaDtata)
  }

  override def write(record: T): Unit = {
    MessageWriter.writeTopLevelMessage(recordConsumer, record)
  }

  override def prepareForWrite(recordConsumer: RecordConsumer): Unit = {
    this.recordConsumer = recordConsumer
  }
}

object ScalaPBWriteSupport {
  val SCALAPB_CLASS_WRITE = "parquet.scalapb.writeClass"

  def setSchema[T <: GeneratedMessage](config: Configuration, protoClass: Class[T]) = {
    config.setClass(SCALAPB_CLASS_WRITE, protoClass, classOf[GeneratedMessage])
  }
} 

//设置package包名称以及导入依赖的类
package com.trueaccord.scalapb.parquet

import java.util

import com.trueaccord.scalapb.{GeneratedMessage, GeneratedMessageCompanion, Message}
import org.apache.hadoop.conf.Configuration
import org.apache.parquet.hadoop.api.{InitContext, ReadSupport}
import org.apache.parquet.hadoop.api.ReadSupport.ReadContext
import org.apache.parquet.io.api.{GroupConverter, RecordMaterializer}
import org.apache.parquet.schema.MessageType

class ScalaPBReadSupport[T <: GeneratedMessage with Message[T]] extends ReadSupport[T] {
  override def prepareForRead(
    configuration: Configuration,
    keyValueMetaData: util.Map[String, String],
    fileSchema: MessageType,
    readContext: ReadContext): RecordMaterializer[T] = {
    val protoClass = Option(keyValueMetaData.get(ScalaPBReadSupport.PB_CLASS)).getOrElse(throw new RuntimeException(s"Value for ${ScalaPBReadSupport.PB_CLASS} not found."))
    val cmp = {
      import scala.reflect.runtime.universe

      val runtimeMirror = universe.runtimeMirror(getClass.getClassLoader)

      val module = runtimeMirror.staticModule(protoClass)

      val obj = runtimeMirror.reflectModule(module)

      obj.instance.asInstanceOf[GeneratedMessageCompanion[T]]
    }

    new RecordMaterializer[T] {
      val root = new ProtoMessageConverter[T](cmp, fileSchema, onEnd = _ => ())

      override def getRootConverter: GroupConverter = root

      override def getCurrentRecord: T = root.getCurrentRecord
    }
  }

  override def init(context: InitContext): ReadContext = {
    new ReadContext(context.getFileSchema)
  }
}

object ScalaPBReadSupport {
  val PB_CLASS = "parquet.scalapb.class"
} 

//设置package包名称以及导入依赖的类
package com.trueaccord.scalapb.parquet

import com.google.protobuf.ByteString
import com.google.protobuf.Descriptors.FieldDescriptor.JavaType
import com.google.protobuf.Descriptors.{EnumValueDescriptor, FieldDescriptor}
import com.trueaccord.scalapb.{GeneratedMessage, Message}
import org.apache.parquet.io.api.Binary
import org.apache.parquet.Log
import org.apache.parquet.io.api.RecordConsumer

object MessageWriter {
  val log = Log.getLog(this.getClass)

  def writeTopLevelMessage[T <: GeneratedMessage with Message[T]](consumer: RecordConsumer, m: T) = {
    consumer.startMessage()
    writeAllFields(consumer, m)
    consumer.endMessage()
  }

  private def writeAllFields[T <: GeneratedMessage](consumer: RecordConsumer, m: T): Unit = {
    m.getAllFields.foreach {
      case (fd, value) =>
        consumer.startField(fd.getName, fd.getIndex)
        if (fd.isRepeated) {
          value.asInstanceOf[Seq[Any]].foreach {
            v =>
              writeSingleField(consumer, fd, v)
          }
        } else {
          writeSingleField(consumer, fd, value)
        }
        consumer.endField(fd.getName, fd.getIndex)
    }
  }

  private def writeSingleField(consumer: RecordConsumer, fd: FieldDescriptor, v: Any) = fd.getJavaType match {
    case JavaType.BOOLEAN => consumer.addBoolean(v.asInstanceOf[Boolean])
    case JavaType.INT => consumer.addInteger(v.asInstanceOf[Int])
    case JavaType.LONG => consumer.addLong(v.asInstanceOf[Long])
    case JavaType.FLOAT => consumer.addFloat(v.asInstanceOf[Float])
    case JavaType.DOUBLE => consumer.addDouble(v.asInstanceOf[Double])
    case JavaType.BYTE_STRING => consumer.addBinary(Binary.fromByteArray(v.asInstanceOf[ByteString].toByteArray))
    case JavaType.STRING => consumer.addBinary(Binary.fromString(v.asInstanceOf[String]))
    case JavaType.MESSAGE =>
      consumer.startGroup()
      writeAllFields(consumer, v.asInstanceOf[GeneratedMessage])
      consumer.endGroup()
    case JavaType.ENUM => consumer.addBinary(Binary.fromString(v.asInstanceOf[EnumValueDescriptor].getName))
    case javaType =>
      throw new UnsupportedOperationException("Cannot convert Protocol Buffer: unknown type " + javaType)
  }
} 

//设置package包名称以及导入依赖的类
package com.trueaccord.scalapb.spark

import com.trueaccord.scalapb.parquet.{ScalaPBInputFormat, ScalaPBOutputFormat, ScalaPBWriteSupport}
import com.trueaccord.scalapb.{GeneratedMessage, Message}
import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession

import scala.reflect.ClassTag

object ProtoParquet {
  def loadParquet[T <: GeneratedMessage with Message[T]](sc: SparkContext, input: String)(implicit vt: ClassTag[T]): RDD[T] = {
    sc.newAPIHadoopRDD(
      conf = sc.hadoopConfiguration,
      fClass = classOf[ScalaPBInputFormat[T]],
      kClass = classOf[Void],
      vClass = vt.runtimeClass.asInstanceOf[Class[T]])
      .map(_._2)
  }

  def loadParquet[T <: GeneratedMessage with Message[T]](sc: SparkSession, input: String)(implicit vt: ClassTag[T]): RDD[T] = {
    sc.sparkContext.newAPIHadoopFile(
      input,
      fClass = classOf[ScalaPBInputFormat[T]],
      kClass = classOf[Void],
      vClass = vt.runtimeClass.asInstanceOf[Class[T]])
      .map(_._2)
  }

  def saveParquet[T <: GeneratedMessage with Message[T]](rdd: RDD[T], path: String)(implicit vt: ClassTag[T]) = {
    val config = rdd.context.hadoopConfiguration
    ScalaPBWriteSupport.setSchema(config, vt.runtimeClass.asInstanceOf[Class[T]])
    rdd.map(t => (null, t))
      .saveAsNewAPIHadoopFile(
        path = path,
        keyClass = classOf[Void],
        valueClass = vt.runtimeClass,
        outputFormatClass = classOf[ScalaPBOutputFormat[T]],
        conf = config)
  }
} 

//设置package包名称以及导入依赖的类
package grpcgateway.handlers

import java.nio.charset.StandardCharsets

import com.trueaccord.scalapb.GeneratedMessage
import com.trueaccord.scalapb.json.JsonFormat
import io.grpc.ManagedChannel
import io.netty.buffer.Unpooled
import io.netty.channel.ChannelHandler.Sharable
import io.netty.channel.{ChannelFutureListener, ChannelHandlerContext, ChannelInboundHandlerAdapter}
import io.netty.handler.codec.http._

import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success}

@Sharable
abstract class GrpcGatewayHandler(channel: ManagedChannel)(implicit ec: ExecutionContext)
  extends ChannelInboundHandlerAdapter {

  def name: String

  def shutdown(): Unit =
    if (!channel.isShutdown) channel.shutdown()

  def unaryCall(method: HttpMethod, uri: String, body: String): Future[GeneratedMessage]

  override def channelRead(ctx: ChannelHandlerContext, msg: scala.Any): Unit = msg match {
    case req: FullHttpRequest =>
      val body = req.content().toString(StandardCharsets.UTF_8)
      unaryCall(req.method(), req.uri(), body)
        .map(JsonFormat.toJsonString)
        .map(_.getBytes(StandardCharsets.UTF_8))
        .onComplete {
          case Success(json) =>
            val res = new DefaultFullHttpResponse(
              req.protocolVersion(),
              HttpResponseStatus.OK,
              Unpooled.copiedBuffer(json)
            )
            res.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json")
            HttpUtil.setContentLength(res, json.length)
            HttpUtil.setKeepAlive(res, HttpUtil.isKeepAlive(req))
            ctx.writeAndFlush(res)

          case Failure(e) =>
            val status = e match {
              case _: UnsupportedOperationException => HttpResponseStatus.NOT_FOUND
              case _: NoSuchElementException => HttpResponseStatus.BAD_REQUEST
              case _ => HttpResponseStatus.INTERNAL_SERVER_ERROR
            }
            val res = new DefaultHttpResponse(req.protocolVersion(), status)
            ctx.writeAndFlush(res).addListener(ChannelFutureListener.CLOSE)
        }

    case _ => super.channelRead(ctx, msg)
  }
}