forked from mongodb/mongo-python-driver
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy paththread_util.py
More file actions
118 lines (93 loc) · 3.28 KB
/
Copy paththread_util.py
File metadata and controls
118 lines (93 loc) · 3.28 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# Copyright 2012 10gen, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities to abstract the differences between threads and greenlets."""
import threading
import weakref
have_greenlet = True
try:
import greenlet
except ImportError:
have_greenlet = False
class Ident(object):
def __init__(self):
self._refs = {}
def watching(self):
"""Is the current thread or greenlet being watched for death?"""
return self.get() in self._refs
def unwatch(self):
self._refs.pop(self.get(), None)
def get(self):
"""An id for this thread or greenlet"""
raise NotImplementedError
def watch(self, callback):
"""Run callback when this thread or greenlet dies. callback takes
one meaningless argument.
"""
raise NotImplementedError
class ThreadIdent(Ident):
def __init__(self):
super(ThreadIdent, self).__init__()
self._local = threading.local()
# We watch for thread-death using a weakref callback to a thread local.
# Weakrefs are permitted on subclasses of object but not object() itself.
class ThreadVigil(object):
pass
def get(self):
if not hasattr(self._local, 'vigil'):
self._local.vigil = ThreadIdent.ThreadVigil()
return id(self._local.vigil)
def watch(self, callback):
tid = self.get()
self._refs[tid] = weakref.ref(self._local.vigil, callback)
class GreenletIdent(Ident):
def get(self):
return id(greenlet.getcurrent())
def watch(self, callback):
current = greenlet.getcurrent()
tid = self.get()
if hasattr(current, 'link'):
# This is a Gevent Greenlet (capital G), which inherits from
# greenlet and provides a 'link' method to detect when the
# Greenlet exits.
current.link(callback)
self._refs[tid] = None
else:
# This is a non-Gevent greenlet (small g), or it's the main
# greenlet.
self._refs[tid] = weakref.ref(current, callback)
def create_ident(use_greenlets):
if use_greenlets:
return GreenletIdent()
else:
return ThreadIdent()
class Counter(object):
"""A thread- or greenlet-local counter.
"""
def __init__(self, use_greenlets):
self.ident = create_ident(use_greenlets)
self._counters = {}
def inc(self):
tid = self.ident.get()
self._counters.setdefault(tid, 0)
self._counters[tid] += 1
return self._counters[tid]
def dec(self):
tid = self.ident.get()
if self._counters.get(tid, 0) > 0:
self._counters[tid] -= 1
return self._counters[tid]
else:
return 0
def get(self):
return self._counters.get(self.ident.get(), 0)