The Partition Key We Thought We Couldn’t Change
Was partitioning by customer really the only option?
Part 3 of 3 | Previously: Reconstructing Identity in AWS CloudTrail with Apache Flink | A hot partition in Flink - how bad could it be?
In the first part, I introduced the requirement to trace actions performed through assumed roles back to their original actors. We considered partitioning by (customerId, accessKeyId), which would help distribute the load more evenly across subtasks. But then role chaining ruined the party, and we circled back to partitioning by customerId, accepting the risk of hot partitions.
In Part 2, we saw what that partitioning decision cost us. All the events of a large customer, along with their accumulated state, had to be handled by a single subtask. The resulting backpressure sent us down a long path of performance optimizations. After trying to address the problem without sacrificing enrichment coverage, we ultimately had to shorten the lookupState TTL for those customers. We could afford that trade-off only because their credentials were typically used soon after issuance. Had they been used hours later, we’d have sacrificed far more enrichment coverage.
But perhaps we’d been trying too hard to make the wrong partitioning strategy work.
We’d chosen customerId because resolving a role chain required access to the identity mappings created by its earlier events. Keeping the entire chain on one subtask guaranteed that access. But what if we could make those mappings available without keeping the events together?
External Identity Mapping
One possibility was to move the identity mappings out of Flink’s local state and into a shared external store, accessible to every subtask.
An AssumeRole processed on one subtask could register a mapping that another subtask would later use to resolve an action. We could then return to partitioning by (customerId, accessKeyId).
The resolution logic would remain similar to what we built in part 1. Whenever an AssumeRole event arrives, we register the newly issued accessKeyId in the external store. If the caller is the original actor, we map the new key to its identity. If the caller is acting through an assumed role, we map it to the accessKeyId used to make the call.
When an action event arrives, we start with its accessKeyId and follow these mappings back through the chain until we reach the original actor. If a mapping is missing along the way, we buffer the action until it becomes available.
But there’s a catch. Consider the three events above arriving in reverse order. The CreateDBSnapshot action reaches subtask 1 first and gets buffered. The second AssumeRole arrives at the same subtask and registers ASIA2 ->ASIA1 in the external store. Finally, Alice’s AssumeRole arrives at subtask 0, registering ASIA1 -> Alice. The store now has everything needed to resolve the buffered CreateDBSnapshot. But nothing notifies subtask 1 that the missing mapping is now available. The action remains buffered, even though its original actor can now be identified.
We could periodically retry the missing lookups, but that means making repeated requests for mappings that might not exist yet. Longer retry intervals would reduce those requests, but at the cost of higher enrichment latency.
Alternatively, we could use notifications to trigger retries whenever a missing mapping becomes available.
If we were already introducing a messaging system, though, why limit it to notifying subtasks about changes in an external store? We could use those messages to carry the resolved mappings themselves, keeping lookupState inside Flink.
toResolve as a stream
In our original design, toResolve was a local queue in each subtask helping us to resolve chains. We would pop an accessKeyId, find the events waiting on it in bufferState, and resolve them. If one of those events was an intermediate AssumeRole, resolving it would register a new accessKeyId and push it into the queue. The process repeated until the queue was empty.
What if we replaced that queue with an external, replayable messaging stream? Here’s what that would look like:
Our Flink job would now consume two streams: the original CloudTrail events and the resolution events it produces itself. Instead of partitioning by customerId, we’d partition by (customerId, accessKeyId), distributing a customer’s events more evenly across subtasks. This time, CloudTrail events would be routed by the caller’s accessKeyId, including intermediate AssumeRole events. An AssumeRole performed by an already identifiable caller can be resolved directly, producing the first mapping in the chain. Resolution messages, on the other hand, would be routed by the newly issued accessKeyId.
Get Snir Zur’s stories in your inbox
Join Medium for free to get updates from this writer.
Each subtask would maintain its own lookupState and bufferState. When an event arrives, we’d look up its caller. If the identity isn’t known yet, we’d buffer the event. Otherwise, we resolve it immediately.
Whenever an AssumeRole is resolved, the enricher would sink the newly issued key and its resolved identity information to the external stream. The message would then be consumed by the same Flink job and routed by its accessKeyId to the responsible subtask.
That subtask would register the mapping in lookupState and resolve any events waiting in bufferState. If the mapping arrived before those events, it would simply be available when they arrived.
Seeing It in Action
Let’s revisit the same four events from part 1, arriving in reverse order. This time, ASIA1 belongs to subtask 0 and ASIA2 to subtask 1.
As before, events 2, 3, and 4 arrive before the AssumeRole that explains their caller, so they get buffered. The difference is that events 2 and 3 wait on subtask 0, while event 4 waits on subtask 1.
Finally, event 1 arrives: Alice assumes role-1 and receives ASIA1. Since Alice’s identity is already known, the enricher sinks a resolution message for ASIA1 to the external stream:
{
"customerId": "customer-123",
"accessKeyId": "ASIA1",
"originalActor": "Alice",
"roleChain": []
}The event flows back into Flink and reaches subtask 0, which registers ASIA1, releases event 2, and resolves event 3, the intermediate AssumeRole that issued ASIA2.
This time, however, ASIA2 belongs to another subtask. So instead of registering it locally, subtask 0 sinks another resolution message to the external stream, now carrying ASIA2 and Alice’s identity.
{
"customerId": "customer-123",
"accessKeyId": "ASIA2",
"originalActor": "Alice",
"roleChain": ["role-1"]
}Subtask 1 consumes that event, registers ASIA2, and releases event 4, which had been waiting there since the beginning.
The same chain that once seemed to force us to partition by customerId could now be resolved across different subtasks.
Looking Back
We’d treated partitioning by customerId as an unavoidable consequence of role chaining. But it turns out, keeping the pieces together isn’t always the only way to put them together.
We started by tracing actions back to their original actors. We ended up tracing our architecture back to its original assumption.
Not bad for a simple data enrichment.
