refactor: bundle the transition to FLUSHCOMPLETE into a single function

modules/core/04-channel/keeper/packet.go

@@ -489,27 +489,46 @@ func (k *Keeper) AcknowledgePacket(
 
 	// if an upgrade is in progress, handling packet flushing and update channel state appropriately
 	if channel.State == types.FLUSHING {
-		// counterparty upgrade is written in the OnChanUpgradeAck step.
-		counterpartyUpgrade, found := k.GetCounterpartyUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel())
-		if found {
-			timeout := counterpartyUpgrade.Timeout
-			selfHeight, selfTimestamp := clienttypes.GetSelfHeight(ctx), uint64(ctx.BlockTime().UnixNano())
-
-			if timeout.Elapsed(selfHeight, selfTimestamp) {
-				// packet flushing timeout has expired, abort the upgrade and return nil,
-				// committing an error receipt to state, restoring the channel and successfully acknowledging the packet.
-				k.MustAbortUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), timeout.ErrTimeoutElapsed(selfHeight, selfTimestamp))
-				return nil
-			}
-
-			// set the channel state to flush complete if all packets have been acknowledged/flushed.
-			if !k.HasInflightPackets(ctx, packet.GetSourcePort(), packet.GetSourceChannel()) {
-				channel.State = types.FLUSHCOMPLETE
-				k.SetChannel(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), channel)
-				emitChannelFlushCompleteEvent(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), channel)
-			}
+		if aborted := k.handleFlushState(ctx, &packet, &channel); aborted {
+			k.Logger(ctx).Info(
+				"upgrade aborted",
+				"port_id", packet.GetSourcePort(),
+				"channel_id", packet.GetSourceChannel(),
+				"upgrade_sequence", channel.UpgradeSequence,
+			)
+			return nil
 		}
 	}
 
 	return nil
 }
+
+// handleFlushState is called when a packet is acknowledged or timed out and the channel is in
+// FLUSHING state. It checks if the upgrade has timed out and if so, aborts the upgrade. If all
+// packets have completed their lifecycle, it sets the channel state to FLUSHCOMPLETE and
+// emits a channel_flush_complete event. Returns true if the upgrade was aborted, false otherwise.
+func (k *Keeper) handleFlushState(ctx sdk.Context, packet *types.Packet, channel *types.Channel) bool {
+	counterpartyUpgrade, found := k.GetCounterpartyUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel())
+	if !found {
+		return false
+	}
+
+	timeout := counterpartyUpgrade.Timeout
+	selfHeight, selfTimestamp := clienttypes.GetSelfHeight(ctx), uint64(ctx.BlockTime().UnixNano())
+
+	if timeout.Elapsed(selfHeight, selfTimestamp) {
+		// packet flushing timeout has expired, abort the upgrade
+		// committing an error receipt to state, deleting upgrade information and restoring the channel.
+		k.MustAbortUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), timeout.ErrTimeoutElapsed(selfHeight, selfTimestamp))
+		return true
+	}
+
+	// set the channel state to flush complete if all packets have been acknowledged/flushed.
+	if !k.HasInflightPackets(ctx, packet.GetSourcePort(), packet.GetSourceChannel()) {
+		channel.State = types.FLUSHCOMPLETE
+		k.SetChannel(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), *channel)
+		emitChannelFlushCompleteEvent(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), *channel)
+	}
+
+	return false
+}

modules/core/04-channel/keeper/timeout.go

@@ -152,23 +152,14 @@ func (k *Keeper) TimeoutExecuted(
 
 	// if an upgrade is in progress, handling packet flushing and update channel state appropriately
 	if channel.State == types.FLUSHING && channel.Ordering == types.UNORDERED {
-		counterpartyUpgrade, found := k.GetCounterpartyUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel())
-		// once we have received the counterparty timeout in the channel UpgradeAck or UpgradeConfirm handshake steps
-		// then we can move to flushing complete if the timeout has not passed and there are no in-flight packets
-		if found {
-			timeout := counterpartyUpgrade.Timeout
-			selfHeight, selfTimestamp := clienttypes.GetSelfHeight(ctx), uint64(ctx.BlockTime().UnixNano())
-
-			if timeout.Elapsed(selfHeight, selfTimestamp) {
-				// packet flushing timeout has expired, abort the upgrade and return nil,
-				// committing an error receipt to state, restoring the channel and successfully timing out the packet.
-				k.MustAbortUpgrade(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), timeout.ErrTimeoutElapsed(selfHeight, selfTimestamp))
-			} else if !k.HasInflightPackets(ctx, packet.GetSourcePort(), packet.GetSourceChannel()) {
-				// set the channel state to flush complete if all packets have been flushed.
-				channel.State = types.FLUSHCOMPLETE
-				k.SetChannel(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), channel)
-				emitChannelFlushCompleteEvent(ctx, packet.GetSourcePort(), packet.GetSourceChannel(), channel)
-			}
+		if aborted := k.handleFlushState(ctx, &packet, &channel); aborted {
+			k.Logger(ctx).Info(
+				"upgrade aborted",
+				"port_id", packet.GetSourcePort(),
+				"channel_id", packet.GetSourceChannel(),
+				"upgrade_sequence", channel.UpgradeSequence,
+			)
+			return nil
 		}
 	}