Use CreateFunctionHandle callback in actions (#36362)

GitOrigin-RevId: 71dadabdf60a1bc8ae06f2be4b761cf70504a71d

Author: goffrie

crates/isolate/src/environment/action/async_syscall.rs

@@ -17,10 +17,7 @@ use errors::{
     ErrorMetadataAnyhowExt,
 };
 use model::{
-    components::{
-        auth::propagate_component_auth,
-        handles::function_handle_not_found,
-    },
+    components::auth::propagate_component_auth,
     file_storage::FileStorageId,
 };
 use serde::{
@@ -462,15 +459,10 @@ impl<RT: Runtime> TaskExecutor<RT> {
                 self.resolve_function(&reference)?
             },
         };
-        // TODO(lee) remove preloaded function handles and call action callback instead,
-        // after the callback is deployed to backend & usher.
-        let handle = {
-            let function_handles = self.function_handles.lock();
-            function_handles.get(&function_path).cloned()
-        };
-        let Some(handle) = handle else {
-            anyhow::bail!(function_handle_not_found());
-        };
+        let handle = self
+            .action_callbacks
+            .create_function_handle(self.identity.clone(), function_path)
+            .await?;
         Ok(serde_json::to_value(String::from(handle))?)
     }
 }

crates/isolate/src/environment/action/mod.rs

@@ -259,7 +259,6 @@ impl<RT: Runtime> ActionEnvironment<RT> {
         let syscall_trace = Arc::new(Mutex::new(SyscallTrace::new()));
         let (task_retval_sender, task_responses) = mpsc::unbounded_channel();
         let resources = Arc::new(Mutex::new(BTreeMap::new()));
-        let function_handles = Arc::new(Mutex::new(BTreeMap::new()));
         let convex_origin_override = Arc::new(Mutex::new(None));
         let task_executor = TaskExecutor {
             rt: rt.clone(),
@@ -276,7 +275,6 @@ impl<RT: Runtime> ActionEnvironment<RT> {
             context,
             resources: resources.clone(),
             component_id: component,
-            function_handles: function_handles.clone(),
             convex_origin_override: convex_origin_override.clone(),
         };
         let (pending_task_sender, pending_task_receiver) = spsc::unbounded_channel();
@@ -300,7 +298,6 @@ impl<RT: Runtime> ActionEnvironment<RT> {
                 module_loader,
                 default_system_env_vars,
                 resources,
-                function_handles,
                 convex_origin_override,
             ),
             syscall_trace,

crates/isolate/src/environment/action/phase.rs

@@ -6,9 +6,7 @@ use std::{
 
 use anyhow::Context;
 use common::{
-    bootstrap_model::components::handles::FunctionHandle,
     components::{
-        CanonicalizedComponentFunctionPath,
         ComponentId,
         Reference,
         Resource,
@@ -30,10 +28,7 @@ use database::{
 use errors::ErrorMetadata;
 use model::{
     canonical_urls::CanonicalUrlsModel,
-    components::{
-        handles::FunctionHandlesModel,
-        ComponentsModel,
-    },
+    components::ComponentsModel,
     config::module_loader::ModuleLoader,
     environment_variables::{
         types::{
@@ -98,7 +93,6 @@ enum ActionPreloaded<RT: Runtime> {
         module_loader: Arc<dyn ModuleLoader<RT>>,
         default_system_env_vars: BTreeMap<EnvVarName, EnvVarValue>,
         resources: Arc<Mutex<BTreeMap<Reference, Resource>>>,
-        function_handles: Arc<Mutex<BTreeMap<CanonicalizedComponentFunctionPath, FunctionHandle>>>,
         convex_origin_override: Arc<Mutex<Option<ConvexOrigin>>>,
     },
     Preloading,
@@ -119,7 +113,6 @@ impl<RT: Runtime> ActionPhase<RT> {
         module_loader: Arc<dyn ModuleLoader<RT>>,
         default_system_env_vars: BTreeMap<EnvVarName, EnvVarValue>,
         resources: Arc<Mutex<BTreeMap<Reference, Resource>>>,
-        function_handles: Arc<Mutex<BTreeMap<CanonicalizedComponentFunctionPath, FunctionHandle>>>,
         convex_origin_override: Arc<Mutex<Option<ConvexOrigin>>>,
     ) -> Self {
         Self {
@@ -131,7 +124,6 @@ impl<RT: Runtime> ActionPhase<RT> {
                 module_loader,
                 default_system_env_vars,
                 resources,
-                function_handles,
                 convex_origin_override,
             },
         }
@@ -151,7 +143,6 @@ impl<RT: Runtime> ActionPhase<RT> {
             module_loader,
             default_system_env_vars,
             resources,
-            function_handles,
             convex_origin_override,
         } = preloaded
         else {
@@ -247,16 +238,6 @@ impl<RT: Runtime> ActionPhase<RT> {
             )
         };
 
-        {
-            let handles = with_release_permit(
-                timeout,
-                permit_slot,
-                FunctionHandlesModel::new(&mut tx).preload(),
-            )
-            .await?;
-            *function_handles.lock() = handles;
-        }
-
         self.preloaded = ActionPreloaded::Ready {
             modules,
             env_vars,

crates/isolate/src/environment/action/task_executor.rs

@@ -5,7 +5,6 @@ use std::{
 };
 
 use common::{
-    bootstrap_model::components::handles::FunctionHandle,
     components::{
         CanonicalizedComponentFunctionPath,
         ComponentId,
@@ -80,7 +79,6 @@ pub struct TaskExecutor<RT: Runtime> {
     pub context: ExecutionContext,
     pub resources: Arc<Mutex<BTreeMap<Reference, Resource>>>,
     pub component_id: ComponentId,
-    pub function_handles: Arc<Mutex<BTreeMap<CanonicalizedComponentFunctionPath, FunctionHandle>>>,
     pub convex_origin_override: Arc<Mutex<Option<ConvexOrigin>>>,
 }
 

crates/model/src/components/handles.rs

@@ -128,26 +128,6 @@ impl<'a, RT: Runtime> FunctionHandlesModel<'a, RT> {
         })
     }
 
-    pub async fn preload(
-        &mut self,
-    ) -> anyhow::Result<BTreeMap<CanonicalizedComponentFunctionPath, FunctionHandle>> {
-        let mut handles = BTreeMap::new();
-        let index_query = Query::full_table_scan(FUNCTION_HANDLES_TABLE.clone(), Order::Asc);
-        let mut query_stream = ResolvedQuery::new(self.tx, TableNamespace::Global, index_query)?;
-        while let Some(doc) = query_stream.next(self.tx, None).await? {
-            let handle: ParsedDocument<FunctionHandleMetadata> = doc.parse()?;
-            if handle.deleted_ts.is_none() {
-                let path = CanonicalizedComponentFunctionPath {
-                    component: BootstrapComponentsModel::new(self.tx)
-                        .must_component_path(handle.component)?,
-                    udf_path: handle.path.clone(),
-                };
-                handles.insert(path, FunctionHandle::new(handle.developer_id()));
-            }
-        }
-        Ok(handles)
-    }
-
     pub async fn get_with_component_path(
         &mut self,
         path: CanonicalizedComponentFunctionPath,