Add WaitingOnPeer/MonitorPersist spans to outbound HTLCs

Author: wvanlint

lightning/src/ln/channel.rs

@@ -593,15 +593,52 @@ impl OutboundHTLCOutput {
 	where
 		L::Target: Logger,
 	{
+		mem::drop(self.state_wrapper.waiting_on_peer_span.take());
+		mem::drop(self.state_wrapper.waiting_on_monitor_persist_span.take());
 		mem::drop(self.state_wrapper.span.take());
 		self.state_wrapper =
 			OutboundHTLCStateWrapper::new(state, self.span.as_user_span_ref::<L>(), logger);
 	}
+
+	fn waiting_on_peer<L: Deref>(&mut self, logger: &L)
+	where
+		L::Target: Logger,
+	{
+		self.state_wrapper.waiting_on_peer_span = Some(BoxedSpan::new(logger.start(
+			Span::WaitingOnPeer,
+			self.state_wrapper.span.as_ref().map(|s| s.as_user_span_ref::<L>()).flatten(),
+		)));
+	}
+
+	fn peer_responded(&mut self) {
+		if self.state_wrapper.waiting_on_peer_span.is_some() {
+			mem::drop(self.state_wrapper.waiting_on_peer_span.take());
+		}
+	}
+
+	fn waiting_on_monitor_persist<L: Deref>(&mut self, logger: &L)
+	where
+		L::Target: Logger,
+	{
+		self.state_wrapper.waiting_on_monitor_persist_span = Some(BoxedSpan::new(logger.start(
+			Span::WaitingOnMonitorPersist,
+			self.state_wrapper.span.as_ref().map(|s| s.as_user_span_ref::<L>()).flatten(),
+		)));
+	}
+
+	fn monitor_persisted(&mut self) {
+		if self.state_wrapper.waiting_on_monitor_persist_span.is_some() {
+			mem::drop(self.state_wrapper.waiting_on_monitor_persist_span.take());
+		}
+	}
 }
 
 #[cfg_attr(test, derive(Clone, Debug, PartialEq))]
 struct OutboundHTLCStateWrapper {
 	state: OutboundHTLCState,
+	waiting_on_peer_span: Option<BoxedSpan>,
+	waiting_on_monitor_persist_span: Option<BoxedSpan>,
+	// Drop full span last.
 	span: Option<BoxedSpan>,
 }
 
@@ -615,7 +652,12 @@ impl OutboundHTLCStateWrapper {
 	{
 		let state_span =
 			logger.start(Span::OutboundHTLCState { state: (&state).into() }, parent_span);
-		OutboundHTLCStateWrapper { state, span: Some(BoxedSpan::new(state_span)) }
+		OutboundHTLCStateWrapper {
+			state,
+			span: Some(BoxedSpan::new(state_span)),
+			waiting_on_peer_span: None,
+			waiting_on_monitor_persist_span: None,
+		}
 	}
 }
 
@@ -6866,6 +6908,7 @@ where
 						return Err(ChannelError::close(format!("Remote tried to fulfill/fail HTLC ({}) before it had been committed", htlc_id))),
 					OutboundHTLCState::Committed => {
 						htlc.set_state(OutboundHTLCState::RemoteRemoved(outcome), logger);
+						htlc.waiting_on_peer(logger);
 					},
 					OutboundHTLCState::AwaitingRemoteRevokeToRemove(_) | OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) | OutboundHTLCState::RemoteRemoved(_) =>
 						return Err(ChannelError::close(format!("Remote tried to fulfill/fail HTLC ({}) that they'd already fulfilled/failed", htlc_id))),
@@ -7284,6 +7327,9 @@ where
 					claimed_htlcs.push((SentHTLCId::from_source(&htlc.source), preimage));
 				}
 				htlc.set_state(OutboundHTLCState::AwaitingRemoteRevokeToRemove(reason), logger);
+				if self.context.channel_state.is_awaiting_remote_revoke() {
+					htlc.waiting_on_peer(logger);
+				}
 				need_commitment = true;
 			}
 		}
@@ -7849,6 +7895,7 @@ where
 					let mut reason = OutboundHTLCOutcome::Success(PaymentPreimage([0u8; 32]), None);
 					mem::swap(outcome, &mut reason);
 					htlc.set_state(OutboundHTLCState::AwaitingRemovedRemoteRevoke(reason), logger);
+					htlc.waiting_on_monitor_persist(logger);
 					require_commitment = true;
 				}
 			}
@@ -8288,6 +8335,16 @@ where
 				_ => {},
 			}
 		}
