113 lines
2.7 KiB
Go
113 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type Queue struct {
|
|
driver Driver
|
|
}
|
|
|
|
func NewQueue(ctx context.Context, driver Driver) (Queue, error) {
|
|
if _, err := driver.ExecContext(ctx, `
|
|
CREATE TABLE IF NOT EXISTS queue (
|
|
id INTEGER PRIMARY KEY,
|
|
topic TEXT NOT NULL,
|
|
updated INTEGER NOT NULL,
|
|
reservation TEXT,
|
|
payload TEXT
|
|
);
|
|
`); err != nil {
|
|
return Queue{}, fmt.Errorf("failed to create table: %w", err)
|
|
}
|
|
return Queue{driver: driver}, nil
|
|
}
|
|
|
|
func (q Queue) Enqueue(ctx context.Context, topic string, b []byte) error {
|
|
_, err := q.driver.ExecContext(ctx, `
|
|
INSERT INTO queue (topic, updated, payload) VALUES (?, ?, ?)
|
|
`,
|
|
topic,
|
|
time.Now().Unix(),
|
|
b,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (q Queue) Syn(ctx context.Context, topic string) (string, []byte, error) {
|
|
for {
|
|
reservation, m, err := q.syn(ctx, topic)
|
|
if reservation != nil || err != nil {
|
|
return string(reservation), m, err
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return "", nil, ctx.Err()
|
|
case <-time.After(time.Second):
|
|
}
|
|
}
|
|
}
|
|
|
|
func (q Queue) syn(ctx context.Context, topic string) ([]byte, []byte, error) {
|
|
now := time.Now().Unix()
|
|
reservation := []byte(uuid.New().String())
|
|
var payload []byte
|
|
if result, err := q.driver.ExecContext(ctx, `
|
|
UPDATE queue
|
|
SET
|
|
updated = ?, reservation = ?
|
|
WHERE
|
|
id IN (
|
|
SELECT id
|
|
FROM queue
|
|
WHERE
|
|
topic == ?
|
|
AND (
|
|
reservation IS NULL
|
|
OR ? - updated > 60
|
|
)
|
|
LIMIT 1
|
|
)
|
|
`, now, reservation, topic, now); err != nil {
|
|
return nil, nil, fmt.Errorf("failed to assign reservation: %w", err)
|
|
} else if n, err := result.RowsAffected(); err != nil {
|
|
return nil, nil, fmt.Errorf("failed to assign reservation: no count: %w", err)
|
|
} else if n == 0 {
|
|
return nil, nil, fmt.Errorf("failed to assign reservation: zero updates")
|
|
}
|
|
|
|
rows, err := q.driver.QueryContext(ctx, `
|
|
SELECT payload
|
|
FROM queue
|
|
WHERE reservation==?
|
|
LIMIT 1
|
|
`, reservation)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to query reservation: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
if err := rows.Scan(&payload); err != nil {
|
|
return nil, nil, fmt.Errorf("failed to parse reservation: %w", err)
|
|
}
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, nil, fmt.Errorf("failed to page reservation: %w", err)
|
|
}
|
|
|
|
return reservation, payload, nil
|
|
}
|
|
|
|
func (q Queue) Ack(ctx context.Context, reservation string) error {
|
|
_, err := q.driver.ExecContext(ctx, `
|
|
DELETE FROM queue
|
|
WHERE reservation==?
|
|
`, reservation)
|
|
return err
|
|
}
|