probe: Ensure all external callers always call end_probe_session()

Rework ProbeSessionHelper's multi_probe_start() and multi_probe_end()
to start_probe_session() and end_probe_session().  Ensure all external
callers always invoke these methods prior to running run_probe().

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

1 files changed, 26 insertions, 21 deletions

diff --git a/klippy/extras/probe.py b/klippy/extras/probe.py
index 12e6109f..ca778cc9 100644
--- a/klippy/extras/probe.py
+++ b/klippy/extras/probe.py
@@ -1,6 +1,6 @@
 # Z-Probe support
 #
-# Copyright (C) 2017-2021  Kevin O'Connor <kevin@koconnor.net>
+# Copyright (C) 2017-2024  Kevin O'Connor <kevin@koconnor.net>
 #
 # This file may be distributed under the terms of the GNU GPLv3 license.
 import logging
@@ -128,7 +128,6 @@ class ProbeCommandHelper:
         fo_gcmd = gcode.create_gcode_command("", "", fo_params)
         # Probe bed sample_count times
         probe_session = self.probe.start_probe_session(fo_gcmd)
-        probe_session.multi_probe_begin()
         positions = []
         while len(positions) < sample_count:
             # Probe position
@@ -137,7 +136,7 @@ class ProbeCommandHelper:
             # Retract
             liftpos = [None, None, pos[2] + params['sample_retract_dist']]
             self._move(liftpos, params['lift_speed'])
-        probe_session.multi_probe_end()
+        probe_session.end_probe_session()
         # Calculate maximum, minimum and average values
         max_value = max([p[2] for p in positions])
         min_value = min([p[2] for p in positions])
@@ -225,7 +224,6 @@ class ProbeSessionHelper:
     def __init__(self, config, mcu_probe):
         self.printer = config.get_printer()
         self.mcu_probe = mcu_probe
-        self.multi_probe_pending = False
         gcode = self.printer.lookup_object('gcode')
         self.dummy_gcode_cmd = gcode.create_gcode_command("", "", {})
         # Infer Z position to move to during a probe
@@ -252,21 +250,31 @@ class ProbeSessionHelper:
                                                  minval=0.)
         self.samples_retries = config.getint('samples_tolerance_retries', 0,
                                              minval=0)
+        # Session state
+        self.multi_probe_pending = False
         # Register event handlers
         self.printer.register_event_handler("gcode:command_error",
                                             self._handle_command_error)
     def _handle_command_error(self):
-        try:
-            self.multi_probe_end()
-        except:
-            logging.exception("Multi-probe end")
-    def multi_probe_begin(self):
+        if self.multi_probe_pending:
+            try:
+                self.end_probe_session()
+            except:
+                logging.exception("Multi-probe end")
+    def _probe_state_error(self):
+        raise self.printer.command_error(
+            "Internal probe error - start/end probe session mismatch")
+    def start_probe_session(self, gcmd):
+        if self.multi_probe_pending:
+            self._probe_state_error()
         self.mcu_probe.multi_probe_begin()
         self.multi_probe_pending = True
-    def multi_probe_end(self):
-        if self.multi_probe_pending:
-            self.multi_probe_pending = False
-            self.mcu_probe.multi_probe_end()
+        return self
+    def end_probe_session(self):
+        if not self.multi_probe_pending:
+            self._probe_state_error()
+        self.multi_probe_pending = False
+        self.mcu_probe.multi_probe_end()
     def get_probe_params(self, gcmd=None):
         if gcmd is None:
             gcmd = self.dummy_gcode_cmd
@@ -315,10 +323,9 @@ class ProbeSessionHelper:
                            % (epos[0], epos[1], epos[2]))
         return epos[:3]
     def run_probe(self, gcmd):
+        if not self.multi_probe_pending:
+            self._probe_state_error()
         params = self.get_probe_params(gcmd)
-        must_notify_multi_probe = not self.multi_probe_pending
-        if must_notify_multi_probe:
-            self.multi_probe_begin()
         toolhead = self.printer.lookup_object('toolhead')
         probexy = toolhead.get_position()[:2]
         retries = 0
@@ -341,8 +348,6 @@ class ProbeSessionHelper:
                 toolhead.manual_move(
                     probexy + [pos[2] + params['sample_retract_dist']],
                     params['lift_speed'])
-        if must_notify_multi_probe:
-            self.multi_probe_end()
         # Calculate and return result
         return calc_probe_z_average(positions, params['samples_result'])
 
@@ -436,14 +441,13 @@ class ProbePointsHelper:
             raise gcmd.error("horizontal_move_z can't be less than"
                              " probe's z_offset")
         probe_session = probe.start_probe_session(gcmd)
-        probe_session.multi_probe_begin()
         while 1:
             done = self._move_next()
             if done:
                 break
             pos = probe_session.run_probe(gcmd)
             self.results.append(pos)
-        probe_session.multi_probe_end()
+        probe_session.end_probe_session()
     def _manual_probe_start(self):
         done = self._move_next()
         if not done:
@@ -460,6 +464,7 @@ class ProbePointsHelper:
 def run_single_probe(probe, gcmd):
     probe_session = probe.start_probe_session(gcmd)
     pos = probe_session.run_probe(gcmd)
+    probe_session.end_probe_session()
     return pos
 
 
@@ -553,7 +558,7 @@ class PrinterProbe:
     def get_status(self, eventtime):
         return self.cmd_helper.get_status(eventtime)
     def start_probe_session(self, gcmd):
-        return self.probe_session
+        return self.probe_session.start_probe_session(gcmd)
 
 def load_config(config):
     return PrinterProbe(config, ProbeEndstopWrapper(config))