Simpy Store Batch Processing

Question

I'm trying to create a producer/consumer simulation where the consumer processes items in batches. The problem is that the Store.get() function removes items from the Store as soon as it is called, but I need it to wait until I've called yield:

import simpy

def producer(env, Q):
    item = 0
    while True:
        yield Q.put(item)
        print('Submit item:%d'%item)
        item += 1

def consumer(env, Q):
    while True:
        yield env.timeout(20)
        events = [Q.get() for i in range(4)]
        items = yield env.all_of(events)
        print([items for items in items.values()])


env = simpy.Environment()
maxQD = 2
Q = simpy.Store(env, capacity=maxQD)

env.process(producer(env, Q))
env.process(consumer(env, Q))
env.run(until=500)

Which produces the following output:

Submit item:0
Submit item:1
Submit item:2
Submit item:3
Submit item:4
[0, 1, 2, 3]
Submit item:5
Submit item:6
Submit item:7
Submit item:8
[4, 5, 6, 7]
Submit item:9
Submit item:10
Submit item:11
Submit item:12
[8, 9, 10, 11]
...

With maxQD set to 2, I would have expected just:

Submit item:0
Submit item:1

with the consumer blocking until it successfully gets 4 items, and the producer unable to add more than 2.

You can sort of fix this problem by checking len(Q.items):

import simpy

def producer(env, Q):
    item = 0
    while True:
        yield Q.put(item)
        print('Submit item:%d'%item)
        item += 1

def consumer(env, Q):
    while True:
        yield env.timeout(20)
        if len(Q.items) >= 4:
            events = [Q.get() for i in range(4)]
            items = yield env.all_of(events)
            print([items for items in items.values()])


env = simpy.Environment()
maxQD = 4
Q = simpy.Store(env, capacity=maxQD)

env.process(producer(env, Q))
env.process(consumer(env, Q))
env.run(until=500)

But you still get the frustrating behavior that the get() removes the items before yielding which makes it look like 5 items were added to the Q (note maxQD changed to 4):

Submit item:0
Submit item:1
Submit item:2
Submit item:3
Submit item:4
[0, 1, 2, 3]

Answer 1

Solved by subclassing Store:

class BatchGet(simpy.resources.base.Get):
    def __init__(self, resource, count):
        self.count = count
        super().__init__(resource)

class BatchStore(simpy.resources.store.Store):
    get = simpy.core.BoundClass(BatchGet)

    def _do_put(self, event):
        if len(self.items) + len(event.item) <= self._capacity:
            self.items.extend(event.item)
            event.succeed()   

    def _do_get(self, event):
        count = event.count
        if len(self.items) >= count:
            ret = self.items[:count]
            self.items = self.items[count:]
            event.succeed(ret)

Puts must take a list (since it's a batch of items), and gets return a list.