init push

backend/mq/mq.go

@@ -0,0 +1,45 @@
+package mq
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/wire"
+
+	"github.com/chaitin/panda-wiki/config"
+	"github.com/chaitin/panda-wiki/log"
+	"github.com/chaitin/panda-wiki/mq/nats"
+	"github.com/chaitin/panda-wiki/mq/types"
+)
+
+// Message represents a generic message that can be from either Kafka or NATS
+type Message interface {
+	GetData() []byte
+	GetTopic() string
+}
+
+type MQConsumer interface {
+	StartConsumerHandlers(ctx context.Context) error
+	RegisterHandler(topic string, handler func(ctx context.Context, msg types.Message) error) error
+	Close() error
+}
+
+type MQProducer interface {
+	Produce(ctx context.Context, topic string, key string, value []byte) error
+}
+
+func NewMQConsumer(config *config.Config, logger *log.Logger) (MQConsumer, error) {
+	if config.MQ.Type == "nats" {
+		return nats.NewMQConsumer(logger, config)
+	}
+	return nil, fmt.Errorf("invalid mq type: %s", config.MQ.Type)
+}
+
+func NewMQProducer(config *config.Config, logger *log.Logger) (MQProducer, error) {
+	if config.MQ.Type == "nats" {
+		return nats.NewMQProducer(config, logger)
+	}
+	return nil, fmt.Errorf("invalid mq type: %s", config.MQ.Type)
+}
+
+var ProviderSet = wire.NewSet(NewMQConsumer, NewMQProducer)

backend/mq/nats/consumer.go

@@ -0,0 +1,156 @@
+package nats
+
+import (
+	"context"
+	"sync"
+
+	"github.com/nats-io/nats.go"
+
+	"github.com/chaitin/panda-wiki/config"
+	"github.com/chaitin/panda-wiki/domain"
+	"github.com/chaitin/panda-wiki/log"
+	"github.com/chaitin/panda-wiki/mq/types"
+)
+
+type MQConsumer struct {
+	conn     *nats.Conn
+	js       nats.JetStreamContext
+	handlers map[string]*nats.Subscription
+	mutex    sync.Mutex
+	logger   *log.Logger
+}
+
+func NewMQConsumer(logger *log.Logger, config *config.Config) (*MQConsumer, error) {
+	opts := []nats.Option{
+		nats.Name("panda-wiki"),
+	}
+
+	// if user and password are configured, add authentication
+	if user := config.MQ.NATS.User; user != "" {
+		opts = append(opts, nats.UserInfo(user, config.MQ.NATS.Password))
+	}
+
+	// connect to nats server
+	conn, err := nats.Connect(config.MQ.NATS.Server, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	// get jetstream context
+	js, err := conn.JetStream()
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	return &MQConsumer{
+		conn:     conn,
+		js:       js,
+		handlers: make(map[string]*nats.Subscription),
+		logger:   logger.WithModule("mq.nats"),
+	}, nil
+}
+
+func (c *MQConsumer) RegisterHandler(topic string, handler func(ctx context.Context, msg types.Message) error) error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	c.logger.Info("registering handler for topic", log.String("topic", topic))
+
+	// 对于 anydoc.persistence.doc.task.export 主题，使用 Core NATS 订阅
+	if topic == domain.AnydocTaskExportTopic {
+		return c.registerCoreNATSHandler(topic, handler)
+	}
+
+	return c.registerJetStreamHandler(topic, handler)
+}
+
+// registerCoreNATSHandler 使用 Core NATS 订阅主题
+func (c *MQConsumer) registerCoreNATSHandler(topic string, handler func(ctx context.Context, msg types.Message) error) error {
+	sub, err := c.conn.Subscribe(topic, func(msg *nats.Msg) {
+		c.logger.Debug("received message via Core NATS",
+			log.String("topic", topic),
+			log.Int("data_size", len(msg.Data)))
+
+		if err := handler(context.Background(), &Message{msg: msg}); err != nil {
+			c.logger.Error("handle message failed",
+				log.String("topic", topic),
+				log.Error(err))
+			return
+		}
+
+	})
+	if err != nil {
+		c.logger.Error("failed to subscribe to topic via Core NATS",
+			log.String("topic", topic),
+			log.Error(err))
+		return err
+	}
+
+	c.logger.Info("successfully subscribed to topic via Core NATS", log.String("topic", topic))
+	c.handlers[topic] = sub
+	return nil
+}
+
+// registerJetStreamHandler 使用 JetStream 订阅主题
+func (c *MQConsumer) registerJetStreamHandler(topic string, handler func(ctx context.Context, msg types.Message) error) error {
+	consumerName := domain.TopicConsumerName[topic]
+
+	// Choose deliver policy based on topic
+	var deliverPolicy nats.SubOpt
+	if topic == domain.VectorTaskTopic {
+		deliverPolicy = nats.DeliverNew()
+	} else {
+		deliverPolicy = nats.DeliverAll()
+	}
+
+	sub, err := c.js.Subscribe(topic, func(msg *nats.Msg) {
+		c.logger.Debug("received message via JetStream",
+			log.String("topic", topic),
+			log.Int("data_size", len(msg.Data)))
+
+		if err := handler(context.Background(), &Message{msg: msg}); err != nil {
+			c.logger.Error("handle message failed",
+				log.String("topic", topic),
+				log.Error(err))
+			return
+		}
+
+		if err := msg.Ack(); err != nil {
+			c.logger.Error("failed to ack message",
+				log.String("topic", topic),
+				log.Error(err))
+		}
+	}, deliverPolicy, nats.AckExplicit(), nats.Durable(consumerName), nats.ConsumerName(consumerName))
+	if err != nil {
+		c.logger.Error("failed to subscribe to topic via JetStream",
+			log.String("topic", topic),
+			log.Error(err))
+		return err
+	}
+
+	c.logger.Info("successfully subscribed to topic via JetStream", log.String("topic", topic))
+	c.handlers[topic] = sub
+	return nil
+}
+
+func (c *MQConsumer) StartConsumerHandlers(ctx context.Context) error {
+	<-ctx.Done()
+	return nil
+}
+
+func (c *MQConsumer) Close() error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	// close all subscriptions
+	for _, sub := range c.handlers {
+		if err := sub.Unsubscribe(); err != nil {
+			c.logger.Error("unsubscribe failed", log.Any("error", err))
+		}
+	}
+
+	// close connection
+	c.conn.Close()
+	return nil
+}

