Implementation:Risingwavelabs Risingwave Load Generation Loop

Metadata

Property	Value
File	integration_tests/datagen/load_gen.go
Package	main
Language	Go
Lines	163
Category	Integration Test Data Generator Engine
Repository	https://github.com/risingwavelabs/risingwave

Overview

The Load Generation Loop is the central execution engine of the datagen tool. It connects data generators to sink outputs with rate control, serving as the bridge between the data production and data delivery layers. The module provides three key functions:

• createSink: Factory function that instantiates the appropriate sink implementation (Postgres, MySQL, Kafka, Pulsar, Kinesis, S3, or NATS) based on configuration.
• newGen: Factory function that instantiates the appropriate data generator (ad-click, ad-ctr, twitter, cdn-metrics, clickstream, ecommerce, delivery, livestream, nexmark, or compatible-data) based on the configured mode.
• generateLoad: The main orchestration loop that creates the sink, spawns the generator in a background goroutine, and consumes generated records at a rate-limited pace.

The loop uses the uber-go/ratelimit library for QPS throttling, periodically flushes the sink every 10 seconds, supports optional topic filtering, and can terminate after a configurable total number of events.

Code Reference

Source Location

integration_tests/datagen/load_gen.go

Signature

func createSink(ctx context.Context, cfg gen.GeneratorConfig) (sink.Sink, error)

Creates a sink instance based on the cfg.Sink field. Supports: "postgres", "mysql", "kafka", "pulsar", "kinesis", "s3", "nats". Returns an error for unrecognized sink types.

func newGen(cfg gen.GeneratorConfig) (gen.LoadGenerator, error)

Creates a generator instance based on the cfg.Mode field. Supports: "ad-click", "ad-ctr", "twitter", "cdn-metrics", "clickstream", "ecommerce", "delivery", "livestream" (alias "superset"), "nexmark", "compatible-data". Returns an error for unrecognized modes.

func spawnGen(ctx context.Context, cfg gen.GeneratorConfig, outCh chan<- sink.SinkRecord) (gen.LoadGenerator, error)

Creates a generator and immediately spawns its Load method in a new goroutine, returning the generator for topic information.

func generateLoad(ctx context.Context, cfg gen.GeneratorConfig) error

The main load generation loop. Orchestrates the full pipeline from sink creation through record consumption with rate limiting.

Import

import (
	"context"
	"datagen/ad_click"
	"datagen/ad_ctr"
	"datagen/cdn_metrics"
	"datagen/clickstream"
	"datagen/compatible_data"
	"datagen/delivery"
	"datagen/ecommerce"
	"datagen/gen"
	"datagen/livestream"
	"datagen/nexmark"
	"datagen/sink"
	"datagen/sink/kafka"
	"datagen/sink/kinesis"
	"datagen/sink/mysql"
	"datagen/sink/nats"
	"datagen/sink/postgres"
	"datagen/sink/pulsar"
	"datagen/sink/s3"
	"datagen/twitter"
	"fmt"
	"log"
	"time"

	"go.uber.org/ratelimit"
)

I/O Contract

Inputs

• ctx (context.Context): Controls the lifecycle of the generation loop. Cancellation stops all generation and consumption.
• cfg (gen.GeneratorConfig): Full configuration including sink type, generator mode, QPS limit, format, topic filter, and total event count.

Internal Data Flow

Step	Description
1	createSink instantiates the configured sink implementation
2	A buffered channel (outCh) of capacity 1000 is created for record passing
3	spawnGen creates the generator and starts Load() in a goroutine
4	sinkImpl.Prepare() is called with the generator's Kafka topics
5	The main select loop consumes records from outCh, applies rate limiting, writes to sink, and periodically flushes

Outputs

• Records are written to the configured sink via sinkImpl.WriteRecord().
• Progress is logged every 10 seconds showing total records sent and elapsed time.
• If cfg.PrintInsert is true, each record's SQL representation is printed to stdout.

Termination Conditions

• Context cancellation (ctx.Done())
• Total event count reached (cfg.TotalEvents > 0 && count >= cfg.TotalEvents)

Error Handling

• Sink creation errors are returned immediately.
• Generator creation errors are returned immediately.
• Sink write and flush errors cause the loop to return with the error.
• The sink is always closed via defer.

Usage Examples

Typical Invocation Flow

The generateLoad function is called from runCommand() in main.go:

func runCommand() error {
	terminateCh := make(chan os.Signal, 1)
	signal.Notify(terminateCh, os.Interrupt, syscall.SIGTERM)

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		<-terminateCh
		log.Println("Cancelled")
		cancel()
	}()
	return generateLoad(ctx, cfg)
}

Rate Limiting and Flushing

// Rate limiter initialized at configured QPS (no burst slack)
rl := ratelimit.New(cfg.Qps, ratelimit.WithoutSlack)

// In the main loop, after writing each record:
_ = rl.Take() // Blocks until the next allowed send time
count++

// Periodic flush every 10 seconds:
if time.Since(prevTime) >= 10*time.Second {
	log.Printf("Sent %d records in total (Elapsed: %s)", count, time.Since(initTime).String())
	prevTime = time.Now()
	if err := sinkImpl.Flush(ctx); err != nil {
		return err
	}
}

Topic Filtering

// Records can be filtered by topic before writing:
if cfg.Topic != "" && record.Topic() != cfg.Topic {
	continue
}

Related Pages

• Risingwavelabs_Risingwave_LoadGenerator_Interface -- The LoadGenerator interface and GeneratorConfig that drive this loop
• Risingwavelabs_Risingwave_Datagen_Sink_Interface -- The Sink and SinkRecord interfaces used for output
• Risingwavelabs_Risingwave_Datagen_CLI -- The CLI entry point that calls generateLoad
• Risingwavelabs_Risingwave_Ad_CTR_Generator -- One of the generator implementations dispatched by newGen