chore: import upstream snapshot with attribution

python/ray/data/_internal/execution/block_ref_counter.py

@@ -0,0 +1,85 @@
+import threading
+from collections import defaultdict
+from typing import Callable, Dict, Optional
+
+import ray
+from ray._private.worker import global_worker
+
+
+class BlockRefCounter:
+    """Tracks object-store memory usage per operator via Ray Core callbacks.
+
+    The callback fires when:
+    - All Python ObjectRefs wrapping the block's ObjectID are garbage-collected, AND
+    - All Ray tasks that received the block as an argument have completed.
+    """
+
+    def __init__(
+        self,
+        add_object_out_of_scope_callback: Optional[
+            Callable[["ray.ObjectRef", Callable[[bytes], None]], bool]
+        ] = None,
+    ):
+        if add_object_out_of_scope_callback is None:
+            add_object_out_of_scope_callback = (
+                global_worker.core_worker.add_object_out_of_scope_callback  # pyrefly: ignore[missing-attribute]
+            )
+        self._add_callback_fn = add_object_out_of_scope_callback
+        # IDs of live blocks. Stale callbacks (fired after clear()) check
+        # membership here and no-op, preventing negative _bytes_by_producer.
+        self._registered_ids: set[bytes] = set()
+        # (producer_id -> total live bytes); maintained incrementally for O(1) reads.
+        self._bytes_by_producer: Dict[str, int] = defaultdict(int)
+        self._lock = threading.Lock()
+
+    def on_block_produced(
+        self,
+        block_ref: "ray.ObjectRef",
+        size_bytes: int,
+        producer_id: str,
+    ) -> None:
+        """Register a block and attribute its memory to producer_id.
+
+        Registers a Ray Core out-of-scope callback so that when all references
+        to block_ref are gone the bytes are automatically removed from the
+        producer's usage.
+
+        Idempotent: calling twice with the same block_ref is a no-op.
+        """
+        id_binary = block_ref.binary()
+        with self._lock:
+            if id_binary in self._registered_ids:
+                return
+            self._registered_ids.add(id_binary)
+            self._bytes_by_producer[producer_id] += size_bytes
+
+        def _on_object_freed(id_bytes: bytes) -> None:
+            with self._lock:
+                if id_bytes not in self._registered_ids:
+                    # Already cleared (e.g. by clear()), nothing to do.
+                    return
+                self._registered_ids.discard(id_bytes)
+                self._bytes_by_producer[producer_id] -= size_bytes
+
+        # TODO(srayhome): This raises ValueError for blocks not owned by this
+        # worker (e.g., materialized dataset passed to streaming_split). We may
+        # need to guard this with RefBundle.owns_blocks or skip registration at
+        # the InputDataBuffer level.
+        registered = self._add_callback_fn(block_ref, _on_object_freed)
+        if not registered:
+            _on_object_freed(id_binary)
+
+    def get_object_store_memory_usage(self, producer_id: str) -> int:
+        """Total bytes of live blocks attributed to producer_id."""
+        with self._lock:
+            return self._bytes_by_producer.get(producer_id, 0)
+
+    def clear(self) -> None:
+        """Reset all accounting, e.g. on executor shutdown.
+
+        Any previously registered Ray Core callbacks firing after clear()
+        will be silently ignored because _registered_ids is empty.
+        """
+        with self._lock:
+            self._registered_ids.clear()
+            self._bytes_by_producer.clear()