Refactor MTProxy checker to support multiple DC probing with enhanced timeout handling. Update README and Docker documentation to include new -dc-ids flag for sequential DC checks, clarify exit codes, and improve error reporting for each DC. Adjust logic in checkOneLine and related functions to aggregate results from multiple DCs, ensuring robust performance and clearer output.

2026-04-11 15:12:08 +07:00
parent a846ef5373
commit 24a7315a16
6 changed files with 245 additions and 53 deletions
@@ -19,6 +19,7 @@ import (
 	"mtproxy_checker/internal/allowlist"
 	"mtproxy_checker/internal/checker"
 	"mtproxy_checker/internal/checkresult"
+	"mtproxy_checker/internal/dclist"
 	"mtproxy_checker/internal/parseurl"
 	"mtproxy_checker/internal/secret"
 )
@@ -35,7 +36,7 @@ type config struct {
 	checkInterval   time.Duration
 	httpAddr        string
 	checkTimeout    time.Duration
-	dcID            int16
+	dcIDs           []int16
 	probe           checker.ProbeMode
 	allowedPrefixes []netip.Prefix
 }
@@ -66,14 +67,22 @@ func loadConfig() (*config, error) {
 	if err != nil || checkTimeout <= 0 {
 		return nil, fmt.Errorf("MTPROXY_CHECK_TIMEOUT: invalid duration %q", timeoutStr)
 	}
-	dcStr := strings.TrimSpace(os.Getenv("MTPROXY_DC_ID"))
-	if dcStr == "" {
-		dcStr = "2"
+	dcIDsRaw := strings.TrimSpace(os.Getenv("MTPROXY_DC_IDS"))
+	dcIDs, err := dclist.ParseList(dcIDsRaw)
+	if err != nil {
+		return nil, fmt.Errorf("MTPROXY_DC_IDS: %w", err)
 	}
-	var dcParsed int64
-	_, err = fmt.Sscanf(dcStr, "%d", &dcParsed)
-	if err != nil || dcParsed < -32768 || dcParsed > 32767 {
-		return nil, fmt.Errorf("MTPROXY_DC_ID: invalid int16 %q", dcStr)
+	if dcIDs == nil {
+		dcStr := strings.TrimSpace(os.Getenv("MTPROXY_DC_ID"))
+		if dcStr == "" {
+			dcStr = "2"
+		}
+		var dcParsed int64
+		_, err = fmt.Sscanf(dcStr, "%d", &dcParsed)
+		if err != nil || dcParsed < -32768 || dcParsed > 32767 {
+			return nil, fmt.Errorf("MTPROXY_DC_ID: invalid int16 %q", dcStr)
+		}
+		dcIDs = []int16{int16(dcParsed)}
 	}
 	allowedRaw := os.Getenv("MTPROXY_ALLOWED_IPS")
 	prefixes, err := allowlist.ParseCommaList(allowedRaw)
@@ -86,20 +95,28 @@ func loadConfig() (*config, error) {
 		checkInterval:   interval,
 		httpAddr:        httpAddr,
 		checkTimeout:    checkTimeout,
-		dcID:            int16(dcParsed),
+		dcIDs:           dcIDs,
 		probe:           probe,
 		allowedPrefixes: prefixes,
 	}, nil
 }

+type dcProbeResult struct {
+	DC       int16  `json:"dc"`
+	OK       bool   `json:"ok"`
+	ExitCode int    `json:"exit_code"`
+	Error    string `json:"error,omitempty"`
+}
+
 type proxyEntry struct {
-	RawLine    string `json:"raw_line"`
-	URL        string `json:"url,omitempty"`
-	OK         bool   `json:"ok"`
-	ExitCode   int    `json:"exit_code"`
-	Error      string `json:"error,omitempty"`
-	ParseError string `json:"parse_error,omitempty"`
-	CheckedAt  string `json:"checked_at,omitempty"`
+	RawLine    string          `json:"raw_line"`
+	URL        string          `json:"url,omitempty"`
+	OK         bool            `json:"ok"`
+	ExitCode   int             `json:"exit_code"`
+	Error      string          `json:"error,omitempty"`
+	ParseError string          `json:"parse_error,omitempty"`
+	CheckedAt  string          `json:"checked_at,omitempty"`
+	DCs        []dcProbeResult `json:"dcs,omitempty"`
 }

 type snapshot struct {
@@ -149,9 +166,33 @@ func readProxyLines(path string) ([]string, error) {
 	return lines, nil
 }

-func checkOneLine(ctx context.Context, line string, dcID int16, probe checker.ProbeMode) proxyEntry {
+func aggregateExitFromDCs(dcs []dcProbeResult) int {
+	has3, has4 := false, false
+	for _, d := range dcs {
+		if d.ExitCode == 3 {
+			has3 = true
+		}
+		if d.ExitCode == 4 {
+			has4 = true
+		}
+	}
+	if has3 {
+		return 3
+	}
+	if has4 {
+		return 4
+	}
+	return 1
+}
+
+func checkOneLine(ctx context.Context, line string, dcIDs []int16, perDCTimeout time.Duration, probe checker.ProbeMode) proxyEntry {
 	now := time.Now().UTC().Format(time.RFC3339)
 	ent := proxyEntry{RawLine: line, CheckedAt: now}
+	if len(dcIDs) == 0 {
+		ent.ExitCode = 2
+		ent.Error = "no DC ids configured"
+		return ent
+	}
 	t, err := parseurl.ParseTGProxy(line)
 	if err != nil {
 		ent.ExitCode = 2
@@ -167,12 +208,46 @@ func checkOneLine(ctx context.Context, line string, dcID int16, probe checker.Pr
 		ent.Error = ent.ParseError
 		return ent
 	}
-	err = checker.Check(ctx, t.Host, t.Port, parsed, dcID, &checker.Options{Probe: probe})
-	code, msg := checkresult.Classify(err)
-	ent.ExitCode = code
-	ent.OK = err == nil
-	if msg != "" {
-		ent.Error = msg
+	if len(dcIDs) == 1 {
+		dcCtx, cancel := context.WithTimeout(ctx, perDCTimeout)
+		defer cancel()
+		err = checker.Check(dcCtx, t.Host, t.Port, parsed, dcIDs[0], &checker.Options{Probe: probe})
+		code, msg := checkresult.Classify(err)
+		ent.ExitCode = code
+		ent.OK = err == nil
+		if msg != "" {
+			ent.Error = msg
+		}
+		return ent
+	}
+	var dcs []dcProbeResult
+	anyOK := false
+	for _, dc := range dcIDs {
+		dcCtx, cancel := context.WithTimeout(ctx, perDCTimeout)
+		err := checker.Check(dcCtx, t.Host, t.Port, parsed, dc, &checker.Options{Probe: probe})
+		cancel()
+		code, msg := checkresult.Classify(err)
+		ok := err == nil
+		if ok {
+			anyOK = true
+		}
+		dcs = append(dcs, dcProbeResult{DC: dc, OK: ok, ExitCode: code, Error: msg})
+	}
+	ent.DCs = dcs
+	ent.OK = anyOK
+	if ent.OK {
+		ent.ExitCode = 0
+	} else {
+		ent.ExitCode = aggregateExitFromDCs(dcs)
+		var parts []string
+		for _, d := range dcs {
+			if d.Error != "" {
+				parts = append(parts, fmt.Sprintf("dc%d: %s", d.DC, d.Error))
+			} else {
+				parts = append(parts, fmt.Sprintf("dc%d: fail", d.DC))
+			}
+		}
+		ent.Error = strings.Join(parts, "; ")
 	}
 	return ent
 }
@@ -193,8 +268,12 @@ func runCycle(cfg *config, st *store) {
 	}
 	entries := make([]proxyEntry, 0, len(lines))
 	for _, line := range lines {
-		ctx, cancel := context.WithTimeout(context.Background(), cfg.checkTimeout)
-		ent := checkOneLine(ctx, line, cfg.dcID, cfg.probe)
+		total := cfg.checkTimeout
+		if len(cfg.dcIDs) > 1 {
+			total = cfg.checkTimeout * time.Duration(len(cfg.dcIDs))
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), total)
+		ent := checkOneLine(ctx, line, cfg.dcIDs, cfg.checkTimeout, cfg.probe)
 		cancel()
 		entries = append(entries, ent)
 	}
@@ -274,7 +353,7 @@ func run() error {

 	errCh := make(chan error, 1)
 	go func() {
-		log.Printf("listening on %s, list=%s interval=%s", cfg.httpAddr, cfg.listFile, cfg.checkInterval)
+		log.Printf("listening on %s, list=%s interval=%s dcs=%v", cfg.httpAddr, cfg.listFile, cfg.checkInterval, cfg.dcIDs)
 		errCh <- srv.ListenAndServe()
 	}()