insights/livestream/events/kafka.go

package events

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
	"github.com/insights/insights/livestream/bot"
	"github.com/insights/insights/livestream/configs"
	"github.com/insights/insights/livestream/geo"
	"github.com/insights/insights/livestream/metrics"
	jlexer "github.com/mailru/easyjson/jlexer"
	jwriter "github.com/mailru/easyjson/jwriter"
	"github.com/prometheus/client_golang/prometheus"
)

// FlexibleString handles JSON values that can be either a string or a number,
// converting them to a string representation. This is needed because some SDKs
// (e.g., insights-ruby) may send distinct_id as an integer instead of a string.
type FlexibleString string

func (f *FlexibleString) UnmarshalEasyJSON(in *jlexer.Lexer) {
	if in.IsNull() {
		in.Skip()
		*f = ""
		return
	}

	// Try to detect the type by looking at the first character
	data := in.Raw()
	if len(data) == 0 {
		*f = ""
		return
	}

	// If it starts with a quote, it's a string
	if data[0] == '"' {
		var s string
		if err := json.Unmarshal(data, &s); err != nil {
			in.AddError(err)
			return
		}
		*f = FlexibleString(s)
		return
	}

	// Otherwise, try to parse as a number and convert to string
	var n json.Number
	if err := json.Unmarshal(data, &n); err != nil {
		// If that fails, just use the raw value as a string
		*f = FlexibleString(string(data))
		return
	}
	*f = FlexibleString(n.String())
}

func (f FlexibleString) MarshalEasyJSON(out *jwriter.Writer) {
	out.String(string(f))
}

func (f *FlexibleString) UnmarshalJSON(data []byte) error {
	if len(data) == 0 || string(data) == "null" {
		*f = ""
		return nil
	}

	// If it starts with a quote, it's a string
	if data[0] == '"' {
		var s string
		if err := json.Unmarshal(data, &s); err != nil {
			return err
		}
		*f = FlexibleString(s)
		return nil
	}

	// Otherwise, try to parse as a number and convert to string
	var n json.Number
	if err := json.Unmarshal(data, &n); err != nil {
		return fmt.Errorf("FlexibleString: cannot unmarshal %s", string(data))
	}
	*f = FlexibleString(n.String())
	return nil
}

func (f FlexibleString) MarshalJSON() ([]byte, error) {
	return json.Marshal(string(f))
}

//easyjson:json
type InsightsEventWrapper struct {
	Uuid       string         `json:"uuid"`
	DistinctId FlexibleString `json:"distinct_id"`
	Ip         string         `json:"ip"`
	Data       string         `json:"data"`
	Token      string         `json:"token"`
	Timestamp  string         `json:"timestamp"`
}

//easyjson:json
type InsightsEvent struct {
	Token      string         `json:"api_key,omitempty"`
	Event      string         `json:"event"`
	Properties map[string]any `json:"properties"`
	Timestamp  any            `json:"timestamp,omitempty"`

	Uuid        string
	DistinctId  string
	Lat         float64
	Lng         float64
	CountryCode string

	// Bot classification (populated by bot.Classifier)
	IsBot           bool
	TrafficType     string
	TrafficCategory string
	BotName         string
}

type KafkaConsumerInterface interface {
	SubscribeTopics(topics []string, rebalanceCb kafka.RebalanceCb) error
	ReadMessage(timeout time.Duration) (*kafka.Message, error)
	Close() error
}

type InsightsKafkaConsumer struct {
	consumer      KafkaConsumerInterface
	topic         string
	geolocator    geo.GeoLocator
	botClassifier *bot.Classifier
	incoming      chan []byte
	outgoingChan  chan InsightsEvent
	statsChan     chan CountEvent
	parallel      int
	Broker        *RedisEventBroker
	// Shared across all runParsing goroutines so the log cadence is global,
	// not per-worker. The Prometheus counter is already global; this keeps
	// the log line consistent with it.
	droppedNoToken atomic.Uint64
}

