marmots/stackless.py

103 lines
3 KiB
Python

# This file is part of MarMOTS.
#
# MarMOTS is free software: you can redistribute it and/or modify it under the terms of the GNU Affero
# General Public License as published by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# MarMOTS is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the
# implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General
# Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License along with MarMOTS. If not,
# see <https://www.gnu.org/licenses/>.
#
# Copyright 2009, 2010, 2011, 2020 Jeremy Penner
from greenlet import greenlet
import traceback
from collections import deque
class Scheduler(object):
def __init__(self):
self.running = None
self.rgtasklet = []
self.g = None
def schedule(self, tasklet, *args):
tasklet.args = args
tasklet.scheduled = True
self.rgtasklet.append(tasklet)
def kill(self, tasklet):
try:
self.rgtasklet.remove(tasklet)
except:
pass
if tasklet.g == greenlet.getcurrent():
raise greenlet.GreenletExit()
elif not tasklet.g.dead:
tasklet.g.parent = greenlet.getcurrent()
tasklet.g.throw()
def run(self):
self.g = greenlet.getcurrent()
while len(self.rgtasklet) > 0:
tasklets = self.rgtasklet
self.rgtasklet = []
for tasklet in tasklets:
try:
self.running = tasklet
args = tasklet.args
tasklet.args = None
tasklet.scheduled = False
tasklet.g.parent = self.g
tasklet.g.switch(*args)
except:
traceback.print_exc()
self.kill(tasklet)
self.running = None
self.g = None
def suspend(self):
self.g.switch()
class Tasklet(object):
def __init__(self, cb, scheduler):
self.g = greenlet(cb)
self.args = None
self.scheduler = scheduler
self.scheduled = False
def __call__(self, *args):
self.scheduler.schedule(self, *args)
def kill(self):
self.scheduler.kill(self)
class Channel(object):
def __init__(self, scheduler):
self.queue = deque()
self.tasklet = None
self.scheduler = scheduler
def receive(self):
while len(self.queue) == 0:
self.tasklet = self.scheduler.running
self.tasklet.scheduled = True
self.scheduler.suspend()
self.tasklet = None
return self.queue.popleft()
def send(self, val):
self.queue.append(val)
if self.tasklet:
self.scheduler.schedule(self.tasklet)
scheduler = Scheduler()
def tasklet(cb):
return Tasklet(cb, scheduler)
def channel():
return Channel(scheduler)
def run():
scheduler.run()