From e8305f6c8c0e6eb69495b9f7d6aa36e8a5b2105b Mon Sep 17 00:00:00 2001
From: Tiexin Guo <tiexin.guo@canonical.com>
Date: Tue, 26 Mar 2024 11:47:57 +0800
Subject: [PATCH] fix(pebble): catch socket.timeout exception in
 pebble.Client.exec() (#1155)

Catch `socket.timeout` exception and continue polling till the end of the deadline, added a test.
---
 CHANGES.md          |  1 +
 ops/pebble.py       | 12 ++++++++++--
 test/test_pebble.py | 13 +++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 07ae6bbe1..cea76b8b0 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -2,6 +2,7 @@
 
 * Updated Pebble Notices `get_notices` parameter name to `users=all` (previously `select=all`).
 * Added `Model.get_cloud_spec` which uses the `credential-get` hook tool to get details of the cloud where the model is deployed.
+* Fixed an issue where `pebble.Client.exec` might leak a `socket.timeout` (`builtins.TimeoutError`) exception.
 * Updated code examples in the docstring of `ops.testing` from unittest to pytest style.
 * Refactored main.py, creating a new `_Manager` class.
 
diff --git a/ops/pebble.py b/ops/pebble.py
index 60bb9d7b5..638885f18 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -18,6 +18,7 @@
 """
 
 import binascii
+import builtins
 import copy
 import dataclasses
 import datetime
@@ -1722,6 +1723,11 @@ class Client:
     connecting to or transferring data with Pebble will raise a :class:`ConnectionError`. When an
     error occurs executing the request, such as trying to add an invalid layer or execute a command
     that does not exist, an :class:`APIError` is raised.
+
+    The ``timeout`` parameter specifies a timeout in seconds for blocking operations like the
+    connection attempt to Pebble; used by ``urllib.request.OpenerDirector.open``. It's not for
+    methods like :meth:`start_services` and :meth:`replan_services` mentioned above, and it's not
+    for the command execution timeout defined in method :meth:`Client.exec`.
     """
 
     _chunk_size = 8192
@@ -2029,8 +2035,10 @@ def _wait_change_using_wait(self, change_id: ChangeID, timeout: Optional[float])
 
             try:
                 return self._wait_change(change_id, this_timeout)
-            except TimeoutError:
-                # Catch timeout from wait endpoint and loop to check deadline
+            except (socket.timeout, builtins.TimeoutError):
+                # NOTE: in Python 3.10 socket.timeout is an alias of TimeoutError,
+                # but we still need to support Python 3.8, so catch both.
+                # Catch timeout from wait endpoint and loop to check deadline.
                 pass
 
         raise TimeoutError(f'timed out waiting for change {change_id} ({timeout} seconds)')
diff --git a/test/test_pebble.py b/test/test_pebble.py
index ccd96d00c..5be22663c 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -18,6 +18,7 @@
 import io
 import json
 import signal
+import socket
 import tempfile
 import test.fake_pebble as fake_pebble
 import typing
@@ -2028,6 +2029,18 @@ def test_wait_change_error(self):
             ('GET', '/v1/changes/70/wait', {'timeout': '4.000s'}, None),
         ])
 
+    def test_wait_change_socket_timeout(self):
+        def timeout_response(n: float):
+            self.time.sleep(n)
+            raise socket.timeout("socket.timeout: timed out")
+
+        self.client.responses.append(lambda: timeout_response(3))
+
+        with self.assertRaises(pebble.TimeoutError) as cm:
+            self.client.wait_change(pebble.ChangeID('70'), timeout=3)
+        self.assertIsInstance(cm.exception, pebble.Error)
+        self.assertIsInstance(cm.exception, TimeoutError)
+
     def test_add_layer(self):
         okay_response = {
             "result": True,