Merge pull request #81 from hengyoush/feature/overview-2

feat. add overview subcommand

agent/analysis/analysis.go

@@ -14,17 +14,20 @@ import (
 type aggregator struct {
 	*analysis_common.AnalysisOptions
 	*ConnStat
+	isSub bool
 }
 
 func createAggregatorWithHumanReadableClassId(humanReadableClassId string,
-	classId analysis_common.ClassId, aggregateOption *analysis_common.AnalysisOptions) *aggregator {
-	aggregator := createAggregator(classId, aggregateOption)
+	classId analysis_common.ClassId,
+	aggregateOption *analysis_common.AnalysisOptions, isSub bool) *aggregator {
+	aggregator := createAggregator(classId, aggregateOption, isSub)
 	aggregator.HumanReadbleClassId = humanReadableClassId
 	return aggregator
 }
 
-func createAggregator(classId analysis_common.ClassId, aggregateOption *analysis_common.AnalysisOptions) *aggregator {
+func createAggregator(classId analysis_common.ClassId, aggregateOption *analysis_common.AnalysisOptions, isSub bool) *aggregator {
 	aggregator := aggregator{}
+	aggregator.isSub = isSub
 	aggregator.reset(classId, aggregateOption)
 	return &aggregator
 }
@@ -34,6 +37,10 @@ func (a *aggregator) reset(classId analysis_common.ClassId, aggregateOption *ana
 	a.ConnStat = &ConnStat{
 		ClassId:       classId,
 		ClassfierType: aggregateOption.ClassfierType,
+		IsSub:         a.isSub,
+	}
+	if a.isSub {
+		a.ConnStat.ClassfierType = aggregateOption.SubClassfierType
 	}
 	a.SamplesMap = make(map[analysis_common.MetricType][]*analysis_common.AnnotatedRecord)
 	a.PercentileCalculators = make(map[analysis_common.MetricType]*PercentileCalculator)
@@ -65,9 +72,14 @@ func (a *aggregator) receive(record *analysis_common.AnnotatedRecord) error {
 		metricType := analysis_common.MetricType(rawMetricType)
 
 		if enabled {
-			samples := a.SamplesMap[metricType]
 			MetricExtract := analysis_common.GetMetricExtractFunc[float64](metricType)
-			a.SamplesMap[metricType] = AddToSamples(samples, record, MetricExtract, a.AnalysisOptions.SampleLimit)
+			samples := a.SamplesMap[metricType]
+			// only sample if aggregator is sub or no sub classfier
+			if a.isSub || a.SubClassfierType == analysis_common.None {
+				a.SamplesMap[metricType] = AddToSamples(samples, record, MetricExtract, a.AnalysisOptions.SampleLimit)
+			} else {
+				a.SamplesMap[metricType] = AddToSamples(samples, record, MetricExtract, 1)
+			}
 
 			metricValue := MetricExtract(record)
 
@@ -102,6 +114,7 @@ func AddToSamples[T analysis_common.MetricValueType](samples []*analysis_common.
 
 type Analyzer struct {
 	Classfier
+	subClassfier Classfier
 	*analysis_common.AnalysisOptions
 	common.SideEnum // 那一边的统计指标TODO 根据参数自动推断
 	Aggregators     map[analysis_common.ClassId]*aggregator
@@ -120,7 +133,7 @@ func CreateAnalyzer(recordsChannel <-chan *analysis_common.AnnotatedRecord, opts
 	// ac.AddToFastStopper(stopper)
 	opts.Init()
 	analyzer := &Analyzer{
-		Classfier:       getClassfier(opts.ClassfierType),
+		Classfier:       getClassfier(opts.ClassfierType, *opts),
 		recordsChannel:  recordsChannel,
 		Aggregators:     make(map[analysis_common.ClassId]*aggregator),
 		AnalysisOptions: opts,
@@ -129,6 +142,9 @@ func CreateAnalyzer(recordsChannel <-chan *analysis_common.AnnotatedRecord, opts
 		renderStopper:   renderStopper,
 		ctx:             ctx,
 	}
+	if opts.SubClassfierType != analysis_common.None {
+		analyzer.subClassfier = getClassfier(opts.SubClassfierType, *opts)
+	}
 	opts.CurrentReceivedSamples = func() int {
 		return analyzer.recordReceived
 	}
@@ -189,14 +205,33 @@ func (a *Analyzer) analyze(record *analysis_common.AnnotatedRecord) {
 			if ok {
 				humanReadableClassId := humanReadableFunc(record)
 				a.Aggregators[class] = createAggregatorWithHumanReadableClassId(humanReadableClassId,
-					class, a.AnalysisOptions)
+					class, a.AnalysisOptions, false)
 			} else {
-				a.Aggregators[class] = createAggregator(class, a.AnalysisOptions)
+				a.Aggregators[class] = createAggregator(class, a.AnalysisOptions, false)
 			}
 
 			aggregator = a.Aggregators[class]
 		}
 		aggregator.receive(record)
+
+		if a.subClassfier != nil {
+			subClassId, err := a.subClassfier(record)
+			if err == nil {
+				fullClassId := class + "||" + subClassId
+				subAggregator, exists := a.Aggregators[fullClassId]
+				if !exists {
+					subHumanReadableFunc, ok := getClassIdHumanReadableFunc(a.AnalysisOptions.SubClassfierType, *a.AnalysisOptions)
+					if ok {
+						subHumanReadableClassId := subHumanReadableFunc(record)
+						a.Aggregators[fullClassId] = createAggregatorWithHumanReadableClassId(subHumanReadableClassId, fullClassId, a.AnalysisOptions, true)
+					} else {
+						a.Aggregators[fullClassId] = createAggregator(fullClassId, a.AnalysisOptions, true)
+					}
+					subAggregator = a.Aggregators[fullClassId]
+				}
+				subAggregator.receive(record)
+			}
+		}
 	} else {
 		common.DefaultLog.Warnf("classify error: %v\n", err)
 	}

agent/analysis/classfier.go

@@ -47,7 +47,25 @@ func init() {
 		}
 	}
 
+	classfierMap[anc.ProtocolAdaptive] = func(ar *anc.AnnotatedRecord) (anc.ClassId, error) {
+		redisReq, ok := ar.Record.Request().(*protocol.RedisMessage)
+		if !ok {
+			return "_not_a_redis_req_", nil
+		} else {
+			return anc.ClassId(redisReq.Command()), nil
+		}
+	}
+
 	classIdHumanReadableMap = make(map[anc.ClassfierType]ClassIdAsHumanReadable)
+	classIdHumanReadableMap[anc.RemoteIp] = func(ar *anc.AnnotatedRecord) string {
+		return ar.ConnDesc.RemoteAddr.String()
+	}
+	classIdHumanReadableMap[anc.RemotePort] = func(ar *anc.AnnotatedRecord) string {
+		return fmt.Sprintf("%d", ar.ConnDesc.RemotePort)
+	}
+	classIdHumanReadableMap[anc.LocalPort] = func(ar *anc.AnnotatedRecord) string {
+		return fmt.Sprintf("%d", ar.ConnDesc.LocalPort)
+	}
 	classIdHumanReadableMap[anc.Conn] = func(ar *anc.AnnotatedRecord) string {
 		return ar.ConnDesc.SimpleString()
 	}
@@ -73,6 +91,33 @@ func init() {
 	}
 }
 
-func getClassfier(classfierType anc.ClassfierType) Classfier {
-	return classfierMap[classfierType]
+func getClassfier(classfierType anc.ClassfierType, options anc.AnalysisOptions) Classfier {
+	if classfierType == anc.ProtocolAdaptive {
+		return func(ar *anc.AnnotatedRecord) (anc.ClassId, error) {
+			c, ok := options.ProtocolSpecificClassfiers[bpf.AgentTrafficProtocolT(ar.Protocol)]
+			if !ok {
+				return classfierMap[anc.RemoteIp](ar)
+			} else {
+				return classfierMap[c](ar)
+			}
+		}
+	} else {
+		return classfierMap[classfierType]
+	}
+}
+
+func getClassIdHumanReadableFunc(classfierType anc.ClassfierType, options anc.AnalysisOptions) (ClassIdAsHumanReadable, bool) {
+	if classfierType == anc.ProtocolAdaptive {
+		return func(ar *anc.AnnotatedRecord) string {
+			c, ok := options.ProtocolSpecificClassfiers[bpf.AgentTrafficProtocolT(ar.Protocol)]
+			if !ok {
+				return classIdHumanReadableMap[anc.RemoteIp](ar)
+			} else {
+				return classIdHumanReadableMap[c](ar)
+			}
+		}, true
+	} else {
+		f, ok := classIdHumanReadableMap[classfierType]
+		return f, ok
+	}
 }

agent/analysis/common/classfier.go

@@ -1,15 +1,16 @@
 package common
 
 var ClassfierTypeNames = map[ClassfierType]string{
-	None:         "none",
-	Conn:         "conn",
-	RemotePort:   "remote-port",
-	LocalPort:    "local-port",
-	RemoteIp:     "remote-ip",
-	Protocol:     "protocol",
-	HttpPath:     "http-path",
-	RedisCommand: "redis-command",
-	Default:      "default",
+	None:             "none",
+	Conn:             "conn",
+	RemotePort:       "remote-port",
+	LocalPort:        "local-port",
+	RemoteIp:         "remote-ip",
+	Protocol:         "protocol",
+	HttpPath:         "http-path",
+	RedisCommand:     "redis-command",
+	ProtocolAdaptive: "protocol-adaptive",
+	Default:          "default",
 }
 
 const (
@@ -26,6 +27,8 @@ const (
 
 	// Redis
 	RedisCommand
+
+	ProtocolAdaptive
 )
 
 type ClassId string

agent/analysis/common/types.go

@@ -3,6 +3,7 @@ package common
 import (
 	"fmt"
 	"kyanos/agent/protocol"
+	"kyanos/bpf"
 	"kyanos/common"
 	ac "kyanos/common"
 
@@ -14,7 +15,9 @@ type AnalysisOptions struct {
 	SampleLimit          int
 	Side                 ac.SideEnum
 	ClassfierType
-	CleanWhenHarvest bool
+	SubClassfierType           ClassfierType
+	ProtocolSpecificClassfiers map[bpf.AgentTrafficProtocolT]ClassfierType
+	CleanWhenHarvest           bool
 
 	// Fast Inspect Options
 	SlowMode               bool
@@ -23,6 +26,9 @@ type AnalysisOptions struct {
 	TargetSamples          int
 	CurrentReceivedSamples func() int
 	HavestSignal           chan struct{}
+
+	// overview mode
+	Overview bool
 }
 
 func (a *AnalysisOptions) Init() {

agent/analysis/types.go

@@ -18,6 +18,7 @@ type ConnStat struct {
 	ClassId             anc.ClassId
 	HumanReadbleClassId string
 	ClassfierType       anc.ClassfierType
+	IsSub               bool
 }
 
 func (c *ConnStat) ClassIdAsHumanReadable(classId anc.ClassId) string {
@@ -31,10 +32,13 @@ func (c *ConnStat) ClassIdAsHumanReadable(classId anc.ClassId) string {
 	case anc.LocalPort:
 		fallthrough
 	case anc.RemoteIp:
-		return string(classId)
+		return c.HumanReadbleClassId
 	case anc.Protocol:
 		return c.HumanReadbleClassId
 	default:
+		if c.HumanReadbleClassId != "" {
+			return c.HumanReadbleClassId
+		}
 		return string(classId)
 	}
 }