routing: add TlvTrafficShaper to bandwidth hints

routing/bandwidth.go

@@ -1,10 +1,14 @@
 package routing
 
 import (
+	"fmt"
+
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/htlcswitch"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/lightningnetwork/lnd/tlv"
 )
 
 // bandwidthHints provides hints about the currently available balance in our
@@ -19,6 +23,42 @@ type bandwidthHints interface {
 	// returned.
 	availableChanBandwidth(channelID uint64,
 		amount lnwire.MilliSatoshi) (lnwire.MilliSatoshi, bool)
+
+	// firstHopCustomBlob returns the custom blob for the first hop of the
+	// payment, if available.
+	firstHopCustomBlob() fn.Option[tlv.Blob]
+}
+
+// TlvTrafficShaper is an interface that allows the sender to determine if a
+// payment should be carried by a channel based on the TLV records that may be
+// present in the `update_add_htlc` message or the channel commitment itself.
+type TlvTrafficShaper interface {
+	AuxHtlcModifier
+
+	// ShouldHandleTraffic is called in order to check if the channel
+	// identified by the provided channel ID may have external mechanisms
+	// that would allow it to carry out the payment.
+	ShouldHandleTraffic(cid lnwire.ShortChannelID,
+		fundingBlob fn.Option[tlv.Blob]) (bool, error)
+
+	// PaymentBandwidth returns the available bandwidth for a custom channel
+	// decided by the given channel aux blob and HTLC blob. A return value
+	// of 0 means there is no bandwidth available. To find out if a channel
+	// is a custom channel that should be handled by the traffic shaper, the
+	// HandleTraffic method should be called first.
+	PaymentBandwidth(htlcBlob, commitmentBlob fn.Option[tlv.Blob],
+		linkBandwidth lnwire.MilliSatoshi) (lnwire.MilliSatoshi, error)
+}
+
+// AuxHtlcModifier is an interface that allows the sender to modify the outgoing
+// HTLC of a payment by changing the amount or the wire message tlv records.
+type AuxHtlcModifier interface {
+	// ProduceHtlcExtraData is a function that, based on the previous extra
+	// data blob of an HTLC, may produce a different blob or modify the
+	// amount of bitcoin this htlc should carry.
+	ProduceHtlcExtraData(totalAmount lnwire.MilliSatoshi,
+		htlcCustomRecords lnwire.CustomRecords) (lnwire.MilliSatoshi,
+		lnwire.CustomRecords, error)
 }
 
 // getLinkQuery is the function signature used to lookup a link.