+		for htlc in self.context.pending_outbound_htlcs.iter_mut() {
+			match htlc.state() {
+				OutboundHTLCState::LocalAnnounced(_) |
+				OutboundHTLCState::AwaitingRemovedRemoteRevoke(_) => {
+					htlc.monitor_persisted();
+					htlc.waiting_on_peer(logger);
+				},
+				_ => {},
+			}
+		}
 
 		// If we're past (or at) the AwaitingChannelReady stage on an outbound (or V2-established) channel,
 		// try to (re-)broadcast the funding transaction as we may have declined to broadcast it when we
@@ -10964,7 +11021,7 @@ where
 		// that are simple to implement, and we do it on the outgoing side because then the failure message that encodes
 		// the hold time still needs to be built in channel manager.
 		let send_timestamp = duration_since_epoch();
-		self.context.pending_outbound_htlcs.push(OutboundHTLCOutput::new(
+		let mut htlc = OutboundHTLCOutput::new(
 			self.context.channel_id(),
 			OutboundHTLCOutputParams {
 				htlc_id: self.context.next_holder_htlc_id,
@@ -10979,7 +11036,9 @@ where
 			},
 			forward_span,
 			logger,
-		));
+		);
+		htlc.waiting_on_monitor_persist(logger);
+		self.context.pending_outbound_htlcs.push(htlc);
 		self.context.next_holder_htlc_id += 1;
 
 		Ok(true)
@@ -11029,6 +11088,7 @@ where
 				let mut reason = OutboundHTLCOutcome::Success(PaymentPreimage([0u8; 32]), None);
 				mem::swap(outcome, &mut reason);
 				htlc.set_state(OutboundHTLCState::AwaitingRemovedRemoteRevoke(reason), logger);
+				htlc.waiting_on_monitor_persist(logger);
 			}
 		}
 		if let Some((feerate, update_state)) = self.context.pending_update_fee {

lightning/src/ln/functional_tests.rs

@@ -11797,6 +11797,20 @@ pub fn test_payment_traces() {
 				},
 				Some(Span::OutboundHTLC { channel_id: channel_id_2, htlc_id: 0 }),
 			),
+			TestSpanBoundary::Start(
+				Span::WaitingOnMonitorPersist,
+				Some(Span::OutboundHTLCState {
+					state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToAdd
+				}),
+			),
+			TestSpanBoundary::End(Span::WaitingOnMonitorPersist,),
+			TestSpanBoundary::Start(
+				Span::WaitingOnPeer,
+				Some(Span::OutboundHTLCState {
+					state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToAdd
+				}),
+			),
+			TestSpanBoundary::End(Span::WaitingOnPeer,),
 			TestSpanBoundary::End(Span::OutboundHTLCState {
 				state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToAdd
 			}),
@@ -11811,6 +11825,10 @@ pub fn test_payment_traces() {
 				Span::OutboundHTLCState { state: OutboundHTLCStateDetails::Committed },
 				Some(Span::OutboundHTLC { channel_id: channel_id_2, htlc_id: 0 }),
 			),
+			TestSpanBoundary::Start(
+				Span::WaitingOnPeer,
+				Some(Span::OutboundHTLCState { state: OutboundHTLCStateDetails::Committed }),
+			),
 			TestSpanBoundary::End(Span::InboundHTLCState {
 				state: Some(InboundHTLCStateDetails::Committed)
 			}),
@@ -11833,6 +11851,7 @@ pub fn test_payment_traces() {
 					state: Some(InboundHTLCStateDetails::AwaitingRemoteRevokeToRemoveFulfill)
 				}),
 			),
+			TestSpanBoundary::End(Span::WaitingOnPeer),
 			TestSpanBoundary::End(Span::OutboundHTLCState {
 				state: OutboundHTLCStateDetails::Committed
 			}),
@@ -11851,6 +11870,20 @@ pub fn test_payment_traces() {
 				},
 				Some(Span::OutboundHTLC { channel_id: channel_id_2, htlc_id: 0 }),
 			),
+			TestSpanBoundary::Start(
+				Span::WaitingOnMonitorPersist,
+				Some(Span::OutboundHTLCState {
+					state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToRemoveSuccess
+				}),
+			),
+			TestSpanBoundary::End(Span::WaitingOnMonitorPersist),
+			TestSpanBoundary::Start(
+				Span::WaitingOnPeer,
+				Some(Span::OutboundHTLCState {
+					state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToRemoveSuccess
+				}),
+			),
+			TestSpanBoundary::End(Span::WaitingOnPeer),
 			TestSpanBoundary::End(Span::OutboundHTLCState {
 				state: OutboundHTLCStateDetails::AwaitingRemoteRevokeToRemoveSuccess
 			}),