remove dict queue

Author: seamile

fastcopy/client.py

@@ -177,9 +177,10 @@ def new_channel(self, sock, user, pkey, password, num=1):
             tp.use_compression(True)
             tp.set_keepalive(60)
             tp.connect(username=user, pkey=pkey, password=password)
-        except paramiko.SSHException:
+        except paramiko.SSHException as e:
+            logging.error(e)
             tp.stop_thread()
-            return
+            return []
 
         try:
             conns = [tp]
@@ -232,8 +233,8 @@ def handshake(self, channel, remote_path: str):
         send_pkt(channel, conn_pkt)
         session_pkt = recv_pkt(channel)
         session_id, = session_pkt.unpack_body()
-        logging.debug(f'[bold cyan]fcp[/bold cyan]: '
-                      f'Channel {channel.get_name()} connected')
+        logging.info(f'[bold cyan]fcp[/bold cyan]: '
+                     f'Channel {channel.get_name()} connected')
 
         return session_id
 
@@ -249,8 +250,8 @@ def _connect(wait):
             for channel in channels:
                 send_pkt(channel, attach_pkt)
                 porter.conn_pool.add(channel)
-                logging.debug(f'[bold cyan]fcp[/bold cyan]: '
-                              f'Channel {channel.get_name()} connected')
+                logging.info(f'[bold cyan]fcp[/bold cyan]: '
+                             f'Channel {channel.get_name()} connected')
 
         for i in range(self.n_channel - 1):
             Thread(target=_connect, args=(0.5 * i,), daemon=True).start()

fastcopy/utils.py

@@ -1,6 +1,5 @@
 import os
 import re
-import sys
 import logging
 from binascii import crc32
 from collections import deque
@@ -14,15 +13,15 @@
 from socket import socket, MSG_WAITALL, error as SocketError
 from struct import pack, unpack
 from threading import Event, Semaphore, Thread
-from typing import (Any, Deque, Dict, Generator, Iterable, List, NamedTuple,
+from typing import (Any, Deque, Dict, Generator, Iterable, List, NamedTuple, Set,
                     Tuple, Union)
 
 from rich.progress import (BarColumn, Progress, TaskID, TextColumn,
                            TransferSpeedColumn)
 
 
 SERVER_ADDR = ('127.0.0.1', 7523)
-CHUNK_SIZE = 8192  # 默认数据块大小 (单位: 字节)
+CHUNK_SIZE = 1024 * 16  # 默认数据块大小 (单位: 字节)
 TIMEOUT = 60 * 5  # 全局超时时间
 LEN_HEAD = 7
 MAX_SEQ = 0xffffffff
@@ -57,28 +56,6 @@ def contains(cls, member: object) -> bool:
         return member in cls.__members__.values()
 
 
-class DictQueue(dict):
-    def __init__(self, size=None):
-        super().__init__()
-        self.size = size or float('inf')
-        self.cur_seq = 0
-
-    def put(self, value):
-        while len(self) >= self.size:
-            self.pop()
-
-        self[self.cur_seq] = value
-        logging.debug(f'[DictQueue] put {self.cur_seq} to {value} into {id(self)}')
-        self.cur_seq = (self.cur_seq + 1) & MAX_SEQ
-
-    def pop(self, seq=None):
-        if seq is None:
-            seq = min(self.keys())
-        value = super().pop(seq)
-        logging.debug(f'[DictQueue] pop {seq} as {value} from {id(self)}')
-        return value
-
-
 class Packet(NamedTuple):
     flag: Flag
     body: bytes
@@ -251,53 +228,39 @@ def __init__(self, size=16):
         self.send_q = Queue(self.size * 16)
         self.recv_q = Queue(self.size * 16)
         self.done = Event()
-        self.connections = {}
+        self.connections: Set[socket] = set()
 
     def send(self, packet: Packet):
         self.send_q.put(packet)
 
     def recv(self, timeout=TIMEOUT) -> Packet:
         return self.recv_q.get(timeout)
 
-    def _send(self, conn: socket, cache: DictQueue):
+    def _send(self, conn: socket):
         conn_name = f'{id(conn):x}'
         while not self.done.is_set():
             packet: Packet = self.send_q.get()
             try:
                 send_pkt(conn, packet)
             except SocketError as e:
-                logging.error(e)
-                sys.exit(1)
-
-            logging.debug(f'[Send] conn-{conn_name}/{cache.cur_seq}: {packet}')
-            cache.put(packet)
+                conn.close()
+                logging.error(f'[Send] conn-{conn_name}: {e}, conn exit.')
+                return
 
-    def _recv(self, conn: socket, cache: DictQueue):
-        n_recv = 0
+    def _recv(self, conn: socket):
         conn_name = f'{id(conn):x}'
         while not self.done.is_set():
             try:
                 packet = recv_pkt(conn)
+                self.recv_q.put(packet)
+                logging.debug(f'[Recv] conn-{conn_name}: {packet}')
             except ConnectionResetError:
+                conn.close()
+                return
+            except PacketError as e:
+                conn.close()
+                logging.error(f'[Recv] conn-{conn_name}: {e}, conn exit.')
                 return
-            except PacketError:
-                resend_pkt = Packet.load(Flag.RESEND, n_recv)
-                send_pkt(conn, resend_pkt)
-                continue
-
-            if packet.flag is Flag.RESEND:
-                seq, = packet.unpack_body()
-                if seq not in cache:
-                    seqs = ','.join(f'{k}' for k in cache.keys())
-                    logging.error(f"{seq} not in {seqs}")
-                    sys.exit(1)
-                pkt = cache[seq]
-                print(f'get {seq} from {id(cache)}')
-                self.send(pkt)
-            else:
-                logging.debug(f'[Recv] conn-{conn_name}/{cache.cur_seq}: {packet}')
-                self.recv_q.put(packet)
-                n_recv = (n_recv + 1) & MAX_SEQ
 
     def add(self, conn: socket):
         '''添加一个连接'''
@@ -308,22 +271,31 @@ def add(self, conn: socket):
         if conn in self.connections:
             return True
 
-        cache = DictQueue(256)
-        t_send = Thread(target=self._send, args=(conn, cache), daemon=True)
+        t_send = Thread(target=self._send, args=(conn,), daemon=True)
         t_send.start()
 
-        t_recv = Thread(target=self._recv, args=(conn, cache), daemon=True)
+        t_recv = Thread(target=self._recv, args=(conn,), daemon=True)
         t_recv.start()
         return True
 
+    def all_closed(self):
+        for conn in self.connections:
+            if not conn.closed:
+                return False
+        return True
+
     def stop(self):
         self.done.set()
         for conn in self.connections:
             conn.close()
 
     def run(self):
         self.done.clear()
-        self.done.wait()
+        while not self.done.wait(1):
+            if not self.connections:
+                continue
+            elif self.all_closed():
+                break
         self.stop()
 
 

setup.cfg

@@ -1,4 +1,4 @@
 [metadata]
-desciption-file = README.md
+desciption_file = README.md
 [bdist_wheel]
 universal=1

setup.py

@@ -4,8 +4,8 @@
     long_description = f_readme.read()
 
 setuptools.setup(
-    name="fast-copy",
-    version="0.1.1",
+    name="fastcopy",
+    version="0.1.3",
     python_requires=">=3.6",
     author="Seamile",
     author_email="[email protected]",