mirror of
				https://github.com/zulip/zulip.git
				synced 2025-10-31 12:03:46 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			97 lines
		
	
	
		
			3.2 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			97 lines
		
	
	
		
			3.2 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
| import ctypes
 | |
| import logging
 | |
| import sys
 | |
| import threading
 | |
| import time
 | |
| from types import TracebackType
 | |
| from typing import Callable, Optional, Tuple, Type, TypeVar
 | |
| 
 | |
| # Based on https://code.activestate.com/recipes/483752/
 | |
| 
 | |
| 
 | |
| class TimeoutExpiredError(Exception):
 | |
|     """Exception raised when a function times out."""
 | |
| 
 | |
|     def __str__(self) -> str:
 | |
|         return "Function call timed out."
 | |
| 
 | |
| 
 | |
| ResultT = TypeVar("ResultT")
 | |
| 
 | |
| 
 | |
| def timeout(timeout: float, func: Callable[[], ResultT]) -> ResultT:
 | |
|     """Call the function in a separate thread.
 | |
|     Return its return value, or raise an exception,
 | |
|     within approximately 'timeout' seconds.
 | |
| 
 | |
|     The function may receive a TimeoutExpiredError exception
 | |
|     anywhere in its code, which could have arbitrary
 | |
|     unsafe effects (resources not released, etc.).
 | |
|     It might also fail to receive the exception and
 | |
|     keep running in the background even though
 | |
|     timeout() has returned.
 | |
| 
 | |
|     This may also fail to interrupt functions which are
 | |
|     stuck in a long-running primitive interpreter
 | |
|     operation."""
 | |
| 
 | |
|     class TimeoutThread(threading.Thread):
 | |
|         def __init__(self) -> None:
 | |
|             threading.Thread.__init__(self)
 | |
|             self.result: Optional[ResultT] = None
 | |
|             self.exc_info: Tuple[
 | |
|                 Optional[Type[BaseException]],
 | |
|                 Optional[BaseException],
 | |
|                 Optional[TracebackType],
 | |
|             ] = (None, None, None)
 | |
| 
 | |
|             # Don't block the whole program from exiting
 | |
|             # if this is the only thread left.
 | |
|             self.daemon = True
 | |
| 
 | |
|         def run(self) -> None:
 | |
|             try:
 | |
|                 self.result = func()
 | |
|             except BaseException:
 | |
|                 self.exc_info = sys.exc_info()
 | |
| 
 | |
|         def raise_async_timeout(self) -> None:
 | |
|             # This function is called from another thread; we attempt
 | |
|             # to raise a TimeoutExpiredError in _this_ thread.
 | |
|             assert self.ident is not None
 | |
|             ctypes.pythonapi.PyThreadState_SetAsyncExc(
 | |
|                 ctypes.c_ulong(self.ident),
 | |
|                 ctypes.py_object(TimeoutExpiredError),
 | |
|             )
 | |
| 
 | |
|     thread = TimeoutThread()
 | |
|     thread.start()
 | |
|     thread.join(timeout)
 | |
| 
 | |
|     if thread.is_alive():
 | |
|         # We need to retry, because an async exception received while
 | |
|         # the thread is in a system call is simply ignored.
 | |
|         for i in range(10):
 | |
|             thread.raise_async_timeout()
 | |
|             time.sleep(0.1)
 | |
|             if not thread.is_alive():
 | |
|                 break
 | |
|         if thread.exc_info[1] is not None:
 | |
|             # Re-raise the exception we sent, if possible, so the
 | |
|             # stacktrace originates in the slow code
 | |
|             raise thread.exc_info[1].with_traceback(thread.exc_info[2])
 | |
|         # If we don't have that for some reason (e.g. we failed to
 | |
|         # kill it), just raise from here; the thread _may still be
 | |
|         # running_ because it failed to see any of our exceptions, and
 | |
|         # we just ignore it.
 | |
|         if thread.is_alive():
 | |
|             logging.warning("Failed to time out backend thread")
 | |
|         raise TimeoutExpiredError
 | |
| 
 | |
|     if thread.exc_info[1] is not None:
 | |
|         # Died with some other exception; re-raise it
 | |
|         raise thread.exc_info[1].with_traceback(thread.exc_info[2])
 | |
| 
 | |
|     assert thread.result is not None
 | |
|     return thread.result
 |