marmots/stackless.py

# 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()