backend/mq/nats/message.go

@@ -0,0 +1,21 @@
+package nats
+
+import (
+	"github.com/nats-io/nats.go"
+
+	"github.com/chaitin/panda-wiki/mq/types"
+)
+
+type Message struct {
+	msg *nats.Msg
+}
+
+func (m *Message) GetData() []byte {
+	return m.msg.Data
+}
+
+func (m *Message) GetTopic() string {
+	return m.msg.Subject
+}
+
+var _ types.Message = (*Message)(nil)

backend/mq/nats/producer.go

@@ -0,0 +1,128 @@
+package nats
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/nats-io/nats.go"
+
+	"github.com/chaitin/panda-wiki/config"
+	"github.com/chaitin/panda-wiki/log"
+)
+
+type MQProducer struct {
+	conn   *nats.Conn
+	js     nats.JetStreamContext
+	logger *log.Logger
+}
+
+func (p *MQProducer) EnsureStreams() error {
+	streams := []struct {
+		name     string
+		subjects []string
+	}{
+		{
+			name:     "task",
+			subjects: []string{"apps.panda-wiki.summary.task", "apps.panda-wiki.vector.task"},
+		},
+		{
+			name:     "scraper",
+			subjects: []string{"apps.panda-wiki.scraper.>"},
+		},
+	}
+
+	for _, stream := range streams {
+		_, err := p.js.StreamInfo(stream.name)
+		if err == nil {
+			p.logger.Debug("stream already exists",
+				log.String("stream", stream.name))
+			continue
+		}
+
+		// Stream doesn't exist, create it
+		_, err = p.js.AddStream(&nats.StreamConfig{
+			Name:       stream.name,
+			Subjects:   stream.subjects,
+			Storage:    nats.FileStorage,
+			Retention:  nats.LimitsPolicy,
+			Discard:    nats.DiscardOld,
+			MaxAge:     7 * 24 * time.Hour,
+			MaxBytes:   1 * 1024 * 1024 * 1024,
+			MaxMsgs:    1000000,
+			MaxMsgSize: 50 * 1024 * 1024,
+			Replicas:   1,
+			Duplicates: 120 * time.Second,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to create stream %s: %w", stream.name, err)
+		}
+
+		p.logger.Info("created stream",
+			log.String("stream", stream.name),
+			log.Any("subjects", stream.subjects))
+	}
+
+	return nil
+}
+
+func NewMQProducer(config *config.Config, logger *log.Logger) (*MQProducer, error) {
+	opts := []nats.Option{
+		nats.Name("panda-wiki"),
+	}
+
+	if user := config.MQ.NATS.User; user != "" {
+		opts = append(opts, nats.UserInfo(user, config.MQ.NATS.Password))
+	}
+
+	server := config.MQ.NATS.Server
+
+	conn, err := nats.Connect(server, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
+	}
+
+	js, err := conn.JetStream()
+	if err != nil {
+		conn.Close()
+		return nil, fmt.Errorf("failed to get JetStream context: %w", err)
+	}
+
+	producer := &MQProducer{
+		conn:   conn,
+		js:     js,
+		logger: logger,
+	}
+
+	// Ensure streams exist
+	if err := producer.EnsureStreams(); err != nil {
+		conn.Close()
+		return nil, fmt.Errorf("failed to ensure streams: %w", err)
+	}
+
+	return producer, nil
+}
+
+func (p *MQProducer) Produce(ctx context.Context, topic string, key string, value []byte) error {
+	p.logger.Debug("publishing message",
+		log.String("topic", topic),
+		log.String("key", key),
+		log.Int("value_size", len(value)))
+
+	_, err := p.js.Publish(topic, value)
+	if err != nil {
+		p.logger.Error("failed to publish message",
+			log.String("topic", topic),
+			log.Error(err))
+		return fmt.Errorf("failed to publish message: %w", err)
+	}
+
+	p.logger.Debug("message published successfully",
+		log.String("topic", topic))
+	return nil
+}
+
+func (p *MQProducer) Close() error {
+	p.conn.Close()
+	return nil
+}

backend/mq/types/message.go

@@ -0,0 +1,7 @@
+package types
+
+// Message represents a generic message that can be from either Kafka or NATS
+type Message interface {
+	GetData() []byte
+	GetTopic() string
+}