Phase 5 Batch 4: analyzer 规则引擎 + 3 个高层 Tool
- internal/analyzer: 纯函数规则引擎
- Finding / Severity / Thresholds / Input / StageMetric / ExecutorMetric
- 3 规则: data_skew (max/min ratio), gc_pressure (GC/CPU ratio),
bottleneck (shuffle read+write GB)
- 严重度阶梯: warning / critical (阈值 ×2)
- Analyze 入口数据驱动 for ... range
- 11 个单测 (3 规则 × normal/warning/critical + 零除防御)
- fetch_spark_metrics: SHS 业务封装, format=summary 触发 analyzer
- fetch_cluster_env: RM /cluster/info + /cluster/metrics 聚合
- analyze_spark_log: RM 日志降级链 + LLM-Ready Prompt 拼装
- prompt 结构: cluster 元信息 + log source + log tail +
heuristic findings + suggested LLM analysis
- findings 是 first-pass filter, LLM 做最终判断
- deps.go: +AnalyzerThresholds
- main.go: 注入 cfg 的 3 个阈值
- 端到端实测: analyze_spark_log 完整 prompt 渲染
完成 11 个 Tool 表面 (list_clusters / spark_submit / 4 RM /
fetch_spark_metrics / fetch_cluster_env / analyze_spark_log /
fetch_url / upload_file)
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,116 @@
|
||||
package analyzer
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Analyze runs all heuristic rules and returns the triggered findings.
|
||||
func Analyze(in Input, t Thresholds) []Finding {
|
||||
var out []Finding
|
||||
out = append(out, checkDataSkew(in, t)...)
|
||||
out = append(out, checkGCPressure(in, t)...)
|
||||
out = append(out, checkBottleneck(in, t)...)
|
||||
return out
|
||||
}
|
||||
|
||||
func checkDataSkew(in Input, t Thresholds) []Finding {
|
||||
var out []Finding
|
||||
for stage, m := range in.StageMetrics {
|
||||
if m.MinPartitionBytes <= 0 || m.MaxPartitionBytes <= 0 {
|
||||
continue
|
||||
}
|
||||
ratio := float64(m.MaxPartitionBytes) / float64(m.MinPartitionBytes)
|
||||
if ratio > t.DataSkewRatio {
|
||||
sev := SeverityWarning
|
||||
if ratio > t.DataSkewRatio*2 {
|
||||
sev = SeverityCritical
|
||||
}
|
||||
out = append(out, Finding{
|
||||
Rule: "data_skew",
|
||||
Severity: sev,
|
||||
Stage: stage,
|
||||
Evidence: map[string]any{
|
||||
"max_partition_bytes": m.MaxPartitionBytes,
|
||||
"min_partition_bytes": m.MinPartitionBytes,
|
||||
"ratio": ratio,
|
||||
"threshold": t.DataSkewRatio,
|
||||
},
|
||||
Message: fmt.Sprintf("stage %q 数据倾斜: max/min = %.1fx (阈值 %.1fx)", stage, ratio, t.DataSkewRatio),
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func checkGCPressure(in Input, t Thresholds) []Finding {
|
||||
var out []Finding
|
||||
var totalGC, totalCPU int64
|
||||
var worstExec string
|
||||
var worstRatio float64
|
||||
for id, m := range in.ExecutorMetrics {
|
||||
if m.CPUTimeMS <= 0 {
|
||||
continue
|
||||
}
|
||||
r := float64(m.GCTimeMS) / float64(m.CPUTimeMS)
|
||||
totalGC += m.GCTimeMS
|
||||
totalCPU += m.CPUTimeMS
|
||||
if r > worstRatio {
|
||||
worstRatio = r
|
||||
worstExec = id
|
||||
}
|
||||
}
|
||||
if totalCPU > 0 {
|
||||
aggRatio := float64(totalGC) / float64(totalCPU)
|
||||
if aggRatio > t.GCPressureRatio {
|
||||
sev := SeverityWarning
|
||||
if aggRatio > t.GCPressureRatio*2 {
|
||||
sev = SeverityCritical
|
||||
}
|
||||
out = append(out, Finding{
|
||||
Rule: "gc_pressure",
|
||||
Severity: sev,
|
||||
Evidence: map[string]any{
|
||||
"aggregate_gc_ratio": aggRatio,
|
||||
"worst_executor": worstExec,
|
||||
"worst_ratio": worstRatio,
|
||||
"threshold": t.GCPressureRatio,
|
||||
},
|
||||
Message: fmt.Sprintf("GC 压力过高: 聚合 GC/CPU 比率 = %.1f%% (阈值 %.1f%%, 最差 executor %s = %.1f%%)",
|
||||
aggRatio*100, t.GCPressureRatio*100, worstExec, worstRatio*100),
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func checkBottleneck(in Input, t Thresholds) []Finding {
|
||||
var out []Finding
|
||||
thresholdBytes := int64(t.BottleneckShuffleGB * (1 << 30))
|
||||
for stage, m := range in.StageMetrics {
|
||||
total := m.ShuffleReadBytes + m.ShuffleWriteBytes
|
||||
if total < 0 {
|
||||
continue // overflow guard
|
||||
}
|
||||
if total > thresholdBytes {
|
||||
sev := SeverityInfo
|
||||
if total > thresholdBytes*2 {
|
||||
sev = SeverityWarning
|
||||
}
|
||||
if total > thresholdBytes*5 {
|
||||
sev = SeverityCritical
|
||||
}
|
||||
out = append(out, Finding{
|
||||
Rule: "bottleneck",
|
||||
Severity: sev,
|
||||
Stage: stage,
|
||||
Evidence: map[string]any{
|
||||
"shuffle_read_bytes": m.ShuffleReadBytes,
|
||||
"shuffle_write_bytes": m.ShuffleWriteBytes,
|
||||
"shuffle_total_gb": float64(total) / (1 << 30),
|
||||
"threshold_gb": t.BottleneckShuffleGB,
|
||||
},
|
||||
Message: fmt.Sprintf("stage %q 是潜在瓶颈: shuffle 总 %.1f GB (阈值 %.1f GB)",
|
||||
stage, float64(total)/(1<<30), t.BottleneckShuffleGB),
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
package analyzer
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAnalyze(t *testing.T) {
|
||||
defaultThresholds := Thresholds{
|
||||
DataSkewRatio: 3.0,
|
||||
GCPressureRatio: 0.1,
|
||||
BottleneckShuffleGB: 50.0,
|
||||
}
|
||||
|
||||
t.Run("data_skew", func(t *testing.T) {
|
||||
t.Run("normal", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 0": {MaxPartitionBytes: 100, MinPartitionBytes: 50},
|
||||
},
|
||||
}
|
||||
if got := Analyze(in, defaultThresholds); len(got) != 0 {
|
||||
t.Fatalf("expected no findings, got %+v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("critical", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 1": {MaxPartitionBytes: 1000, MinPartitionBytes: 10},
|
||||
},
|
||||
}
|
||||
got := Analyze(in, defaultThresholds)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 finding, got %+v", got)
|
||||
}
|
||||
if got[0].Rule != "data_skew" || got[0].Severity != SeverityCritical || got[0].Stage != "stage 1" {
|
||||
t.Fatalf("unexpected finding: %+v", got[0])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("zero_min_partition_bytes", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 2": {MaxPartitionBytes: 1000, MinPartitionBytes: 0},
|
||||
},
|
||||
}
|
||||
if got := Analyze(in, defaultThresholds); len(got) != 0 {
|
||||
t.Fatalf("expected no findings for zero min, got %+v", got)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("gc_pressure", func(t *testing.T) {
|
||||
t.Run("normal", func(t *testing.T) {
|
||||
in := Input{
|
||||
ExecutorMetrics: map[string]ExecutorMetric{
|
||||
"1": {GCTimeMS: 100, CPUTimeMS: 10000},
|
||||
},
|
||||
}
|
||||
if got := Analyze(in, defaultThresholds); len(got) != 0 {
|
||||
t.Fatalf("expected no findings, got %+v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("warning", func(t *testing.T) {
|
||||
in := Input{
|
||||
ExecutorMetrics: map[string]ExecutorMetric{
|
||||
"1": {GCTimeMS: 2000, CPUTimeMS: 10000},
|
||||
},
|
||||
}
|
||||
got := Analyze(in, defaultThresholds)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 finding, got %+v", got)
|
||||
}
|
||||
if got[0].Rule != "gc_pressure" || got[0].Severity != SeverityWarning {
|
||||
t.Fatalf("unexpected finding: %+v", got[0])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("critical", func(t *testing.T) {
|
||||
in := Input{
|
||||
ExecutorMetrics: map[string]ExecutorMetric{
|
||||
"1": {GCTimeMS: 5000, CPUTimeMS: 10000},
|
||||
},
|
||||
}
|
||||
got := Analyze(in, defaultThresholds)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 finding, got %+v", got)
|
||||
}
|
||||
if got[0].Rule != "gc_pressure" || got[0].Severity != SeverityCritical {
|
||||
t.Fatalf("unexpected finding: %+v", got[0])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("zero_cpu_time", func(t *testing.T) {
|
||||
in := Input{
|
||||
ExecutorMetrics: map[string]ExecutorMetric{
|
||||
"1": {GCTimeMS: 5000, CPUTimeMS: 0},
|
||||
},
|
||||
}
|
||||
if got := Analyze(in, defaultThresholds); len(got) != 0 {
|
||||
t.Fatalf("expected no findings for zero cpu, got %+v", got)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("bottleneck", func(t *testing.T) {
|
||||
t.Run("normal", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 0": {ShuffleReadBytes: 1 << 30, ShuffleWriteBytes: 0},
|
||||
},
|
||||
}
|
||||
if got := Analyze(in, defaultThresholds); len(got) != 0 {
|
||||
t.Fatalf("expected no findings, got %+v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("warning", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 1": {ShuffleReadBytes: 120 << 30, ShuffleWriteBytes: 0},
|
||||
},
|
||||
}
|
||||
got := Analyze(in, defaultThresholds)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 finding, got %+v", got)
|
||||
}
|
||||
if got[0].Rule != "bottleneck" || got[0].Severity != SeverityWarning {
|
||||
t.Fatalf("unexpected finding: %+v", got[0])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("critical", func(t *testing.T) {
|
||||
in := Input{
|
||||
StageMetrics: map[string]StageMetric{
|
||||
"stage 2": {ShuffleReadBytes: 300 << 30, ShuffleWriteBytes: 0},
|
||||
},
|
||||
}
|
||||
got := Analyze(in, defaultThresholds)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 finding, got %+v", got)
|
||||
}
|
||||
if got[0].Rule != "bottleneck" || got[0].Severity != SeverityCritical {
|
||||
t.Fatalf("unexpected finding: %+v", got[0])
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package analyzer
|
||||
|
||||
// Severity is the heuristic finding severity.
|
||||
type Severity string
|
||||
|
||||
const (
|
||||
SeverityInfo Severity = "info"
|
||||
SeverityWarning Severity = "warning"
|
||||
SeverityCritical Severity = "critical"
|
||||
)
|
||||
|
||||
// Finding is a single triggered heuristic rule.
|
||||
type Finding struct {
|
||||
Rule string `json:"rule"`
|
||||
Severity Severity `json:"severity"`
|
||||
Stage string `json:"stage,omitempty"`
|
||||
Evidence map[string]any `json:"evidence"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// Thresholds holds analyzer thresholds injected at runtime.
|
||||
type Thresholds struct {
|
||||
DataSkewRatio float64
|
||||
GCPressureRatio float64
|
||||
BottleneckShuffleGB float64
|
||||
}
|
||||
|
||||
// Input is the data passed to the heuristic rules.
|
||||
type Input struct {
|
||||
StageMetrics map[string]StageMetric `json:"stage_metrics"`
|
||||
ExecutorMetrics map[string]ExecutorMetric `json:"executor_metrics"`
|
||||
}
|
||||
|
||||
// StageMetric contains per-stage measurements.
|
||||
type StageMetric struct {
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
ShuffleReadBytes int64 `json:"shuffle_read_bytes"`
|
||||
ShuffleWriteBytes int64 `json:"shuffle_write_bytes"`
|
||||
MaxPartitionBytes int64 `json:"max_partition_bytes"`
|
||||
MinPartitionBytes int64 `json:"min_partition_bytes"`
|
||||
}
|
||||
|
||||
// ExecutorMetric contains per-executor measurements.
|
||||
type ExecutorMetric struct {
|
||||
GCTimeMS int64 `json:"gc_time_ms"`
|
||||
CPUTimeMS int64 `json:"cpu_time_ms"`
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/analyzer"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/rm"
|
||||
)
|
||||
|
||||
const AnalyzeSparkLogName = "analyze_spark_log"
|
||||
|
||||
const defaultAnalyzeMaxLogBytes = 10240
|
||||
|
||||
// NewAnalyzeSparkLogTool returns the schema for the analyze_spark_log MCP Tool.
|
||||
func NewAnalyzeSparkLogTool() mcp.Tool {
|
||||
return mcp.NewTool(AnalyzeSparkLogName,
|
||||
mcp.WithDescription("Fetch YARN application logs and produce an LLM-ready analysis prompt with heuristic findings."),
|
||||
mcp.WithString("cluster_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("ID of the configured cluster (from list_clusters)"),
|
||||
),
|
||||
mcp.WithString("app_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("YARN application ID, e.g. application_1234567890_0001"),
|
||||
),
|
||||
mcp.WithString("container",
|
||||
mcp.Description("Container ID used for the aggregated-logs fallback"),
|
||||
),
|
||||
mcp.WithNumber("max_log_bytes",
|
||||
mcp.Description("Maximum bytes of the log tail to include in the prompt"),
|
||||
mcp.DefaultNumber(10240),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// AnalyzeSparkLogHandler retrieves logs and builds an LLM-ready prompt.
|
||||
func (d *Deps) AnalyzeSparkLogHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
clusterID, err := req.RequireString("cluster_id")
|
||||
if err != nil {
|
||||
return errResult("analyze_spark_log: " + err.Error()), nil
|
||||
}
|
||||
appID, err := req.RequireString("app_id")
|
||||
if err != nil {
|
||||
return errResult("analyze_spark_log: " + err.Error()), nil
|
||||
}
|
||||
|
||||
args := req.GetArguments()
|
||||
container := ""
|
||||
if v, ok := args["container"].(string); ok {
|
||||
container = v
|
||||
}
|
||||
if container == "" {
|
||||
container = fmt.Sprintf("container_%s_01", appID)
|
||||
}
|
||||
|
||||
maxLogBytes := defaultAnalyzeMaxLogBytes
|
||||
if v, ok := args["max_log_bytes"].(float64); ok {
|
||||
maxLogBytes = int(v)
|
||||
}
|
||||
if maxLogBytes <= 0 {
|
||||
maxLogBytes = defaultAnalyzeMaxLogBytes
|
||||
}
|
||||
|
||||
callLog := startToolCall(ctx, d.Logger, AnalyzeSparkLogName, map[string]any{
|
||||
"cluster_id": clusterID,
|
||||
"app_id": appID,
|
||||
"container": container,
|
||||
"max_log_bytes": maxLogBytes,
|
||||
})
|
||||
defer callLog.End()
|
||||
|
||||
cl, err := d.ClusterRepo.Get(ctx, clusterID)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("analyze_spark_log: cluster " + clusterID + ": " + err.Error()), nil
|
||||
}
|
||||
|
||||
rmc := rm.New(d.HTTPClient, cl)
|
||||
body, source, err := rmc.GetLogs(ctx, appID, container)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("analyze_spark_log: " + err.Error()), nil
|
||||
}
|
||||
|
||||
logTail := truncateTail(string(body), maxLogBytes)
|
||||
|
||||
var findings []analyzer.Finding
|
||||
input, metricsErr := d.sparkMetricsInput(ctx, cl, appID)
|
||||
if metricsErr == nil {
|
||||
findings = analyzer.Analyze(input, d.AnalyzerThresholds)
|
||||
}
|
||||
|
||||
prompt := buildSparkLogPrompt(appID, cl, source, len(body), maxLogBytes, logTail, findings)
|
||||
|
||||
result := map[string]any{
|
||||
"findings": findings,
|
||||
"log_source": source,
|
||||
"log_tail": logTail,
|
||||
"prompt": prompt,
|
||||
}
|
||||
callLog.WithResult(map[string]any{"source": source, "bytes": len(body), "findings": len(findings)})
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
|
||||
func buildSparkLogPrompt(appID string, cl *cluster.Cluster, source string, totalBytes, maxBytes int, logTail string, findings []analyzer.Finding) string {
|
||||
findingsBlock := "no heuristic findings; rely on raw log + LLM analysis"
|
||||
if len(findings) > 0 {
|
||||
findingsBlock = encodeJSON(findings)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("# Spark Log Analysis for %s\n\n## Cluster\n%s (%s) — RM=%s, SHS=%s\n\n## Log Source\n%s (%d bytes, last %d shown below)\n\n## Log Tail\n"+"```"+"\n%s\n"+"```"+"\n\n## Heuristic Findings\n%s\n\n## Suggested LLM Analysis\n1. Check for ERROR/Exception stack traces\n2. Look for OOM/Timeout/Shuffle fetch failures\n3. Identify slow stages (compare to median)\n4. Suggest mitigations\n", appID, cl.Name, cl.ID, cl.RMURL, cl.SHSURL, source, totalBytes, maxBytes, logTail, findingsBlock)
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
const FetchClusterEnvName = "fetch_cluster_env"
|
||||
|
||||
// NewFetchClusterEnvTool returns the schema for the fetch_cluster_env MCP Tool.
|
||||
func NewFetchClusterEnvTool() mcp.Tool {
|
||||
return mcp.NewTool(FetchClusterEnvName,
|
||||
mcp.WithDescription("Fetch YARN cluster environment information and metrics from the ResourceManager."),
|
||||
mcp.WithString("cluster_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("ID of the configured cluster (from list_clusters)"),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// FetchClusterEnvHandler fetches RM cluster info and metrics.
|
||||
func (d *Deps) FetchClusterEnvHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
clusterID, err := req.RequireString("cluster_id")
|
||||
if err != nil {
|
||||
return errResult("fetch_cluster_env: " + err.Error()), nil
|
||||
}
|
||||
|
||||
callLog := startToolCall(ctx, d.Logger, FetchClusterEnvName, map[string]any{
|
||||
"cluster_id": clusterID,
|
||||
})
|
||||
defer callLog.End()
|
||||
|
||||
cl, err := d.ClusterRepo.Get(ctx, clusterID)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_cluster_env: cluster " + clusterID + ": " + err.Error()), nil
|
||||
}
|
||||
|
||||
base := strings.TrimRight(cl.RMURL, "/")
|
||||
infoBody, err := d.fetchInternal(ctx, cl, base+"/ws/v1/cluster/info")
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_cluster_env: info: " + err.Error()), nil
|
||||
}
|
||||
metricsBody, err := d.fetchInternal(ctx, cl, base+"/ws/v1/cluster/metrics")
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_cluster_env: metrics: " + err.Error()), nil
|
||||
}
|
||||
|
||||
var info, metrics any
|
||||
_ = json.Unmarshal(infoBody, &info)
|
||||
_ = json.Unmarshal(metricsBody, &metrics)
|
||||
|
||||
result := map[string]any{
|
||||
"cluster_info": info,
|
||||
"metrics": metrics,
|
||||
"fetched_at": time.Now().UTC().Format(time.RFC3339),
|
||||
}
|
||||
callLog.WithResult(map[string]any{"info_bytes": len(infoBody), "metrics_bytes": len(metricsBody)})
|
||||
return textResult(encodeJSON(result)), nil
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"spark-mcp-go/internal/analyzer"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
)
|
||||
|
||||
const FetchSparkMetricsName = "fetch_spark_metrics"
|
||||
|
||||
type shsExecutor struct {
|
||||
ID string `json:"id"`
|
||||
GCTimeMS int64 `json:"gc_time_ms"`
|
||||
CPUTimeMS int64 `json:"cpu_time_ms"`
|
||||
}
|
||||
|
||||
type shsStage struct {
|
||||
StageID int `json:"stageId"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
ShuffleReadBytes int64 `json:"shuffle_read_bytes"`
|
||||
ShuffleWriteBytes int64 `json:"shuffle_write_bytes"`
|
||||
MaxPartitionBytes int64 `json:"max_partition_bytes"`
|
||||
MinPartitionBytes int64 `json:"min_partition_bytes"`
|
||||
}
|
||||
|
||||
// NewFetchSparkMetricsTool returns the schema for the fetch_spark_metrics MCP Tool.
|
||||
func NewFetchSparkMetricsTool() mcp.Tool {
|
||||
return mcp.NewTool(FetchSparkMetricsName,
|
||||
mcp.WithDescription("Fetch Spark History Server metrics for an application. Returns SHS executor/stage data plus optional heuristic findings."),
|
||||
mcp.WithString("cluster_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("ID of the configured cluster whose SHS is used"),
|
||||
),
|
||||
mcp.WithString("app_id",
|
||||
mcp.Required(),
|
||||
mcp.Description("Spark application ID, e.g. application_1234567890_0001"),
|
||||
),
|
||||
mcp.WithString("format",
|
||||
mcp.Description("Output format: raw SHS JSON or summary with heuristic findings"),
|
||||
mcp.Enum("summary", "raw"),
|
||||
mcp.DefaultString("summary"),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
// FetchSparkMetricsHandler fetches SHS executor and stage data and optionally runs analysis.
|
||||
func (d *Deps) FetchSparkMetricsHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
clusterID, err := req.RequireString("cluster_id")
|
||||
if err != nil {
|
||||
return errResult("fetch_spark_metrics: " + err.Error()), nil
|
||||
}
|
||||
appID, err := req.RequireString("app_id")
|
||||
if err != nil {
|
||||
return errResult("fetch_spark_metrics: " + err.Error()), nil
|
||||
}
|
||||
|
||||
format := "summary"
|
||||
if v, ok := req.GetArguments()["format"].(string); ok && v != "" {
|
||||
format = v
|
||||
}
|
||||
|
||||
callLog := startToolCall(ctx, d.Logger, FetchSparkMetricsName, map[string]any{
|
||||
"cluster_id": clusterID,
|
||||
"app_id": appID,
|
||||
"format": format,
|
||||
})
|
||||
defer callLog.End()
|
||||
|
||||
cl, err := d.ClusterRepo.Get(ctx, clusterID)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_spark_metrics: cluster " + clusterID + ": " + err.Error()), nil
|
||||
}
|
||||
|
||||
base := strings.TrimRight(cl.SHSURL, "/")
|
||||
execURL := base + "/api/v1/applications/" + appID + "/executors"
|
||||
stageURL := base + "/api/v1/applications/" + appID + "/stages"
|
||||
|
||||
execBody, err := d.fetchInternal(ctx, cl, execURL)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_spark_metrics: executors: " + err.Error()), nil
|
||||
}
|
||||
stageBody, err := d.fetchInternal(ctx, cl, stageURL)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_spark_metrics: stages: " + err.Error()), nil
|
||||
}
|
||||
|
||||
var rawExec, rawStage any
|
||||
_ = json.Unmarshal(execBody, &rawExec)
|
||||
_ = json.Unmarshal(stageBody, &rawStage)
|
||||
|
||||
if format != "summary" {
|
||||
return textResult(encodeJSON(map[string]any{
|
||||
"executors": rawExec,
|
||||
"stages": rawStage,
|
||||
})), nil
|
||||
}
|
||||
|
||||
input, err := d.sparkMetricsInput(ctx, cl, appID)
|
||||
if err != nil {
|
||||
callLog.WithError(err)
|
||||
return errResult("fetch_spark_metrics: analyze: " + err.Error()), nil
|
||||
}
|
||||
findings := analyzer.Analyze(input, d.AnalyzerThresholds)
|
||||
|
||||
return textResult(encodeJSON(map[string]any{
|
||||
"executors": rawExec,
|
||||
"stages": rawStage,
|
||||
"findings": findings,
|
||||
})), nil
|
||||
}
|
||||
|
||||
// fetchInternal performs an SSRF-aware, authenticated GET to a cluster URL.
|
||||
func (d *Deps) fetchInternal(ctx context.Context, cl *cluster.Cluster, rawURL string) ([]byte, error) {
|
||||
allowedHosts := buildAllowedHosts(cl)
|
||||
|
||||
targetURL, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse url: %w", err)
|
||||
}
|
||||
targetHost := strings.ToLower(targetURL.Hostname())
|
||||
if !hostAllowed(targetHost, allowedHosts) {
|
||||
return nil, fmt.Errorf("host %q is not in the cluster allowlist", targetHost)
|
||||
}
|
||||
|
||||
if err := httpclient.CheckURL(rawURL, allowedHosts...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build request: %w", err)
|
||||
}
|
||||
if err := httpclient.ApplyAuth(httpReq, "", cl); err != nil {
|
||||
return nil, fmt.Errorf("auth: %w", err)
|
||||
}
|
||||
|
||||
resp, err := d.HTTPClient.DoWithRedirect(ctx, httpReq, 5, allowedHosts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
return nil, fmt.Errorf("%s returned %d", rawURL, resp.StatusCode)
|
||||
}
|
||||
return io.ReadAll(resp.Body)
|
||||
}
|
||||
|
||||
// sparkMetricsInput converts SHS executor/stage JSON into analyzer input.
|
||||
func (d *Deps) sparkMetricsInput(ctx context.Context, cl *cluster.Cluster, appID string) (analyzer.Input, error) {
|
||||
base := strings.TrimRight(cl.SHSURL, "/")
|
||||
execURL := base + "/api/v1/applications/" + appID + "/executors"
|
||||
stageURL := base + "/api/v1/applications/" + appID + "/stages"
|
||||
|
||||
execBody, err := d.fetchInternal(ctx, cl, execURL)
|
||||
if err != nil {
|
||||
return analyzer.Input{}, err
|
||||
}
|
||||
stageBody, err := d.fetchInternal(ctx, cl, stageURL)
|
||||
if err != nil {
|
||||
return analyzer.Input{}, err
|
||||
}
|
||||
|
||||
var execs []shsExecutor
|
||||
if err := json.Unmarshal(execBody, &execs); err != nil {
|
||||
return analyzer.Input{}, fmt.Errorf("parse executors: %w", err)
|
||||
}
|
||||
var stages []shsStage
|
||||
if err := json.Unmarshal(stageBody, &stages); err != nil {
|
||||
return analyzer.Input{}, fmt.Errorf("parse stages: %w", err)
|
||||
}
|
||||
|
||||
input := analyzer.Input{
|
||||
StageMetrics: make(map[string]analyzer.StageMetric),
|
||||
ExecutorMetrics: make(map[string]analyzer.ExecutorMetric),
|
||||
}
|
||||
for _, e := range execs {
|
||||
input.ExecutorMetrics[e.ID] = analyzer.ExecutorMetric{
|
||||
GCTimeMS: e.GCTimeMS,
|
||||
CPUTimeMS: e.CPUTimeMS,
|
||||
}
|
||||
}
|
||||
for _, s := range stages {
|
||||
key := fmt.Sprintf("stage %d", s.StageID)
|
||||
input.StageMetrics[key] = analyzer.StageMetric{
|
||||
DurationMS: s.DurationMS,
|
||||
ShuffleReadBytes: s.ShuffleReadBytes,
|
||||
ShuffleWriteBytes: s.ShuffleWriteBytes,
|
||||
MaxPartitionBytes: s.MaxPartitionBytes,
|
||||
MinPartitionBytes: s.MinPartitionBytes,
|
||||
}
|
||||
}
|
||||
return input, nil
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"spark-mcp-go/internal/analyzer"
|
||||
"spark-mcp-go/internal/cluster"
|
||||
"spark-mcp-go/internal/httpclient"
|
||||
"spark-mcp-go/internal/storage"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func testDepsWithCluster(t *testing.T, rmURL, shsURL string) *Deps {
|
||||
t.Helper()
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
cl := &cluster.Cluster{ID: "cluster-a", Name: "Cluster A", RMURL: rmURL, SHSURL: shsURL, SparkSubmitExecuteBin: "/usr/bin/spark-submit", IsActive: true, AuthType: cluster.AuthNone}
|
||||
if err := db.Clusters().Create(context.Background(), cl); err != nil {
|
||||
t.Fatalf("create cluster: %v", err)
|
||||
}
|
||||
return &Deps{HTTPClient: httpclient.New(httpclient.Config{Timeout: 5 * time.Second, MaxResponseBytes: 1 << 20}), ClusterRepo: db.Clusters(), MaxResponseBytes: 1 << 20, DataDir: t.TempDir(), AnalyzerThresholds: analyzer.Thresholds{DataSkewRatio: 3.0, GCPressureRatio: 0.1, BottleneckShuffleGB: 50.0}}
|
||||
}
|
||||
|
||||
func TestFetchSparkMetrics_Summary(t *testing.T) {
|
||||
shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/v1/applications/app_123/executors":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`[{"id":"1","gc_time_ms":5000,"cpu_time_ms":10000}]`))
|
||||
case "/api/v1/applications/app_123/stages":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`[{"stageId":1,"duration_ms":30000,"shuffle_read_bytes":0,"shuffle_write_bytes":0,"max_partition_bytes":100,"min_partition_bytes":10}]`))
|
||||
default:
|
||||
t.Errorf("unexpected SHS path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer shsSrv.Close()
|
||||
deps := testDepsWithCluster(t, "http://unused", shsSrv.URL)
|
||||
req := newToolRequest("fetch_spark_metrics", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123", "format": "summary"})
|
||||
res, err := deps.FetchSparkMetricsHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
findings, ok := payload["findings"].([]any)
|
||||
if !ok || len(findings) == 0 {
|
||||
t.Fatalf("expected findings, got %+v", payload["findings"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchClusterEnv(t *testing.T) {
|
||||
rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/ws/v1/cluster/info":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"clusterInfo":{"id":"rm1","name":"test"}}`))
|
||||
case "/ws/v1/cluster/metrics":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(`{"clusterMetrics":{"appsSubmitted":5}}`))
|
||||
default:
|
||||
t.Errorf("unexpected RM path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer rmSrv.Close()
|
||||
deps := testDepsWithCluster(t, rmSrv.URL, "http://unused")
|
||||
req := newToolRequest("fetch_cluster_env", map[string]any{"cluster_id": "cluster-a"})
|
||||
res, err := deps.FetchClusterEnvHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
if payload["cluster_info"] == nil {
|
||||
t.Errorf("cluster_info missing")
|
||||
}
|
||||
if payload["metrics"] == nil {
|
||||
t.Errorf("metrics missing")
|
||||
}
|
||||
if _, ok := payload["fetched_at"].(string); !ok {
|
||||
t.Errorf("fetched_at missing or not string")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnalyzeSparkLog(t *testing.T) {
|
||||
rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/ws/v1/cluster/apps/app_123/amContainerLogs":
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
case "/ws/v1/cluster/apps/app_123/aggregated-logs":
|
||||
w.Write([]byte("first line\nsecond line\nERROR: something"))
|
||||
default:
|
||||
t.Errorf("unexpected RM path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer rmSrv.Close()
|
||||
shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/v1/applications/app_123/executors":
|
||||
w.Write([]byte(`[]`))
|
||||
case "/api/v1/applications/app_123/stages":
|
||||
w.Write([]byte(`[]`))
|
||||
default:
|
||||
t.Errorf("unexpected SHS path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer shsSrv.Close()
|
||||
deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL)
|
||||
req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"})
|
||||
res, err := deps.AnalyzeSparkLogHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
text, ok := mcp.AsTextContent(res.Content[0])
|
||||
if !ok {
|
||||
t.Fatalf("content is not text: %T", res.Content[0])
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
prompt := payload["prompt"].(string)
|
||||
if !strings.Contains(prompt, "Spark Log Analysis") {
|
||||
t.Errorf("prompt missing header")
|
||||
}
|
||||
if !strings.Contains(payload["log_tail"].(string), "ERROR: something") {
|
||||
t.Errorf("log tail missing content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnalyzeSparkLog_LogSourceReported(t *testing.T) {
|
||||
rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/ws/v1/cluster/apps/app_123/amContainerLogs":
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
case "/ws/v1/cluster/apps/app_123/aggregated-logs":
|
||||
w.Write([]byte("log body"))
|
||||
default:
|
||||
t.Errorf("unexpected RM path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer rmSrv.Close()
|
||||
shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte(`[]`))
|
||||
}))
|
||||
defer shsSrv.Close()
|
||||
deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL)
|
||||
req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"})
|
||||
res, err := deps.AnalyzeSparkLogHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
text, _ := mcp.AsTextContent(res.Content[0])
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
if payload["log_source"] == "" {
|
||||
t.Errorf("log_source empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnalyzeSparkLog_PromptIncludesFindings(t *testing.T) {
|
||||
rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/ws/v1/cluster/apps/app_123/amContainerLogs":
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
case "/ws/v1/cluster/apps/app_123/aggregated-logs":
|
||||
w.Write([]byte("log body"))
|
||||
default:
|
||||
t.Errorf("unexpected RM path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer rmSrv.Close()
|
||||
shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/v1/applications/app_123/executors":
|
||||
w.Write([]byte(`[{"id":"1","gc_time_ms":5000,"cpu_time_ms":10000}]`))
|
||||
case "/api/v1/applications/app_123/stages":
|
||||
w.Write([]byte(`[]`))
|
||||
default:
|
||||
t.Errorf("unexpected SHS path %q", r.URL.Path)
|
||||
}
|
||||
}))
|
||||
defer shsSrv.Close()
|
||||
deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL)
|
||||
req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"})
|
||||
res, err := deps.AnalyzeSparkLogHandler(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("handler error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("unexpected error result: %v", res.Content)
|
||||
}
|
||||
text, _ := mcp.AsTextContent(res.Content[0])
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
prompt := payload["prompt"].(string)
|
||||
if !strings.Contains(prompt, "GC 压力") {
|
||||
t.Errorf("prompt missing GC finding; prompt:\n%s", prompt)
|
||||
}
|
||||
findings, ok := payload["findings"].([]any)
|
||||
if !ok || len(findings) == 0 {
|
||||
t.Fatalf("expected findings in result")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user