From 14949e6587f7204d4e5cc9f23f54715101f08a8e Mon Sep 17 00:00:00 2001
From: Matt Keeter <matt@oxide.computer>
Date: Wed, 20 Nov 2024 11:24:24 -0500
Subject: [PATCH] Move clear_client_state call to a single place

---
 task/control-plane-agent/src/dump.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/task/control-plane-agent/src/dump.rs b/task/control-plane-agent/src/dump.rs
index 400931dfc..095fd3a45 100644
--- a/task/control-plane-agent/src/dump.rs
+++ b/task/control-plane-agent/src/dump.rs
@@ -165,6 +165,19 @@ impl DumpState {
         key: [u8; 16],
         seq: u32,
         buf: &mut [u8],
+    ) -> Result<Option<DumpSegment>, SpError> {
+        let r = self.task_dump_read_continue_inner(key, seq, buf);
+        if matches!(r, Ok(None) | Err(..)) {
+            self.clear_client_state(key);
+        }
+        r
+    }
+
+    pub(crate) fn task_dump_read_continue_inner(
+        &mut self,
+        key: [u8; 16],
+        seq: u32,
+        buf: &mut [u8],
     ) -> Result<Option<DumpSegment>, SpError> {
         let Some(state) =
             self.clients.iter_mut().flatten().find(|c| c.key == key)
@@ -202,7 +215,6 @@ impl DumpState {
         if pos.offset + SEGMENT_DATA_SIZE > header.written {
             if header.next == 0 {
                 // we're done, because there's no more dump areas
-                self.clear_client_state(key);
                 return Ok(None);
             }
 
@@ -217,7 +229,6 @@ impl DumpState {
             if header.contents != humpty::DumpContents::SingleTask.into()
                 || header.nsegments != 0
             {
-                self.clear_client_state(key);
                 return Ok(None);
             }
             // Skip the area header