Linux和Windows文件高速傳輸

以下代碼使用socket實現了文件傳輸:

server,運行在目標服務端上:

import os
import signal
import psutil
import socket
import threading
from multiprocessing import Queue, Process


class AtomicInteger:
    def __init__(self, value=0):
        self._value = value
        self._lock = threading.Lock()

    def inc(self):
        with self._lock:
            self._value += 1
            return self._value

    def dec(self):
        with self._lock:
            self._value -= 1
            return self._value

    @property
    def value(self):
        with self._lock:
            return self._value

    @value.setter
    def value(self, v):
        with self._lock:
            self._value = v


def sendall(con, data):
    count = 0
    with memoryview(data) as view, view.cast("B") as byte_view:
        amount = len(byte_view)
        while count < amount:
            v = con.send(byte_view[count:])
            count += v


class ConnectionHandler(threading.Thread):

    def __init__(self, con, amount: AtomicInteger):
        super().__init__()
        amount.inc()
        self.amount = amount
        self.con = con

    def run(self) -> None:
        con = self.con
        con.sendall(b'connected')
        res: bytes = con.recv(409600)
        """download|ab.txt"""
        """upload|ab.txt"""
        temp= res.decode('utf8').split('|')
        if temp.__len__() != 2:
            return
        mode, file = temp
        if mode == 'download':
            with open(file, 'rb') as f:
                sendall(con, f.read())
        elif mode == 'upload':
            with open(file, 'wb') as f:
                fragment = f.read(1024)
                while fragment:
                    fragment = f.read(1024)
                    f.write(fragment)
        self.amount.dec()
        con.shutdown(1)
        con.close()
        print("mode: %s\tfile: %s ok....." % (mode, file))


class ClientHandler(Process):

    def __init__(self, client_queue: Queue, thread_num_update_queue: Queue, amount_queue: Queue):
        super().__init__()
        self.client_queue = client_queue
        self.thread_num_update_queue = thread_num_update_queue
        self.amount_queue = amount_queue
        self.thread_amount = None

    def run(self):
        threading.Thread(target=self.update_amount_server).start()
        self.thread_amount = AtomicInteger(0)
        while True:
            connect, addr = self.client_queue.get()
            ConnectionHandler(connect, self.thread_amount).start()

    def update_amount_server(self):
        while True:
            cmd = self.thread_num_update_queue.get()
            if cmd == 1:
                self.amount_queue.put(self.thread_amount.value)
            elif cmd == 2:
                return


def server(queue_list: list):
    service = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    service.bind(('127.0.0.1', 7777))
    service.listen(9999)
    print('server start')
    min_queue = None
    while True:
        con = service.accept()
        min_amount = 999999
        for client_queue, thread_num_update_queue, amount_queue in queue_list:
            thread_num_update_queue.put(1)
            amount = amount_queue.get()
            if amount < min_amount:
                min_amount = amount
                min_queue = client_queue
        min_queue.put(con)
        print(min_amount, 'min_amount')


if __name__ == '__main__':
    core_num = psutil.cpu_count()
    print(core_num - 1)
    q_list = []
    ps = []
    for i in range(core_num - 1):
        _client_queue = Queue()
        _thread_num_update_queue = Queue(maxsize=1)
        _amount_queue = Queue(maxsize=1)
        q_list.append((_client_queue, _thread_num_update_queue, _amount_queue))
        tp = ClientHandler(_client_queue, _thread_num_update_queue, _amount_queue)
        ps.append(tp)
        tp.start()
    p_server = Process(target=server, args=(q_list,))
    p_server.start()
    while True:
        _service_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        _service_.bind(('127.0.0.1', 6767))
        _service_.listen(2)
        _con, addr = _service_.accept()
        content = _con.recv(4096)
        if content == b'exit':
            for p in ps:
                os.kill(p.pid,signal.SIGABRT)
            os.kill(p_server.pid,signal.SIGABRT)
            break
    print('exit ok')


client客戶端:

import os
import socket
import threading


def download(target, local):
    if not os.path.isfile(target):
        return
    print(target, local)
    client = socket.socket()
    client.connect(('127.0.0.1', 7777))
    res = client.recv(40960)
    if res != b'connected':
        print('連接失敗')
        return
    print(res)
    client.sendall(bytes("download|%s" % target, encoding='utf8'))
    dir_name, _ = os.path.split(local)
    if not os.path.exists(dir_name):
        os.makedirs(dir_name, exist_ok=True)
    with open(local, 'wb')as f:
        fragment = client.recv(1024)
        while fragment:
            f.write(fragment)
            fragment = client.recv(1024)
        print('退出了')


def get_file_list(dir_name):
    res = []
    for i in os.walk(dir_name):
        print(i)
        for filename in i[2]:
            res.append("%s/%s" % (i[0], filename))
    print(res)
    return res


file_list = get_file_list('G:/Driver驅動')
file_list = get_file_list('G:/test000')
save_list = map(lambda x: "F" + x[1:], file_list)
threads = []
for i in zip(file_list, save_list):
    threads.append(threading.Thread(target=download, args=i))
for t in threads:
    t.setDaemon(False)
    t.start()

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章