func NewInsightsKafkaConsumer(
	consumerConfig configs.ConsumerConfig,
	geolocator geo.GeoLocator,
	outgoingChan chan InsightsEvent, statsChan chan CountEvent, parallel int) (*InsightsKafkaConsumer, error) {

	config := &kafka.ConfigMap{
		"bootstrap.servers":          consumerConfig.Brokers,
		"group.id":                   consumerConfig.GroupID,
		"auto.offset.reset":          "latest",
		"enable.auto.commit":         false,
		"security.protocol":          consumerConfig.SecurityProtocol,
		"fetch.message.max.bytes":    1_000_000_000,
		"fetch.max.bytes":            1_000_000_000,
		"queued.max.messages.kbytes": 2_000_000,
	}
	applyKafkaConfigOverrides(config, consumerConfig)

	consumer, err := kafka.NewConsumer(config)
	if err != nil {
		return nil, err
	}

	return &InsightsKafkaConsumer{
		consumer:      consumer,
		topic:         consumerConfig.Topic,
		geolocator:    geolocator,
		botClassifier: bot.NewClassifier(),
		incoming:      make(chan []byte, (1+parallel)*100),
		outgoingChan:  outgoingChan,
		statsChan:     statsChan,
		parallel:      parallel,
	}, nil
}

func (c *InsightsKafkaConsumer) Consume(ctx context.Context) {
	rebalanceCallback := func(consumer *kafka.Consumer, event kafka.Event) error {
		if _, ok := event.(kafka.AssignedPartitions); ok {
			log.Printf("✅ Livestream service ready")
		}
		return nil
	}

	topics := splitTopics(c.topic)
	if err := c.consumer.SubscribeTopics(topics, rebalanceCallback); err != nil {
		log.Fatalf("Failed to subscribe to topics: %v", err)
	}

	for i := 0; i < c.parallel; i++ {
		go c.runParsing(ctx)
	}

	var msgCount, timeoutCount uint64

	for {
		msg, err := c.consumer.ReadMessage(15 * time.Second)
		if err != nil {
			if inErr, ok := errors.AsType[kafka.Error](err); ok {
				if inErr.Code() == kafka.ErrTransport {
					metrics.ConnectFailure.Inc()
				} else if inErr.IsTimeout() {
					metrics.TimeoutConsume.Inc()
					timeoutCount++
					if timeoutCount == 1 || timeoutCount%60 == 0 {
						log.Printf("Events consumer: %d timeouts so far (topic: %s)", timeoutCount, c.topic)
					}
					continue
				}
			}
			log.Printf("Error consuming message: %v", err)
			// TODO capture error to Insights
			continue
		}

		msgCount++
		metrics.MsgConsumed.With(prometheus.Labels{"partition": strconv.Itoa(int(msg.TopicPartition.Partition))}).Inc()
		if msgCount <= 5 || msgCount%10000 == 0 {
			log.Printf("Events message #%d: partition=%d, offset=%d",
				msgCount, msg.TopicPartition.Partition, msg.TopicPartition.Offset)
		}
		c.incoming <- msg.Value
	}
}

func (c *InsightsKafkaConsumer) runParsing(ctx context.Context) {
	for {
		value, ok := <-c.incoming
		if !ok {
			return
		}
		phEvent := parse(c.geolocator, c.botClassifier, value)
		if phEvent.Token == "" {
			metrics.EventsDroppedNoToken.Inc()
			n := c.droppedNoToken.Add(1)
			if n <= 5 || n%10000 == 0 {
				log.Printf("Events dropped (no token): count=%d", n)
			}
			continue
		}
		// Blocking send: matches session_recording_consumer's policy so
		// stats stay in sync with delivered events. statsChan is buffered
		// at 10K (see main.go); a sustained full channel signals a real
		// problem with the stats keeper that we want to back-pressure on,
		// not silently swallow.
		select {
		case c.statsChan <- CountEvent{Token: phEvent.Token, DistinctID: phEvent.DistinctId}:
		case <-ctx.Done():
			return
		}
		if c.Broker != nil {
			c.Broker.Publish(ctx, phEvent)
		} else {
			select {
			case c.outgoingChan <- phEvent:
			case <-ctx.Done():
				return
			}
		}
	}
}

