Actor.reboot method uses the new reboot endpoint (#109)

Author: vdusek

CHANGELOG.md

@@ -1,6 +1,17 @@
 Changelog
 =========
 
+[1.2.0](../../releases/tag/v1.2.0) - 2023-08-23
+-----------------------------------------------
+
+### Added
+
+- Update the `Actor.reboot` method to use the new reboot endpoint
+
+### Internal changes
+
+- Unify indentation in configuration files
+
 [1.1.2](../../releases/tag/v1.1.2) - 2023-08-02
 -----------------------------------------------
 

docs/docs.md

@@ -557,14 +557,16 @@ Get the actor input value from the default key-value store associated with the c
 
 ***
 
-#### [](#actor-get_value) `async Actor.get_value(key)`
+#### [](#actor-get_value) `async Actor.get_value(key, default_value=None)`
 
 Get a value from the default key-value store associated with the current actor run.
 
 * **Parameters**
 
   * **key** (`str`) – The key of the record which to retrieve.
 
+  * **default_value** (`Any`, *optional*) – Default value returned in case the record does not exist.
+
 * **Return type**
 
   `Any`
@@ -859,7 +861,7 @@ and the new input is stored under the INPUT-METAMORPH-1 key in the same default
 
 ***
 
-#### [](#actor-reboot) `async Actor.reboot(*, event_listeners_timeout_secs=5)`
+#### [](#actor-reboot) `async Actor.reboot(*, event_listeners_timeout_secs=5, custom_after_sleep_millis=None)`
 
 Internally reboot this actor.
 
@@ -869,6 +871,8 @@ The system stops the current container and starts a new one, with the same run I
 
   * **event_listeners_timeout_secs** (`int`, *optional*) – How long should the actor wait for actor event listeners to finish before exiting
 
+  * **custom_after_sleep_millis** (`int`, *optional*) – How long to sleep for after the reboot, to wait for the container to be stopped.
+
 * **Return type**
 
   `None`

pyproject.toml

@@ -25,7 +25,7 @@ requires-python = ">=3.8"
 dependencies = [
     "aiofiles >= 22.1.0",
     "aioshutil >= 1.0",
-    "apify-client == 1.3.1",
+    "apify-client == 1.4.0",
     "apify-shared == 1.0.2",
     "colorama >= 0.4.6",
     "cryptography >= 39.0.0",

src/apify/actor.py

@@ -31,6 +31,7 @@
 from .proxy_configuration import ProxyConfiguration
 from .storages import Dataset, KeyValueStore, RequestQueue, StorageClientManager
 
+T = TypeVar('T')
 MainReturnType = TypeVar('MainReturnType')
 
 # This metaclass is needed so you can do `async with Actor: ...` instead of `async with Actor() as a: ...`
@@ -629,19 +630,20 @@ async def _get_input_internal(self) -> Any:
         return input_value
 
     @classmethod
-    async def get_value(cls, key: str) -> Any:
+    async def get_value(cls, key: str, default_value: Optional[T] = None) -> Any:
         """Get a value from the default key-value store associated with the current actor run.
 
         Args:
             key (str): The key of the record which to retrieve.
+            default_value (Any, optional): Default value returned in case the record does not exist.
         """
-        return await cls._get_default_instance().get_value(key=key)
+        return await cls._get_default_instance().get_value(key=key, default_value=default_value)
 
-    async def _get_value_internal(self, key: str) -> Any:
+    async def _get_value_internal(self, key: str, default_value: Optional[T] = None) -> Any:
         self._raise_if_not_initialized()
 
         key_value_store = await self.open_key_value_store()
-        value = await key_value_store.get_value(key)
+        value = await key_value_store.get_value(key, default_value)
         return value
 
     @classmethod
@@ -1110,38 +1112,48 @@ async def reboot(
         cls,
         *,
         event_listeners_timeout_secs: Optional[int] = EVENT_LISTENERS_TIMEOUT_SECS,
+        custom_after_sleep_millis: Optional[int] = None,
     ) -> None:
         """Internally reboot this actor.
 
         The system stops the current container and starts a new one, with the same run ID and default storages.
 
         Args:
             event_listeners_timeout_secs (int, optional): How long should the actor wait for actor event listeners to finish before exiting
+            custom_after_sleep_millis (int, optional): How long to sleep for after the reboot, to wait for the container to be stopped.
         """
-        return await cls._get_default_instance().reboot(event_listeners_timeout_secs=event_listeners_timeout_secs)
+        return await cls._get_default_instance().reboot(
+            event_listeners_timeout_secs=event_listeners_timeout_secs,
+            custom_after_sleep_millis=custom_after_sleep_millis,
+        )
 
     async def _reboot_internal(
         self,
         *,
         event_listeners_timeout_secs: Optional[int] = EVENT_LISTENERS_TIMEOUT_SECS,
+        custom_after_sleep_millis: Optional[int] = None,
     ) -> None:
         self._raise_if_not_initialized()
 
         if not self.is_at_home():
             self.log.error('Actor.reboot() is only supported when running on the Apify platform.')
             return
 
+        if not custom_after_sleep_millis:
+            custom_after_sleep_millis = self._config.metamorph_after_sleep_millis
+
         await self._cancel_event_emitting_intervals()
 
         self._event_manager.emit(ActorEventTypes.PERSIST_STATE, {'isMigrating': True})
         self._was_final_persist_state_emitted = True
 
         await self._event_manager.close(event_listeners_timeout_secs=event_listeners_timeout_secs)
 
-        # If is_at_home() is True, config.actor_id is always set
-        assert self._config.actor_id is not None
+        assert self._config.actor_run_id is not None
+        await self._apify_client.run(self._config.actor_run_id).reboot()
 
-        await self.metamorph(self._config.actor_id)
+        if custom_after_sleep_millis:
+            await asyncio.sleep(custom_after_sleep_millis / 1000)
 
     @classmethod
     async def add_webhook(

tests/integration/conftest.py

@@ -15,7 +15,7 @@
 from apify.storages import Dataset, KeyValueStore, RequestQueue, StorageClientManager
 from apify_client import ApifyClientAsync
 from apify_client.clients.resource_clients import ActorClientAsync
-from apify_client.consts import ActorJobStatus, ActorSourceType
+from apify_shared.consts import ActorJobStatus, ActorSourceType
 
 from ._utils import generate_unique_resource_name
 

tests/integration/test_actor_api_helpers.py

@@ -323,6 +323,35 @@ async def main_outer() -> None:
         assert await inner_actor.last_run().get() is None
 
 
+class TestActorReboot:
+    async def test_actor_reboot(self, make_actor: ActorFactory) -> None:
+        async def main() -> None:
+            async with Actor:
+                print('Starting...')
+                cnt = await Actor.get_value('reboot_counter', 0)
+
+                if cnt < 2:
+                    print(f'Rebooting (cnt = {cnt})...')
+                    await Actor.set_value('reboot_counter', cnt + 1)
+                    await Actor.reboot()
+                    await Actor.set_value('THIS_KEY_SHOULD_NOT_BE_WRITTEN', 'XXX')
+
+                print('Finishing...')
+
+        actor = await make_actor('actor_rebooter', main_func=main)
+        run_result = await actor.call(run_input={'counter_key': 'reboot_counter'})
+
+        assert run_result is not None
+        assert run_result['status'] == 'SUCCEEDED'
+
+        not_written_value = await actor.last_run().key_value_store().get_record('THIS_KEY_SHOULD_NOT_BE_WRITTEN')
+        assert not_written_value is None
+
+        reboot_counter = await actor.last_run().key_value_store().get_record('reboot_counter')
+        assert reboot_counter is not None
+        assert reboot_counter['value'] == 2
+
+
 class TestActorAddWebhook:
     async def test_actor_add_webhook(self, make_actor: ActorFactory) -> None:
         async def main_server() -> None: