|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 | import random |
| 15 | +import threading |
| 16 | +import time |
| 17 | +import weakref |
15 | 18 |
|
16 | 19 | from collections import namedtuple |
| 20 | +from enum import Enum |
17 | 21 | from functools import lru_cache |
18 | 22 | from itertools import islice, cycle, groupby, repeat |
19 | 23 | import logging |
@@ -778,6 +782,14 @@ def new_schedule(self): |
778 | 782 | raise NotImplementedError() |
779 | 783 |
|
780 | 784 |
|
| 785 | +class NoDelayReconnectionPolicy(ReconnectionPolicy): |
| 786 | + """ |
| 787 | + A :class:`.ReconnectionPolicy` subclass which does not sleep. |
| 788 | + """ |
| 789 | + def new_schedule(self): |
| 790 | + return repeat(0) |
| 791 | + |
| 792 | + |
781 | 793 | class ConstantReconnectionPolicy(ReconnectionPolicy): |
782 | 794 | """ |
783 | 795 | A :class:`.ReconnectionPolicy` subclass which sleeps for a fixed delay |
@@ -864,6 +876,172 @@ def _add_jitter(self, value): |
864 | 876 | return min(max(self.base_delay, delay), self.max_delay) |
865 | 877 |
|
866 | 878 |
|
| 879 | +class _ShardReconnectionScheduler(object): |
| 880 | + def schedule(self, host_id, shard_id, method, *args, **kwargs): |
| 881 | + raise NotImplementedError() |
| 882 | + |
| 883 | +class ShardReconnectionPolicy(object): |
| 884 | + """ |
| 885 | + Base class for shard reconnection policies. |
| 886 | +
|
| 887 | + On `new_scheduler` instantiate a scheduler that behaves according to the policy |
| 888 | + """ |
| 889 | + def new_scheduler(self, session) -> _ShardReconnectionScheduler: |
| 890 | + raise NotImplementedError() |
| 891 | + |
| 892 | + |
| 893 | +class NoDelayShardReconnectionPolicy(ShardReconnectionPolicy): |
| 894 | + """ |
| 895 | + Shard reconnection policy that does not have delay. |
| 896 | + Does not allow schedule multiple reconnections for same host+shard, silently ignores attempts to do that. |
| 897 | +
|
| 898 | + On `new_scheduler` instantiate a scheduler that behaves according to the policy |
| 899 | + """ |
| 900 | + def new_scheduler(self, session) -> _ShardReconnectionScheduler: |
| 901 | + return _NoDelayShardReconnectionScheduler(session) |
| 902 | + |
| 903 | + |
| 904 | +class _NoDelayShardReconnectionScheduler(_ShardReconnectionScheduler): |
| 905 | + def __init__(self, session): |
| 906 | + self.session = weakref.proxy(session) |
| 907 | + self.already_scheduled = {} |
| 908 | + |
| 909 | + def _execute(self, scheduled_key, method, *args, **kwargs): |
| 910 | + try: |
| 911 | + method(*args, **kwargs) |
| 912 | + finally: |
| 913 | + self.already_scheduled[scheduled_key] = False |
| 914 | + |
| 915 | + def schedule(self, host_id, shard_id, method, *args, **kwargs): |
| 916 | + scheduled_key = f'{host_id}-{shard_id}' |
| 917 | + if self.already_scheduled.get(scheduled_key): |
| 918 | + return |
| 919 | + |
| 920 | + self.already_scheduled[scheduled_key] = True |
| 921 | + if not self.session.is_shutdown: |
| 922 | + self.session.submit(self._execute, scheduled_key, method, *args, **kwargs) |
| 923 | + |
| 924 | + |
| 925 | +class ShardReconnectionPolicyScope(Enum): |
| 926 | + """ |
| 927 | + A scope for `ShardReconnectionPolicy`, in particular `NoConcurrentShardReconnectionPolicy` |
| 928 | + """ |
| 929 | + Cluster = 0 |
| 930 | + Host = 1 |
| 931 | + |
| 932 | + |
| 933 | +class NoConcurrentShardReconnectionPolicy(ShardReconnectionPolicy): |
| 934 | + """ |
| 935 | + A shard reconnection policy that allows only one pending connection per scope, where scope could be `Host`, `Cluster` |
| 936 | + For backoff it uses `ReconnectionPolicy`, when there is no more reconnections to scheduled backoff policy is reminded |
| 937 | + For all scopes does not allow schedule multiple reconnections for same host+shard, it silently ignores attempts to do that. |
| 938 | +
|
| 939 | + On `new_scheduler` instantiate a scheduler that behaves according to the policy |
| 940 | + """ |
| 941 | + def __init__(self, shard_reconnection_scope, reconnection_policy): |
| 942 | + if not isinstance(shard_reconnection_scope, ShardReconnectionPolicyScope): |
| 943 | + raise ValueError("shard_reconnection_scope must be a ShardReconnectionPolicyScope") |
| 944 | + if not isinstance(reconnection_policy, ReconnectionPolicy): |
| 945 | + raise ValueError("reconnection_policy must be a ReconnectionPolicy") |
| 946 | + self.shard_reconnection_scope = shard_reconnection_scope |
| 947 | + self.reconnection_policy = reconnection_policy |
| 948 | + |
| 949 | + def new_scheduler(self, session) -> _ShardReconnectionScheduler: |
| 950 | + return _NoConcurrentShardReconnectionScheduler(session, self.shard_reconnection_scope, self.reconnection_policy) |
| 951 | + |
| 952 | + |
| 953 | +class _ScopeBucket(object): |
| 954 | + """ |
| 955 | + Holds information for a shard reconnection scope, schedules and executes reconnections. |
| 956 | + """ |
| 957 | + def __init__(self, session, reconnection_policy): |
| 958 | + self._items = [] |
| 959 | + self.session = session |
| 960 | + self.reconnection_policy = reconnection_policy |
| 961 | + self.lock = threading.Lock() |
| 962 | + self.running = False |
| 963 | + self.schedule = self.reconnection_policy.new_schedule() |
| 964 | + |
| 965 | + def _get_delay(self): |
| 966 | + if self.schedule is None: |
| 967 | + self.schedule = self.reconnection_policy.new_schedule() |
| 968 | + try: |
| 969 | + return next(self.schedule) |
| 970 | + except StopIteration: |
| 971 | + self.schedule = self.reconnection_policy.new_schedule() |
| 972 | + return next(self.schedule) |
| 973 | + |
| 974 | + def _schedule(self): |
| 975 | + if self.session.is_shutdown: |
| 976 | + return |
| 977 | + delay = self._get_delay() |
| 978 | + if delay: |
| 979 | + self.session.cluster.scheduler.schedule(delay, self._run) |
| 980 | + else: |
| 981 | + self.session.submit(self._run) |
| 982 | + |
| 983 | + def _run(self): |
| 984 | + if self.session.is_shutdown: |
| 985 | + return |
| 986 | + |
| 987 | + with self.lock: |
| 988 | + try: |
| 989 | + item = self._items.pop() |
| 990 | + except IndexError: |
| 991 | + self.running = False |
| 992 | + self.schedule = None |
| 993 | + return |
| 994 | + |
| 995 | + method, args, kwargs = item |
| 996 | + try: |
| 997 | + method(*args, **kwargs) |
| 998 | + finally: |
| 999 | + self._schedule() |
| 1000 | + |
| 1001 | + def add(self, method, *args, **kwargs): |
| 1002 | + with self.lock: |
| 1003 | + self._items.append([method, args, kwargs]) |
| 1004 | + if not self.running: |
| 1005 | + self.running = True |
| 1006 | + self._schedule() |
| 1007 | + |
| 1008 | + |
| 1009 | +class _NoConcurrentShardReconnectionScheduler(_ShardReconnectionScheduler): |
| 1010 | + def __init__(self, session, shard_reconnection_scope, reconnection_policy): |
| 1011 | + self.already_scheduled = {} |
| 1012 | + self.scopes = {} |
| 1013 | + self.shard_reconnection_scope = shard_reconnection_scope |
| 1014 | + self.reconnection_policy = reconnection_policy |
| 1015 | + self.session = session |
| 1016 | + self.lock = threading.Lock() |
| 1017 | + |
| 1018 | + def _execute(self, scheduled_key, method, *args, **kwargs): |
| 1019 | + try: |
| 1020 | + method(*args, **kwargs) |
| 1021 | + finally: |
| 1022 | + with self.lock: |
| 1023 | + self.already_scheduled[scheduled_key] = False |
| 1024 | + |
| 1025 | + def schedule(self, host_id, shard_id, method, *args, **kwargs): |
| 1026 | + if self.shard_reconnection_scope == ShardReconnectionPolicyScope.Cluster: |
| 1027 | + scope_hash = "global-cluster-scope" |
| 1028 | + else: |
| 1029 | + scope_hash = host_id |
| 1030 | + scheduled_key = f'{host_id}-{shard_id}' |
| 1031 | + |
| 1032 | + with self.lock: |
| 1033 | + if self.already_scheduled.get(scheduled_key): |
| 1034 | + return False |
| 1035 | + self.already_scheduled[scheduled_key] = True |
| 1036 | + |
| 1037 | + scope_info = self.scopes.get(scope_hash, 0) |
| 1038 | + if not scope_info: |
| 1039 | + scope_info = _ScopeBucket(self.session, self.reconnection_policy) |
| 1040 | + self.scopes[scope_hash] = scope_info |
| 1041 | + scope_info.add(self._execute, scheduled_key, method,*args, **kwargs) |
| 1042 | + return True |
| 1043 | + |
| 1044 | + |
867 | 1045 | class RetryPolicy(object): |
868 | 1046 | """ |
869 | 1047 | A policy that describes whether to retry, rethrow, or ignore coordinator |
|
0 commit comments