How can I elagently organize this multithreaded python code?

Question

I am working on a Python service that subscribes real-time streaming data from one messaging broker and publishes to another broker, in some situations I also need to get snapshot data from other data source on network disconnection or system recovery. While the streaming data comes from one thread, and some service events happen in another thread, I decided to create a data processing thread to just pop the queue one by one. I got it to work but later I tried to keep the snapshot fetching logic in a separate thread and that's where things get messy.

I know this is a long question with a lot of specific nuances but I tried to make the example here as clear as I can.

So here is what the 1st attempt looks like, and it works well:

import queue
import threading

def process_data(data_queue, data_store):
    # data_store is my internal cache data structure. 
    # so for simplicity and demonstration purpose, I assume the following:
    # if its type is dict, it's snapshot data
    # if its type is tuple, it's a key/value pair and that's an incremental update data
    # if it is -1, we terminate the queue processing
    # if it is -2, we need to retrieve a snapshot
    while True:
        x = data_queue.get()
        if isinstance(x, dict):
            data_store.on_snapshot(x)
        elif isinstance(x, tuple):
            k, v = x
            data_store.on_update(k, v)
        elif isinstance(x, int):
            if x == -1:
                data_queue.task_done()
                break
            elif x == -2:
                get_snapshot() # this is potentially a long blocking call
            else:
                print('unknown int', x)
        else: 
            print('unknown data', x)

        data_queue.task_done()


if __name__ == '__main__':
    data_store = DataStore()
    data_queue = queue.Queue()

    # start other threads that write data to the queue
    start_data_writer1(data_queue)
    start_data_writer2(data_queue)
    start_thread_for_some_event(data_queue) # may put -2 in the queue for snapshot
    
    process_thread = threading.Thread(
        target=process_data, 
        args=(data_queue, data_store))

    process_thread.start()
    data_queue.put(-2) # signal a snapshot fetching    
    do_something_else()

    try:
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print('terminating...')
    finally:
        # to break out of the infinite loop in process_data()
        data_queue.put(-1) 
        process_thread.join()
        data_queue.join()

This way works however I don't particularly like the fact that I am calling get_snapshot() function in the processing thread. I thought the idea of the processing thread is to be busy at popping data off the queue all the time, unless there's nothing to pop. In the above implementation, during the get_snapshot call, it is possible that queue could be built up due to other writer threads.

So I tried something else and I also wanted to be able to exit the program gracefully. That's where things get really ugly. I created a new thread for occasionally fetching the snapshot and used condition object for thread communication. This is what I did on top of the existing code:

snapshot_lock = threading.Lock()
cond = threading.Condition(snapshot_lock)
need_for_snapshot = False # used to trigger snapshots
keep_snapshot_thread = True # flag if the snapshot thread is done

# then I need to add this new function to run snapshot fetching
def fetch_snapshot(data_queue):
    global need_for_snapshot
    global keep_snapshot_thread
    
    def _need_snapshot():
        global need_for_snapshot
        return need_for_snapshot 
    
    while True:
        with cond:
            cond.wait_for(_need_snapshot)
            if not keep_snapshot_thread:
                break
            data_queue.put(get_snapshot()) # the long blocking function
            need_for_snapshot = False


# in process_data() function, everything stays the same except for the `if` statement for handling `x == 2`

def process_data(data_queue, data_store):
   global need_for_snapshot
   while True:
       x = data_queue.get()
       # omitting some old code
       elif isinstance(x, int):
           if x == -1:
               data_queue.task_done()
               break
           elif x == -2: 
               **with cond:
                   need_for_snapshot = True
                   cond.notify()**
    # more code omitted    

if __name__ == '__main__':
    # same code as before except for the finally part
    try:
        # start other threads...omitting some code
        # when a snapshot is needed in these threads
        # do the following
        # with cond:
        #     need_for_snapshot = True
        #     cond.notify()

        # start snapshot worker thread 
        snapshot_thread = threading.Thread(
            target=fetch_snapshot, args=(data_queue,))
        process_thread = threading.Thread(
            target=process_data,
            args=(data_queue, data_store))

        snapshot_thread.start()
        process_thread.start()
        data_queue.put(-2) # signal fetching a snapshot

        # omitting more code here...

    finally: 
        keep_snapshot_thread = False
        # we don't technically need to trigger another snapshot now
        # but the code below is to unblock the cond.wait_for() part
        # since keep_snapshot_thread flag is just flipped, we can use 
        # it to break out of the infinite loop in fetch_snapshot thread.
        # This is the part that I feel hacky...
        with cond: 
            need_for_snapshot = True
            cond.notify()
        snapshot_t.join()

        data_queue.put(-1) # signal the termination of process_thread
        process_t.join()
        data_queue.join()

I think I got this to work, especially that the program can exit gracefully when I hit ctrl-c but it is so ugly and tricky that I had to play with it quick a bit to get it to work correctly.

Is there some way I can write it more elegantly? Is there some sort of pattern that we generally use to solve this type of problem? Thank you so much for your help.

Answer 1

The standard technique for handling multiple producers and multiple consumers is to use an Event is_done and a joinable Queue work .

The worker queues do nothing but:

    while not event.is_set():
         grab something to do from the work queue
         do it
         work.task_done()

Your main worker does the following:

    start the jobs that produce work
    wait for them to be done
    work.join() # wait for the queue to be empty
    event.set() # let the workers exit