Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# Copyright (C) 2019 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16import functools 

17 

18import grpc 

19 

20 

21class ExecContext: 

22 

23 __max_workers = None 

24 __busy_workers = None 

25 

26 @classmethod 

27 def init(cls, max_workers): 

28 # Max 4/5 workers for blocking requests: 

29 cls.__max_workers = 4 * max_workers // 5 

30 cls.__busy_workers = 0 

31 

32 @classmethod 

33 def request_worker(cls): 

34 if cls.__max_workers is None: 

35 return True 

36 if cls.__busy_workers >= cls.__max_workers: 

37 return False 

38 cls.__busy_workers += 1 

39 return True 

40 

41 @classmethod 

42 def release_worker(cls): 

43 if cls.__max_workers is None: 

44 return 

45 cls.__busy_workers -= 1 

46 

47 

48def limit(exec_context): 

49 """RPC method decorator for execution resource control. 

50 

51 This decorator is design to be used together with an :class:`ExecContext` 

52 execution context holder:: 

53 

54 @limit(ExecContext) 

55 def Execute(self, request, context): 

56 

57 Args: 

58 exec_context(ExecContext): Execution context holder. 

59 """ 

60 def __limit_decorator(behavior): 

61 """RPC resource control method decorator.""" 

62 @functools.wraps(behavior) 

63 def __limit_wrapper(self, request, context): 

64 """RPC resource control method wrapper.""" 

65 if not exec_context.request_worker(): 

66 context.abort(grpc.StatusCode.RESOURCE_EXHAUSTED, 

67 "Max. number of simultaneous request reached") 

68 return None 

69 

70 context.add_callback(__limit_callback) 

71 

72 return behavior(self, request, context) 

73 

74 def __limit_callback(): 

75 exec_context.release_worker() 

76 

77 return __limit_wrapper 

78 

79 return __limit_decorator