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

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/segmentio/kafka-go"

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

const identityBatchOperationTimeout = 30 * time.Second

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

	cfg := loadConfig()
	if err := cfg.Validate(); err != nil {
		logger.Error("invalid identity writer config", "error", err)
		os.Exit(1)
	}
	db, err := sql.Open("mysql", cfg.MySQLDSN)
	if err != nil {
		logger.Error("mysql open failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()
	db.SetMaxOpenConns(cfg.MySQLMaxOpenConns)
	db.SetMaxIdleConns(cfg.MySQLMaxIdleConns)
	db.SetConnMaxLifetime(cfg.MySQLConnMaxLifetime)
	if err := db.PingContext(ctx); err != nil {
		logger.Error("mysql ping failed", "error", err)
		os.Exit(1)
	}
	if cfg.EnsureSchema {
		if err := identity.EnsureJT808RegistrationSchema(ctx, db); err != nil {
			logger.Error("jt808 registration schema bootstrap failed", "error", err)
			os.Exit(1)
		}
	}

	registry := metrics.NewRegistry()
	metrics.RegisterKafkaConsumerInfo(registry, "vehicle-identity-writer", cfg.KafkaGroup, []string{cfg.KafkaTopic})
	registry.SetGauge("vehicle_identity_writer_config", metrics.Labels{"setting": "batch_size"}, float64(cfg.BatchSize))
	registry.SetGauge("vehicle_identity_writer_config", metrics.Labels{"setting": "batch_wait_ms"}, float64(cfg.BatchWait.Milliseconds()))
	registry.SetGauge("vehicle_identity_writer_config", metrics.Labels{"setting": "location_touch_interval_seconds"}, cfg.LocationTouchInterval.Seconds())
	registry.SetGauge("vehicle_identity_writer_config", metrics.Labels{"setting": "workers"}, float64(cfg.Workers))
	health.Start(ctx, logger, health.NewServer(cfg.HealthAddr, "vehicle-identity-writer", []health.Check{
		{Name: "mysql", Check: db.PingContext},
	}, registry))

	store := identity.NewJT808RegistrationStore(db)

	logger.Info("identity writer started",
		"group", cfg.KafkaGroup,
		"topic", cfg.KafkaTopic,
		"workers", cfg.Workers,
		"batch_size", cfg.BatchSize,
		"batch_wait_ms", cfg.BatchWait.Milliseconds(),
		"location_touch_interval_seconds", cfg.LocationTouchInterval.Seconds())
	var workers sync.WaitGroup
	for workerID := 1; workerID <= cfg.Workers; workerID++ {
		workers.Add(1)
		go func(id int) {
			defer workers.Done()
			runIdentityConsumer(ctx, logger, registry, store, cfg, id)
		}(workerID)
	}
	workers.Wait()
}

func runIdentityConsumer(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	store registrationBatchStore,
	cfg config,
	workerID int,
) {
	reader := kafka.NewReader(kafka.ReaderConfig{
		Brokers:     cfg.KafkaBrokers,
		GroupID:     cfg.KafkaGroup,
		GroupTopics: []string{cfg.KafkaTopic},
		MinBytes:    1,
		MaxBytes:    10e6,
		StartOffset: cfg.StartOffset,
	})
	defer reader.Close()

	// Kafka keeps one phone on one partition. Per-worker throttling avoids a
	// global hot lock; a rebalance can only cause one harmless idempotent touch.
	projector := identity.NewJT808RegistrationProjector(cfg.Location, cfg.LocationTouchInterval)
	workerLabels := metrics.Labels{"worker": strconv.Itoa(workerID)}
	registry.SetGauge("vehicle_identity_writer_worker_active", workerLabels, 1)
	defer registry.SetGauge("vehicle_identity_writer_worker_active", workerLabels, 0)

	logger.Info("identity kafka consumer started", "worker", workerID, "topic", cfg.KafkaTopic)
	for {
		first, err := reader.FetchMessage(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			logger.Error("kafka fetch failed", "error", err)
			if !waitForRetry(ctx, cfg.RetryDelay) {
				return
			}
			continue
		}
		batch := collectIdentityBatch(ctx, reader, first, cfg.BatchSize, cfg.BatchWait)
		if !processIdentityBatchReliablyForWorker(ctx, logger, registry, projector, store, reader, batch, cfg.RetryDelay, workerLabels) {
			return
		}
	}
}

type registrationBatchStore interface {
	UpsertBatch(context.Context, []identity.JT808RegistrationFact) error
}

type registrationProjector interface {
	ProjectBatch([]envelope.FrameEnvelope) []identity.JT808RegistrationFact
	MarkPersisted([]identity.JT808RegistrationFact)
}

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

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

const (
	identityFailureWrite  = "write_error"
	identityFailureCommit = "commit_error"
)

type identityBatchFailure struct {
	reason string
	err    error
}

func (e *identityBatchFailure) Error() string { return e.err.Error() }
func (e *identityBatchFailure) Unwrap() error { return e.err }

func collectIdentityBatch(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 = 20 * 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 processIdentityBatch(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	projector registrationProjector,
	store registrationBatchStore,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
) error {
	return processIdentityBatchAttemptForWorker(ctx, logger, registry, projector, store, committer, messages, true, nil)
}

func processIdentityBatchAttempt(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	projector registrationProjector,
	store registrationBatchStore,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
	recordReceived bool,
) error {
	return processIdentityBatchAttemptForWorker(ctx, logger, registry, projector, store, committer, messages, recordReceived, nil)
}

func processIdentityBatchAttemptForWorker(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	projector registrationProjector,
	store registrationBatchStore,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
	recordReceived bool,
	workerLabels metrics.Labels,
) error {
	if len(messages) == 0 {
		return nil
	}
	operationCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), identityBatchOperationTimeout)
	defer cancel()

	registry.SetGauge("vehicle_identity_writer_batch_pending_messages", workerLabels, float64(len(messages)))
	defer registry.SetGauge("vehicle_identity_writer_batch_pending_messages", workerLabels, 0)
	envelopes := make([]envelope.FrameEnvelope, 0, len(messages))
	for _, message := range messages {
		if recordReceived {
			recordIdentityMessage(registry, message, "received")
			registry.SetKafkaLag("vehicle_identity_writer_kafka_lag", message.Topic, message.Partition, message.Offset, message.HighWaterMark)
		}
		var env envelope.FrameEnvelope
		if err := json.Unmarshal(message.Value, &env); err != nil {
			if recordReceived {
				recordIdentityMessage(registry, message, "invalid_json")
				logger.Warn("skip invalid identity raw envelope", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "error", err)
			}
			continue
		}
		if status, err := topics.ValidateRawEnvelope(message.Topic, env); err != nil {
			if recordReceived {
				recordIdentityMessage(registry, message, status)
				logger.Warn("skip mismatched identity raw envelope", "topic", message.Topic, "partition", message.Partition, "offset", message.Offset, "protocol", env.Protocol, "event_id", env.StableEventID(), "error", err)
			}
			continue
		}
		envelopes = append(envelopes, env)
	}
	facts := projector.ProjectBatch(envelopes)
	registry.SetGauge("vehicle_identity_writer_batch_pending_facts", workerLabels, float64(len(facts)))
	defer registry.SetGauge("vehicle_identity_writer_batch_pending_facts", workerLabels, 0)
	if len(facts) > 0 {
		started := time.Now()
		if err := store.UpsertBatch(operationCtx, facts); err != nil {
			recordIdentityWrite(registry, "error", len(facts), time.Since(started))
			return &identityBatchFailure{reason: identityFailureWrite, err: fmt.Errorf("upsert jt808 registration facts: %w", err)}
		}
		projector.MarkPersisted(facts)
		recordIdentityWrite(registry, "ok", len(facts), time.Since(started))
	}
	if err := committer.CommitMessages(operationCtx, messages...); err != nil {
		for _, message := range messages {
			recordIdentityCommit(registry, message, "error")
		}
		return &identityBatchFailure{reason: identityFailureCommit, err: fmt.Errorf("commit identity kafka batch: %w", err)}
	}
	for _, message := range messages {
		recordIdentityCommit(registry, message, "ok")
	}
	return nil
}

func processIdentityBatchReliably(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	projector registrationProjector,
	store registrationBatchStore,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
	retryDelay time.Duration,
) bool {
	return processIdentityBatchReliablyForWorker(ctx, logger, registry, projector, store, committer, messages, retryDelay, nil)
}

func processIdentityBatchReliablyForWorker(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	projector registrationProjector,
	store registrationBatchStore,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
	retryDelay time.Duration,
	workerLabels metrics.Labels,
) bool {
	defer registry.SetGauge("vehicle_identity_writer_retry_pending_messages", workerLabels, 0)
	recordReceived := true
	for {
		err := processIdentityBatchAttemptForWorker(ctx, logger, registry, projector, store, committer, messages, recordReceived, workerLabels)
		if err == nil {
			return true
		}
		if ctx.Err() != nil {
			return false
		}
		reason := identityFailureReason(err)
		registry.SetGauge("vehicle_identity_writer_retry_pending_messages", workerLabels, float64(len(messages)))
		registry.IncCounter("vehicle_identity_writer_batch_retries_total", metrics.Labels{"reason": reason})
		logger.Error("identity batch failed; retrying without fetching newer offsets", "messages", len(messages), "reason", reason, "error", err)
		if reason == identityFailureCommit {
			return retryIdentityCommit(ctx, logger, registry, committer, messages, retryDelay)
		}
		if !waitForRetry(ctx, retryDelay) {
			return false
		}
		recordReceived = false
	}
}

