i6000-pusher/cmd/i6000pusher/config.go

package main

import (
	"fmt"
	"gopkg.in/yaml.v3"
	"os"
	"time"
)

type AggregateConfig struct {
	ApdexGoodMS   int `yaml:"apdexGoodMS"`
	ApdexFairMS   int `yaml:"apdexFairMS"`
	LoadTimeoutMS int `yaml:"loadTimeoutMS"`
}

func (c AggregateConfig) LoadTimeout() time.Duration {
	return time.Duration(c.LoadTimeoutMS) * time.Millisecond
}

type ClickhouseConfig struct {
	Servers  []string `yaml:"servers"`
	Database string   `yaml:"database"`
	User     string   `yaml:"user"`
	Pwd      string   `yaml:"pwd"`
}

type JobConfig struct {
	TraceIDChannelSize   int `yaml:"traceIDChannelSize"`
	FlatTraceChannelSize int `yaml:"flatTraceChannelSize"`
	TimeRangeChannelSize int `yaml:"timeRangeChannelSize"`
}

type KafkaSinkConfig struct {
	Parallelism int      `yaml:"parallelism"`
	Brokers     []string `yaml:"brokers"`
	Topic       string
}

type TraceLoaderConfig struct {
	Parallelism int `yaml:"parallelism"`
}

type PushTaskConfig struct {
	RunIntervalSeconds        int `yaml:"runIntervalSeconds"`
	LatencyMinutes            int `yaml:"latencyMinutes"`
	DebugQueryDurationMinutes int `yaml:"debugQueryDurationMinutes"`
}

func (c PushTaskConfig) RunInterval() time.Duration {
	return time.Duration(c.RunIntervalSeconds) * time.Second
}

func (c PushTaskConfig) LatencyDuration() time.Duration {
	return time.Duration(c.LatencyMinutes) * time.Minute
}

func (c PushTaskConfig) GetDebugQueryDuration() time.Duration {
	return time.Duration(c.DebugQueryDurationMinutes) * time.Minute
}

type RocketMQConfig struct {
	Parallelism   int    `yaml:"parallelism"`
	TopoTopic     string `yaml:"topoTopic"`
	MetricTopic   string `yaml:"metricTopic"`
	TopoGroupID   string `yaml:"topoGroupID"`
	MetricGroupID string `yaml:"metricGroupID"`
}

type Config struct {
	Aggregate       AggregateConfig   `yaml:"aggregate"`
	TraceLoader     TraceLoaderConfig `yaml:"traceLoader"`
	Job             JobConfig         `yaml:"job"`
	KafkaSink       KafkaSinkConfig   `yaml:"kafkaSink"`
	Clickhouse      ClickhouseConfig  `yaml:"clickhouse"`
	AppAlias2CmdbID map[string]string `yaml:"appAlias2CmdbID"`
	PushTaskConfig  PushTaskConfig    `yaml:"pushTaskConfig"`
	IncludeApps     []string          `yaml:"includeApps"`
	RocketMq        RocketMQConfig    `yaml:"rocketMq"`
}

func (c Config) findServiceCmdbID(serviceName string, serviceName2AppAlias map[string]string) string {
	alias, ok := serviceName2AppAlias[serviceName]
	if !ok {
		return ""
	}
	cmdbID, findCmdb := c.AppAlias2CmdbID[alias]
	if !findCmdb {
		return ""
	}
	return cmdbID
}

func (c Config) GetSNRelation() *ServiceNameAppAliasCmdbIDRelation {
	return &ServiceNameAppAliasCmdbIDRelation{
		AppAlias2CmdbID: c.AppAlias2CmdbID,
	}
}

func parseConfig(cp string) (Config, error) {
	var c Config
	configContents, err := os.ReadFile(cp)
	if err != nil {
		return c, fmt.Errorf("read %s: %w", cp, err)
	}
	if errUn := yaml.Unmarshal(configContents, &c); errUn != nil {
		return c, fmt.Errorf("unmarshal %s: %w", string(configContents), errUn)
	}
	return c, nil
}

type ServiceNameAppAliasCmdbIDRelation struct {
	AppAlias2CmdbID map[string]string
	sn2AA           map[string]string
}

func (sn ServiceNameAppAliasCmdbIDRelation) GetCmdbIDFromAppAlias(appAlias string) (string, bool) {
	aa, findAA := sn.AppAlias2CmdbID[appAlias]
	return aa, findAA
}