test_ok2/contrib/pygreen/greensock2.py

491 lines
14 KiB
Python

""" This is a base implementation of thread-like network programming
on top of greenlets. From API available here it's quite unlikely
that you would like to use anything except wait(). Higher level interface
is available in pipe directory
"""
import os, sys
try:
from stackless import greenlet
except ImportError:
import py
greenlet = py.magic.greenlet
from collections import deque
from select import select as _select
from time import time as _time
from heapq import heappush, heappop, heapify
TRACE = False
def meetingpoint():
senders = deque() # list of senders, or [None] if Giver closed
receivers = deque() # list of receivers, or [None] if Receiver closed
return (MeetingPointGiver(senders, receivers),
MeetingPointAccepter(senders, receivers))
def producer(func, *args, **kwds):
iterable = func(*args, **kwds)
giver, accepter = meetingpoint()
def autoproducer():
try:
giver.wait()
for obj in iterable:
giver.give(obj)
giver.wait()
finally:
giver.close()
autogreenlet(autoproducer)
return accepter
class MeetingPointBase(object):
def __init__(self, senders, receivers):
self.senders = senders
self.receivers = receivers
self.g_active = g_active
def close(self):
while self.senders:
if self.senders[0] is None:
break
packet = self.senders.popleft()
if packet.g_from is not None:
self.g_active.append(packet.g_from)
else:
self.senders.append(None)
while self.receivers:
if self.receivers[0] is None:
break
other = self.receivers.popleft()
self.g_active.append(other)
else:
self.receivers.append(None)
__del__ = close
def closed(self):
return self.receivers and self.receivers[0] is None
class MeetingPointGiver(MeetingPointBase):
def give(self, obj):
if self.receivers:
if self.receivers[0] is None:
raise MeetingPointClosed
other = self.receivers.popleft()
g_active.append(g_getcurrent())
packet = _Packet()
packet.payload = obj
other.switch(packet)
if not packet.accepted:
raise Interrupted("packet not accepted")
else:
packet = _Packet()
packet.g_from = g_getcurrent()
packet.payload = obj
try:
self.senders.append(packet)
g_dispatcher.switch()
if not packet.accepted:
raise Interrupted("packet not accepted")
except:
remove_by_id(self.senders, packet)
raise
def give_queued(self, obj):
if self.receivers:
self.give(obj)
else:
packet = _Packet()
packet.g_from = None
packet.payload = obj
self.senders.append(packet)
def ready(self):
return self.receivers and self.receivers[0] is not None
def wait(self):
if self.receivers:
if self.receivers[0] is None:
raise MeetingPointClosed
else:
packet = _Packet()
packet.g_from = g_getcurrent()
packet.empty = True
self.senders.append(packet)
try:
g_dispatcher.switch()
if not packet.accepted:
raise Interrupted("no accepter found")
except:
remove_by_id(self.senders, packet)
raise
def trigger(self):
if self.ready():
self.give(None)
class MeetingPointAccepter(MeetingPointBase):
def accept(self):
while self.senders:
if self.senders[0] is None:
raise MeetingPointClosed
packet = self.senders.popleft()
packet.accepted = True
if packet.g_from is not None:
g_active.append(packet.g_from)
if not packet.empty:
return packet.payload
g = g_getcurrent()
self.receivers.append(g)
try:
packet = g_dispatcher.switch()
except:
remove_by_id(self.receivers, g)
raise
if type(packet) is not _Packet:
remove_by_id(self.receivers, g)
raise Interrupted("no packet")
packet.accepted = True
return packet.payload
def ready(self):
for packet in self.senders:
if packet is None:
return False
if not packet.empty:
return True
return False
def wait_trigger(self, timeout=None, default=None):
if timeout is None:
return self.accept()
else:
timer = Timer(timeout)
try:
try:
return self.accept()
finally:
timer.stop()
except Interrupted:
if timer.finished:
return default
raise
class MeetingPointClosed(greenlet.GreenletExit):
pass
class Interrupted(greenlet.GreenletExit):
pass
class ConnexionClosed(greenlet.GreenletExit):
pass
class _Packet(object):
empty = False
accepted = False
def remove_by_id(d, obj):
lst = [x for x in d if x is not obj]
d.clear()
d.extend(lst)
def wait_input(sock):
_register(g_iwtd, sock)
def recv(sock, bufsize):
wait_input(sock)
buf = sock.recv(bufsize)
if not buf:
raise ConnexionClosed("inbound connexion closed")
return buf
def recvall(sock, bufsize):
in_front = False
data = []
while bufsize > 0:
_register(g_iwtd, sock, in_front=in_front)
buf = sock.recv(bufsize)
if not buf:
raise ConnexionClosed("inbound connexion closed")
data.append(buf)
bufsize -= len(buf)
in_front = True
return ''.join(data)
def read(fd, bufsize):
assert fd >= 0
wait_input(fd)
buf = os.read(fd, bufsize)
if not buf:
raise ConnexionClosed("inbound connexion closed")
return buf
def readall(fd, bufsize):
assert fd >= 0
in_front = False
data = []
while bufsize > 0:
_register(g_iwtd, fd, in_front=in_front)
buf = os.read(fd, bufsize)
if not buf:
raise ConnexionClosed("inbound connexion closed")
data.append(buf)
bufsize -= len(buf)
in_front = True
return ''.join(data)
def wait_output(sock):
_register(g_owtd, sock)
def sendall(sock, buffer):
in_front = False
while buffer:
_register(g_owtd, sock, in_front=in_front)
count = sock.send(buffer)
buffer = buffer[count:]
in_front = True
def writeall(fd, buffer):
assert fd >= 0
in_front = False
while buffer:
_register(g_owtd, fd, in_front=in_front)
count = os.write(fd, buffer)
if not count:
raise ConnexionClosed("outbound connexion closed")
buffer = buffer[count:]
in_front = True
def sleep(duration):
timer = Timer(duration)
try:
_suspend_forever()
finally:
ok = timer.finished
timer.stop()
if not ok:
raise Interrupted
def _suspend_forever():
g_dispatcher.switch()
def oneof(*callables):
assert callables
for c in callables:
assert callable(c)
greenlets = [tracinggreenlet(c) for c in callables]
g_active.extend(greenlets)
res = g_dispatcher.switch()
for g in greenlets:
g.interrupt()
return res
def allof(*callables):
for c in callables:
assert callable(c)
greenlets = [tracinggreenlet(lambda i=i, c=c: (i, c()))
for i, c in enumerate(callables)]
g_active.extend(greenlets)
result = [None] * len(callables)
for _ in callables:
num, res = g_dispatcher.switch()
result[num] = res
return tuple(result)
class Timer(object):
started = False
finished = False
def __init__(self, timeout):
self.g = g_getcurrent()
entry = (_time() + timeout, self)
if g_timers_mixed:
g_timers.append(entry)
else:
heappush(g_timers, entry)
def stop(self):
global g_timers_mixed
if not self.finished:
for i, (activationtime, timer) in enumerate(g_timers):
if timer is self:
g_timers[i] = g_timers[-1]
g_timers.pop()
g_timers_mixed = True
break
self.finished = True
# ____________________________________________________________
class tracinggreenlet(greenlet):
def __init__(self, function, *args, **kwds):
self.function = function
self.args = args
self.kwds = kwds
def __repr__(self):
## args = ', '.join([repr(s) for s in self.args] +
## ['%s=%r' % keyvalue for keyvalue in self.kwds.items()])
## return '<autogreenlet %s(%s)>' % (self.function.__name__, args)
return '<%s %s at %s>' % (self.__class__.__name__,
self.function.__name__,
hex(id(self)))
def run(self):
self.trace("start")
try:
res = self.function(*self.args, **self.kwds)
except Exception, e:
self.trace("stop (%s%s)", e.__class__.__name__,
str(e) and (': '+str(e)))
raise
else:
self.trace("done")
return res
def trace(self, msg, *args):
if TRACE:
print self, msg % args
def interrupt(self):
self.throw(Interrupted)
class autogreenlet(tracinggreenlet):
def __init__(self, *args, **kwargs):
super(autogreenlet, self).__init__(*args, **kwargs)
self.parent = g_dispatcher
g_active.append(self)
g_active = deque()
g_iwtd = {}
g_owtd = {}
g_timers = []
g_timers_mixed = False
g_getcurrent = greenlet.getcurrent
def _register(g_wtd, sock, in_front=False):
d = g_wtd.setdefault(sock, deque())
g = g_getcurrent()
if in_front:
d.appendleft(g)
else:
d.append(g)
try:
if g_dispatcher.switch() is not g_wtd:
raise Interrupted
except:
remove_by_id(d, g)
raise
##def _unregister_timer():
## ...
def check_dead_greenlets(mapping):
to_remove = [i for i, v in mapping.items() if not v]
for k in to_remove:
del mapping[k]
#def check_waiters(active):
# if active in g_waiters:
# for g in g_waiters[active]:
# g.switch()
# del g_waiters[active]
def dispatcher_mainloop():
global g_timers_mixed
GreenletExit = greenlet.GreenletExit
while 1:
try:
while g_active:
#print 'active:', g_active[0]
g_active.popleft().switch()
# active.switch()
# if active.dead:
# check_waiters(active)
# del active
if g_timers:
if g_timers_mixed:
heapify(g_timers)
g_timers_mixed = False
activationtime, timer = g_timers[0]
delay = activationtime - _time()
if delay <= 0.0:
if timer.started:
heappop(g_timers)
#print 'timeout:', g
timer.finished = True
timer.g.switch()
# if timer.g.dead:
# check_waiters(timer.g)
continue
delay = 0.0
timer.started = True
else:
check_dead_greenlets(g_iwtd)
check_dead_greenlets(g_owtd)
if not (g_iwtd or g_owtd):
# nothing to do, switch to the main greenlet
g_dispatcher.parent.switch()
continue
delay = None
#print 'selecting...', g_iwtd.keys(), g_owtd.keys(), delay
iwtd, owtd, _ = _select(g_iwtd.keys(), g_owtd.keys(), [], delay)
#print 'done'
for s in owtd:
if s in g_owtd:
d = g_owtd[s]
#print 'owtd:', d[0]
# XXX: Check if d is non-empty
try:
g = d.popleft()
except IndexError:
g = None
if not d:
try:
del g_owtd[s]
except KeyError:
pass
if g:
g.switch(g_owtd)
# if g.dead:
# check_waiters(g)
for s in iwtd:
if s in g_iwtd:
d = g_iwtd[s]
#print 'iwtd:', d[0]
# XXX: Check if d is non-empty
try:
g = d.popleft()
except IndexError:
g = None
if not d:
try:
del g_iwtd[s]
except KeyError:
pass
if g:
g.switch(g_iwtd)
# if g.dead:
# check_waiters(g)
except GreenletExit:
raise
except:
import sys
g_dispatcher.parent.throw(*sys.exc_info())
g_dispatcher = greenlet(dispatcher_mainloop)
#g_waiters = {}