river
Contents
River
简单来说，River 是一个专为 Go 设计的作业队列库，但它不是那种「又一个轮子」的存在。它最大的特点是什么？把队列和 PostgreSQL深度绑定。
传统方案里，我们通常会用 Redis、RabbitMQ 这些专门的消息中间件来做队列。这没问题，但也带来了一个老大难问题：分布式事物
想象一个场景：你在处理订单支付，需要先更新数据库的订单状态，再往队列里扔一个「发送支付成功短信」的任务。如果数据库更新成功了，但往队列里塞任务失败了呢？或者反过来？这就是分布式事务的经典困境。
River 的解决方案很直接：既然都用 PostgreSQL 存数据，那队列也放 PostgreSQL 里不就完了？ 一个事务搞定所有事情，要么全成功，要么全回滚。是不是有种「大道至简」的感觉？
package main
import (
"context"
"database/sql"
"errors"
"fmt"
"log"
"os"
"os/signal"
"syscall"
"time"
"github.com/riverqueue/river"
"github.com/riverqueue/river/riverdriver/riverdatabasesql"
"github.com/riverqueue/river/rivermigrate"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
const (
OrderStatusPending = "pending"
OrderStatusProcessing = "processing"
OrderStatusCompleted = "completed"
OrderQueue = "order"
)
// Order 是业务表，由 GORM 管理。
type Order struct {
ID int64 `gorm:"primaryKey;autoIncrement"`
OrderNo string `gorm:"size:64;not null;uniqueIndex"`
Product string `gorm:"size:128;not null"`
Quantity int `gorm:"not null"`
Status string `gorm:"size:32;not null;index"`
CreatedAt time.Time `gorm:"not null"`
UpdatedAt time.Time `gorm:"not null"`
}
func (Order) TableName() string {
return "orders"
}
// ProcessOrderArgs 是 River 任务参数。
// River 会将其序列化成 JSON 存入 river_job 表。
type ProcessOrderArgs struct {
OrderID int64 `json:"order_id"`
}
// Kind 必须保持稳定，不能随意修改。
// River 根据 Kind 找到对应的 Worker。
func (ProcessOrderArgs) Kind() string {
return "process_order"
}
// InsertOpts 设置任务默认参数。
func (ProcessOrderArgs) InsertOpts() river.InsertOpts {
return river.InsertOpts{
Queue: OrderQueue,
MaxAttempts: 5,
}
}
// ProcessOrderWorker 负责处理订单任务。
type ProcessOrderWorker struct {
river.WorkerDefaults[ProcessOrderArgs]
db *gorm.DB
}
func NewProcessOrderWorker(db *gorm.DB) *ProcessOrderWorker {
return &ProcessOrderWorker{
db: db,
}
}
// Work 执行具体任务。
// 返回 nil 表示任务成功；返回 error 表示任务失败并触发重试。
func (w *ProcessOrderWorker) Work(
ctx context.Context,
job *river.Job[ProcessOrderArgs],
) error {
orderID := job.Args.OrderID
log.Printf(
"开始处理订单任务: job_id=%d order_id=%d attempt=%d",
job.ID,
orderID,
job.Attempt,
)
var order Order
err := w.db.WithContext(ctx).
First(&order, orderID).
Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
// 订单已经不存在，继续重试通常没有意义。
// 返回 nil，将任务标记为完成。
log.Printf("订单不存在，忽略任务: order_id=%d", orderID)
return nil
}
return fmt.Errorf("查询订单失败: %w", err)
}
// Worker 必须尽量保证幂等。
// River 任务可能因为进程退出、网络异常等原因被再次执行。
if order.Status == OrderStatusCompleted {
log.Printf("订单已经完成，跳过重复处理: order_id=%d", orderID)
return nil
}
result := w.db.WithContext(ctx).
Model(&Order{}).
Where(
"id = ? AND status != ?",
orderID,
OrderStatusCompleted,
).
Update("status", OrderStatusProcessing)
if result.Error != nil {
return fmt.Errorf("更新订单为处理中失败: %w", result.Error)
}
// 模拟耗时业务，并响应任务取消。
select {
case <-time.After(2 * time.Second):
case <-ctx.Done():
return ctx.Err()
}
result = w.db.WithContext(ctx).
Model(&Order{}).
Where("id = ?", orderID).
Update("status", OrderStatusCompleted)
if result.Error != nil {
return fmt.Errorf("更新订单为已完成失败: %w", result.Error)
}
log.Printf(
"订单处理完成: job_id=%d order_id=%d",
job.ID,
orderID,
)
return nil
}
// CreateOrderAndEnqueue 在同一个事务中：
// 1. 使用 GORM 创建订单；
// 2. 使用 River 创建异步任务；
// 3. 一起提交或一起回滚。
func CreateOrderAndEnqueue(
ctx context.Context,
db *gorm.DB,
riverClient *river.Client[*sql.Tx],
product string,
quantity int,
) (*Order, error) {
tx := db.WithContext(ctx).Begin()
if tx.Error != nil {
return nil, fmt.Errorf("开启事务失败: %w", tx.Error)
}
committed := false
defer func() {
if !committed {
_ = tx.Rollback().Error
}
}()
order := &Order{
OrderNo: fmt.Sprintf("ORD-%d", time.Now().UnixNano()),
Product: product,
Quantity: quantity,
Status: OrderStatusPending,
}
if err := tx.Create(order).Error; err != nil {
return nil, fmt.Errorf("创建订单失败: %w", err)
}
// GORM 开启事务后，底层连接是 *sql.Tx。
// 将同一个 *sql.Tx 交给 River，实现事务性投递任务。
sqlTx, ok := tx.Statement.ConnPool.(*sql.Tx)
if !ok {
return nil, errors.New("无法从 GORM 事务获取 *sql.Tx")
}
insertResult, err := riverClient.InsertTx(
ctx,
sqlTx,
ProcessOrderArgs{
OrderID: order.ID,
},
nil,
)
if err != nil {
return nil, fmt.Errorf("创建 River 任务失败: %w", err)
}
if err := tx.Commit().Error; err != nil {
return nil, fmt.Errorf("提交事务失败: %w", err)
}
committed = true
log.Printf(
"订单及任务创建成功: order_id=%d job_id=%d",
order.ID,
insertResult.Job.ID,
)
return order, nil
}
func migrateRiver(ctx context.Context, sqlDB *sql.DB) error {
migrator, err := rivermigrate.New(
riverdatabasesql.New(sqlDB),
nil,
)
if err != nil {
return fmt.Errorf("创建 River Migrator 失败: %w", err)
}
result, err := migrator.Migrate(
ctx,
rivermigrate.DirectionUp,
nil,
)
if err != nil {
return fmt.Errorf("执行 River 数据库迁移失败: %w", err)
}
if len(result.Versions) == 0 {
log.Println("River 数据库结构已经是最新版本")
return nil
}
for _, version := range result.Versions {
log.Printf(
"River 数据库迁移完成: version=%d name=%s duration=%s",
version.Version,
version.Name,
version.Duration,
)
}
return nil
}
func openDatabase(
ctx context.Context,
databaseURL string,
) (sqlDB *sql.DB, gormDB *gorm.DB, err error) {
// sqlDB.SetMaxOpenConns(30)
// sqlDB.SetMaxIdleConns(10)
// sqlDB.SetConnMaxLifetime(30 * time.Minute)
// sqlDB.SetConnMaxIdleTime(10 * time.Minute)
// if err := sqlDB.PingContext(ctx); err != nil {
// _ = sqlDB.Close()
// return nil, nil, fmt.Errorf("连接数据库失败: %w", err)
// }
gormDB, err = gorm.Open(
postgres.Open(databaseURL),
&gorm.Config{
Logger: logger.Default.LogMode(logger.Info),
},
)
if err != nil {
return
}
sqlDB, err = gormDB.DB()
if err != nil {
return
}
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(50)
sqlDB.SetConnMaxLifetime(time.Hour)
if err := sqlDB.PingContext(ctx); err != nil {
_ = sqlDB.Close()
return nil, nil, fmt.Errorf("连接数据库失败: %w", err)
}
return sqlDB, gormDB, nil
}
func main() {
databaseURL := os.Getenv("DATABASE_URL")
if databaseURL == "" {
databaseURL = "postgres://kk:kkm@localhost:5432/river_demo?sslmode=disable"
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
sqlDB, gormDB, err := openDatabase(ctx, databaseURL)
if err != nil {
log.Fatal(err)
}
defer sqlDB.Close()
// 必须先创建和升级 River 所需的数据库表。
if err := migrateRiver(ctx, sqlDB); err != nil {
log.Fatalf("River 数据库迁移失败: %v", err)
}
// 这里只迁移业务表。
// River 自己的表应通过 river migrate-up 创建。
if err := gormDB.AutoMigrate(&Order{}); err != nil {
log.Fatalf("迁移业务表失败: %v", err)
}
workers := river.NewWorkers()
// 注册任务 Worker。
river.AddWorker(
workers,
NewProcessOrderWorker(gormDB),
)
// riverdatabasesql 允许 River 与 GORM 共用 sql.DB/sql.Tx。
riverClient, err := river.NewClient(
riverdatabasesql.New(sqlDB),
&river.Config{
Queues: map[string]river.QueueConfig{
OrderQueue: {
MaxWorkers: 5,
},
},
Workers: workers,
SoftStopTimeout: 15 * time.Second,
},
)
if err != nil {
log.Fatalf("创建 River Client 失败: %v", err)
}
appCtx, stop := signal.NotifyContext(
ctx,
syscall.SIGINT,
syscall.SIGTERM,
)
defer stop()
// 启动 River Worker。
if err := riverClient.Start(appCtx); err != nil {
log.Fatalf("启动 River Client 失败: %v", err)
}
log.Println("River Worker 已启动")
// 创建一条示例订单，并在同一事务中投递任务。
order, err := CreateOrderAndEnqueue(
appCtx,
gormDB,
riverClient,
"MacBook Pro",
1,
)
if err != nil {
log.Fatalf("创建示例订单失败: %v", err)
}
log.Printf(
"示例订单已创建: id=%d order_no=%s status=%s",
order.ID,
order.OrderNo,
order.Status,
)
// 等待 SIGINT 或 SIGTERM。
<-appCtx.Done()
log.Println("正在停止 River Worker")
// Start 使用的 context 被取消后，River 会进入优雅停止流程。
<-riverClient.Stopped()
log.Println("River Worker 已停止")
}