func retryIdentityCommit(
	ctx context.Context,
	logger *slog.Logger,
	registry *metrics.Registry,
	committer kafkaMessageCommitter,
	messages []kafka.Message,
	retryDelay time.Duration,
) bool {
	for {
		if !waitForRetry(ctx, retryDelay) {
			return false
		}
		operationCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), identityBatchOperationTimeout)
		err := committer.CommitMessages(operationCtx, messages...)
		cancel()
		if err == nil {
			for _, message := range messages {
				recordIdentityCommit(registry, message, "ok")
			}
			return true
		}
		for _, message := range messages {
			recordIdentityCommit(registry, message, "error")
		}
		registry.IncCounter("vehicle_identity_writer_batch_retries_total", metrics.Labels{"reason": identityFailureCommit})
		logger.Error("identity kafka commit retry failed", "messages", len(messages), "error", err)
	}
}

func identityFailureReason(err error) string {
	var failure *identityBatchFailure
	if errors.As(err, &failure) && failure.reason != "" {
		return failure.reason
	}
	return identityFailureWrite
}

var identityWriteDurationBucketsMS = []float64{1, 5, 10, 25, 50, 100, 250, 500, 1000, 5000}

func recordIdentityMessage(registry *metrics.Registry, message kafka.Message, status string) {
	labels := metrics.Labels{"topic": message.Topic, "status": status}
	registry.IncCounter("vehicle_identity_writer_kafka_messages_total", labels)
	metrics.RecordLastActivity(registry, "vehicle_identity_writer_last_message_unix_seconds", labels)
}

func recordIdentityWrite(registry *metrics.Registry, status string, facts int, elapsed time.Duration) {
	labels := metrics.Labels{"status": status}
	registry.IncCounter("vehicle_identity_writer_batches_total", labels)
	registry.AddCounter("vehicle_identity_writer_facts_total", labels, float64(facts))
	registry.ObserveHistogram("vehicle_identity_writer_write_duration_ms_histogram", labels, identityWriteDurationBucketsMS, float64(elapsed.Milliseconds()))
	metrics.RecordLastActivity(registry, "vehicle_identity_writer_last_write_unix_seconds", labels)
}

func recordIdentityCommit(registry *metrics.Registry, message kafka.Message, status string) {
	labels := metrics.Labels{"topic": message.Topic, "status": status}
	registry.IncCounter("vehicle_identity_writer_kafka_commits_total", labels)
	metrics.RecordLastActivity(registry, "vehicle_identity_writer_last_commit_unix_seconds", labels)
}

