lingniu-vehicle-ingest/go/vehicle-gateway/cmd/history-writer/main.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/segmentio/kafka-go"
	_ "github.com/taosdata/driver-go/v3/taosWS"

	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/health"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/history"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/metrics"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/observability"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/topics"
)

func main() {
	logger := observability.NewLogger("vehicle-history-writer")
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	cfg := loadConfig()
	db, err := sql.Open(cfg.TDengineDriver, cfg.TDengineDSN)
	if err != nil {
		logger.Error("tdengine open failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()
	if err := db.PingContext(ctx); err != nil {
		logger.Error("tdengine ping failed", "error", err)
		os.Exit(1)
	}
	registry := metrics.NewRegistry()
	health.Start(ctx, logger, health.NewServer(env("HEALTH_ADDR", ""), "vehicle-history-writer", []health.Check{
		{Name: "tdengine", Check: db.PingContext},
	}, registry))

	writer := history.NewWriter(db)
	if cfg.EnsureSchema {
		if err := writer.EnsureSchema(ctx, cfg.TDengineDatabase); err != nil {
			logger.Error("tdengine schema bootstrap failed", "error", err)
			os.Exit(1)
		}
	}

	reader := kafka.NewReader(kafka.ReaderConfig{
		Brokers:     cfg.KafkaBrokers,
		GroupID:     cfg.KafkaGroup,
		GroupTopics: cfg.KafkaTopics,
		MinBytes:    1,
		MaxBytes:    10e6,
	})
	defer reader.Close()

	logger.Info("history writer started",
		"driver", cfg.TDengineDriver,
		"group", cfg.KafkaGroup,
		"topics", strings.Join(cfg.KafkaTopics, ","),
		"batch_size", cfg.BatchSize,
		"batch_wait_ms", cfg.BatchWait)

	for {
		message, err := reader.FetchMessage(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			logger.Error("kafka fetch failed", "error", err)
			continue
		}
		batch := collectHistoryBatch(ctx, reader, message, cfg.BatchSize, time.Duration(cfg.BatchWait)*time.Millisecond)
		processHistoryBatch(ctx, logger, registry, writer, reader, batch)
	}
}

const kafkaMessageOperationTimeout = 30 * time.Second

type historyAppender interface {
	AppendAll(context.Context, envelope.FrameEnvelope) error
	AppendAllBatch(context.Context, []envelope.FrameEnvelope) error
}

type kafkaMessageCommitter interface {
	CommitMessages(context.Context, ...kafka.Message) error
}

type kafkaMessageFetcher interface {
	FetchMessage(context.Context) (kafka.Message, error)
}

func collectHistoryBatch(ctx context.Context, fetcher kafkaMessageFetcher, first kafka.Message, maxSize int, maxWait time.Duration) []kafka.Message {
	if maxSize <= 1 {
		return []kafka.Message{first}
	}
	if maxWait <= 0 {
		maxWait = 100 * time.Millisecond
	}
	batch := []kafka.Message{first}
	deadline := time.Now().Add(maxWait)
	for len(batch) < maxSize {
		remaining := time.Until(deadline)
		if remaining <= 0 {
			break
		}
		fetchCtx, cancel := context.WithTimeout(ctx, remaining)
		message, err := fetcher.FetchMessage(fetchCtx)
		cancel()
		if err != nil {
			break
		}
		batch = append(batch, message)
	}
	return batch
}

func processHistoryMessage(ctx context.Context, logger interface {
	Error(string, ...any)
	Warn(string, ...any)
}, registry *metrics.Registry, appender historyAppender, committer kafkaMessageCommitter, message kafka.Message) {
	messageCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), kafkaMessageOperationTimeout)
	defer cancel()

	addWriterMetric(registry, "vehicle_history_kafka_messages_total", message, "received")
	addWriterLagMetric(registry, message)
	var env envelope.FrameEnvelope
	if err := json.Unmarshal(message.Value, &env); err != nil {
		addWriterMetric(registry, "vehicle_history_kafka_messages_total", message, "invalid_json")
		logger.Warn("skip invalid envelope json", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "error", err)
		_ = committer.CommitMessages(messageCtx, message)
		return
	}
	if err := appender.AppendAll(messageCtx, env); err != nil {
		addWriterMetric(registry, "vehicle_history_writes_total", message, "error")
		logger.Error("tdengine append failed", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "event_id", env.StableEventID(), "error", err)
		return
	}
	addWriterMetric(registry, "vehicle_history_writes_total", message, "ok")
	if err := committer.CommitMessages(messageCtx, message); err != nil {
		addWriterMetric(registry, "vehicle_history_kafka_commits_total", message, "error")
		logger.Error("kafka commit failed", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "error", err)
		return
	}
	addWriterMetric(registry, "vehicle_history_kafka_commits_total", message, "ok")
}

