brisonus_data_collector/serial_recorder.py

import serial
import time
import threading
from queue import Queue
from collections import deque
from typing import Optional
import logging
import os
from datetime import datetime
from frame_processor import FrameProcessor
from frame_finder import FrameFinder

class SerialRecorder:
    """串口数据记录器 - 记录原始数据并进行帧处理"""
    def __init__(self, port: str = None, baudrate: int = None, buffer_size: int = 4*1024*1024,
                 output_dir: str = 'dumps', max_frames_per_file: int = 10000, num_channels: int = 8):
        """
        初始化串口记录器
        Args:
            port: 串口名称
            baudrate: 波特率
            buffer_size: 缓冲区大小（字节）
            output_dir: 输出目录
            max_frames_per_file: 每个文件最大帧数
            num_channels: 通道数量，默认为8
        """
        self.port = port
        self.baudrate = baudrate
        self.buffer_size = buffer_size
        self.output_dir = output_dir
        self.max_frames_per_file = max_frames_per_file
        self.num_channels = num_channels
        
        # 初始化主缓冲区
        self.buffer = deque(maxlen=buffer_size)  # 主接收缓冲区
        self.buffer_lock = threading.Lock()  # 主缓冲区锁
        
        # 串口和线程相关
        self.serial: Optional[serial.Serial] = None
        self.is_running = False
        self.read_thread: Optional[threading.Thread] = None
        self.save_thread: Optional[threading.Thread] = None
        self.process_thread: Optional[threading.Thread] = None
        self.bytes_received = 0
        self.start_time = None
        self.last_data_time = None
        self.error_count = 0
        self.max_errors = 5
        
        # Ping-pong buffer 设置
        self.save_buffer_size = 4 * 1024 * 1024  # 4MB
        self.save_buffer_a = deque(maxlen=self.save_buffer_size)
        self.save_buffer_b = deque(maxlen=self.save_buffer_size)
        self.current_save_buffer = self.save_buffer_a
        self.save_buffer_full = False
        self.save_buffer_event = threading.Event()
        self.save_buffer_lock = threading.Lock()  # 保存缓冲区锁
        
        # 帧处理器设置
        self.frame_processor = FrameProcessor(num_channels=self.num_channels)
        self.process_buffer = bytearray()  # 用于帧处理的缓冲区
        self.process_buffer_lock = threading.Lock()
        self.process_interval = 0.001  # 1ms处理间隔
        
        # 帧保存相关
        self.last_save_time = time.time()
        self.save_interval = 1.0  # 每秒检查一次
        
        # 配置日志
        log_dir = "logs"
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
            
        log_file = os.path.join(log_dir, f"recorder_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
        self.logger.info(f"初始化完成，主缓冲区大小: {self.buffer_size} 字节")

    def set_port(self, port: str, baudrate: int):
        """设置串口参数"""
        self.port = port
        self.baudrate = baudrate

    def test_connection(self) -> bool:
        """测试串口连接"""
        try:
            if self.serial and self.serial.is_open:
                self.serial.close()
                
            self.serial = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                timeout=0.001,  # 减小超时时间
                write_timeout=1,  # 添加写入超时
                bytesize=serial.EIGHTBITS,  # 8位数据位
                parity=serial.PARITY_NONE,  # 无校验
                stopbits=serial.STOPBITS_ONE,  # 1位停止位
                rtscts=True  # 启用硬件流控
            )
            # 设置接收缓冲区大小
            self.serial.set_buffer_size(rx_size=1024*1024)  # 1MB接收缓冲区
            self.logger.info(f"成功连接到串口 {self.port}")
            return True
        except Exception as e:
            self.logger.error(f"连接串口失败: {e}")
            return False

    def disconnect(self) -> bool:
        """断开串口连接"""
        try:
            if self.serial and self.serial.is_open:
                self.serial.close()
                self.logger.info("已断开串口连接")
            return True
        except Exception as e:
            self.logger.error(f"断开连接失败: {e}")
            return False

    def start(self) -> bool:
        """启动数据记录"""
        if not self.serial or not self.serial.is_open:
            if not self.test_connection():
                return False

        self.start_time = time.time()
        self.last_data_time = time.time()
        self.bytes_received = 0
        self.buffer.clear()
        self.save_buffer_a.clear()
        self.save_buffer_b.clear()
        self.current_save_buffer = self.save_buffer_a
        self.save_buffer_full = False
        self.save_buffer_event.clear()
        self.process_buffer = bytearray()

        self.is_running = True
        self.read_thread = threading.Thread(target=self._read_loop)
        self.save_thread = threading.Thread(target=self._save_loop)
        self.process_thread = threading.Thread(target=self._process_loop)
        self.read_thread.daemon = True
        self.save_thread.daemon = True
        self.process_thread.daemon = True
        self.read_thread.start()
        self.save_thread.start()
        self.process_thread.start()
        
        self.logger.info("开始记录数据")
        return True

    def _process_loop(self):
        """帧处理循环"""
        while self.is_running:
            try:
                with self.buffer_lock:
                    if len(self.buffer) > 0:
                        # 复制数据到处理缓冲区
                        with self.process_buffer_lock:
                            self.process_buffer.extend(self.buffer)
                            self.buffer.clear()
                
                # 处理数据
                if len(self.process_buffer) > 0:
                    with self.process_buffer_lock:
                        self.frame_processor.process_buffer(self.process_buffer)
                        self.process_buffer = bytearray()
                
                # 检查是否需要保存帧数据
                current_time = time.time()
                if current_time - self.last_save_time >= self.save_interval:
                    self._check_and_save_frames()
                    self.last_save_time = current_time
                
                # 等待处理间隔
                time.sleep(self.process_interval)
                
            except Exception as e:
                self.logger.error(f"帧处理错误: {e}")
                time.sleep(0.1)

    def _check_and_save_frames(self):
        """检查并保存帧数据"""
        try:
            queue_size = self.frame_processor.frame_queue.qsize()
            if queue_size >= self.max_frames_per_file:
                self._save_frames(self.max_frames_per_file)
        except Exception as e:
            self.logger.error(f"检查并保存帧数据时出错: {e}")
            
    def _save_frames(self, max_frames: int = None):
        """保存指定数量的帧"""
        try:
            stats = self.frame_processor.dump_frames(
                output_dir=self.output_dir,
                format='dat',
                max_frames=max_frames
            )
            self.logger.info(f"已保存 {stats['exported_frames']} 帧数据")
        except Exception as e:
            self.logger.error(f"保存帧数据时出错: {e}")
            
    def _save_remaining_frames(self):
        """保存剩余的帧数据"""
        try:
            if not self.frame_processor.frame_queue.empty():
                stats = self.frame_processor.dump_frames(
                    output_dir=self.output_dir,
                    format='dat'
                )
                self.logger.info(f"已保存剩余 {stats['exported_frames']} 帧数据")
        except Exception as e:
            self.logger.error(f"保存剩余帧数据时出错: {e}")

    def _save_loop(self):
        """数据保存循环"""
        while self.is_running or self.save_buffer_event.is_set():
            try:
                # 等待缓冲区满或程序结束
                if not self.save_buffer_full and self.is_running:
                    time.sleep(0.1)
                    continue

                with self.save_buffer_lock:
                    # 确定要保存的缓冲区
                    buffer_to_save = self.save_buffer_a if self.current_save_buffer == self.save_buffer_b else self.save_buffer_b
                    
                    if len(buffer_to_save) > 0:
                        # 创建保存目录
                        records_dir = "records"
                        if not os.path.exists(records_dir):
                            os.makedirs(records_dir)

                        # 生成文件名
                        filename = f"received_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.bin"
                        filepath = os.path.join(records_dir, filename)

                        # 保存数据
                        with open(filepath, 'wb') as f:
                            data = bytes(buffer_to_save)
                            f.write(data)
                        
                        self.logger.info(f"数据已保存到文件: {filepath}")
                        self.logger.info(f"保存数据大小: {len(data)} 字节")
                        
                        # 清空已保存的缓冲区
                        buffer_to_save.clear()
                        self.logger.info("已清空保存缓冲区")
                    
                    # 重置状态
                    self.save_buffer_full = False
                    self.save_buffer_event.clear()

            except Exception as e:
                self.logger.error(f"保存数据错误: {e}")
                time.sleep(0.1)

    def _read_loop(self):
        """数据读取循环"""
        CHUNK_SIZE = 4096
        last_progress_time = time.time()
        PROGRESS_INTERVAL = 1.0
        self.error_count = 0
        
        while self.is_running:
            try:
                if self.serial.in_waiting:
                    # 读取数据
                    chunk = self.serial.read(min(CHUNK_SIZE, self.serial.in_waiting))
                    chunk_size = len(chunk)
                    
                    if chunk_size > 0:
                        self.bytes_received += chunk_size
                        self.last_data_time = time.time()
                        self.error_count = 0
                        
                        # 更新主缓冲区
                        with self.buffer_lock:
                            # 如果主缓冲区接近满，清空它
                            if len(self.buffer) > self.buffer_size * 0.9:
                                self.logger.warning("主缓冲区接近满，执行清空操作")
                                self.buffer.clear()
                            self.buffer.extend(chunk)
                        
                        # 更新保存缓冲区
                        with self.save_buffer_lock:
                            # 检查当前缓冲区是否已满
                            if len(self.current_save_buffer) >= self.save_buffer_size:
                                # 等待保存线程处理完上一个缓冲区
                                while self.save_buffer_full and self.is_running:
                                    time.sleep(0.01)
                                
                                if not self.is_running:
                                    break
                                
                                # 切换到另一个缓冲区
                                self.current_save_buffer = self.save_buffer_b if self.current_save_buffer == self.save_buffer_a else self.save_buffer_a
                                self.save_buffer_full = True
                                self.save_buffer_event.set()
                                self.logger.info(f"切换到新的保存缓冲区，当前缓冲区大小: {len(self.current_save_buffer)}")
                            
                            self.current_save_buffer.extend(chunk)
                        
                        # 定期打印进度
                        current_time = time.time()
                        if current_time - last_progress_time >= PROGRESS_INTERVAL:
                            stats = self.frame_processor.get_statistics()
                            self.logger.info(f"接收速度: {chunk_size/PROGRESS_INTERVAL/1024:.1f} KB/s")
                            self.logger.info(f"总帧数: {stats['total_frames']}")
                            self.logger.info(f"有效帧数: {stats['valid_frames']}")
                            self.logger.info(f"帧处理速度: {stats['frames_per_second']:.1f} 帧/秒")
                            last_progress_time = current_time
                    
                    # 检查数据超时
                    if time.time() - self.last_data_time > 5.0:  # 5秒超时
                        self.error_count += 1
                        self.logger.warning(f"数据接收超时 ({self.error_count}/{self.max_errors})")
                        if self.error_count >= self.max_errors:
                            self.logger.error("达到最大错误次数，停止记录")
                            self.is_running = False
                            break
                        time.sleep(0.1)
                        
                else:
                    time.sleep(0.001)  # 无数据时短暂休眠
                    
            except Exception as e:
                self.logger.error(f"读取数据错误: {e}")
                self.error_count += 1
                if self.error_count >= self.max_errors:
                    self.logger.error("达到最大错误次数，停止记录")
                    self.is_running = False
                    break
                time.sleep(0.1)

    def stop(self) -> bool:
        """停止数据记录"""
        try:
            self.is_running = False
            self.save_buffer_event.set()  # 触发保存线程处理剩余数据
            
            if self.save_thread:
                self.save_thread.join(timeout=1.0)
            if self.read_thread:
                self.read_thread.join(timeout=1.0)
            if self.process_thread:
                self.process_thread.join(timeout=1.0)
                
            # 保存剩余的帧数据
            self._save_remaining_frames()
            
            # 记录统计信息
            if self.start_time is not None:
                total_time = time.time() - self.start_time
                speed = self.bytes_received / total_time / 1024 if total_time > 0 else 0
                stats = self.frame_processor.get_statistics()
                
                self.logger.info(f"数据记录完成")
                self.logger.info(f"总时间: {total_time:.2f} 秒")
                self.logger.info(f"平均速度: {speed:.1f} KB/s")
                self.logger.info(f"接收字节数: {self.bytes_received}")
                self.logger.info(f"缓冲区使用: {len(self.buffer)}/{self.buffer_size}")
                self.logger.info(f"总帧数: {stats['total_frames']}")
                self.logger.info(f"有效帧数: {stats['valid_frames']}")
                self.logger.info(f"损坏帧数: {stats['corrupt_frames']}")
                self.logger.info(f"无效字节数: {stats['invalid_bytes']}")
                self.logger.info(f"帧处理速度: {stats['frames_per_second']:.1f} 帧/秒")
            else:
                self.logger.info("数据记录未正常启动")
                
            return True
        except Exception as e:
            self.logger.error(f"停止记录失败: {e}")
            return False

    def set_channel_nums(self, num_channels: int) -> bool:
        """
        设置通道数量
        Args:
            num_channels: 新的通道数量
        Returns:
            bool: 设置是否成功
        """
        try:
            if num_channels < 1 or num_channels > 32:
                self.logger.error(f"通道数量必须在1-32之间，当前值: {num_channels}")
                return False
                
            self.num_channels = num_channels
            # 更新帧处理器的通道数
            self.frame_processor.num_channels = num_channels
            # 更新帧查找器的通道数
            self.frame_processor.finder = FrameFinder(num_channels)
            
            self.logger.info(f"通道数量已更新为: {num_channels}")
            return True
            
        except Exception as e:
            self.logger.error(f"设置通道数量失败: {e}")
            return False

if __name__ == "__main__":
    # 使用示例
    recorder = SerialRecorder(port="COM51", baudrate=4000000)
    
    try:
        if recorder.start():
            print("开始记录数据...")
            
            last_stats_time = time.time()
            STATS_INTERVAL = 1.0  # 统计信息显示间隔
            
            while True:
                current_time = time.time()
                
                # 每秒显示一次统计信息
                if current_time - last_stats_time >= STATS_INTERVAL:
                    stats = recorder.frame_processor.get_statistics()
                    print(f"\n当前状态:")
                    print(f"已接收字节: {stats['total_bytes']}")
                    print(f"主缓冲区: {len(recorder.buffer)}/{recorder.buffer_size} 字节")
                    print(f"总帧数: {stats['total_frames']}")
                    print(f"有效帧数: {stats['valid_frames']}")
                    print(f"损坏帧数: {stats['corrupt_frames']}")
                    print(f"无效字节数: {stats['invalid_bytes']}")
                    print(f"帧处理速度: {stats['frames_per_second']:.1f} 帧/秒")
                    
                    last_stats_time = current_time
                
                time.sleep(0.1)  # 降低主循环频率
                
    except KeyboardInterrupt:
        print("\n停止记录...")
    finally:
        recorder.stop()