prnvbn

Reusable and Testable Go and Rust Protobuf Kafka Consumers

Introduction

I have had to do this many times: consume messages from a Kafka1 topic, decode protobuf payloads, and run some processing on top of them. The use cases vary. Sometimes I need to track the latest value. Sometimes I am writing an event log to a database. Sometimes I am enriching data before sending it somewhere else.

Across microservices, the pattern repeats itself. Each service ends up rewriting the same consumer setup. New topic, new protobuf type, and a bit of service specific business logic.

Did somebody say over abstraction?

This can easily become an over abstraction if you try to share it across multiple microservices. If that works for you, great but my goal here is simpler.

I wanted to achieve two things:

  • reduce duplication within a single service that needs to consume multiple topics
  • keep business logic separate from the underlying consumption logic (which also makes testing easier)

Even if the abstraction never gets shared beyond one service, it still pays off.

Rust Consumer

The Rust implementation uses rdkafka for Kafka consumption and prost for protobuf decoding. The Consumer struct is generic over the protobuf message type T and the processor P, which keeps things type safe without relying on dynamic traits.

use anyhow::{Context, Result};
use futures::StreamExt;
use prost::Message as ProstMessage;
use rdkafka::{
    consumer::{Consumer as KafkaConsumer, StreamConsumer},
    ClientConfig, Message as KafkaMessage,
};
use std::env;
use std::marker::PhantomData;
use tracing::{debug, error, info};

pub struct Consumer<T, P>
where
    T: ProstMessage + Default + Send + Sync + 'static,
    P: MessageProcessor<T>,
{
    consumer: StreamConsumer,
    processor: P,
    _pd: PhantomData<T>,
}

impl<T, P> Consumer<T, P>
where
    T: ProstMessage + Default + Send + Sync + 'static,
    P: MessageProcessor<T>,
{
    pub fn new(processor: P) -> Result<Self> {
        let mut ccfg = ClientConfig::new();
        ccfg.set("bootstrap.servers", env::var("KAFKA_SERVER")?)
            .set("group.id", env::var("KAFKA_GROUP")?);

        let consumer = ccfg.create().context("Failed to create stream consumer")?;

        Ok(Self {
            consumer,
            processor,
            _pd: PhantomData,
        })
    }

    pub async fn start(self) -> Result<()> {
        info!("Starting consumer");

        let topic = env::var("KAFKA_TOPIC").context("KAFKA_TOPIC not set")?;

        self.consumer.subscribe(&[&topic])?;

        let mut stream = self.consumer.stream();
        while let Some(result) = stream.next().await {
            let message = match result {
                Ok(message) => message,
                Err(e) => {
                    error!(error = %e, "Error polling consumer");
                    continue;
                }
            };

            let bytes = match message.payload() {
                Some(bytes) => bytes,
                None => {
                    error!("No payload found");
                    continue;
                }
            };

            let message = match T::decode(bytes) {
                Ok(message) => message,
                Err(e) => {
                    error!(error = %e, "Error decoding message");
                    continue;
                }
            };

            match self.processor.process(message) {
                Ok(_) => {
                    debug!("Message processed successfully");
                }
                Err(e) => {
                    error!(error = %e, "Error processing message");
                }
            }
        }

        Ok(())
    }
}
pub trait MessageProcessor<M>
where
    M: ProstMessage + Default + Send + Sync + 'static,
{
    fn process(&self, message: M) -> Result<()>;
}

Go Consumer

The Go implementation uses franz-go for Kafka consumption and the standard google.golang.org/protobuf package for protobuf decoding. The Poller struct uses generics to keep message handling type safe. The Msg constraint makes sure only protobuf types are used, and the Processor interface keeps the business logic separate and easy to plug in.

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"

	"github.com/twmb/franz-go/pkg/kgo"
	"google.golang.org/protobuf/proto"
)

type Poller[T any, M Msg[T]] struct {
	cl        *kgo.Client
	processor Processor[T, M]
}

type Msg[T any] interface {
	*T
	proto.Message
}

func NewPoller[T any, M Msg[T]](processor Processor[T, M]) (*Poller[T, M], error) {
	kafkaOpts := []kgo.Opt{
		kgo.SeedBrokers(os.Getenv("KAFKA_SERVER")),
		kgo.ConsumerGroup(os.Getenv("KAFKA_GROUP_ID")),
		kgo.ConsumeTopics(os.Getenv("KAFKA_TOPIC")),
	}

	cl, err := kgo.NewClient(kafkaOpts...)
	if err != nil {
		return nil, fmt.Errorf("failed to create kafka client: %w", err)
	}

	return &Poller[T, M]{cl: cl, processor: processor}, nil
}

func (p *Poller[T, M]) Run(ctx context.Context) error {
	defer p.cl.Close()
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		fetches := p.cl.PollFetches(ctx)

		iter := fetches.RecordIter()
		for !iter.Done() {
			record := iter.Next()

			data := M(new(T))
			err := proto.Unmarshal(record.Value, data)
			if err != nil {
				slog.Error("failed to unmarshal record", "error", err)
				continue
			}

			if p.processor != nil {
				slog.Debug("processor is nil")
				continue
			}

			err = p.processor.Process(ctx, data)
			if err != nil {
				slog.Error("failed to process data", "error", err)
				continue
			}
		}
	}
}

type Processor[T any, M Msg[T]] interface {
	Process(ctx context.Context, data M) error
}

Notes


  1. I prefer using Redpanda, not Apache Kafka ↩︎