feat: 基于KCP的Ping程序，利用目前的传输栈
This commit is contained in:
246
cmd/kcpping/platform_linux.go
Normal file
246
cmd/kcpping/platform_linux.go
Normal file
@@ -0,0 +1,246 @@
|
||||
//go:build linux
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"omnisocketgo/cmd/internal/latencylog"
|
||||
peerpkg "omnisocketgo/cmd/internal/peer"
|
||||
"omnisocketgo/cmd/internal/protocol"
|
||||
)
|
||||
|
||||
func runPlatform(cfg config, stdout, stderr io.Writer, now func() time.Time) error {
|
||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
||||
defer stop()
|
||||
|
||||
client, closeLogger, err := dialKCPClient(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer closeLogger()
|
||||
defer client.Close()
|
||||
|
||||
if cfg.echo {
|
||||
return runEchoMode(ctx, client, stderr)
|
||||
}
|
||||
return runPingMode(ctx, client, cfg, stdout, stderr, now)
|
||||
}
|
||||
|
||||
func dialKCPClient(cfg config) (*peerpkg.KCPClient, func(), error) {
|
||||
options := make([]peerpkg.Option, 0, 3)
|
||||
closeLogger := func() {}
|
||||
|
||||
if cfg.latencyLog != "" {
|
||||
logger, err := latencylog.NewJSONLLogger(cfg.latencyLog)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("create latency logger %s: %w", cfg.latencyLog, err)
|
||||
}
|
||||
options = append(options, peerpkg.WithLogger(logger))
|
||||
closeLogger = func() {
|
||||
_ = logger.Close()
|
||||
}
|
||||
}
|
||||
if cfg.bindIP != "" {
|
||||
options = append(options, peerpkg.WithBindIP(cfg.bindIP))
|
||||
}
|
||||
if cfg.bindDevice != "" {
|
||||
options = append(options, peerpkg.WithBindDevice(cfg.bindDevice))
|
||||
}
|
||||
|
||||
client, err := peerpkg.DialKCP(cfg.server, cfg.id, options...)
|
||||
if err != nil {
|
||||
closeLogger()
|
||||
return nil, nil, fmt.Errorf("dial kcp server %s: %w", cfg.server, err)
|
||||
}
|
||||
return client, closeLogger, nil
|
||||
}
|
||||
|
||||
func runPingMode(ctx context.Context, client *peerpkg.KCPClient, cfg config, stdout, stderr io.Writer, now func() time.Time) error {
|
||||
if err := writePingHeader(stdout, cfg); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
receiveCh := make(chan protocol.Message, 32)
|
||||
receiveErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
for {
|
||||
msg, err := client.Receive()
|
||||
if err != nil {
|
||||
receiveErrCh <- err
|
||||
return
|
||||
}
|
||||
receiveCh <- msg
|
||||
}
|
||||
}()
|
||||
|
||||
tracker := newPingTracker(cfg.timeout)
|
||||
expiryTicker := time.NewTicker(expiryPollInterval(cfg.timeout))
|
||||
defer expiryTicker.Stop()
|
||||
|
||||
var sendTicker *time.Ticker
|
||||
if cfg.count == 0 || cfg.count > 1 {
|
||||
sendTicker = time.NewTicker(cfg.interval)
|
||||
defer sendTicker.Stop()
|
||||
}
|
||||
|
||||
nextSeq := uint64(1)
|
||||
if err := sendPing(client, tracker, cfg, nextSeq, now); err != nil {
|
||||
return err
|
||||
}
|
||||
nextSeq++
|
||||
|
||||
stopSending := cfg.count == 1
|
||||
receiveErrSeen := false
|
||||
|
||||
for {
|
||||
if stopSending && tracker.pendingCount() == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
var sendTick <-chan time.Time
|
||||
if !stopSending && sendTicker != nil {
|
||||
sendTick = sendTicker.C
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
stopSending = true
|
||||
case <-expiryTicker.C:
|
||||
for _, seq := range tracker.expire(now()) {
|
||||
if err := writeTimeout(stdout, seq); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case msg := <-receiveCh:
|
||||
if err := handlePingMessage(tracker, msg, stdout, stderr, now); err != nil {
|
||||
return err
|
||||
}
|
||||
case err := <-receiveErrCh:
|
||||
receiveErrSeen = true
|
||||
if ctx.Err() != nil && isExpectedCloseError(err) {
|
||||
break
|
||||
}
|
||||
if stopSending && tracker.pendingCount() == 0 && isExpectedCloseError(err) {
|
||||
break
|
||||
}
|
||||
return fmt.Errorf("receive reply: %w", err)
|
||||
case <-sendTick:
|
||||
if cfg.count > 0 && tracker.sent >= cfg.count {
|
||||
stopSending = true
|
||||
continue
|
||||
}
|
||||
if err := sendPing(client, tracker, cfg, nextSeq, now); err != nil {
|
||||
return err
|
||||
}
|
||||
nextSeq++
|
||||
if cfg.count > 0 && tracker.sent >= cfg.count {
|
||||
stopSending = true
|
||||
}
|
||||
}
|
||||
|
||||
if receiveErrSeen && stopSending && tracker.pendingCount() == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
for _, seq := range tracker.expire(now()) {
|
||||
if err := writeTimeout(stdout, seq); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return writeSummary(stdout, cfg.to, tracker.summary())
|
||||
}
|
||||
|
||||
func sendPing(client *peerpkg.KCPClient, tracker *pingTracker, cfg config, seq uint64, now func() time.Time) error {
|
||||
sentAt := now()
|
||||
payload, err := buildPingPayload(seq, sentAt.UnixNano(), cfg.size)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := client.SendText(cfg.to, string(payload)); err != nil {
|
||||
return fmt.Errorf("send ping seq=%d: %w", seq, err)
|
||||
}
|
||||
tracker.markSent(seq, sentAt)
|
||||
return nil
|
||||
}
|
||||
|
||||
func handlePingMessage(tracker *pingTracker, msg protocol.Message, stdout, stderr io.Writer, now func() time.Time) error {
|
||||
switch msg.Type {
|
||||
case protocol.MessageTypeText:
|
||||
payload, err := parsePingPayload(msg.Body)
|
||||
if err != nil {
|
||||
_, writeErr := fmt.Fprintf(stderr, "ignore non-ping text message from %s: %v\n", msg.From, err)
|
||||
if writeErr != nil {
|
||||
return writeErr
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
result := tracker.observeReply(payload, now())
|
||||
switch result.disposition {
|
||||
case replyMatched:
|
||||
return writeMatchedReply(stdout, payload.Seq, result.rtt)
|
||||
case replyDuplicate:
|
||||
_, err := fmt.Fprintf(stderr, "seq=%d duplicate or late reply ignored\n", payload.Seq)
|
||||
return err
|
||||
case replyUnexpected:
|
||||
_, err := fmt.Fprintf(stderr, "seq=%d unexpected reply ignored\n", payload.Seq)
|
||||
return err
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
case protocol.MessageTypeError:
|
||||
_, err := fmt.Fprintf(stderr, "server error: %s\n", strings.TrimSpace(string(msg.Body)))
|
||||
return err
|
||||
default:
|
||||
_, err := fmt.Fprintf(stderr, "unexpected message type %s from %s ignored\n", msg.Type, msg.From)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
func runEchoMode(ctx context.Context, client *peerpkg.KCPClient, stderr io.Writer) error {
|
||||
receiveErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
receiveErrCh <- client.ReceiveLoop(func(msg protocol.Message) error {
|
||||
switch msg.Type {
|
||||
case protocol.MessageTypeText:
|
||||
return client.SendText(msg.From, string(msg.Body))
|
||||
case protocol.MessageTypeError:
|
||||
_, err := fmt.Fprintf(stderr, "server error: %s\n", strings.TrimSpace(string(msg.Body)))
|
||||
return err
|
||||
default:
|
||||
_, err := fmt.Fprintf(stderr, "unexpected message type %s from %s ignored\n", msg.Type, msg.From)
|
||||
return err
|
||||
}
|
||||
})
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
case err := <-receiveErrCh:
|
||||
if err == nil || (ctx.Err() != nil && isExpectedCloseError(err)) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("echo receive loop: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
func isExpectedCloseError(err error) bool {
|
||||
if err == nil {
|
||||
return true
|
||||
}
|
||||
message := err.Error()
|
||||
return errors.Is(err, context.Canceled) ||
|
||||
strings.Contains(message, "closed") ||
|
||||
strings.Contains(message, "broken pipe") ||
|
||||
strings.Contains(message, "io: read/write on closed pipe")
|
||||
}
|
||||
Reference in New Issue
Block a user