Skip to content

Commit 768bea4

Browse files
andrew-nowakandrew-nowak
committed
restructure the thrall stream graph to pass through source completion futures
1 parent bc4fbc1 commit 768bea4

2 files changed

Lines changed: 58 additions & 53 deletions

File tree

‎thrall/app/ThrallComponents.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ class ThrallComponents(context: Context) extends GridComponents(context, new Thr
7777
applicationLifecycle.addStopHook(() => {
7878
logger.info("thrall stream is closing; saying byebye")
7979
thrallStreamProcessor.killSwitch.shutdown()
80+
8081
streamRunning
8182
})
8283

‎thrall/app/lib/ThrallStreamProcessor.scala‎

Lines changed: 57 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,16 @@
11
package lib
22

3-
import java.time.Instant
4-
import org.apache.pekko.actor.ActorSystem
5-
import org.apache.pekko.stream.scaladsl.{GraphDSL, Keep, MergePreferred, MergePrioritized, Sink, Source}
6-
import org.apache.pekko.stream.{KillSwitch, KillSwitches, Materializer, SourceShape}
7-
import org.apache.pekko.{Done, NotUsed}
83
import com.gu.kinesis.KinesisRecord
94
import com.gu.mediaservice.lib.DateTimeUtils
10-
import com.gu.mediaservice.lib.aws.UpdateMessage
115
import com.gu.mediaservice.lib.logging._
12-
import com.gu.mediaservice.model.{ExternalThrallMessage, InternalThrallMessage, ThrallMessage}
13-
import lib.kinesis.{MessageTranslator, ThrallEventConsumer}
6+
import com.gu.mediaservice.model.{ExternalThrallMessage, ThrallMessage}
7+
import lib.kinesis.ThrallEventConsumer
8+
import org.apache.pekko.Done
9+
import org.apache.pekko.actor.ActorSystem
10+
import org.apache.pekko.stream.scaladsl.{GraphDSL, MergePreferred, Source}
11+
import org.apache.pekko.stream.{FlowShape, KillSwitches, Materializer}
1412

13+
import java.time.Instant
1514
import scala.concurrent.{ExecutionContextExecutor, Future}
1615
import scala.util.{Failure, Success}
1716

@@ -55,67 +54,72 @@ class ThrallStreamProcessor(
5554
actorSystem: ActorSystem
5655
) extends GridLogging {
5756

57+
import GraphDSL.Implicits._
58+
59+
private def waitForBoth(f1: Future[Done], f2: Future[Done]): Future[Done] = {
60+
f1.zip(f2).map(_ => Done)
61+
}
62+
5863
val killSwitch = KillSwitches.shared("thrall-kill-switch")
5964

6065
implicit val mat: Materializer = Materializer.matFromSystem(actorSystem)
6166
implicit val dispatcher: ExecutionContextExecutor = actorSystem.getDispatcher
6267

63-
val mergedKinesisSource: Source[TaggedRecord[ThrallMessage], NotUsed] = Source.fromGraph(GraphDSL.create() { implicit graphBuilder =>
64-
import GraphDSL.Implicits._
68+
val uiRecordSource = uiSource.via(killSwitch.flow).map(kinesisRecord =>
69+
TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, UiPriority, kinesisRecord.markProcessed))
6570

66-
val uiRecordSource = uiSource.via(killSwitch.flow).map(kinesisRecord =>
67-
TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, UiPriority, kinesisRecord.markProcessed))
71+
val automationRecordSource = automationSource.via(killSwitch.flow).map(kinesisRecord =>
72+
TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, AutomationPriority, kinesisRecord.markProcessed))
6873

69-
val automationRecordSource = automationSource.via(killSwitch.flow).map(kinesisRecord =>
70-
TaggedRecord(kinesisRecord.data.toArray, kinesisRecord.approximateArrivalTimestamp, AutomationPriority, kinesisRecord.markProcessed))
74+
val migrationMessagesSource = migrationSource.via(killSwitch.flow).map { case MigrationRecord(internalThrallMessage, time) =>
75+
TaggedRecord(internalThrallMessage, time, MigrationPriority, () => {})
76+
}
7177

72-
val migrationMessagesSource = migrationSource.via(killSwitch.flow).map { case MigrationRecord(internalThrallMessage, time) =>
73-
TaggedRecord(internalThrallMessage, time, MigrationPriority, () => {})
78+
val uiAndAutomationRecordsMerge: Source[TaggedRecord[Array[Byte]], Future[Done]] = {
79+
val mergeGraph = GraphDSL.createGraph(automationRecordSource) { implicit b => r =>
80+
val merge = b.add(MergePreferred[TaggedRecord[Array[Byte]]](1, eagerComplete = false))
81+
r ~> merge.in(0)
82+
FlowShape(merge.in(1), merge.out)
7483
}
84+
uiRecordSource.viaMat(mergeGraph)(waitForBoth)
85+
}
7586

76-
// merge together ui and automation kinesis records
77-
val uiAndAutomationRecordsMerge = graphBuilder.add(MergePreferred[TaggedRecord[Array[Byte]]](1))
78-
uiRecordSource ~> uiAndAutomationRecordsMerge.preferred
79-
automationRecordSource ~> uiAndAutomationRecordsMerge.in(0)
80-
81-
// parse the kinesis records into thrall update messages (dropping those that fail)
82-
val uiAndAutomationMessagesSource: PortOps[TaggedRecord[ExternalThrallMessage]] =
83-
uiAndAutomationRecordsMerge.out
84-
.map { taggedRecord =>
85-
val parsedRecord = ThrallEventConsumer
86-
.parseRecord(taggedRecord.payload, taggedRecord.arrivalTimestamp)
87-
.map(
88-
message => taggedRecord.copy(payload = message)
89-
)
90-
// If we failed to parse the record (Left), we'll drop it below because we can't process it.
91-
// However we still need to mark the record as processed, otherwise the kinesis stream can't progress
92-
// and checkpoint will be stuck at this message forevermore.
93-
parsedRecord.left.foreach(_ => taggedRecord.markProcessed())
94-
parsedRecord
95-
}
96-
// drop unparseable records
97-
.collect {
98-
case Right(taggedRecord) => taggedRecord
99-
}
100-
101-
// merge in the re-ingestion source (preferring ui/automation)
102-
val mergePreferred = graphBuilder.add(MergePreferred[TaggedRecord[ThrallMessage]](1))
103-
uiAndAutomationMessagesSource ~> mergePreferred.preferred
104-
migrationMessagesSource ~> mergePreferred.in(0)
105-
106-
SourceShape(mergePreferred.out)
107-
})
108-
109-
def createStream(): Source[(TaggedRecord[ThrallMessage], Stopwatch, ThrallMessage), NotUsed] = {
110-
mergedKinesisSource.mapAsync(1) { result =>
87+
val uiAndAutomationMessagesSource: Source[TaggedRecord[ExternalThrallMessage], Future[Done]] =
88+
uiAndAutomationRecordsMerge
89+
.map { taggedRecord =>
90+
val parsedRecord = ThrallEventConsumer
91+
.parseRecord(taggedRecord.payload, taggedRecord.arrivalTimestamp)
92+
.map(
93+
message => taggedRecord.copy(payload = message)
94+
)
95+
// If we failed to parse the record (Left), we'll drop it below because we can't process it.
96+
// However we still need to mark the record as processed, otherwise the kinesis stream can't progress
97+
// and checkpoint will be stuck at this message forevermore.
98+
parsedRecord.left.foreach(_ => taggedRecord.markProcessed())
99+
parsedRecord
100+
}.collect {
101+
// drop unparseable records
102+
case Right(taggedRecord) => taggedRecord
103+
}
104+
105+
val allSourcesMerged: Source[TaggedRecord[ThrallMessage], Future[Done]] = {
106+
val mergeGraph = GraphDSL.createGraph(migrationMessagesSource) { implicit b => r =>
107+
val merge = b.add(MergePreferred[TaggedRecord[ThrallMessage]](1, eagerComplete = false))
108+
r ~> merge.in(0)
109+
FlowShape(merge.in(1), merge.out)
110+
}
111+
uiAndAutomationMessagesSource.viaMat(mergeGraph)(waitForBoth)
112+
}
113+
114+
def createStream(): Source[(TaggedRecord[ThrallMessage], Stopwatch, ThrallMessage), Future[Done]] = {
115+
allSourcesMerged.mapAsync(1) { result =>
111116
val stopwatch = Stopwatch.start
112117
consumer.processMessage(result.payload)
113118
.recover { case _ => () }
114119
.map(_ => (result, stopwatch, result.payload))
115120
}
116-
117-
118121
}
122+
119123
def run(): Future[Done] = {
120124
val stream = this.createStream().runForeach {
121125
case (taggedRecord, stopwatch, _) =>

0 commit comments

Comments
 (0)