Adding RequestContext

Support creating the RequestContext in ServerChannelHandler

network/src/main/scala/com/linkedin/norbert/network/common/RequestContext.scala

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2009-2010 LinkedIn, Inc
+ *
+ * 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.
+ */
+package com.linkedin.norbert.network.common
+
+import java.util.UUID
+import com.google.protobuf.Message
+
+case class RequestContext(requestId: UUID, receivedAt: Long, message: Message)

network/src/main/scala/com/linkedin/norbert/network/netty/ServerChannelHandler.scala

@@ -20,8 +20,10 @@ import com.linkedin.norbert.logging.Logging
 import com.linkedin.norbert.protos.NorbertProtos
 import com.linkedin.norbert.network.InvalidMessageException
 import org.jboss.netty.channel._
-import com.linkedin.norbert.network.server.{MessageExecutor, MessageHandlerRegistry}
 import com.google.protobuf.{InvalidProtocolBufferException, Message}
+import com.linkedin.norbert.network.server.{MessageHandlerRegistry, MessageExecutor, MessageExecutor, MessageHandlerRegistry}
+import java.util.UUID
+import com.linkedin.norbert.network.common.RequestContext
 
 @ChannelPipelineCoverage("all")
 class ServerChannelHandler(channelGroup: ChannelGroup, messageHandlerRegistry: MessageHandlerRegistry, messageExecutor: MessageExecutor) extends SimpleChannelHandler with Logging {
@@ -36,15 +38,18 @@ class ServerChannelHandler(channelGroup: ChannelGroup, messageHandlerRegistry: M
     val norbertMessage = e.getMessage.asInstanceOf[NorbertProtos.NorbertMessage]
     log.ifTrace("messageRecieved [%s]: %s", channel, norbertMessage)
 
+    val requestId = new UUID(norbertMessage.getRequestIdMsb, norbertMessage.getRequestIdLsb)
+
     if (norbertMessage.getStatus != NorbertProtos.NorbertMessage.Status.OK) {
       log.warn("Received invalid message: %s", norbertMessage)
-      channel.write(newErrorMessage(norbertMessage, new InvalidMessageException("Recieved a request in the error state")))
+      channel.write(ResponseHelper.errorResponse(requestId, new InvalidMessageException("Recieved a request in the error state")))
     } else {
       try {
         val di = messageHandlerRegistry.requestMessageDefaultInstanceFor(norbertMessage.getMessageName)
         val message = di.newBuilderForType.mergeFrom(norbertMessage.getMessage).build
         log.ifDebug("Queuing to MessageExecutor: %s", message)
-        messageExecutor.executeMessage(message, either => responseHandler(norbertMessage, channel, either))
+        val context = RequestContext(requestId, System.currentTimeMillis, message)
+        messageExecutor.executeMessage(message, either => responseHandler(context, channel, either))
       } catch {
         case ex: InvalidMessageException => log.error(ex, "Recieved invalid message")
         case ex: InvalidProtocolBufferException => log.error(ex, "Error deserializing message")
@@ -54,31 +59,31 @@ class ServerChannelHandler(channelGroup: ChannelGroup, messageHandlerRegistry: M
 
   override def exceptionCaught(ctx: ChannelHandlerContext, e: ExceptionEvent) = log.info(e.getCause, "Caught exception in network layer")
 
-  def responseHandler(norbertMessage: NorbertProtos.NorbertMessage, channel: Channel, either: Either[Exception, Message]) {
+  def responseHandler(context: RequestContext, channel: Channel, either: Either[Exception, Message]) {
     val message = either match {
-      case Left(ex) => newErrorMessage(norbertMessage, ex)
-      case Right(message) => newBuilder(norbertMessage).
-          setMessageName(message.getDescriptorForType.getFullName).
-          setMessage(message.toByteString).
-          build
+      case Left(ex) => ResponseHelper.errorResponse(context.requestId, ex)
+      case Right(message) => ResponseHelper.responseBuilder(context.requestId)
+              .setMessageName(message.getDescriptorForType.getFullName)
+              .setMessage(message.toByteString)
+              .build
     }
 
     log.ifDebug("Sending response: %s", message)
 
     channel.write(message)
   }
+}
 
-  private def newBuilder(norbertMessage: NorbertProtos.NorbertMessage) = {
-    val builder = NorbertProtos.NorbertMessage.newBuilder()
-    builder.setRequestIdMsb(norbertMessage.getRequestIdMsb)
-    builder.setRequestIdLsb(norbertMessage.getRequestIdLsb)
-    builder
+private[netty] object ResponseHelper {
+  def responseBuilder(requestId: UUID) = {
+    NorbertProtos.NorbertMessage.newBuilder.setRequestIdMsb(requestId.getMostSignificantBits).setRequestIdLsb(requestId.getLeastSignificantBits)
   }
 
-  private def newErrorMessage(norbertMessage: NorbertProtos.NorbertMessage, ex: Exception) = {
-    newBuilder(norbertMessage).setMessageName(ex.getClass.getName).
-        setStatus(NorbertProtos.NorbertMessage.Status.ERROR).
-        setErrorMessage(if (ex.getMessage == null) "" else ex.getMessage).
-        build
+  def errorResponse(requestId: UUID, ex: Exception) = {
+    responseBuilder(requestId)
+            .setMessageName(ex.getClass.getName)
+            .setStatus(NorbertProtos.NorbertMessage.Status.ERROR)
+            .setErrorMessage(if (ex.getMessage == null) "" else ex.getMessage)
+            .build
   }
 }