Implementation:Risingwavelabs Risingwave LoadGenerator Interface

Metadata

Property	Value
File	integration_tests/datagen/gen/generator.go
Package	gen
Language	Go
Lines	108
Category	Integration Test Data Generator Framework
Repository	https://github.com/risingwavelabs/risingwave

Overview

The LoadGenerator Interface module defines the core data generation framework for RisingWave's integration test datagen tool. It provides:

• The LoadGenerator interface that all specific data generators (ad-click, clickstream, ecommerce, etc.) must implement.
• The GeneratorConfig struct that holds all configuration parameters including sink settings, QPS throttling, data generation mode, and format options.
• Timestamp layout constants (RwTimestampNaiveLayout and RwTimestamptzLayout) that standardize timestamp formatting across all generators.
• The RandDist interface and two implementations (UniformDist and PoissonDist) that provide configurable random distributions for realistic data generation.

The Poisson distribution (default) produces more realistic data with lower tail probabilities, while the Uniform distribution (activated via HeavyTail config flag) provides equal probability across the range.

Code Reference

Source Location

integration_tests/datagen/gen/generator.go

Signature

type LoadGenerator interface {
	KafkaTopics() []string
	Load(ctx context.Context, outCh chan<- sink.SinkRecord)
}

LoadGenerator is the core interface that all data generators must satisfy:

• KafkaTopics(): Returns the list of Kafka topic names that this generator writes to.
• Load(): Generates data records and sends them to the output channel. The method should respect context cancellation for graceful shutdown.

type GeneratorConfig struct {
	Postgres postgres.PostgresConfig
	Mysql    mysql.MysqlConfig
	Kafka    kafka.KafkaConfig
	Pulsar   pulsar.PulsarConfig
	Kinesis  kinesis.KinesisConfig
	S3       s3.S3Config
	Nats     nats.NatsConfig

	PrintInsert bool
	Mode        string
	Sink        string
	Qps         int
	HeavyTail   bool
	Format      string
	Topic       string
	TotalEvents int64
}

type RandDist interface {
	Rand(max float64) float64
}

func NewRandDist(cfg GeneratorConfig) RandDist

Import

import (
	"context"
	"datagen/sink"
	"datagen/sink/kafka"
	"datagen/sink/kinesis"
	"datagen/sink/mysql"
	"datagen/sink/nats"
	"datagen/sink/postgres"
	"datagen/sink/pulsar"
	"datagen/sink/s3"
	"time"

	"gonum.org/v1/gonum/stat/distuv"
)

I/O Contract

GeneratorConfig Fields

Field	Type	Description
Postgres	postgres.PostgresConfig	PostgreSQL sink connection configuration
Mysql	mysql.MysqlConfig	MySQL sink connection configuration
Kafka	kafka.KafkaConfig	Kafka sink configuration (brokers, topic recreation)
Pulsar	pulsar.PulsarConfig	Pulsar sink configuration
Kinesis	kinesis.KinesisConfig	Kinesis sink configuration (region, endpoint, stream)
S3	s3.S3Config	S3 sink configuration (region, bucket, endpoint)
Nats	nats.NatsConfig	NATS sink configuration (URL, JetStream)
PrintInsert	bool	Whether to print every event's SQL insert to stdout
Mode	string	Data generation mode (e.g., "ad-ctr", "clickstream", "ecommerce")
Sink	string	Target sink type (e.g., "kafka", "postgres", "s3")
Qps	int	Throttled requests per second
HeavyTail	bool	If true, use Uniform distribution; otherwise use Poisson
Format	string	Record format for message queues ("json", "protobuf", "avro")
Topic	string	Optional topic filter; if set, only matching records are emitted
TotalEvents	int64	Total events to generate; 0 means run indefinitely

Timestamp Constants

const RwTimestampNaiveLayout = time.DateTime      // "2006-01-02 15:04:05"
const RwTimestamptzLayout = time.RFC3339           // "2006-01-02T15:04:05Z07:00"

Random Distribution Behavior

• PoissonDist (default): Uses Poisson distribution with lambda = max/2. Produces values concentrated around the center with lower tail probability, simulating more realistic data patterns.
• UniformDist: Uses Uniform distribution over [0, max]. All values are equally likely. Activated when HeavyTail is true.

Both distributions lazily cache their internal distribution objects keyed by the max parameter.

Usage Examples

Implementing a LoadGenerator

package mygen

import (
	"context"
	"datagen/gen"
	"datagen/sink"
)

type myGenerator struct{}

func NewMyGenerator() gen.LoadGenerator {
	return &myGenerator{}
}

func (g *myGenerator) KafkaTopics() []string {
	return []string{"my_topic"}
}

func (g *myGenerator) Load(ctx context.Context, outCh chan<- sink.SinkRecord) {
	for {
		select {
		case <-ctx.Done():
			return
		case outCh <- generateMyRecord():
		}
	}
}

Using the Random Distribution

cfg := gen.GeneratorConfig{HeavyTail: false}
dist := gen.NewRandDist(cfg) // Returns PoissonDist

value := dist.Rand(100.0) // Random value in [0, 100] with Poisson distribution

Formatting Timestamps

import "datagen/gen"
import "time"

// Timestamp with timezone (RFC3339)
tstz := time.Now().Format(gen.RwTimestamptzLayout)
// e.g., "2024-01-15T10:30:00Z"

// Naive timestamp (no timezone)
ts := time.Now().Format(gen.RwTimestampNaiveLayout)
// e.g., "2024-01-15 10:30:00"

Related Pages

• Risingwavelabs_Risingwave_Ad_CTR_Generator -- An example generator implementing the LoadGenerator interface
• Risingwavelabs_Risingwave_Datagen_Sink_Interface -- The SinkRecord and Sink interfaces consumed by generators
• Risingwavelabs_Risingwave_Load_Generation_Loop -- The orchestration loop that uses GeneratorConfig and spawns generators
• Risingwavelabs_Risingwave_Datagen_CLI -- The CLI that populates GeneratorConfig from command-line arguments