Python+zmq实现负载均衡

# encoding: utf-8

import zmq
import time
import random
from threading import Thread
from vbsl.config import *
from vbsl.get_logs import *
from vbsl.parse_logs import *

NBR_WORKERS = 30


def worker_thread(context=None):
    context = context or zmq.Context.instance()
    worker = context.socket(zmq.REQ)
    worker.connect("tcp://localhost:5671")

    total = 0
    while True:
        # Tell the router we're ready for work
        worker.send(b"ready")

        # Get workload from router, until finished
        workload = worker.recv_string()
        finished = workload == "END"
        if finished:
            print("Processed: %d tasks" % total)
            break
        total += 1

        # Do the work
        station = workload
        ip = get_ip(station)
        if ip is None:
            continue
        print(station, ip)
        file_name = f"/tmp/{station}-{get_date()}.log"
        if os.path.isfile(file_name):
            print(f"{file_name} exist.")
        else:
            print("start script")
            os.system(F"ssh jumper "/home/bsp/IP/get_logs.sh {ip} {station} {get_date()}"")


def main():
    context = zmq.Context.instance()
    client = context.socket(zmq.ROUTER)
    client.bind("tcp://*:5671")

    for _ in range(NBR_WORKERS):
        Thread(target=worker_thread).start()

    random.shuffle(STATIONS)
    print(len(STATIONS))
    print(STATIONS)

    for station in STATIONS:
        # LRU worker is next waiting in the queue
        address, empty, ready = client.recv_multipart()
        client.send_multipart([
            address,
            b'',
            bytes(station, encoding="utf-8"),
        ])
        time.sleep(1)

    # Now ask mama to shut down and report their results
    for _ in range(NBR_WORKERS):
        address, empty, ready = client.recv_multipart()
        client.send_multipart([
            address,
            b'',
            b'END',
        ])


if __name__ == "__main__":
    main()