func processHistoryBatch(ctx context.Context, logger interface {
	Error(string, ...any)
	Warn(string, ...any)
}, registry *metrics.Registry, appender historyAppender, committer kafkaMessageCommitter, messages []kafka.Message) {
	if len(messages) == 0 {
		return
	}
	messageCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), kafkaMessageOperationTimeout)
	defer cancel()

	envelopes := make([]envelope.FrameEnvelope, 0, len(messages))
	for _, message := range messages {
		addWriterMetric(registry, "vehicle_history_kafka_messages_total", message, "received")
		addWriterLagMetric(registry, message)
		var env envelope.FrameEnvelope
		if err := json.Unmarshal(message.Value, &env); err != nil {
			addWriterMetric(registry, "vehicle_history_kafka_messages_total", message, "invalid_json")
			logger.Warn("skip invalid envelope json", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "error", err)
			continue
		}
		envelopes = append(envelopes, env)
	}
	if len(envelopes) > 0 {
		started := time.Now()
		err := appender.AppendAllBatch(messageCtx, envelopes)
		elapsed := time.Since(started)
		status := "ok"
		if err != nil {
			status = "error"
		}
		addBatchMetric(registry, "vehicle_history_batch_flush_total", status, 1)
		addBatchMetric(registry, "vehicle_history_batch_rows_total", status, float64(len(envelopes)))
		setBatchDuration(registry, status, elapsed)
		if err != nil {
			for _, message := range messages {
				addWriterMetric(registry, "vehicle_history_writes_total", message, "error")
			}
			first := messages[0]
			logger.Error("tdengine batch append failed", "topic", first.Topic, "partition", first.Partition, "offset", first.Offset, "rows", len(envelopes), "error", err)
			return
		}
		for _, message := range messages {
			addWriterMetric(registry, "vehicle_history_writes_total", message, "ok")
		}
	}
	if err := committer.CommitMessages(messageCtx, messages...); err != nil {
		for _, message := range messages {
			addWriterMetric(registry, "vehicle_history_kafka_commits_total", message, "error")
		}
		first := messages[0]
		logger.Error("kafka batch commit failed", "topic", first.Topic, "partition", first.Partition, "offset", first.Offset, "messages", len(messages), "error", err)
		return
	}
	for _, message := range messages {
		addWriterMetric(registry, "vehicle_history_kafka_commits_total", message, "ok")
	}
}

func addWriterMetric(registry *metrics.Registry, name string, message kafka.Message, status string) {
	if registry == nil {
		return
	}
	registry.IncCounter(name, metrics.Labels{"topic": message.Topic, "status": status})
}

func addBatchMetric(registry *metrics.Registry, name string, status string, value float64) {
	if registry == nil {
		return
	}
	registry.AddCounter(name, metrics.Labels{"status": status}, value)
}

func setBatchDuration(registry *metrics.Registry, status string, elapsed time.Duration) {
	if registry == nil {
		return
	}
	registry.SetGauge("vehicle_history_batch_flush_duration_ms", metrics.Labels{"status": status}, float64(elapsed.Milliseconds()))
}

func addWriterLagMetric(registry *metrics.Registry, message kafka.Message) {
	if registry == nil {
		return
	}
	registry.SetKafkaLag("vehicle_history_kafka_lag", message.Topic, message.Partition, message.Offset, message.HighWaterMark)
}

type config struct {
	KafkaBrokers     []string
	KafkaTopics      []string
	KafkaGroup       string
	TDengineDriver   string
	TDengineDSN      string
	TDengineDatabase string
	EnsureSchema     bool
	BatchSize        int
	BatchWait        int
}

func loadConfig() config {
	return config{
		KafkaBrokers:     splitCSV(env("KAFKA_BROKERS", "127.0.0.1:9092")),
		KafkaTopics:      splitCSV(env("KAFKA_TOPICS", strings.Join([]string{topics.RawGB32960, topics.RawJT808, topics.RawYutongMQTT}, ","))),
		KafkaGroup:       env("KAFKA_GROUP", "go-history-writer"),
		TDengineDriver:   env("TDENGINE_DRIVER", "taosWS"),
		TDengineDSN:      env("TDENGINE_DSN", ""),
		TDengineDatabase: env("TDENGINE_DATABASE", history.DefaultDatabase),
		EnsureSchema:     env("TDENGINE_ENSURE_SCHEMA", "true") != "false",
		BatchSize:        envInt("HISTORY_BATCH_SIZE", 200),
		BatchWait:        envInt("HISTORY_BATCH_WAIT_MS", 100),
	}
}

func env(key string, fallback string) string {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return fallback
	}
	return value
}

func envInt(key string, fallback int) int {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return fallback
	}
	parsed, err := strconv.Atoi(value)
	if err != nil {
		return fallback
	}
	return parsed
}

func splitCSV(value string) []string {
	var out []string
	for _, item := range strings.Split(value, ",") {
		item = strings.TrimSpace(item)
		if item != "" {
			out = append(out, item)
		}
	}
	return out
}