func parse(geolocator geo.GeoLocator, classifier *bot.Classifier, kafkaMessage []byte) InsightsEvent {
	var wrapperMessage InsightsEventWrapper
	if err := json.Unmarshal(kafkaMessage, &wrapperMessage); err != nil {
		log.Printf("Error decoding JSON %s: %v", err, string(kafkaMessage))
	}

	if wrapperMessage.Timestamp != "" {
		if eventTime, err := time.Parse(time.RFC3339Nano, wrapperMessage.Timestamp); err == nil {
			if lag := time.Since(eventTime).Seconds(); lag >= 0 {
				metrics.EventLagHistogram.Observe(lag)
			}
		}
	}

	phEvent := InsightsEvent{
		Timestamp:  wrapperMessage.Timestamp,
		Token:      "",
		Event:      "",
		Properties: make(map[string]any),
	}

	data := []byte(wrapperMessage.Data)
	if err := json.Unmarshal(data, &phEvent); err != nil {
		log.Printf("Error decoding JSON %s: %v", err, string(data))
	}

	phEvent.Uuid = wrapperMessage.Uuid
	phEvent.DistinctId = string(wrapperMessage.DistinctId)

	if wrapperMessage.Token != "" {
		phEvent.Token = wrapperMessage.Token
	} else if phEvent.Token == "" {
		if tokenValue, ok := phEvent.Properties["token"].(string); ok {
			phEvent.Token = tokenValue
		} else {
			log.Printf("No valid token found in event with UUID: %s", wrapperMessage.Uuid)
		}
	}

	var ipStr = ""
	if ipValue, ok := phEvent.Properties["$ip"]; ok {
		if ipProp, ok := ipValue.(string); ok && ipProp != "" {
			ipStr = ipProp
		}
	} else if wrapperMessage.Ip != "" {
		ipStr = wrapperMessage.Ip
	}

	if ipStr != "" {
		geoResult, err := geolocator.Lookup(ipStr)
		if err != nil {
			metrics.GeoIPLookupFailures.Inc()
		}
		phEvent.Lat = geoResult.Latitude
		phEvent.Lng = geoResult.Longitude
		phEvent.CountryCode = geoResult.CountryCode
	}

	if classifier != nil && shouldClassifyBot(phEvent.Event) {
		userAgent := extractUserAgent(phEvent.Properties)
		if userAgent != "" {
			result := classifier.Classify(userAgent)
			phEvent.IsBot = result.IsBot
			phEvent.TrafficType = result.TrafficType
			phEvent.TrafficCategory = result.TrafficCategory
			phEvent.BotName = result.BotName
			// Inject $virt_* properties so they flow through both the
			// in-memory filter and the Redis pub/sub path.
			if result.TrafficType != "" {
				phEvent.Properties["$virt_is_bot"] = result.IsBot
				phEvent.Properties["$virt_traffic_type"] = result.TrafficType
				phEvent.Properties["$virt_traffic_category"] = result.TrafficCategory
				if result.BotName != "" {
					phEvent.Properties["$virt_bot_name"] = result.BotName
				}
			}
		}
	}

	return phEvent
}

var botClassifyEvents = map[string]bool{
	"$pageview":    true,
	"$pageleave":   true,
	"$screen":      true,
	"$http_log":    true,
	"$autocapture": true,
}

func shouldClassifyBot(event string) bool {
	return botClassifyEvents[event]
}

func extractUserAgent(props map[string]any) string {
	if uaValue, ok := props["$user_agent"]; ok {
		if ua, ok := uaValue.(string); ok && ua != "" {
			return ua
		}
	}
	if rawUA, ok := props["$raw_user_agent"]; ok {
		if ua, ok := rawUA.(string); ok && ua != "" {
			return ua
		}
	}
	return ""
}

func (c *InsightsKafkaConsumer) Close() {
	if err := c.consumer.Close(); err != nil {
		// TODO capture error to Insights
		log.Printf("Failed to close consumer: %v", err)
	}
	close(c.incoming)
}

func (c *InsightsKafkaConsumer) IncomingRatio() float64 {
	return float64(len(c.incoming)) / float64(cap(c.incoming))
}

func applyKafkaConfigOverrides(config *kafka.ConfigMap, consumerConfig configs.ConsumerConfig) {
	if consumerConfig.ClientID != "" {
		_ = config.SetKey("client.id", consumerConfig.ClientID)
	}
	if consumerConfig.SessionTimeoutMs > 0 {
		_ = config.SetKey("session.timeout.ms", consumerConfig.SessionTimeoutMs)
	}
	if consumerConfig.HeartbeatIntervalMs > 0 {
		_ = config.SetKey("heartbeat.interval.ms", consumerConfig.HeartbeatIntervalMs)
	}
	if consumerConfig.MaxPollIntervalMs > 0 {
		_ = config.SetKey("max.poll.interval.ms", consumerConfig.MaxPollIntervalMs)
	}
}

func splitTopics(topic string) []string {
	out := make([]string, 0, 2)
	for p := range strings.SplitSeq(topic, ",") {
		if t := strings.TrimSpace(p); t != "" {
			out = append(out, t)
		}
	}
	return out
}