ensure safety of insert using entry api in single lock hold

Author: adriangb

datafusion/proto/src/physical_plan/from_proto.rs

@@ -424,10 +424,10 @@ pub fn parse_physical_expr(
 
     // Insert into cache if an ID is present
     if let Some(id) = proto.id {
-        decode_ctx.insert_cached_expr(id, Arc::clone(&pexpr));
+        Ok(decode_ctx.insert_cached_expr(id, Arc::clone(&pexpr)))
+    } else {
+        Ok(pexpr)
     }
-
-    Ok(pexpr)
 }
 
 fn parse_required_physical_expr(

datafusion/proto/src/physical_plan/mod.rs

@@ -133,8 +133,15 @@ impl<'a> DecodeContext<'a> {
     }
 
     /// Insert a physical expression into the cache with the given ID.
-    pub fn insert_cached_expr(&self, id: u64, expr: Arc<dyn PhysicalExpr>) {
-        self.cache.lock().unwrap().insert(id, expr);
+    /// Returns the inserted expression or an existing one if the ID was already present.
+    #[must_use]
+    pub fn insert_cached_expr(
+        &self,
+        id: u64,
+        expr: Arc<dyn PhysicalExpr>,
+    ) -> Arc<dyn PhysicalExpr> {
+        let mut cache = self.cache.lock().unwrap();
+        Arc::clone(cache.entry(id).or_insert(expr))
     }
 }