the-algorithm/tweetypie/common/src/scala/com/twitter/tweetypie/storage/UpdateTweetHandler.scala

65 lines
2.4 KiB
Scala

package com.twitter.tweetypie.storage
import com.twitter.finagle.stats.StatsReceiver
import com.twitter.stitch.Stitch
import com.twitter.storage.client.manhattan.kv.DeniedManhattanException
import com.twitter.storage.client.manhattan.kv.ManhattanValue
import com.twitter.tweetypie.storage.TweetUtils._
import com.twitter.tweetypie.thriftscala.Tweet
import com.twitter.util.Throw
import com.twitter.util.Time
object UpdateTweetHandler {
def apply(
insert: ManhattanOperations.Insert,
stats: StatsReceiver
): TweetStorageClient.UpdateTweet = { (tpTweet: Tweet, fields: Seq[Field]) =>
require(
fields.forall(!TweetFields.coreFieldIds.contains(_)),
"Core fields cannot be modified by calling updateTweet; use addTweet instead."
)
require(
areAllFieldsDefined(tpTweet, fields),
s"Input tweet $tpTweet does not have specified fields $fields set"
)
val now = Time.now
val storedTweet = StorageConversions.toStoredTweetForFields(tpTweet, fields.toSet)
val tweetId = storedTweet.id
Stats.updatePerFieldQpsCounters("updateTweet", fields.map(_.id), 1, stats)
val (fieldIds, stitchesPerTweet) =
fields.map { field =>
val fieldId = field.id
val tweetKey = TweetKey.fieldKey(tweetId, fieldId)
val blob = storedTweet.getFieldBlob(fieldId).get
val value = ManhattanValue(TFieldBlobCodec.toByteBuffer(blob), Some(now))
val record = TweetManhattanRecord(tweetKey, value)
(fieldId, insert(record).liftToTry)
}.unzip
Stitch.collect(stitchesPerTweet).map { seqOfTries =>
val fieldkeyAndMhResults = fieldIds.zip(seqOfTries).toMap
// If even a single field was rate limited, we will send an overall OverCapacity TweetResponse
val wasRateLimited = fieldkeyAndMhResults.exists { keyAndResult =>
keyAndResult._2 match {
case Throw(e: DeniedManhattanException) => true
case _ => false
}
}
if (wasRateLimited) {
buildTweetOverCapacityResponse("updateTweets", tweetId, fieldkeyAndMhResults)
} else {
buildTweetResponse("updateTweets", tweetId, fieldkeyAndMhResults)
}
}
}
private def areAllFieldsDefined(tpTweet: Tweet, fields: Seq[Field]) = {
val storedTweet = StorageConversions.toStoredTweetForFields(tpTweet, fields.toSet)
fields.map(_.id).forall(storedTweet.getFieldBlob(_).isDefined)
}
}