64 lines
1.5 KiB
Go
64 lines
1.5 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/emil/deepres/internal/models"
|
|
"github.com/nats-io/nats.go"
|
|
)
|
|
|
|
type Queue struct {
|
|
conn *nats.Conn
|
|
}
|
|
|
|
func New(url string) (*Queue, error) {
|
|
conn, err := nats.Connect(url, nats.Timeout(5*time.Second))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("connect to nats: %w", err)
|
|
}
|
|
return &Queue{conn: conn}, nil
|
|
}
|
|
|
|
func (q *Queue) Close() { q.conn.Close() }
|
|
|
|
func (q *Queue) PublishFast(ctx context.Context, req *models.ResearchRequest) error {
|
|
data, err := json.Marshal(req)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal fast request: %w", err)
|
|
}
|
|
return q.conn.Publish("research.fast", data)
|
|
}
|
|
|
|
func (q *Queue) PublishDeep(ctx context.Context, req *models.ResearchRequest) error {
|
|
data, err := json.Marshal(req)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal deep request: %w", err)
|
|
}
|
|
return q.conn.Publish("research.deep", data)
|
|
}
|
|
|
|
func (q *Queue) SubscribeFast(handler func(*models.ResearchRequest)) (*nats.Subscription, error) {
|
|
sub, err := q.conn.QueueSubscribe("research.fast", "workers", func(msg *nats.Msg) {
|
|
var req models.ResearchRequest
|
|
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
|
return
|
|
}
|
|
handler(&req)
|
|
})
|
|
return sub, err
|
|
}
|
|
|
|
func (q *Queue) SubscribeDeep(handler func(*models.ResearchRequest)) (*nats.Subscription, error) {
|
|
sub, err := q.conn.QueueSubscribe("research.deep", "workers", func(msg *nats.Msg) {
|
|
var req models.ResearchRequest
|
|
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
|
return
|
|
}
|
|
handler(&req)
|
|
})
|
|
return sub, err
|
|
}
|