func waitForRetry(ctx context.Context, delay time.Duration) bool {
	if delay <= 0 {
		delay = time.Second
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}

type config struct {
	KafkaBrokers          []string
	KafkaTopic            string
	KafkaGroup            string
	MySQLDSN              string
	MySQLMaxOpenConns     int
	MySQLMaxIdleConns     int
	MySQLConnMaxLifetime  time.Duration
	EnsureSchema          bool
	HealthAddr            string
	Location              *time.Location
	LocationTouchInterval time.Duration
	BatchSize             int
	BatchWait             time.Duration
	RetryDelay            time.Duration
	StartOffset           int64
	Workers               int
}

func loadConfig() config {
	location, err := time.LoadLocation(env("LOCAL_TZ", "Asia/Shanghai"))
	if err != nil {
		location = time.FixedZone("Asia/Shanghai", 8*60*60)
	}
	startOffset := kafka.LastOffset
	if strings.EqualFold(env("KAFKA_START_OFFSET", "last"), "first") {
		startOffset = kafka.FirstOffset
	}
	return config{
		KafkaBrokers:          splitCSV(env("KAFKA_BROKERS", "127.0.0.1:9092")),
		KafkaTopic:            env("KAFKA_TOPIC", topics.RawJT808),
		KafkaGroup:            env("KAFKA_GROUP", "go-identity-writer"),
		MySQLDSN:              env("MYSQL_DSN", ""),
		MySQLMaxOpenConns:     envInt("MYSQL_MAX_OPEN_CONNS", 8),
		MySQLMaxIdleConns:     envInt("MYSQL_MAX_IDLE_CONNS", 4),
		MySQLConnMaxLifetime:  time.Duration(envInt("MYSQL_CONN_MAX_LIFETIME_SECONDS", 300)) * time.Second,
		EnsureSchema:          envBool("MYSQL_ENSURE_SCHEMA", true),
		HealthAddr:            env("HEALTH_ADDR", "127.0.0.1:20217"),
		Location:              location,
		LocationTouchInterval: time.Duration(envInt("JT808_REGISTRATION_LOCATION_TOUCH_INTERVAL_SECONDS", 600)) * time.Second,
		BatchSize:             envInt("IDENTITY_WRITER_BATCH_SIZE", 500),
		BatchWait:             time.Duration(envInt("IDENTITY_WRITER_BATCH_WAIT_MS", 20)) * time.Millisecond,
		RetryDelay:            time.Duration(envInt("IDENTITY_WRITER_RETRY_DELAY_MS", 1000)) * time.Millisecond,
		StartOffset:           startOffset,
		Workers:               envInt("IDENTITY_WRITER_WORKERS", 3),
	}
}

func (c config) Validate() error {
	if len(c.KafkaBrokers) == 0 {
		return fmt.Errorf("KAFKA_BROKERS is required")
	}
	if strings.TrimSpace(c.MySQLDSN) == "" {
		return fmt.Errorf("MYSQL_DSN is required")
	}
	protocol, ok := topics.ProtocolForKnownRawTopic(c.KafkaTopic)
	if !ok || protocol != string(envelope.ProtocolJT808) {
		return fmt.Errorf("identity writer consumes JT808 raw topic only, got %q", c.KafkaTopic)
	}
	if c.BatchSize <= 0 {
		return fmt.Errorf("IDENTITY_WRITER_BATCH_SIZE must be positive")
	}
	if c.Workers <= 0 {
		return fmt.Errorf("IDENTITY_WRITER_WORKERS must be positive")
	}
	return nil
}

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 envBool(key string, fallback bool) bool {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return fallback
	}
	parsed, err := strconv.ParseBool(value)
	if err != nil {
		return fallback
	}
	return parsed
}

func splitCSV(value string) []string {
	parts := strings.Split(value, ",")
	out := make([]string, 0, len(parts))
	for _, item := range parts {
		if trimmed := strings.TrimSpace(item); trimmed != "" {
			out = append(out, trimmed)
		}
	}
	return out
}