-
Notifications
You must be signed in to change notification settings - Fork 1
/
titan_server.py
254 lines (213 loc) · 8.73 KB
/
titan_server.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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
import asyncio
import uuid
import pulsar
from pulsar import get_actor
from pulsar.apps import wsgi, ws
from pulsar.apps.wsgi.utils import LOGGER
from gremlinclient import aiohttp_client, connection
from models import Person
from proto import titan_pb2
from wsrpc import WSRPC
# Simple example. This will have to be more carefully handled -
# mananging serializer names etc....
@asyncio.coroutine
def create_person(request_id, blob, pool):
person = titan_pb2.Person()
person.ParseFromString(blob)
name = person.name
email = person.email
url = person.url
return (yield from Person().create(name=name,
email=email,
url=url,
request_id=request_id,
pool=pool,
future_class=asyncio.Future))
def serialize_person(response):
try:
person = titan_pb2.Person()
person.name = response.name
person.url = response.url
person.email = response.email
person.id = response.id
person.label = response.get_label()
response = person.SerializeToString()
finally:
return response
# Simple example. This will have to be more carefully handled.
command_map = {
"create_person": create_person,
"serialize_person": serialize_person}
# `API` commands
@pulsar.command(ack=True)
@asyncio.coroutine
def add_task(request, method, blob):
"""Called by the rpc service to add data to the incoming task queue"""
self = request.actor.app
request_id = str(uuid.uuid4())
yield from self.add_task(request_id, method, blob)
return request_id
@pulsar.command(ack=True)
@asyncio.coroutine
def read_response(request, request_id):
"""Called by rpc service to stream results from the goblin app"""
self = request.actor.app
blob = yield from self.read_response(request_id)
return blob
#Internal app commands - how workers communicate with the monitor process
# read from queues etc.
@pulsar.command(ack=True)
@asyncio.coroutine
def get_task(request):
"""Called by worker process to try to get a task from the main incoming
task queue (runs in monitor context)"""
self = request.actor.app
return self.get_task()
@pulsar.command(ack=True)
@asyncio.coroutine
def enqueue_response(request, request_id, response):
"""Called by a worker to put a response on the monitor controlled response
queue (to be read by the rpc service calling read_response)"""
LOGGER.info("Enqueued by monitor: {}\naid: {}".format(
request.actor.is_monitor(), request.actor.aid))
LOGGER.info(response)
self = request.actor.app
yield from self.response_queues[request_id].put(response)
@pulsar.command(ack=True)
@asyncio.coroutine
def process_task(request, request_id, method, blob):
"""This implements streaming data processing using user defined
functions to perform tasks like serialization/parsing. Enqueues
processed responses on the monitor controlled response_queues"""
LOGGER.info("Processed by monitor: {}\naid: {}".format(
request.actor.is_monitor(), request.actor.aid))
# Simple example. This will have to be more carefully handled
pool = request.actor.pool
serializer = "serialize_{}".format(method.split('_')[-1])
try:
resp = yield from command_map[method](request_id, blob, pool)
except KeyError:
raise KeyError("Unknown command issued")
else:
if isinstance(resp, connection.Stream):
while True:
msg = yield from resp.read()
if msg:
msg = command_map[serializer](msg)
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, msg)
if msg is None:
break
else:
try:
msg = command_map[serializer](resp)
except KeyError:
raise KeyError("Unregistered serializer")
else:
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, msg)
# Demonstrate streaming
yield from asyncio.sleep(1)
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, "Hello")
yield from asyncio.sleep(1)
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, "Streaming")
yield from asyncio.sleep(1)
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, "World")
# Message for client to terminate
yield from request.actor.send(
'monitor', 'enqueue_response', request_id, None)
class Goblin(pulsar.Application):
cfg = pulsar.Config(workers=2)
def monitor_start(self, monitor):
"""Setup message queues"""
# This lives in the monitor context
# Queue incoming messages from rpc service
self.incoming_queue = asyncio.Queue(maxsize=250)
# These queues hold response data that can be asynchronously read
# by the rpc service
self.response_queues = {}
@asyncio.coroutine
def add_task(self, request_id, method, blob):
"""Adzd a task to the incoming task queue"""
self.response_queues[request_id] = asyncio.Queue()
yield from self.incoming_queue.put((request_id, method, blob))
@asyncio.coroutine
def read_response(self, request_id):
"""This method allows the rpc service to read from the response queues
maintained by the app."""
try:
queue = self.response_queues[request_id]
except KeyError:
raise KeyError("Bad request id")
else:
resp = yield from queue.get()
if resp is None:
del self.response_queues[request_id]
return resp
def worker_start(self, worker, exc=None):
"""Setup the global goblin variables, then start asking the monitor
for tasks..."""
worker.pool = aiohttp_client.Pool("ws://localhost:8182",
future_class=asyncio.Future,
loop=worker._loop,
force_release=True)
# check the queue periodically for tasks...
worker._loop.call_soon(self.start_working, worker)
def worker_stopping(self, worker, exc=None):
"""Close the connection pool for this process"""
worker._loop.call_soon(pulsar.ensure_future, worker.pool.close())
def start_working(self, worker):
"""Don't be lazy"""
pulsar.ensure_future(self.run(worker))
@asyncio.coroutine
def run(self, worker):
"""Try to get tasks from the monitor. If tasks are available process
using same worker, if not, wait a second, and ask again...
BE PERSISTENT!"""
request_id, method, blob = yield from worker.send(
worker.monitor, 'get_task')
if request_id and method and blob:
yield from pulsar.send(
worker.aid, 'process_task', request_id, method, blob)
worker._loop.call_later(1, self.start_working, worker)
def get_task(self):
"""Check for tasks, if available, pass data to calling worker for
processing..."""
try:
request_id, method, blob = self.incoming_queue.get_nowait()
except asyncio.QueueEmpty:
LOGGER.info("No tasks available :( :( :(")
return None, None, None
else:
return request_id, method, blob
# def actorparams(self, monitor, params):
# pool = aiohttp_client.Pool()
# params.update({"pool": })
class TitanRPC(WSRPC):
def rpc_dispatch(self, websocket, method, blob):
"""A generic dispatch function that sends commands to the
Goblin application"""
request_id = yield from pulsar.send(
'goblin_server', 'add_task', method, blob)
while True:
blob = yield from pulsar.send(
'goblin_server', 'read_response', request_id)
if blob is None:
break
websocket.write(blob)
websocket.write_close()
class TitanRPCSite(wsgi.LazyWsgi):
"""Handler for the RPCServer"""
def setup(self, environ):
wm = ws.WebSocket('/', TitanRPC())
return wsgi.WsgiHandler(middleware=[wm], async=True)
class Server(pulsar.MultiApp):
@asyncio.coroutine
def build(self):
yield self.new_app(wsgi.WSGIServer, callable=TitanRPCSite())
yield self.new_app(Goblin, name='goblin')
if __name__ == "__main__":
Server().start()