/
taskdist.py
154 lines (133 loc) · 5.09 KB
/
taskdist.py
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# taskdist.py
#
# Support for distributed tasks
import time
import redis
# Adjust these settings as appropriate
REDIS_HOST = "localhost"
REDIS_PORT = 6379
REDIS_DB = 0
from multiprocessing.connection import Client, Listener
import tasklib
class ProxyTask(tasklib.Task):
def __init__(self,proxyname,target,conn,export=True):
super().__init__(name="proxy", export=export)
self.proxyname = proxyname
self.target = target
self.conn = conn
def run(self):
try:
while True:
msg = self.recv()
self.conn.send((self.target,msg))
finally:
self.conn.close()
if self.export:
tasklib.unregister(self.proxyname)
# Establish a proxy connection to another machine
def proxy(proxyname,target,address,authkey):
conn = Client(address,authkey=authkey)
pxy = ProxyTask(proxyname,target,conn)
pxy.start()
tasklib.register(proxyname,pxy)
# Task that receives messages from clients and sends them locally
class DispatchClientTask(tasklib.Task):
def __init__(self,conn):
super().__init__(name="dispatchclient")
self.conn = conn
def run(self):
try:
while True:
target,msg = self.conn.recv()
tasklib.send(target,msg)
finally:
self.conn.close()
# Task that listens for incoming connections
class DispatcherTask(tasklib.Task):
def __init__(self,address,authkey):
super().__init__(name="dispatcher")
self.address = address
self.authkey = authkey
def run(self):
serv = Listener(self.address,authkey=self.authkey)
while True:
try:
client = serv.accept()
DispatchClientTask(client).start()
except Exception as e:
self.log.info("Error : %s", e, exc_info=True)
_dispatcher = None
def start_dispatcher(address,authkey):
global _dispatcher
if _dispatcher:
return
_dispatcher = DispatcherTask(address,authkey)
_dispatcher.start()
...
# Set of all task names that we registered
_registered = set()
# Global task registration.
# Stores a dispatcher (address, authkey) tuple in the registry
def global_register(target):
if not _dispatcher:
return
print("Registering", target)
r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB)
r.set(target,"%s:%d" % _dispatcher.address)
_registered.add(target)
# Global task unregistration
def global_unregister(*targets):
if _dispatcher and targets:
print("Unregistering", targets)
r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB)
for name in (set(targets) & _registered):
r.delete(name)
_registered.remove(name)
# Unregister all tasks at interpreter exit
import atexit
atexit.register(lambda : global_unregister(*_registered))
# Task that receives messages for all unknown targets and creates proxies.
class ResolverTask(tasklib.Task):
def __init__(self,host,port,db,authkey):
super().__init__(name="resolver")
self.host = host
self.port = port
self.db = db
self.authkey = authkey
def run(self):
# Establish a connection with the redis server
r = redis.Redis(host=self.host,port=self.port,db=self.db)
tasklib.register("resolver",self)
try:
while True:
# Get an outgoing message
target, msg = self.recv()
print("RESOLVER:", target, msg)
# Check if a proxy was already registered
if tasklib.lookup(target):
tasklib.send(target,msg)
continue
# See if redis knows anything about the target
try:
target_info = r.get(target)
if target_info:
host, port = target_info.decode('utf-8').split(":")
port = int(port)
# Create a proxy task
proxy(target,target,(host,port),self.authkey)
self.log.info("Connected to %s", (host,port))
# Send the message to the proxy
tasklib.send(target,msg)
else:
self.log.info("Nothing known about target '%s'", target)
except Exception as e:
self.log.info("Couldn't resolve '%s' : %s:%s", target, type(e),e)
finally:
tasklib.unregister("resolver")
_resolver = None
def start_resolver(authkey):
global _resolver
if _resolver:
return
_resolver = ResolverTask(REDIS_HOST, REDIS_PORT, REDIS_DB,authkey)
_resolver.start()