From 33dbd1c11eee87c9c652c387dbc765bcc52a1086 Mon Sep 17 00:00:00 2001
From: Nikolay Petrov <sungon@gmail.com>
Date: Sat, 18 Jan 2025 12:46:44 -0500
Subject: [PATCH] relay stats

---
 client/peer_relay.go | 9 ++++++++-
 relay/clients.go     | 5 +++++
 relay/server.go      | 2 ++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/client/peer_relay.go b/client/peer_relay.go
index de62c94..9df25e7 100644
--- a/client/peer_relay.go
+++ b/client/peer_relay.go
@@ -13,6 +13,7 @@ import (
 	"github.com/connet-dev/connet/netc"
 	"github.com/connet-dev/connet/pb"
 	"github.com/connet-dev/connet/pbc"
+	"github.com/connet-dev/connet/quicc"
 	"github.com/quic-go/quic-go"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -88,7 +89,7 @@ func (r *relayPeer) connect(ctx context.Context) (quic.Connection, error) {
 
 	cfg := r.serverConf.Load()
 	r.logger.Debug("dialing relay", "relay", r.serverHostport, "addr", addr, "server", cfg.name, "cert", cfg.key)
-	return r.local.direct.transport.Dial(ctx, addr, &tls.Config{
+	return r.local.direct.transport.Dial(quicc.RTTContext(ctx), addr, &tls.Config{
 		Certificates: []tls.Certificate{r.local.clientCert},
 		RootCAs:      cfg.cas,
 		ServerName:   cfg.name,
@@ -96,6 +97,7 @@ func (r *relayPeer) connect(ctx context.Context) (quic.Connection, error) {
 	}, &quic.Config{
 		MaxIdleTimeout:  20 * time.Second,
 		KeepAlivePeriod: 10 * time.Second,
+		Tracer:          quicc.RTTTracer,
 	})
 }
 
@@ -120,6 +122,11 @@ func (r *relayPeer) keepalive(ctx context.Context, conn quic.Connection) error {
 		if err := r.heartbeat(ctx, stream); err != nil {
 			return err
 		}
+		if rttStats := quicc.RTTStats(conn); rttStats != nil {
+			r.logger.Debug("rtt-pro", "last", rttStats.LatestRTT(), "smoothed", rttStats.SmoothedRTT())
+		} else {
+			r.logger.Debug("no rtt")
+		}
 	}
 }
 
diff --git a/relay/clients.go b/relay/clients.go
index 241b5f4..46c03de 100644
--- a/relay/clients.go
+++ b/relay/clients.go
@@ -15,6 +15,7 @@ import (
 	"github.com/connet-dev/connet/netc"
 	"github.com/connet-dev/connet/pb"
 	"github.com/connet-dev/connet/pbc"
+	"github.com/connet-dev/connet/quicc"
 	"github.com/klev-dev/kleverr"
 	"github.com/quic-go/quic-go"
 	"golang.org/x/sync/errgroup"
@@ -169,6 +170,7 @@ func (s *clientsServer) run(ctx context.Context, transport *quic.Transport) erro
 	l, err := transport.Listen(s.tlsConf, &quic.Config{
 		MaxIdleTimeout:  20 * time.Second,
 		KeepAlivePeriod: 10 * time.Second,
+		Tracer:          quicc.RTTTracer,
 	})
 	if err != nil {
 		return kleverr.Ret(err)
@@ -369,6 +371,9 @@ func (c *clientConn) heartbeat(ctx context.Context, stream quic.Stream, hbt *pbc
 			if err := pb.Write(stream, &pbc.Response{Heartbeat: req.Heartbeat}); err != nil {
 				return err
 			}
+			if rttStats := quicc.RTTStats(c.conn); rttStats != nil {
+				c.logger.Debug("rtt", "last", rttStats.LatestRTT(), "smoothed", rttStats.SmoothedRTT())
+			}
 		}
 	})
 
diff --git a/relay/server.go b/relay/server.go
index cc791b3..d0e8fcd 100644
--- a/relay/server.go
+++ b/relay/server.go
@@ -11,6 +11,7 @@ import (
 	"slices"
 
 	"github.com/connet-dev/connet/model"
+	"github.com/connet-dev/connet/quicc"
 	"github.com/connet-dev/connet/statusc"
 	"github.com/klev-dev/kleverr"
 	"github.com/quic-go/quic-go"
@@ -94,6 +95,7 @@ func (s *Server) Run(ctx context.Context) error {
 		Conn:               udpConn,
 		ConnectionIDLength: 8,
 		StatelessResetKey:  s.statelessResetKey,
+		ConnContext:        quicc.RTTContext,
 		// TODO review other options
 	}
 	defer transport.Close()