Add initial MQTT scrubber service scaffold
This commit is contained in:
@@ -0,0 +1,152 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"mqqt-scrubber/internal/config"
|
||||
"mqqt-scrubber/internal/model"
|
||||
"mqqt-scrubber/internal/parser"
|
||||
)
|
||||
|
||||
type writer interface {
|
||||
Write(ctx context.Context, records []model.Record) error
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
config config.Config
|
||||
influxClient writer
|
||||
input chan model.RawMessage
|
||||
received atomic.Uint64
|
||||
parsed atomic.Uint64
|
||||
written atomic.Uint64
|
||||
dropped atomic.Uint64
|
||||
failed atomic.Uint64
|
||||
}
|
||||
|
||||
type Snapshot struct {
|
||||
Received uint64 `json:"received"`
|
||||
Parsed uint64 `json:"parsed"`
|
||||
Written uint64 `json:"written"`
|
||||
Dropped uint64 `json:"dropped"`
|
||||
Failed uint64 `json:"failed"`
|
||||
}
|
||||
|
||||
func NewService(cfg config.Config, influxClient writer) *Service {
|
||||
return &Service{
|
||||
config: cfg,
|
||||
influxClient: influxClient,
|
||||
input: make(chan model.RawMessage, cfg.App.BufferSize),
|
||||
}
|
||||
}
|
||||
|
||||
func (service *Service) Enqueue(message model.RawMessage) {
|
||||
service.received.Add(1)
|
||||
|
||||
select {
|
||||
case service.input <- message:
|
||||
default:
|
||||
service.dropped.Add(1)
|
||||
slog.Warn("dropping message because buffer is full", "topic", message.Topic)
|
||||
}
|
||||
}
|
||||
|
||||
func (service *Service) Run(ctx context.Context) error {
|
||||
ticker := time.NewTicker(service.config.App.FlushInterval.Duration)
|
||||
defer ticker.Stop()
|
||||
|
||||
batch := make([]model.Record, 0, service.config.App.BatchSize)
|
||||
var input <-chan model.RawMessage = service.input
|
||||
|
||||
for {
|
||||
if len(batch) >= service.config.App.BatchSize {
|
||||
input = nil
|
||||
} else {
|
||||
input = service.input
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
flushCtx, cancel := context.WithTimeout(context.Background(), service.config.App.FlushTimeout.Duration)
|
||||
err := service.flush(flushCtx, batch)
|
||||
cancel()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
service.logCounters()
|
||||
return nil
|
||||
case message := <-input:
|
||||
records, err := parser.ParseTasmota(message)
|
||||
if err != nil {
|
||||
service.failed.Add(1)
|
||||
slog.Warn("failed to parse message", "topic", message.Topic, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
service.parsed.Add(uint64(len(records)))
|
||||
batch = append(batch, records...)
|
||||
|
||||
if len(batch) >= service.config.App.BatchSize {
|
||||
flushCtx, cancel := context.WithTimeout(ctx, service.config.App.FlushTimeout.Duration)
|
||||
err := service.flush(flushCtx, batch)
|
||||
cancel()
|
||||
if err != nil {
|
||||
slog.Error("failed to flush full batch to influx; keeping batch in memory", "count", len(batch), "error", err)
|
||||
continue
|
||||
}
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
flushCtx, cancel := context.WithTimeout(ctx, service.config.App.FlushTimeout.Duration)
|
||||
err := service.flush(flushCtx, batch)
|
||||
cancel()
|
||||
if err != nil {
|
||||
slog.Error("failed to flush batch to influx; will retry on next interval", "count", len(batch), "error", err)
|
||||
continue
|
||||
}
|
||||
batch = batch[:0]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (service *Service) flush(ctx context.Context, batch []model.Record) error {
|
||||
if len(batch) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := service.influxClient.Write(ctx, batch); err != nil {
|
||||
service.failed.Add(uint64(len(batch)))
|
||||
return err
|
||||
}
|
||||
|
||||
service.written.Add(uint64(len(batch)))
|
||||
slog.Info("flushed records to influx", "count", len(batch))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (service *Service) logCounters() {
|
||||
slog.Info("service counters",
|
||||
"received", service.received.Load(),
|
||||
"parsed", service.parsed.Load(),
|
||||
"written", service.written.Load(),
|
||||
"dropped", service.dropped.Load(),
|
||||
"failed", service.failed.Load(),
|
||||
)
|
||||
}
|
||||
|
||||
func (service *Service) Snapshot() Snapshot {
|
||||
return Snapshot{
|
||||
Received: service.received.Load(),
|
||||
Parsed: service.parsed.Load(),
|
||||
Written: service.written.Load(),
|
||||
Dropped: service.dropped.Load(),
|
||||
Failed: service.failed.Load(),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"mqqt-scrubber/internal/config"
|
||||
"mqqt-scrubber/internal/model"
|
||||
)
|
||||
|
||||
type fakeWriter struct {
|
||||
mu sync.Mutex
|
||||
batches [][]model.Record
|
||||
flushed chan struct{}
|
||||
}
|
||||
|
||||
func newFakeWriter() *fakeWriter {
|
||||
return &fakeWriter{flushed: make(chan struct{}, 1)}
|
||||
}
|
||||
|
||||
func (writer *fakeWriter) Write(_ context.Context, records []model.Record) error {
|
||||
writer.mu.Lock()
|
||||
copyBatch := append([]model.Record(nil), records...)
|
||||
writer.batches = append(writer.batches, copyBatch)
|
||||
writer.mu.Unlock()
|
||||
|
||||
select {
|
||||
case writer.flushed <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (writer *fakeWriter) firstBatch() []model.Record {
|
||||
writer.mu.Lock()
|
||||
defer writer.mu.Unlock()
|
||||
if len(writer.batches) == 0 {
|
||||
return nil
|
||||
}
|
||||
return writer.batches[0]
|
||||
}
|
||||
|
||||
func TestServiceFlushesParsedRecords(t *testing.T) {
|
||||
fake := newFakeWriter()
|
||||
service := NewService(config.Config{
|
||||
App: config.AppConfig{
|
||||
BatchSize: 2,
|
||||
BufferSize: 8,
|
||||
FlushInterval: config.DurationValue{Duration: time.Hour},
|
||||
FlushTimeout: config.DurationValue{Duration: time.Second},
|
||||
},
|
||||
}, fake)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- service.Run(ctx)
|
||||
}()
|
||||
|
||||
service.Enqueue(model.RawMessage{
|
||||
Topic: "tele/tasmota_896001/LWT",
|
||||
Payload: []byte("Online"),
|
||||
ReceivedAt: time.Date(2026, time.March, 12, 15, 21, 39, 0, time.UTC),
|
||||
})
|
||||
service.Enqueue(model.RawMessage{
|
||||
Topic: "tele/tasmota_C88994/SENSOR",
|
||||
Payload: []byte(`{"Time":"2026-03-12T16:23:13","ENERGY":{"TotalStartTime":"2026-02-04T19:13:40","Total":41.385,"Yesterday":1.124,"Today":0.799,"Period":0,"Power":1,"ApparentPower":4,"ReactivePower":4,"Factor":0.22,"Voltage":231,"Current":0.016}}`),
|
||||
ReceivedAt: time.Date(2026, time.March, 12, 16, 23, 13, 0, time.UTC),
|
||||
})
|
||||
|
||||
select {
|
||||
case <-fake.flushed:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for pipeline flush")
|
||||
}
|
||||
|
||||
batch := fake.firstBatch()
|
||||
if len(batch) != 2 {
|
||||
t.Fatalf("expected 2 records in flushed batch, got %d", len(batch))
|
||||
}
|
||||
|
||||
if batch[0].Measurement != "tasmota_lwt" {
|
||||
t.Fatalf("unexpected first measurement: %s", batch[0].Measurement)
|
||||
}
|
||||
if batch[0].Fields["online"] != true {
|
||||
t.Fatalf("unexpected lwt online field: %#v", batch[0].Fields["online"])
|
||||
}
|
||||
|
||||
if batch[1].Measurement != "tasmota_sensor" {
|
||||
t.Fatalf("unexpected second measurement: %s", batch[1].Measurement)
|
||||
}
|
||||
if batch[1].Fields["energy_total"] != float64(41.385) {
|
||||
t.Fatalf("unexpected sensor energy_total field: %#v", batch[1].Fields["energy_total"])
|
||||
}
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case err := <-errCh:
|
||||
if err != nil {
|
||||
t.Fatalf("service returned error: %v", err)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for service shutdown")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user