110 lines
2.5 KiB
Go
110 lines
2.5 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/nats-io/nats.go"
|
|
"github.com/nats-io/nats.go/jetstream"
|
|
"github.com/pys/pdf-count/internal/domain"
|
|
)
|
|
|
|
type Publisher struct {
|
|
js jetstream.JetStream
|
|
stream string
|
|
subject string
|
|
}
|
|
|
|
func NewPublisher(natsURL, streamName, subjectPrefix string) (*Publisher, error) {
|
|
nc, err := nats.Connect(natsURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("nats connect: %w", err)
|
|
}
|
|
js, err := jetstream.New(nc)
|
|
if err != nil {
|
|
nc.Close()
|
|
return nil, err
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
|
|
Name: streamName,
|
|
Subjects: []string{subjectPrefix + ".>"},
|
|
Retention: jetstream.WorkQueuePolicy,
|
|
MaxMsgs: 200000,
|
|
})
|
|
if err != nil {
|
|
nc.Close()
|
|
return nil, fmt.Errorf("create stream: %w", err)
|
|
}
|
|
|
|
return &Publisher{js: js, stream: streamName, subject: subjectPrefix}, nil
|
|
}
|
|
|
|
func (p *Publisher) Publish(ctx context.Context, msg domain.QueueMessage) error {
|
|
data, err := json.Marshal(msg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
subject := fmt.Sprintf("%s.%s", p.subject, msg.JobID)
|
|
_, err = p.js.Publish(ctx, subject, data)
|
|
return err
|
|
}
|
|
|
|
type Consumer struct {
|
|
js jetstream.JetStream
|
|
subject string
|
|
consumerName string
|
|
}
|
|
|
|
func NewConsumer(natsURL, streamName, subjectPrefix, consumerName string) (*Consumer, jetstream.Consumer, *nats.Conn, error) {
|
|
nc, err := nats.Connect(natsURL)
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
js, err := jetstream.New(nc)
|
|
if err != nil {
|
|
nc.Close()
|
|
return nil, nil, nil, err
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
|
|
Name: streamName,
|
|
Subjects: []string{subjectPrefix + ".>"},
|
|
Retention: jetstream.WorkQueuePolicy,
|
|
MaxMsgs: 200000,
|
|
})
|
|
if err != nil {
|
|
nc.Close()
|
|
return nil, nil, nil, err
|
|
}
|
|
|
|
cons, err := js.CreateOrUpdateConsumer(ctx, streamName, jetstream.ConsumerConfig{
|
|
Name: consumerName,
|
|
FilterSubject: subjectPrefix + ".>",
|
|
AckPolicy: jetstream.AckExplicitPolicy,
|
|
MaxDeliver: 3,
|
|
AckWait: 2 * time.Minute,
|
|
MaxAckPending: 256,
|
|
})
|
|
if err != nil {
|
|
nc.Close()
|
|
return nil, nil, nil, err
|
|
}
|
|
|
|
return &Consumer{js: js, subject: subjectPrefix, consumerName: consumerName}, cons, nc, nil
|
|
}
|
|
|
|
func DecodeMessage(data []byte) (domain.QueueMessage, error) {
|
|
var msg domain.QueueMessage
|
|
err := json.Unmarshal(data, &msg)
|
|
return msg, err
|
|
}
|