python3 得到一个可用未绑定端口

太过爱你忘了你带给我的痛 2021-07-28 16:20 438阅读 0赞

### 单线程 单进程环境 ###

def getPort():
        pscmd = "netstat -ntl |grep -v Active| grep -v Proto|awk '{print $4}'|awk -F: '{print $NF}'"
        procs = os.popen(pscmd).read()
        procarr = procs.split("\n")
        tt= random.randint(15000,20000)
        if tt not in procarr:
            return tt
        else:
            return getPort()

### 多线程 多进程环境 ###

写本文的目的是介绍如何在 Python 中实现一个函数：`get_free_port` 返回一个未使用的端口号，并且这个函数支持在多线程和多进程环境中使用。

这也就意味着 `get_free_port` 可以在任何时候任意地方调用返回的端口号都能被绑定。

事实上，这个目标在逻辑上并不能实现，因为这个函数并不知道它返回的端口号是否已经被使用了，所以最佳的实践方式就是记住所有返回的端口号，每次调用都检查端口号是否已经输出过。

如果我们再增加一个函数用来释放端口就显得更符合逻辑了。

在 Python 中调用 `socket.bind(('', 0))` 会自动绑定一个端口号，所以我们可以借助这一特性来实现：

import socket 
    
    def get_free_port():  
        sock = socket.socket()
        sock.bind(('', 0))
        ip, port = sock.getnameinfo()
        sock.close()
        return port

这个函数看上去能工作，但是还远远不够。采用这个方法只能让端口在非常短的时间内不被绑定，难以满足在竞态场景中使用。

由此不难看出，获得一个未被绑定的端口号不难，难的是如何把这个端口号安全正确地返回给调用者。鉴于此，我们需要找到一种机制可以首先保证这个未绑定的端口号不被肆意绑定：

> get a free port -> look at dictionary (and lock file) -> bind a free port -> write a dictionary (and lock file) -> release port -> return the port

使用 `lock file` 能够保证即使端口号未绑定，在未拿到锁之前是不会被其他进程绑定的，在这里使用了。全部代码如下：

# freeport.py
    import fasteners
    import threading
    
    class BindFreePort(object):
        def __init__(self, start, stop):
            self.port = None
    
            import random, socket
    
            self.sock = socket.socket()
    
            while True:
                port = random.randint(start, stop)
                try:
                    self.sock.bind(('', port))
                    self.port = port
                    break
                except Exception:
                    continue
    
        def release(self):
            assert self.port is not None
            self.sock.close()
    
    
    class FreePort(object):
        used_ports = set()
    
        def __init__(self, start=4000, stop=6000):
            self.lock = None
            self.bind = None
            self.port = None
    
            from fasteners.process_lock import InterProcessLock
            import time
            while True:
                bind = BindFreePort(start, stop)
    
                if bind.port in self.used_ports:
                    bind.release()
                    continue
    
                '''
                Since we cannot be certain the user will bind the port 'immediately' (actually it is not possible using
                this flow. We must ensure that the port will not be reacquired even it is not bound to anything
                '''
                lock = InterProcessLock(path='/tmp/socialdna/port_{}_lock'.format(bind.port))
                success = lock.acquire(blocking=False)
    
                if success:
                    self.lock = lock
                    self.port = bind.port
                    self.used_ports.add(bind.port)
                    bind.release()
                    break
    
                bind.release()
                time.sleep(0.01)
    
        def release(self):
            assert self.lock is not None
            assert self.port is not None
            self.used_ports.remove(self.port)
            self.lock.release()

测试如下：

#freeport_test.py
    from freeport import FreePort
    
    def get_and_bind_freeport(*args):
        freeport = FreePort(start=4000, stop=4009)
        import time
        time.sleep(1)
        return freeport.port
    
    class FreePortClassTest(unittest.TestCase):
        def test_one_port(self):
            freeport = FreePort(start=4000, stop=4000)
            self.assertEqual(freeport.port, 4000)
            freeport.release()
    
        def test_many_ports(self):
            freeport = FreePort(start=4000, stop=4000)
            self.assertEqual(freeport.port, 4000)
            freeport.release()
            freeport = FreePort(start=4000, stop=4000)
            self.assertEqual(freeport.port, 4000)
            freeport.release()
    
        def test_many_ports_conflict(self):
            def get_port():
                freeport = FreePort(start=4000, stop=4000)
                return freeport.port
    
            def run():
                self.assertEqual(get_port(), 4000)
    
            freeport = FreePort(start=4000, stop=4000)
            self.assertEqual(freeport.port, 4000)
    
            from multiprocessing import Process
            p = Process(target=run)
            p.start()
            p.join(0.1)
    
            self.assertTrue(p.is_alive(), 'the process should find it hard to acquire a free port')
    
            p.terminate()
            p.join()
    
            freeport.release()
    
        def test_multithread_race_condition(self):
            from multiprocessing.pool import ThreadPool
            jobs = 100
            def get_and_bind_freeport(*args):
                freeport = FreePort(start=4000, stop=4000 + jobs - 1)
                import time
                time.sleep(1)
                freeport.release() # needed because thread will not turn back the file descriptor
                return freeport.port
            p = ThreadPool(jobs)
            ports = p.map(get_and_bind_freeport, range(jobs))
            self.assertEqual(len(ports), len(set(ports)))
    
        def test_multiprocess_race_condition(self):
            from multiprocessing.pool import Pool
            p = Pool(10)
            ports = p.map(get_and_bind_freeport, range(10))
            self.assertEqual(len(ports), len(set(ports)))