mengyastore/mengyastore-backend-go/internal/mq/client.go

package mq

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"

	"mengyastore-backend/internal/storage"
)

// Client manages a single AMQP connection, a publish channel, and naming for one env.
type Client struct {
	conn      *amqp.Connection
	pubCh     *amqp.Channel
	exchange  string
	queue     string
	env       string
	amqpURL   string
	mu        sync.Mutex
	closing   sync.Once
	connected bool

	// 用于重连后重启消费协程（与 StartConsumer 注入的一致）
	consumerCtx  context.Context
	consumerSite *storage.SiteStore
}

// New connects to RabbitMQ and declares exchange + queue + binding (idempotent).
func New(amqpURL, env string) (*Client, error) {
	if amqpURL == "" {
		return nil, fmt.Errorf("empty amqp url")
	}
	conn, err := amqp.DialConfig(amqpURL, amqp.Config{
		Heartbeat: 30 * time.Second,
		Locale:    "en_US",
	})
	if err != nil {
		return nil, fmt.Errorf("amqp dial: %w", err)
	}

	pubCh, err := conn.Channel()
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("amqp channel: %w", err)
	}

	exchange, queue, err := declareTopology(pubCh, env)
	if err != nil {
		pubCh.Close()
		conn.Close()
		return nil, err
	}

	return &Client{
		conn:      conn,
		pubCh:     pubCh,
		exchange:  exchange,
		queue:     queue,
		env:       sanitizeEnv(env),
		amqpURL:   amqpURL,
		connected: true,
	}, nil
}

func isRecoverableAMQP(err error) bool {
	if err == nil {
		return false
	}
	var e *amqp.Error
	if errors.As(err, &e) {
		// 504 channel/connection not open、320 连接被服务端关闭等，通过重连恢复
		if e.Code == amqp.ChannelError || e.Code == amqp.ConnectionForced {
			return true
		}
		return strings.Contains(strings.ToLower(e.Reason), "not open")
	}
	s := strings.ToLower(err.Error())
	return strings.Contains(s, "channel/connection is not open") ||
		strings.Contains(s, "connection closed") ||
		strings.Contains(s, "use of closed network connection") ||
		strings.Contains(s, "eof")
}

// reconnectLocked 在持有 mu 时调用：关闭旧连接并重新拨号、声明拓扑。
func (c *Client) reconnectLocked() error {
	if !c.connected || c.amqpURL == "" {
		return fmt.Errorf("mq client closed")
	}

	if c.pubCh != nil {
		_ = c.pubCh.Close()
		c.pubCh = nil
	}
	if c.conn != nil && !c.conn.IsClosed() {
		_ = c.conn.Close()
	}
	c.conn = nil

	conn, err := amqp.DialConfig(c.amqpURL, amqp.Config{
		Heartbeat: 30 * time.Second,
		Locale:    "en_US",
	})
	if err != nil {
		return fmt.Errorf("amqp reconnect dial: %w", err)
	}
	pubCh, err := conn.Channel()
	if err != nil {
		conn.Close()
		return fmt.Errorf("amqp reconnect channel: %w", err)
	}
	exchange, queue, err := declareTopology(pubCh, c.env)
	if err != nil {
		pubCh.Close()
		conn.Close()
		return err
	}
	c.conn = conn
	c.pubCh = pubCh
	c.exchange = exchange
	c.queue = queue

	if c.consumerCtx != nil && c.consumerSite != nil && c.consumerCtx.Err() == nil {
		go RunOrderEmailConsumer(c.consumerCtx, c.conn, c.env, c.consumerSite)
		log.Printf("[MQ] consumer restarted after reconnect (queue=%s)", c.queue)
	}
	return nil
}

func (c *Client) passiveQueueLocked() error {
	if c.pubCh == nil || !c.connected {
		return fmt.Errorf("mq client closed")
	}
	_, err := c.pubCh.QueueDeclarePassive(c.queue, true, false, false, false, nil)
	return err
}

// PublishOrderEmail publishes a persistent JSON message to the order-email routing key.
func (c *Client) PublishOrderEmail(ctx context.Context, p OrderEmailPayload) error {
	body, err := json.Marshal(p)
	if err != nil {
		return fmt.Errorf("marshal payload: %w", err)
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.pubCh == nil || !c.connected {
		return fmt.Errorf("mq client closed")
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	tryPublish := func() error {
		return c.pubCh.PublishWithContext(ctx,
			c.exchange,
			routingKeyOrderEmail,
			false,
			false,
			amqp.Publishing{
				ContentType:  "application/json",
				DeliveryMode: amqp.Persistent,
				Timestamp:    time.Now(),
				Body:         body,
			},
		)
	}

	err = tryPublish()
	if err != nil && isRecoverableAMQP(err) {
		if rerr := c.reconnectLocked(); rerr == nil {
			err = tryPublish()
		} else {
			err = fmt.Errorf("%w (reconnect: %v)", err, rerr)
		}
	}
	if err != nil {
		return fmt.Errorf("publish: %w", err)
	}
	return nil
}

// QueueInspectInfo returns current queue depth and consumer count (passive declare).
func (c *Client) QueueInspectInfo() (messages int, consumers int, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	try := func() (int, int, error) {
		if c.pubCh == nil || !c.connected {
			return 0, 0, fmt.Errorf("mq client closed")
		}
		q, e := c.pubCh.QueueDeclarePassive(c.queue, true, false, false, false, nil)
		if e != nil {
			return 0, 0, e
		}
		return q.Messages, q.Consumers, nil
	}

	msgs, cons, err := try()
	if err != nil && isRecoverableAMQP(err) {
		if rerr := c.reconnectLocked(); rerr != nil {
			return 0, 0, fmt.Errorf("%w (reconnect: %v)", err, rerr)
		}
		return try()
	}
	return msgs, cons, err
}

// Ping checks that the publish channel can query the declared queue (liveness for /api/health).
func (c *Client) Ping() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	err := c.passiveQueueLocked()
	if err != nil && isRecoverableAMQP(err) {
		if rerr := c.reconnectLocked(); rerr != nil {
			return fmt.Errorf("%w (reconnect: %v)", err, rerr)
		}
		return c.passiveQueueLocked()
	}
	return err
}

// Env returns the sanitized environment suffix used in exchange/queue names.
func (c *Client) Env() string { return c.env }

// StartConsumer runs the order-email consumer until ctx is cancelled (run in a goroutine).
func (c *Client) StartConsumer(ctx context.Context, site *storage.SiteStore) {
	c.mu.Lock()
	c.consumerCtx = ctx
	c.consumerSite = site
	conn := c.conn
	env := c.env
	c.mu.Unlock()
	if conn == nil || site == nil {
		return
	}
	RunOrderEmailConsumer(ctx, conn, env, site)
}

// Close releases the publish channel and connection.
func (c *Client) Close() {
	c.closing.Do(func() {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.connected = false
		c.consumerCtx = nil
		c.consumerSite = nil
		if c.pubCh != nil {
			_ = c.pubCh.Close()
			c.pubCh = nil
		}
		if c.conn != nil {
			_ = c.conn.Close()
			c.conn = nil
		}
	})
}