chore: 初始开源发布
iotlcm-worker-svr:异步执行中枢(RabbitMQ 消费者 + 全量同步编排)。 - Go module:iotlcm-worker-svr → git.5iot.com/iotlcm-public/iotlcm-worker-svr(214 文件) - 凭据与内部标识脱敏:MySQL/Redis 凭据、私网与生产 IP、内网域名、 库名 xiujie_admin → iotlcm、哨兵名占位 - 159 处真实手机号 → 合法 11 位编造号(保留号段与后 4 位,维持互异, 不破坏「校验 MSISDN 格式」等测试断言) - 移除内部资产:CI 触发器、Jenkins 触发测试文档、预演门禁与镜像 digest 记录 - Dockerfile 改用公共基础镜像(golang:1.25-alpine → alpine:3.23), 移除写死的 git 凭据 - README 重写(架构/全量同步模型/配置/已知失败),补 Apache-2.0 LICENSE,LF 行尾 + .gitattributes 已知:3 个测试包在源仓同一提交上同样失败(client/upstream、 service/allocate、service/inbound),属既有缺陷,本仓未引入亦未修复。
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
package rabbitmq
|
||||
|
||||
import "context"
|
||||
|
||||
type DeliveryMetadata struct {
|
||||
DeliveryTag uint64
|
||||
Redelivered bool
|
||||
Consumer string
|
||||
Queue string
|
||||
Headers map[string]any
|
||||
}
|
||||
|
||||
type deliveryMetadataContextKey struct{}
|
||||
|
||||
func WithDeliveryMetadata(ctx context.Context, metadata DeliveryMetadata) context.Context {
|
||||
return context.WithValue(ctx, deliveryMetadataContextKey{}, metadata)
|
||||
}
|
||||
|
||||
func DeliveryMetadataFromContext(ctx context.Context) (DeliveryMetadata, bool) {
|
||||
metadata, ok := ctx.Value(deliveryMetadataContextKey{}).(DeliveryMetadata)
|
||||
return metadata, ok
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
package rabbitmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gogf/gf/v2/errors/gerror"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
amqp "github.com/rabbitmq/amqp091-go"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
Dsn string
|
||||
}
|
||||
|
||||
type channelLifecycle interface {
|
||||
Cancel(consumer string, noWait bool) error
|
||||
Close() error
|
||||
IsClosed() bool
|
||||
}
|
||||
|
||||
type managedConsumer struct {
|
||||
channel channelLifecycle
|
||||
tag string
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// Handler 是 RabbitMQ 消费处理函数。
|
||||
// 返回 nil 时消息 Ack;返回 error 时消息 Nack 并重新入队,等待后续重试。
|
||||
type Handler func(ctx context.Context, body []byte) error
|
||||
|
||||
// JSONHandler 是 RabbitMQ JSON 消息的业务处理函数。
|
||||
type JSONHandler[T any] func(ctx context.Context, message T) error
|
||||
|
||||
var (
|
||||
conn *amqp.Connection
|
||||
connMu sync.Mutex
|
||||
consumers []*managedConsumer
|
||||
consumerMu sync.Mutex
|
||||
closeMu sync.Mutex
|
||||
)
|
||||
|
||||
func GetConfig(ctx context.Context) Config {
|
||||
return Config{
|
||||
Enabled: g.Cfg().MustGet(ctx, "rabbitmq.enabled", false).Bool(),
|
||||
Dsn: g.Cfg().MustGet(ctx, "rabbitmq.dsn", "").String(),
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled 判断 RabbitMQ 是否启用且 DSN 是否已配置。
|
||||
func Enabled(ctx context.Context) bool {
|
||||
cfg := GetConfig(ctx)
|
||||
return cfg.Enabled && cfg.Dsn != ""
|
||||
}
|
||||
|
||||
// Conn 懒加载 RabbitMQ 连接,供多个生产者/消费者复用。
|
||||
func Conn(ctx context.Context) (*amqp.Connection, error) {
|
||||
cfg := GetConfig(ctx)
|
||||
if !cfg.Enabled {
|
||||
return nil, gerror.New("rabbitmq is disabled")
|
||||
}
|
||||
if cfg.Dsn == "" {
|
||||
return nil, gerror.New("rabbitmq dsn is empty")
|
||||
}
|
||||
|
||||
connMu.Lock()
|
||||
defer connMu.Unlock()
|
||||
if conn != nil && !conn.IsClosed() {
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
c, err := amqp.Dial(cfg.Dsn)
|
||||
if err != nil {
|
||||
return nil, gerror.Wrap(err, "rabbitmq connect failed")
|
||||
}
|
||||
conn = c
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// Channel 基于共享连接创建独立 channel。
|
||||
func Channel(ctx context.Context) (*amqp.Channel, error) {
|
||||
c, err := Conn(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ch, err := c.Channel()
|
||||
if err != nil {
|
||||
return nil, gerror.Wrap(err, "rabbitmq open channel failed")
|
||||
}
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// DeclareQueue 声明一个持久化队列。
|
||||
func DeclareQueue(ctx context.Context, queue string) error {
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ch.Close()
|
||||
return declareQueue(ch, queue)
|
||||
}
|
||||
|
||||
// Publish 将结构体序列化为 JSON 后发布到指定队列。
|
||||
func Publish(ctx context.Context, queue string, data any) error {
|
||||
body, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return gerror.Wrap(err, "rabbitmq marshal message failed")
|
||||
}
|
||||
return PublishBytes(ctx, queue, body)
|
||||
}
|
||||
|
||||
// PublishDelayed 将结构体序列化为 JSON 后发布到延迟队列。
|
||||
// 消息在 delayQueue 中等待 delay 后,通过死信路由进入 targetQueue。
|
||||
func PublishDelayed(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, data any) error {
|
||||
body, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return gerror.Wrap(err, "rabbitmq marshal delayed message failed")
|
||||
}
|
||||
return PublishDelayedBytes(ctx, delayQueue, targetQueue, delay, body)
|
||||
}
|
||||
|
||||
// PublishBytes 发布原始消息体到指定队列。
|
||||
func PublishBytes(ctx context.Context, queue string, body []byte) error {
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ch.Close()
|
||||
|
||||
// 发布前也声明一次队列,避免生产者先启动时队列不存在。
|
||||
if err = declareQueue(ch, queue); err != nil {
|
||||
return err
|
||||
}
|
||||
return ch.PublishWithContext(ctx, "", queue, false, false, amqp.Publishing{
|
||||
ContentType: "application/json",
|
||||
DeliveryMode: amqp.Persistent,
|
||||
Timestamp: time.Now(),
|
||||
Body: body,
|
||||
})
|
||||
}
|
||||
|
||||
// PublishDelayedBytes 发布原始消息体到延迟队列。
|
||||
func PublishDelayedBytes(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte) error {
|
||||
if delay <= 0 {
|
||||
return gerror.New("rabbitmq delayed message delay must be positive")
|
||||
}
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ch.Close()
|
||||
|
||||
if err = declareQueue(ch, targetQueue); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = declareDelayQueue(ch, delayQueue, targetQueue); err != nil {
|
||||
return err
|
||||
}
|
||||
return ch.PublishWithContext(ctx, "", delayQueue, false, false, amqp.Publishing{
|
||||
ContentType: "application/json",
|
||||
DeliveryMode: amqp.Persistent,
|
||||
Expiration: fmt.Sprintf("%d", delay.Milliseconds()),
|
||||
Timestamp: time.Now(),
|
||||
Body: body,
|
||||
})
|
||||
}
|
||||
|
||||
// PublishDelayedBytesWithHeaders 发布带消息头的原始消息体到延迟队列。
|
||||
func PublishDelayedBytesWithHeaders(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte, headers amqp.Table) error {
|
||||
if delay <= 0 {
|
||||
return gerror.New("rabbitmq delayed message delay must be positive")
|
||||
}
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ch.Close()
|
||||
|
||||
if err = declareQueue(ch, targetQueue); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = declareDelayQueue(ch, delayQueue, targetQueue); err != nil {
|
||||
return err
|
||||
}
|
||||
return publishWithConfirm(ctx, ch, delayQueue, amqp.Publishing{
|
||||
ContentType: "application/json",
|
||||
DeliveryMode: amqp.Persistent,
|
||||
Expiration: fmt.Sprintf("%d", delay.Milliseconds()),
|
||||
Timestamp: time.Now(),
|
||||
Headers: headers,
|
||||
Body: body,
|
||||
})
|
||||
}
|
||||
|
||||
// PublishBytesWithHeaders 发布带消息头的原始消息体到指定队列。
|
||||
func PublishBytesWithHeaders(ctx context.Context, queue string, body []byte, headers amqp.Table) error {
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer ch.Close()
|
||||
|
||||
if err = declareQueue(ch, queue); err != nil {
|
||||
return err
|
||||
}
|
||||
return publishWithConfirm(ctx, ch, queue, amqp.Publishing{
|
||||
ContentType: "application/json",
|
||||
DeliveryMode: amqp.Persistent,
|
||||
Timestamp: time.Now(),
|
||||
Headers: headers,
|
||||
Body: body,
|
||||
})
|
||||
}
|
||||
|
||||
func publishWithConfirm(ctx context.Context, ch *amqp.Channel, queue string, message amqp.Publishing) error {
|
||||
if err := ch.Confirm(false); err != nil {
|
||||
return gerror.Wrap(err, "rabbitmq enable publisher confirm failed")
|
||||
}
|
||||
confirms := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
|
||||
if err := ch.PublishWithContext(ctx, "", queue, false, false, message); err != nil {
|
||||
return gerror.Wrapf(err, "rabbitmq publish failed, queue:%s", queue)
|
||||
}
|
||||
select {
|
||||
case confirmation, ok := <-confirms:
|
||||
if !ok {
|
||||
return gerror.Newf("rabbitmq publisher confirm channel closed, queue:%s", queue)
|
||||
}
|
||||
if !confirmation.Ack {
|
||||
return gerror.Newf("rabbitmq message was not confirmed, queue:%s", queue)
|
||||
}
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return gerror.Wrapf(ctx.Err(), "wait rabbitmq publisher confirm failed, queue:%s", queue)
|
||||
}
|
||||
}
|
||||
|
||||
// Consume 启动消费者,使用手动 Ack,handler 成功才确认消息。
|
||||
func Consume(ctx context.Context, queue, consumer string, handler Handler) error {
|
||||
ch, err := Channel(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = declareQueue(ch, queue); err != nil {
|
||||
_ = ch.Close()
|
||||
return err
|
||||
}
|
||||
if err = ch.Qos(1, 0, false); err != nil {
|
||||
_ = ch.Close()
|
||||
return gerror.Wrap(err, "rabbitmq set qos failed")
|
||||
}
|
||||
|
||||
// autoAck=false,业务处理成功后手动 Ack,处理失败时 Nack 并重新入队。
|
||||
deliveries, err := ch.Consume(queue, consumer, false, false, false, false, nil)
|
||||
if err != nil {
|
||||
_ = ch.Close()
|
||||
return gerror.Wrap(err, "rabbitmq consume failed")
|
||||
}
|
||||
|
||||
consumerState := &managedConsumer{
|
||||
channel: ch,
|
||||
tag: consumer,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
registerConsumer(consumerState)
|
||||
go func() {
|
||||
defer close(consumerState.done)
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case delivery, ok := <-deliveries:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
consumerTag := delivery.ConsumerTag
|
||||
if consumerTag == "" {
|
||||
consumerTag = consumer
|
||||
}
|
||||
deliveryCtx := WithDeliveryMetadata(ctx, DeliveryMetadata{
|
||||
DeliveryTag: delivery.DeliveryTag,
|
||||
Redelivered: delivery.Redelivered,
|
||||
Consumer: consumerTag,
|
||||
Queue: queue,
|
||||
Headers: cloneHeaders(delivery.Headers),
|
||||
})
|
||||
if handlerErr := handler(deliveryCtx, delivery.Body); handlerErr != nil {
|
||||
g.Log().Errorf(deliveryCtx, "rabbitmq message Nack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t, err:%+v", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered, handlerErr)
|
||||
_ = delivery.Nack(false, true)
|
||||
continue
|
||||
}
|
||||
// handler 内部已经完成业务处理或重新发布下一次重试消息,这里只确认当前消息。
|
||||
g.Log().Infof(deliveryCtx, "rabbitmq message Ack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered)
|
||||
_ = delivery.Ack(false)
|
||||
}
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ConsumeJSON 启动 JSON 消费者,统一完成反序列化和 Ack/Nack 处理。
|
||||
func ConsumeJSON[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error {
|
||||
return Consume(ctx, queue, consumer, func(ctx context.Context, body []byte) error {
|
||||
var message T
|
||||
if err := json.Unmarshal(body, &message); err != nil {
|
||||
g.Log().Errorf(ctx, "rabbitmq message dropped, invalid json, queue:%s, body:%s, err:%+v", queue, string(body), err)
|
||||
return nil
|
||||
}
|
||||
return handler(ctx, message)
|
||||
})
|
||||
}
|
||||
|
||||
// Close 关闭消费者 channel 和共享连接。
|
||||
func Close(ctx context.Context) error {
|
||||
closeMu.Lock()
|
||||
defer closeMu.Unlock()
|
||||
|
||||
consumerMu.Lock()
|
||||
registeredConsumers := consumers
|
||||
consumers = nil
|
||||
consumerMu.Unlock()
|
||||
|
||||
var closeErrors []error
|
||||
for _, consumer := range registeredConsumers {
|
||||
if consumer.channel.IsClosed() {
|
||||
continue
|
||||
}
|
||||
if err := consumer.channel.Cancel(consumer.tag, false); !isClosedError(err) {
|
||||
closeErrors = append(closeErrors, gerror.Wrapf(err, "cancel RabbitMQ consumer failed, consumer:%s", consumer.tag))
|
||||
}
|
||||
}
|
||||
|
||||
if err := waitForConsumers(ctx, registeredConsumers); err != nil {
|
||||
closeErrors = append(closeErrors, err)
|
||||
}
|
||||
|
||||
for _, consumer := range registeredConsumers {
|
||||
if consumer.channel.IsClosed() {
|
||||
continue
|
||||
}
|
||||
if err := consumer.channel.Close(); !isClosedError(err) {
|
||||
closeErrors = append(closeErrors, gerror.Wrapf(err, "close RabbitMQ channel failed, consumer:%s", consumer.tag))
|
||||
}
|
||||
}
|
||||
|
||||
if err := closeConnection(); err != nil {
|
||||
closeErrors = append(closeErrors, err)
|
||||
}
|
||||
return errors.Join(closeErrors...)
|
||||
}
|
||||
|
||||
// declareQueue 统一声明持久化队列,生产者和消费者都会调用,保证队列存在。
|
||||
func declareQueue(ch *amqp.Channel, queue string) error {
|
||||
_, err := ch.QueueDeclare(queue, true, false, false, false, nil)
|
||||
if err != nil {
|
||||
return gerror.Wrapf(err, "rabbitmq declare queue failed, queue:%s", queue)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// declareDelayQueue 声明 TTL 延迟队列,到期后死信转发到目标队列。
|
||||
func declareDelayQueue(ch *amqp.Channel, delayQueue, targetQueue string) error {
|
||||
_, err := ch.QueueDeclare(delayQueue, true, false, false, false, amqp.Table{
|
||||
"x-dead-letter-exchange": "",
|
||||
"x-dead-letter-routing-key": targetQueue,
|
||||
})
|
||||
if err != nil {
|
||||
return gerror.Wrapf(err, "rabbitmq declare delay queue failed, queue:%s", delayQueue)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// registerCloser 记录消费者 channel,服务退出时统一关闭。
|
||||
func registerConsumer(consumer *managedConsumer) {
|
||||
consumerMu.Lock()
|
||||
defer consumerMu.Unlock()
|
||||
consumers = append(consumers, consumer)
|
||||
}
|
||||
|
||||
func waitForConsumers(ctx context.Context, registeredConsumers []*managedConsumer) error {
|
||||
for _, consumer := range registeredConsumers {
|
||||
select {
|
||||
case <-consumer.done:
|
||||
case <-ctx.Done():
|
||||
return gerror.Wrap(ctx.Err(), "wait for RabbitMQ consumers failed")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func closeConnection() error {
|
||||
connMu.Lock()
|
||||
defer connMu.Unlock()
|
||||
if conn == nil || conn.IsClosed() {
|
||||
conn = nil
|
||||
return nil
|
||||
}
|
||||
if err := conn.Close(); !isClosedError(err) {
|
||||
return gerror.Wrap(err, "close RabbitMQ connection failed")
|
||||
}
|
||||
conn = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
func cloneHeaders(headers amqp.Table) map[string]any {
|
||||
if len(headers) == 0 {
|
||||
return nil
|
||||
}
|
||||
cloned := make(map[string]any, len(headers))
|
||||
for key, value := range headers {
|
||||
cloned[key] = value
|
||||
}
|
||||
return cloned
|
||||
}
|
||||
|
||||
func isClosedError(err error) bool {
|
||||
return err == nil || errors.Is(err, amqp.ErrClosed)
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package rabbitmq
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type fakeChannel struct {
|
||||
mu sync.Mutex
|
||||
cancelCalls int
|
||||
closeCalls int
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (f *fakeChannel) Cancel(string, bool) error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.cancelCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeChannel) Close() error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.closeCalls++
|
||||
f.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeChannel) IsClosed() bool {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.closed
|
||||
}
|
||||
|
||||
func TestCloseIsIdempotent(t *testing.T) {
|
||||
resetRabbitMQState(t)
|
||||
|
||||
channel := &fakeChannel{}
|
||||
done := make(chan struct{})
|
||||
close(done)
|
||||
registerConsumer(&managedConsumer{
|
||||
channel: channel,
|
||||
tag: "test-consumer",
|
||||
done: done,
|
||||
})
|
||||
|
||||
if err := Close(context.Background()); err != nil {
|
||||
t.Fatalf("first Close() error = %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if err := Close(ctx); err != nil {
|
||||
t.Fatalf("second Close() error = %v", err)
|
||||
}
|
||||
|
||||
channel.mu.Lock()
|
||||
defer channel.mu.Unlock()
|
||||
if channel.cancelCalls != 1 {
|
||||
t.Fatalf("Cancel() calls = %d, want 1", channel.cancelCalls)
|
||||
}
|
||||
if channel.closeCalls != 1 {
|
||||
t.Fatalf("Close() calls = %d, want 1", channel.closeCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseSkipsAlreadyClosedChannel(t *testing.T) {
|
||||
resetRabbitMQState(t)
|
||||
|
||||
channel := &fakeChannel{closed: true}
|
||||
done := make(chan struct{})
|
||||
close(done)
|
||||
registerConsumer(&managedConsumer{
|
||||
channel: channel,
|
||||
tag: "closed-consumer",
|
||||
done: done,
|
||||
})
|
||||
|
||||
if err := Close(context.Background()); err != nil {
|
||||
t.Fatalf("Close() error = %v", err)
|
||||
}
|
||||
|
||||
channel.mu.Lock()
|
||||
defer channel.mu.Unlock()
|
||||
if channel.cancelCalls != 0 || channel.closeCalls != 0 {
|
||||
t.Fatalf("closed channel calls: Cancel=%d Close=%d, want 0", channel.cancelCalls, channel.closeCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseHonorsContextWhileWaiting(t *testing.T) {
|
||||
resetRabbitMQState(t)
|
||||
|
||||
channel := &fakeChannel{}
|
||||
registerConsumer(&managedConsumer{
|
||||
channel: channel,
|
||||
tag: "blocked-consumer",
|
||||
done: make(chan struct{}),
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if err := Close(ctx); err == nil {
|
||||
t.Fatal("Close() error = nil, want context cancellation error")
|
||||
}
|
||||
|
||||
channel.mu.Lock()
|
||||
defer channel.mu.Unlock()
|
||||
if channel.closeCalls != 1 {
|
||||
t.Fatalf("Close() calls = %d, want forced cleanup", channel.closeCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeliveryMetadataContext(t *testing.T) {
|
||||
expected := DeliveryMetadata{
|
||||
DeliveryTag: 42,
|
||||
Redelivered: true,
|
||||
Consumer: "batch-consumer",
|
||||
Queue: "iotlcm.batch.execute",
|
||||
}
|
||||
ctx := WithDeliveryMetadata(context.Background(), expected)
|
||||
actual, ok := DeliveryMetadataFromContext(ctx)
|
||||
if !ok {
|
||||
t.Fatal("delivery metadata not found in context")
|
||||
}
|
||||
if !reflect.DeepEqual(actual, expected) {
|
||||
t.Fatalf("delivery metadata = %+v, want %+v", actual, expected)
|
||||
}
|
||||
if _, ok = DeliveryMetadataFromContext(context.Background()); ok {
|
||||
t.Fatal("unexpected delivery metadata in empty context")
|
||||
}
|
||||
}
|
||||
|
||||
func resetRabbitMQState(t *testing.T) {
|
||||
t.Helper()
|
||||
consumerMu.Lock()
|
||||
consumers = nil
|
||||
consumerMu.Unlock()
|
||||
connMu.Lock()
|
||||
conn = nil
|
||||
connMu.Unlock()
|
||||
t.Cleanup(func() {
|
||||
consumerMu.Lock()
|
||||
consumers = nil
|
||||
consumerMu.Unlock()
|
||||
connMu.Lock()
|
||||
conn = nil
|
||||
connMu.Unlock()
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
package tdengine
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
|
||||
"github.com/gogf/gf/v2/errors/gerror"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
_ "github.com/taosdata/driver-go/v3/taosRestful"
|
||||
_ "github.com/taosdata/driver-go/v3/taosWS"
|
||||
)
|
||||
|
||||
const (
|
||||
// dailyStable 和 monthlyStable 是 TDengine 中已经提前创建好的超级表名称。
|
||||
// 业务启动时只连接和读写,不负责自动建库建表。
|
||||
dailyStable = "t_sim_card_usage_daily"
|
||||
monthlyStable = "t_sim_card_usage_monthly"
|
||||
poolDailyStable = "t_flow_pool_usage_daily"
|
||||
poolMonthlyStable = "t_flow_pool_usage_monthly"
|
||||
poolCardCountDailyStable = "t_flow_pool_card_count_daily"
|
||||
poolCardCountMonthlyStable = "t_flow_pool_card_count_monthly"
|
||||
defaultBatchSize = 200
|
||||
// defaultDailyDbName / defaultMonthlyDbName 是配置缺失时的兑底库名。
|
||||
// 日流量与月流量分库(duration 不同,而 duration 只能建库时确定)。
|
||||
defaultDailyDbName = "iotlcm_daily"
|
||||
defaultMonthlyDbName = "iotlcm_monthly"
|
||||
)
|
||||
|
||||
// Config 是项目内 TDengine 连接和流量超级表配置。
|
||||
//
|
||||
// 日流量与月流量分库存储(iotlcm_daily / iotlcm_monthly):两库的 duration 需要
|
||||
// 不同——日流量按天分片、月流量按月分片——而 TDengine 的 duration 只能建库时确定、
|
||||
// 事后无法 alter。DbName 保留为兼容位,指向不支持分库的流量池等表的落库位置。
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
DriverName string
|
||||
Dsn string
|
||||
DbName string
|
||||
DailyDbName string
|
||||
MonthlyDbName string
|
||||
DailyStable string
|
||||
MonthlyStable string
|
||||
PoolDailyStable string
|
||||
PoolMonthlyStable string
|
||||
PoolCardCountDailyStable string
|
||||
PoolCardCountMonthlyStable string
|
||||
MaxOpenConns int
|
||||
MaxIdleConns int
|
||||
BatchSize int
|
||||
}
|
||||
|
||||
// pools 按库名缓存连接池。DSN 不携带库名(形如 ws(host:6041)/),
|
||||
// 因此同一个 DSN 可以服务多个 database,日/月分库无需两套连接串。
|
||||
var (
|
||||
poolsMu sync.Mutex
|
||||
pools = map[string]*sql.DB{}
|
||||
)
|
||||
|
||||
// GetConfig 从 GoFrame 配置中读取 TDengine 配置。
|
||||
// 表名保持在代码中固定,避免业务配置文件里暴露过多 TDengine 内部结构。
|
||||
func GetConfig(ctx context.Context) Config {
|
||||
database := g.Cfg().MustGet(ctx, "tsdb.database", "TdEngine").String()
|
||||
// dbName 只是连接池的键名与兑底值,不决定实际访问哪个库:
|
||||
// DSN 末尾留空、库名由 SQL 里的 库名.表名 指定。
|
||||
// 默认值取日库,避免兑底到一个不存在的库名。
|
||||
dbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.dbName", defaultDailyDbName).String()
|
||||
// dailyDbName / monthlyDbName 未配置时回退到 dbName。
|
||||
dailyDbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.dailyDbName", dbName).String()
|
||||
monthlyDbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.monthlyDbName", dbName).String()
|
||||
return Config{
|
||||
Enabled: database == "TdEngine",
|
||||
DriverName: g.Cfg().MustGet(ctx, "tsdb.tdengine.type", "taosWS").String(),
|
||||
Dsn: g.Cfg().MustGet(ctx, "tsdb.tdengine.dsn", "").String(),
|
||||
DbName: dbName,
|
||||
DailyDbName: dailyDbName,
|
||||
MonthlyDbName: monthlyDbName,
|
||||
DailyStable: dailyStable,
|
||||
MonthlyStable: monthlyStable,
|
||||
PoolDailyStable: poolDailyStable,
|
||||
PoolMonthlyStable: poolMonthlyStable,
|
||||
PoolCardCountDailyStable: poolCardCountDailyStable,
|
||||
PoolCardCountMonthlyStable: poolCardCountMonthlyStable,
|
||||
MaxOpenConns: g.Cfg().MustGet(ctx, "tsdb.tdengine.maxOpenConns", 100).Int(),
|
||||
MaxIdleConns: g.Cfg().MustGet(ctx, "tsdb.tdengine.maxIdleConns", 10).Int(),
|
||||
BatchSize: g.Cfg().MustGet(ctx, "tsdb.tdengine.batchSize", defaultBatchSize).Int(),
|
||||
}
|
||||
}
|
||||
|
||||
// DbNameFor 返回指定超级表应该写入的 database 名。
|
||||
//
|
||||
// 日流量与月流量分库:两库的 TDengine duration 不同(日库按天分片、月库按月分片),
|
||||
// 而 duration 只能建库时确定。未单独配置 dailyDbName/monthlyDbName 时回退到 DbName。
|
||||
func (c Config) DbNameFor(stable string) string {
|
||||
switch stable {
|
||||
case c.DailyStable, c.PoolDailyStable, c.PoolCardCountDailyStable:
|
||||
if c.DailyDbName != "" {
|
||||
return c.DailyDbName
|
||||
}
|
||||
case c.MonthlyStable, c.PoolMonthlyStable, c.PoolCardCountMonthlyStable:
|
||||
if c.MonthlyDbName != "" {
|
||||
return c.MonthlyDbName
|
||||
}
|
||||
}
|
||||
return c.DbName
|
||||
}
|
||||
|
||||
// Enabled 只判断当前项目是否启用 TDengine,并且 DSN 是否已经配置。
|
||||
func Enabled(ctx context.Context) bool {
|
||||
cfg := GetConfig(ctx)
|
||||
return cfg.Enabled && cfg.Dsn != ""
|
||||
}
|
||||
|
||||
// DB 懒加载并复用默认库(DbName)的连接池。
|
||||
// 注意:这里不会创建 database/stable,库表需要先通过 SQL 脚本准备好。
|
||||
func DB(ctx context.Context) (*sql.DB, error) {
|
||||
return DBForDb(ctx, GetConfig(ctx).DbName)
|
||||
}
|
||||
|
||||
// DBForDb 返回指定 database 的连接池;同一库名复用同一个池。
|
||||
func DBForDb(ctx context.Context, dbName string) (*sql.DB, error) {
|
||||
cfg := GetConfig(ctx)
|
||||
if !cfg.Enabled {
|
||||
return nil, gerror.New("tdengine is disabled")
|
||||
}
|
||||
if cfg.Dsn == "" {
|
||||
return nil, gerror.New("tdengine dsn is empty")
|
||||
}
|
||||
if dbName == "" {
|
||||
dbName = cfg.DbName
|
||||
}
|
||||
|
||||
poolsMu.Lock()
|
||||
defer poolsMu.Unlock()
|
||||
if pool, ok := pools[dbName]; ok {
|
||||
return pool, nil
|
||||
}
|
||||
|
||||
driverName := cfg.DriverName
|
||||
if driverName == "" {
|
||||
driverName = "taosWS"
|
||||
}
|
||||
pool, err := sql.Open(driverName, cfg.Dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pool.SetMaxIdleConns(cfg.MaxIdleConns)
|
||||
pool.SetMaxOpenConns(cfg.MaxOpenConns)
|
||||
if err := pool.PingContext(ctx); err != nil {
|
||||
_ = pool.Close()
|
||||
return nil, err
|
||||
}
|
||||
pools[dbName] = pool
|
||||
return pool, nil
|
||||
}
|
||||
|
||||
// Exec 执行 TDengine 默认库的写入、更新、删除类 SQL。
|
||||
func Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
||||
conn, err := DB(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conn.ExecContext(ctx, query, args...)
|
||||
}
|
||||
|
||||
// ExecForDb 在指定 database 的连接池上执行写入类 SQL。
|
||||
func ExecForDb(ctx context.Context, dbName, query string, args ...any) (sql.Result, error) {
|
||||
conn, err := DBForDb(ctx, dbName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conn.ExecContext(ctx, query, args...)
|
||||
}
|
||||
|
||||
// Query 执行 TDengine 查询 SQL,调用方负责关闭返回的 rows。
|
||||
func Query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
|
||||
conn, err := DB(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conn.QueryContext(ctx, query, args...)
|
||||
}
|
||||
|
||||
// QueryRow 执行 TDengine 单行查询 SQL。
|
||||
func QueryRow(ctx context.Context, query string, args ...any) (*sql.Row, error) {
|
||||
conn, err := DB(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conn.QueryRowContext(ctx, query, args...), nil
|
||||
}
|
||||
|
||||
// Close 在应用退出时关闭所有 TDengine 连接池。
|
||||
func Close() error {
|
||||
poolsMu.Lock()
|
||||
defer poolsMu.Unlock()
|
||||
var firstErr error
|
||||
for name, pool := range pools {
|
||||
if err := pool.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
delete(pools, name)
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
Reference in New Issue
Block a user