Kafka Partition Keys

[JonathanWillis]

Hi, a few years back I rolled my own Kafka Producer/Consumer library for pulling messages from a topic and dispatching to Command/Event Handlers, and its all setup with 1 line, with the rest done by reflection at start up. I've been considering to swap it out for a more battle tested library, but as mass transit has gone commercial, I looked at Brighter and Wolverine. I've found that Brighter doesn't support multiple message types per topic, so its left me to investigate wolverine.

builder.Host.UseWolverine(opts =>
{
    // Kafka connection
    opts.UseKafka("localhost:29092")
        .ConfigureConsumers(c =>
        {
            c.AutoOffsetReset = Confluent.Kafka.AutoOffsetReset.Earliest;
        });

    opts.PublishAllMessages()
        .ToKafkaTopics();

    opts.ListenToKafkaTopic("order.commands")
        .ProcessInline();

    opts.ListenToKafkaTopic("order.events")
        .ProcessInline();
});

The issue I came across is the only way I could see that you can set partition key is via the Delivery options, such as

app.MapPost("/orders", async (
    CreateOrder command,
    IMessageBus bus) =>
{
    await bus.PublishAsync(command, new DeliveryOptions { PartitionKey = command.OrderId });
    return Results.Accepted($"/orders/{command.OrderId}");
});

I felt this was not ideal, and more suited to use case specific overrides.

I did try the MessagePartitioning but found this was internal, and had no effect on Kafka producers

opts.MessagePartitioning
        .ByMessage<CreateOrder>(x => x.OrderId.ToString())
        .ByMessage<OrderCreated>(x => x.OrderId.ToString())
        .ByMessage<CompleteOrder>(x => x.OrderId.ToString())
        .ByMessage<OrderCompleted>(x => x.OrderId.ToString());

I see we can add outgoing rules, so as an idea I put these in.

1. Custom class to build the partition
2. Add an interface marker and implement it.
3. Fluent interface like MessagePartitioning using lambda.

I'm torn between preferring 2 and 3, but throwing it out there as an idea.

    opts.PublishAllMessages()
        .ToKafkaTopics()
        .AddOutgoingRule(new PartitionRoutingByTypeRule());
        
public class PartitionRoutingByTypeRule : IEnvelopeRule
{
    public void Modify(Envelope envelope)
    {
        if (envelope.Message == null)
        {
            return;
        }

        envelope.PartitionKey ??= envelope.Message switch
        {
            CreateOrder createOrder => createOrder.OrderId.ToString(),
            OrderCreated orderCreated => orderCreated.OrderId.ToString(),
            CompleteOrder completeOrder => completeOrder.OrderId.ToString(),
            OrderCompleted orderCompleted => orderCompleted.OrderId.ToString(),
        };
    }
}

    opts.PublishAllMessages()
        .ToKafkaTopics()
        .AddOutgoingRule(new PartitionRoutingByInterfaceRule());

public class PartitionRoutingByInterfaceRule : IEnvelopeRule
{
    public void Modify(Envelope envelope)
    {
        if (envelope.Message is IPartitionedMessage partitioned)
        {
            envelope.PartitionKey ??= partitioned.PartitionKey;
        }
    }
}

public interface IPartitionedMessage
{
    string PartitionKey { get; }
}

[Topic("order.commands")]
public record CreateOrder(
    string OrderId,
    string CustomerId,
    decimal TotalAmount
) : IPartitionedMessage
{
    public string PartitionKey => OrderId;
}

    opts.PublishAllMessages()
        .ToKafkaTopics()
        .AddOutgoingRule(new PartitionRoutingByInterfaceRuleBuilder()
            .ByMessage<CreateOrder>(x => x.OrderId.ToString())
            .ByMessage<OrderCreated>(x => x.OrderId.ToString())
            .ByMessage<CompleteOrder>(x => x.OrderId.ToString())
            .ByMessage<OrderCompleted>(x => x.OrderId.ToString())
            .Build());

public class PartitionRoutingByInterfaceRuleBuilder
{
    private readonly Dictionary<Type, Func<object, string>> _rules = new();

    public PartitionRoutingByInterfaceRuleBuilder ByMessage<T>(Func<T, string> keySelector)
    {
        _rules[typeof(T)] = msg => keySelector((T)msg!);
        return this;
    }

    public IEnvelopeRule Build()
    {
        return new PartitionRoutingRule(_rules);
    }

    private class PartitionRoutingRule : IEnvelopeRule
    {
        private readonly Dictionary<Type, Func<object, string>> _rules;

        public PartitionRoutingRule(Dictionary<Type, Func<object, string>> rules)
        {
            _rules = rules;
        }

        public void Modify(Envelope envelope)
        {
            if (envelope.Message == null) return;
            if (envelope.PartitionKey != null) return;

            var type = envelope.Message.GetType();
            if (_rules.TryGetValue(type, out var func))
            {
                envelope.PartitionKey = func(envelope.Message);
            }
        }
    }
}