fix: use a directory pool

Signed-off-by: Frost Ming <[email protected]>

src/bentoml/_internal/context.py

@@ -3,8 +3,6 @@
 import contextlib
 import contextvars
 import os
-import shutil
-import tempfile
 import typing as t
 from abc import ABC
 from abc import abstractmethod
@@ -14,6 +12,7 @@
 import starlette.datastructures
 
 from .utils.http import Cookie
+from .utils.temp import TempfilePool
 
 if TYPE_CHECKING:
     import starlette.requests
@@ -26,12 +25,14 @@
     contextvars.ContextVar("response")
 )
 
+request_tempdir_pool = TempfilePool(prefix="bentoml-request-")
+
 
 def request_temp_dir() -> str:
     """A request-unique directory for storing temporary files"""
     request = _request_var.get()
     if not hasattr(request.state, "temp_dir"):
-        request.state.temp_dir = tempfile.mkdtemp(prefix="bentoml-request-")
+        request.state.temp_dir = request_tempdir_pool.acquire()
     return request.state.temp_dir
 
 
@@ -106,7 +107,7 @@ def in_request(
             yield self
         finally:
             if hasattr(request.state, "temp_dir"):
-                shutil.rmtree(request.state.temp_dir, ignore_errors=True)
+                request_tempdir_pool.release(request.state.temp_dir)
             _request_var.reset(request_token)
             _response_var.reset(response_token)
 

src/bentoml/_internal/server/base_app.py

@@ -43,7 +43,9 @@ def on_startup(self) -> list[LifecycleHook]:
 
     @property
     def on_shutdown(self) -> list[LifecycleHook]:
-        return []
+        from ..context import request_tempdir_pool
+
+        return [request_tempdir_pool.cleanup]
 
     def mark_as_ready(self) -> None:
         self._is_ready = True

src/bentoml/_internal/utils/temp.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import shutil
+import tempfile
+from collections import deque
+from functools import partial
+from pathlib import Path
+from threading import Lock
+
+
+class TempfilePool:
+    """A simple pool to get temp directories
+    so they are reused as most as possible.
+    """
+
+    def __init__(
+        self,
+        suffix: str | None = None,
+        prefix: str | None = None,
+        dir: str | None = None,
+    ) -> None:
+        self._pool: deque[str] = deque([])
+        self._lock = Lock()
+        self._new = partial(tempfile.mkdtemp, suffix=suffix, prefix=prefix, dir=dir)
+
+    def cleanup(self) -> None:
+        while len(self._pool):
+            dir = self._pool.popleft()
+            shutil.rmtree(dir, ignore_errors=True)
+
+    def acquire(self) -> str:
+        with self._lock:
+            if not len(self._pool):
+                return self._new()
+            else:
+                return self._pool.popleft()
+
+    def release(self, dir: str) -> None:
+        with self._lock:
+            for child in Path(dir).iterdir():
+                if child.is_dir():
+                    shutil.rmtree(child)
+                else:
+                    child.unlink()
+            self._pool.append(dir)