From c088475dd19599e2643abd8a438125ed95bacccb Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 4 Dec 2024 12:03:55 +0100
Subject: [PATCH] multi: thread through and use AuxTrafficShaper

---
 htlcswitch/link.go | 36 +++++++++++++++++++++++++++++++++++-
 peer/brontide.go   |  5 +++++
 server.go          |  1 +
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/htlcswitch/link.go b/htlcswitch/link.go
index 75a35302cb..44e8f07188 100644
--- a/htlcswitch/link.go
+++ b/htlcswitch/link.go
@@ -293,6 +293,10 @@ type ChannelLinkConfig struct {
 	// ShouldFwdExpEndorsement is a closure that indicates whether the link
 	// should forward experimental endorsement signals.
 	ShouldFwdExpEndorsement func() bool
+
+	// AuxTrafficShaper is an optional auxiliary traffic shaper that can be
+	// used to manage the bandwidth of the link.
+	AuxTrafficShaper fn.Option[AuxTrafficShaper]
 }
 
 // channelLink is the service which drives a channel's commitment update
@@ -3401,8 +3405,38 @@ func (l *channelLink) canSendHtlc(policy models.ForwardingPolicy,
 		return NewLinkError(&lnwire.FailExpiryTooFar{})
 	}
 
+	// We now check the available bandwidth to see if this HTLC can be
+	// forwarded.
+	availableBandwidth := l.Bandwidth()
+	auxBandwidth, err := fn.MapOptionZ(
+		l.cfg.AuxTrafficShaper,
+		func(ts AuxTrafficShaper) fn.Result[OptionalBandwidth] {
+			var htlcBlob fn.Option[tlv.Blob]
+			blob, err := customRecords.Serialize()
+			if err != nil {
+				return fn.Err[OptionalBandwidth](
+					fmt.Errorf("unable to serialize "+
+						"custom records: %w", err))
+			}
+
+			if len(blob) > 0 {
+				htlcBlob = fn.Some(blob)
+			}
+
+			return l.AuxBandwidth(amt, originalScid, htlcBlob, ts)
+		},
+	).Unpack()
+	if err != nil {
+		l.log.Warnf("Unable to determine aux bandwidth: %v", err)
+		return NewLinkError(&lnwire.FailTemporaryNodeFailure{})
+	}
+
+	auxBandwidth.WhenSome(func(bandwidth lnwire.MilliSatoshi) {
+		availableBandwidth = bandwidth
+	})
+
 	// Check to see if there is enough balance in this channel.
-	if amt > l.Bandwidth() {
+	if amt > availableBandwidth {
 		l.log.Warnf("insufficient bandwidth to route htlc: %v is "+
 			"larger than %v", amt, l.Bandwidth())
 		cb := func(upd *lnwire.ChannelUpdate1) lnwire.FailureMessage {
diff --git a/peer/brontide.go b/peer/brontide.go
index 6bc49445ee..7074b10071 100644
--- a/peer/brontide.go
+++ b/peer/brontide.go
@@ -400,6 +400,10 @@ type Config struct {
 	// way contracts are resolved.
 	AuxResolver fn.Option[lnwallet.AuxContractResolver]
 
+	// AuxTrafficShaper is an optional auxiliary traffic shaper that can be
+	// used to manage the bandwidth of peer links.
+	AuxTrafficShaper fn.Option[htlcswitch.AuxTrafficShaper]
+
 	// PongBuf is a slice we'll reuse instead of allocating memory on the
 	// heap. Since only reads will occur and no writes, there is no need
 	// for any synchronization primitives. As a result, it's safe to share
@@ -1330,6 +1334,7 @@ func (p *Brontide) addLink(chanPoint *wire.OutPoint,
 		ShouldFwdExpEndorsement: p.cfg.ShouldFwdExpEndorsement,
 		DisallowQuiescence: p.cfg.DisallowQuiescence ||
 			!p.remoteFeatures.HasFeature(lnwire.QuiescenceOptional),
+		AuxTrafficShaper: p.cfg.AuxTrafficShaper,
 	}
 
 	// Before adding our new link, purge the switch of any pending or live
diff --git a/server.go b/server.go
index f8f8239ed6..7f725553a3 100644
--- a/server.go
+++ b/server.go
@@ -4222,6 +4222,7 @@ func (s *server) peerConnected(conn net.Conn, connReq *connmgr.ConnReq,
 		MsgRouter:              s.implCfg.MsgRouter,
 		AuxChanCloser:          s.implCfg.AuxChanCloser,
 		AuxResolver:            s.implCfg.AuxContractResolver,
+		AuxTrafficShaper:       s.implCfg.TrafficShaper,
 		ShouldFwdExpEndorsement: func() bool {
 			if s.cfg.ProtocolOptions.NoExperimentalEndorsement() {
 				return false