@@ -29,8 +69,10 @@ type getLinkQuery func(lnwire.ShortChannelID) (
 // uses the link lookup provided to query the link for our latest local channel
 // balances.
 type bandwidthManager struct {
-	getLink    getLinkQuery
-	localChans map[lnwire.ShortChannelID]struct{}
+	getLink       getLinkQuery
+	localChans    map[lnwire.ShortChannelID]struct{}
+	firstHopBlob  fn.Option[tlv.Blob]
+	trafficShaper fn.Option[TlvTrafficShaper]
 }
 
 // newBandwidthManager creates a bandwidth manager for the source node provided
@@ -40,11 +82,14 @@ type bandwidthManager struct {
 // allows us to reduce the number of extraneous attempts as we can skip channels
 // that are inactive, or just don't have enough bandwidth to carry the payment.
 func newBandwidthManager(graph Graph, sourceNode route.Vertex,
-	linkQuery getLinkQuery) (*bandwidthManager, error) {
+	linkQuery getLinkQuery, firstHopBlob fn.Option[tlv.Blob],
+	trafficShaper fn.Option[TlvTrafficShaper]) (*bandwidthManager, error) {
 
 	manager := &bandwidthManager{
-		getLink:    linkQuery,
-		localChans: make(map[lnwire.ShortChannelID]struct{}),
+		getLink:       linkQuery,
+		localChans:    make(map[lnwire.ShortChannelID]struct{}),
+		firstHopBlob:  firstHopBlob,
+		trafficShaper: trafficShaper,
 	}
 
 	// First, we'll collect the set of outbound edges from the target
@@ -89,17 +134,111 @@ func (b *bandwidthManager) getBandwidth(cid lnwire.ShortChannelID,
 		return 0
 	}
 
-	// If our link isn't currently in a state where it can  add another
-	// outgoing htlc, treat the link as unusable.
-	if err := link.MayAddOutgoingHtlc(amount); err != nil {
-		log.Warnf("ShortChannelID=%v: cannot add outgoing htlc: %v",
-			cid, err)
+	// bandwidthResult is an inline type that we'll use to pass the
+	// bandwidth result from the external traffic shaper to the main logic
+	// below.
+	type bandwidthResult struct {
+		// bandwidth is the available bandwidth for the channel as
+		// reported by the external traffic shaper. If the external
+		// traffic shaper is not handling the channel, this value will
+		// be fn.None
+		bandwidth fn.Option[lnwire.MilliSatoshi]
+
+		// htlcAmount is the amount we're going to use to check if we
+		// can add another HTLC to the channel. If the external traffic
+		// shaper is handling the channel, we'll use 0 to just sanity
+		// check the number of HTLCs on the channel, since we don't know
+		// the actual HTLC amount that will be sent.
+		htlcAmount fn.Option[lnwire.MilliSatoshi]
+	}
+
+	var (
+		// We will pass the link bandwidth to the external traffic
+		// shaper. This is the current best estimate for the available
+		// bandwidth for the link.
+		linkBandwidth = link.Bandwidth()
+
+		bandwidthErr = func(err error) fn.Result[bandwidthResult] {
+			return fn.Err[bandwidthResult](err)
+		}
+	)
+
+	result, err := fn.MapOptionZ(
+		b.trafficShaper,
+		func(ts TlvTrafficShaper) fn.Result[bandwidthResult] {
+			fundingBlob := link.FundingCustomBlob()
+			shouldHandle, err := ts.ShouldHandleTraffic(
+				cid, fundingBlob,
+			)
+			if err != nil {
+				return bandwidthErr(fmt.Errorf("traffic "+
+					"shaper failed to decide whether to "+
+					"handle traffic: %w", err))
+			}
+
+			log.Debugf("ShortChannelID=%v: external traffic "+
+				"shaper is handling traffic: %v", cid,
+				shouldHandle)
+
+			// If this channel isn't handled by the external traffic
+			// shaper, we'll return early.
+			if !shouldHandle {
+				return fn.Ok(bandwidthResult{})
+			}
+
+			// Ask for a specific bandwidth to be used for the
+			// channel.
+			commitmentBlob := link.CommitmentCustomBlob()
+			auxBandwidth, err := ts.PaymentBandwidth(
+				b.firstHopBlob, commitmentBlob, linkBandwidth,
+			)
+			if err != nil {
+				return bandwidthErr(fmt.Errorf("failed to get "+
+					"bandwidth from external traffic "+
+					"shaper: %w", err))
+			}
+
+			log.Debugf("ShortChannelID=%v: external traffic "+
+				"shaper reported available bandwidth: %v", cid,
+				auxBandwidth)
+
+			// We don't know the actual HTLC amount that will be
+			// sent using the custom channel. But we'll still want
+			// to make sure we can add another HTLC, using the
+			// MayAddOutgoingHtlc method below. Passing 0 into that
+			// method will use the minimum HTLC value for the
+			// channel, which is okay to just check we don't exceed
+			// the max number of HTLCs on the channel. A proper
+			// balance check is done elsewhere.
+			return fn.Ok(bandwidthResult{
+				bandwidth:  fn.Some(auxBandwidth),
+				htlcAmount: fn.Some[lnwire.MilliSatoshi](0),
+			})
+		},
+	).Unpack()
+	if err != nil {
+		log.Errorf("ShortChannelID=%v: failed to get bandwidth from "+
+			"external traffic shaper: %v", cid, err)
+
 		return 0
 	}
 
-	// Otherwise, we'll return the current best estimate for the available
-	// bandwidth for the link.
-	return link.Bandwidth()
+	htlcAmount := result.htlcAmount.UnwrapOr(amount)
+
+	// If our link isn't currently in a state where it can add another
+	// outgoing htlc, treat the link as unusable.
+	if err := link.MayAddOutgoingHtlc(htlcAmount); err != nil {
+		log.Warnf("ShortChannelID=%v: cannot add outgoing "+
+			"htlc with amount %v: %v", cid, htlcAmount, err)
+		return 0
+	}
+
+	// If the external traffic shaper determined the bandwidth, we'll return
+	// that value, even if it is zero (which would mean no bandwidth is
+	// available on that channel).
+	reportedBandwidth := result.bandwidth.UnwrapOr(linkBandwidth)
+
+	return reportedBandwidth
 }
 
 // availableChanBandwidth returns the total available bandwidth for a channel
@@ -116,3 +255,9 @@ func (b *bandwidthManager) availableChanBandwidth(channelID uint64,
 
 	return b.getBandwidth(shortID, amount), true
 }
+
+// firstHopCustomBlob returns the custom blob for the first hop of the payment,
+// if available.
+func (b *bandwidthManager) firstHopCustomBlob() fn.Option[tlv.Blob] {
+	return b.firstHopBlob
+}