243 lines
11 KiB
Scala
243 lines
11 KiB
Scala
package com.twitter.tweetypie.storage
|
|
|
|
import com.twitter.bijection.Conversion.asMethod
|
|
import com.twitter.bijection.Injection
|
|
import com.twitter.scrooge.TFieldBlob
|
|
import com.twitter.storage.client.manhattan.kv._
|
|
import com.twitter.tweetypie.storage.Response.FieldResponse
|
|
import com.twitter.tweetypie.storage.Response.FieldResponseCode
|
|
import com.twitter.tweetypie.storage_internal.thriftscala.CoreFields
|
|
import com.twitter.tweetypie.storage_internal.thriftscala.InternalTweet
|
|
import com.twitter.tweetypie.storage_internal.thriftscala.StoredTweet
|
|
import java.io.ByteArrayOutputStream
|
|
import java.nio.ByteBuffer
|
|
import org.apache.thrift.protocol.TBinaryProtocol
|
|
import org.apache.thrift.transport.TIOStreamTransport
|
|
import org.apache.thrift.transport.TMemoryInputTransport
|
|
import scala.collection.immutable
|
|
import scala.util.control.NoStackTrace
|
|
|
|
// NOTE: All field ids and Tweet structure in this file correspond to the StoredTweet struct ONLY
|
|
|
|
object ByteArrayCodec {
|
|
def toByteBuffer(byteArray: Array[Byte]): ByteBuffer = byteArray.as[ByteBuffer]
|
|
def fromByteBuffer(buffer: ByteBuffer): Array[Byte] = buffer.as[Array[Byte]]
|
|
}
|
|
|
|
object StringCodec {
|
|
private val string2ByteBuffer = Injection.connect[String, Array[Byte], ByteBuffer]
|
|
def toByteBuffer(strValue: String): ByteBuffer = string2ByteBuffer(strValue)
|
|
def fromByteBuffer(buffer: ByteBuffer): String = string2ByteBuffer.invert(buffer).get
|
|
}
|
|
|
|
/**
|
|
* Terminology
|
|
* -----------
|
|
* Tweet id field : The field number of 'tweetId' in the 'Tweet' thrift structure (i.e "1")
|
|
*
|
|
* First AdditionalField id : The ID if the first additional field in 'Tweet' thrift structure. All field Ids less than this are
|
|
* considered internal and all the ids greater than or equal to this field id are considered 'Additional fields'.
|
|
* This is set to 100.
|
|
*
|
|
* Internal Fields : Fields with ids [1 to firstAdditionalFieldid) (excluding firstAdditionalFieldId)
|
|
*
|
|
* Core fields : (Subset of Internal fields)- Fields with ids [1 to 8, 19]. These fields are "packed" together and stored
|
|
* under a single key. This key is referred to as "CoreFieldsKey" (see @TweetKeyType.CoreFieldsKey).
|
|
* Note: Actually field 1 is skipped when packing as this field is the tweet id and it need not be
|
|
* explicitly stored since the pkey already contains the tweet Id)
|
|
*
|
|
* Root Core field id : The field id under which the packed core fields are stored in Manhattan. (This is field Id "1")
|
|
*
|
|
* Required fields : (Subset of Core fields) - Fields with ids [1 to 5] that MUST be present on every tweet.
|
|
*
|
|
* Additional Fields : All fields with field ids >= 'firstAdditionalFieldId'
|
|
*
|
|
* Compiled Additional fields : (Subset of Additional Fields) - All fields that the storage library knows about
|
|
* (i.e present on the latest storage_internal.thrift that is compiled-in).
|
|
*
|
|
* Passthrough fields : (Subset of Additional Fields) - The fields on storage_internal.thrift that the storage library is NOT aware of
|
|
* These field ids are is obtained looking at the "_passThroughFields" member of the scrooge-generated
|
|
* 'Tweet' object.
|
|
*
|
|
* coreFieldsIdInInternalTweet: This is the field id of the core fields (the only field) in the Internal Tweet struct
|
|
*/
|
|
object TweetFields {
|
|
val firstAdditionalFieldId: Short = 100
|
|
val tweetIdField: Short = 1
|
|
val geoFieldId: Short = 9
|
|
|
|
// The field under which all the core field values are stored (in serialized form).
|
|
val rootCoreFieldId: Short = 1
|
|
|
|
val coreFieldIds: immutable.IndexedSeq[FieldId] = {
|
|
val quotedTweetFieldId: Short = 19
|
|
(1 to 8).map(_.toShort) ++ Seq(quotedTweetFieldId)
|
|
}
|
|
val requiredFieldIds: immutable.IndexedSeq[FieldId] = (1 to 5).map(_.toShort)
|
|
|
|
val coreFieldsIdInInternalTweet: Short = 1
|
|
|
|
val compiledAdditionalFieldIds: Seq[FieldId] =
|
|
StoredTweet.metaData.fields.filter(_.id >= firstAdditionalFieldId).map(_.id)
|
|
val internalFieldIds: Seq[FieldId] =
|
|
StoredTweet.metaData.fields.filter(_.id < firstAdditionalFieldId).map(_.id)
|
|
val nonCoreInternalFields: Seq[FieldId] =
|
|
(internalFieldIds.toSet -- coreFieldIds.toSet).toSeq
|
|
def getAdditionalFieldIds(tweet: StoredTweet): Seq[FieldId] =
|
|
compiledAdditionalFieldIds ++ tweet._passthroughFields.keys.toSeq
|
|
}
|
|
|
|
/**
|
|
* Helper object to convert TFieldBlob to ByteBuffer that gets stored in Manhattan.
|
|
*
|
|
* The following is the format in which the TFieldBlob gets stored:
|
|
* [Version][TField][TFieldBlob]
|
|
*/
|
|
object TFieldBlobCodec {
|
|
val BinaryProtocolFactory: TBinaryProtocol.Factory = new TBinaryProtocol.Factory()
|
|
val FormatVersion = 1.0
|
|
|
|
def toByteBuffer(tFieldBlob: TFieldBlob): ByteBuffer = {
|
|
val baos = new ByteArrayOutputStream()
|
|
val prot = BinaryProtocolFactory.getProtocol(new TIOStreamTransport(baos))
|
|
|
|
prot.writeDouble(FormatVersion)
|
|
prot.writeFieldBegin(tFieldBlob.field)
|
|
prot.writeBinary(ByteArrayCodec.toByteBuffer(tFieldBlob.data))
|
|
|
|
ByteArrayCodec.toByteBuffer(baos.toByteArray)
|
|
}
|
|
|
|
def fromByteBuffer(buffer: ByteBuffer): TFieldBlob = {
|
|
val byteArray = ByteArrayCodec.fromByteBuffer(buffer)
|
|
val prot = BinaryProtocolFactory.getProtocol(new TMemoryInputTransport(byteArray))
|
|
|
|
val version = prot.readDouble()
|
|
if (version != FormatVersion) {
|
|
throw new VersionMismatchError(
|
|
"Version mismatch in decoding ByteBuffer to TFieldBlob. " +
|
|
"Actual version: " + version + ". Expected version: " + FormatVersion
|
|
)
|
|
}
|
|
|
|
val tField = prot.readFieldBegin()
|
|
val dataBuffer = prot.readBinary()
|
|
val data = ByteArrayCodec.fromByteBuffer(dataBuffer)
|
|
|
|
TFieldBlob(tField, data)
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Helper object to help convert 'CoreFields' object to/from TFieldBlob (and also to construct
|
|
* 'CoreFields' object from a 'StoredTweet' object)
|
|
*
|
|
* More details:
|
|
* - A subset of fields on the 'StoredTweet' thrift structure (2-8,19) are 'packaged' and stored
|
|
* together as a serialized TFieldBlob object under a single key in Manhattan (see TweetKeyCodec
|
|
* helper object above for more details).
|
|
*
|
|
* - To make the packing/unpacking the fields to/from TFieldBlob object, we created the following
|
|
* two helper thrift structures 'CoreFields' and 'InternalTweet'
|
|
*
|
|
* // The field Ids and types here MUST exactly match field Ids on 'StoredTweet' thrift structure.
|
|
* struct CoreFields {
|
|
* 2: optional i64 user_id
|
|
* ...
|
|
* 8: optional i64 contributor_id
|
|
* ...
|
|
* 19: optional StoredQuotedTweet stored_quoted_tweet
|
|
*
|
|
* }
|
|
*
|
|
* // The field id of core fields MUST be "1"
|
|
* struct InternalTweet {
|
|
* 1: CoreFields coreFields
|
|
* }
|
|
*
|
|
* - Given the above two structures, packing/unpacking fields (2-8,19) on StoredTweet object into a TFieldBlob
|
|
* becomes very trivial:
|
|
* For packing:
|
|
* (i) Copy fields (2-8,19) from StoredTweet object to a new CoreFields object
|
|
* (ii) Create a new InternalTweet object with the 'CoreFields' object constructed in step (i) above
|
|
* (iii) Extract field "1" as a TFieldBlob from InternalField (by calling the scrooge generated "getFieldBlob(1)"
|
|
* function on the InternalField objecton
|
|
*
|
|
* For unpacking:
|
|
* (i) Create an empty 'InternalField' object
|
|
* (ii) Call scrooge-generated 'setField' by passing the tFieldBlob blob (created by packing steps above)
|
|
* (iii) Doing step (ii) above will create a hydrated 'CoreField' object that can be accessed by 'coreFields'
|
|
* member of 'InternalTweet' object.
|
|
*/
|
|
object CoreFieldsCodec {
|
|
val coreFieldIds: Seq[FieldId] = CoreFields.metaData.fields.map(_.id)
|
|
|
|
// "Pack" the core fields i.e converts 'CoreFields' object to "packed" tFieldBlob (See description
|
|
// above for more details)
|
|
def toTFieldBlob(coreFields: CoreFields): TFieldBlob = {
|
|
InternalTweet(Some(coreFields)).getFieldBlob(TweetFields.coreFieldsIdInInternalTweet).get
|
|
}
|
|
|
|
// "Unpack" the core fields from a packed TFieldBlob into a CoreFields object (see description above for
|
|
// more details)
|
|
def fromTFieldBlob(tFieldBlob: TFieldBlob): CoreFields = {
|
|
InternalTweet().setField(tFieldBlob).coreFields.get
|
|
}
|
|
|
|
// "Unpack" the core fields from a packed TFieldBlob into a Map of core-fieldId-> TFieldBlob
|
|
def unpackFields(tFieldBlob: TFieldBlob): Map[Short, TFieldBlob] =
|
|
fromTFieldBlob(tFieldBlob).getFieldBlobs(coreFieldIds)
|
|
|
|
// Create a 'CoreFields' thrift object from 'Tweet' thrift object.
|
|
def fromTweet(tweet: StoredTweet): CoreFields = {
|
|
// As mentioned above, the field ids and types on the 'CoreFields' struct exactly match the
|
|
// corresponding fields on StoredTweet structure. So it is safe to call .getField() on Tweet object and
|
|
// and pass the returned tFleldBlob a 'setField' on 'CoreFields' object.
|
|
coreFieldIds.foldLeft(CoreFields()) {
|
|
case (core, fieldId) =>
|
|
tweet.getFieldBlob(fieldId) match {
|
|
case None => core
|
|
case Some(tFieldBlob) => core.setField(tFieldBlob)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Helper object to convert ManhattanException to FieldResponseCode thrift object
|
|
*/
|
|
object FieldResponseCodeCodec {
|
|
import FieldResponseCodec.ValueNotFoundException
|
|
|
|
def fromManhattanException(mhException: ManhattanException): FieldResponseCode = {
|
|
mhException match {
|
|
case _: ValueNotFoundException => FieldResponseCode.ValueNotFound
|
|
case _: InternalErrorManhattanException => FieldResponseCode.Error
|
|
case _: InvalidRequestManhattanException => FieldResponseCode.InvalidRequest
|
|
case _: DeniedManhattanException => FieldResponseCode.Error
|
|
case _: UnsatisfiableManhattanException => FieldResponseCode.Error
|
|
case _: TimeoutManhattanException => FieldResponseCode.Timeout
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Helper object to construct FieldResponse thrift object from an Exception.
|
|
* This is typically called to convert 'ManhattanException' object to 'FieldResponse' thrift object
|
|
*/
|
|
object FieldResponseCodec {
|
|
class ValueNotFoundException extends ManhattanException("Value not found!") with NoStackTrace
|
|
private[storage] val NotFound = new ValueNotFoundException
|
|
|
|
def fromThrowable(e: Throwable, additionalMsg: Option[String] = None): FieldResponse = {
|
|
val (respCode, errMsg) = e match {
|
|
case mhException: ManhattanException =>
|
|
(FieldResponseCodeCodec.fromManhattanException(mhException), mhException.getMessage)
|
|
case _ => (FieldResponseCode.Error, e.getMessage)
|
|
}
|
|
|
|
val respMsg = additionalMsg.map(_ + ". " + errMsg).orElse(Some(errMsg.toString))
|
|
FieldResponse(respCode, respMsg)
|
|
}
|
|
}
|