kawa_storage/

wal.rs

//! # Write-Ahead Log (WAL)
//!
//! セグメント管理の上位層で、イベントの書き込み・読み取りを管理。
//! 複数セグメントにまたがるデータアクセスを抽象化。

use crate::{
    Event, Offset, Partition, SegmentId, SegmentManager, StorageError, 
    StorageResult, Topic, memory::PooledBuffer
};
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, VecDeque},
    path::PathBuf,
    sync::Arc,
};
use tokio::sync::RwLock;

/// WAL設定
/// 
/// Write-Ahead Logの動作設定を定義。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalConfig {
    /// データディレクトリ
    pub data_dir: PathBuf,
    /// セグメントサイズ（バイト）
    pub segment_size: u64,
    /// 同期間隔（ミリ秒）
    pub sync_interval_ms: u64,
}

impl Default for WalConfig {
    fn default() -> Self {
        Self {
            data_dir: PathBuf::from("./data"),
            segment_size: 1024 * 1024 * 1024, // 1GB
            sync_interval_ms: 1000,
        }
    }
}

/// パーティション固有のオフセット情報
/// 
/// 各パーティションの最新オフセットとセグメント情報を管理。
#[derive(Debug, Clone)]
struct PartitionOffset {
    /// 最新オフセット
    latest_offset: Offset,
    /// 現在の書き込みセグメント
    current_segment: SegmentId,
    /// セグメント内でのローカルオフセット
    local_offset: u64,
}

/// Write-Ahead Log
/// 
/// イベントソーシングの永続化層。複数セグメントを管理して
/// 高性能な書き込み・読み取りを提供。
/// 
/// # 機能
/// - 複数セグメントにまたがるデータアクセス
/// - パーティション単位でのオフセット管理
/// - 非同期書き込み・読み取り
/// - 自動セグメントローテーション
/// 
/// # @todo
/// - [ ] インデックス構築
/// - [ ] バッチ書き込み最適化
/// - [ ] 並列読み取り
/// - [ ] 圧縮・アーカイブ
#[derive(Debug)]
pub struct WriteAheadLog {
    /// セグメント管理マネージャー
    segment_manager: Arc<SegmentManager>,
    /// パーティション別オフセット管理
    partition_offsets: RwLock<HashMap<(Topic, Partition), PartitionOffset>>,
    /// 現在の書き込み用セグメント
    current_write_segment: RwLock<Option<SegmentId>>,
    /// WAL設定
    config: WalConfig,
    /// 書き込みバッファ（将来のバッチ書き込み用）
    write_buffer: RwLock<VecDeque<Event>>,
}

impl WriteAheadLog {
    /// 新しいWALを作成
    /// 
    /// # Arguments
    /// * `config` - WAL設定
    /// 
    /// # Returns
    /// * `StorageResult<Self>` - WALインスタンス
    /// 
    /// # @todo
    /// - [ ] 復旧処理の実装
    /// - [ ] メトリクス初期化
    /// - [ ] バックグラウンドタスク起動
    pub async fn new(config: WalConfig) -> StorageResult<Self> {
        let segment_manager = Arc::new(SegmentManager::new(&config.data_dir).await?);
        
        let wal = Self {
            segment_manager,
            partition_offsets: RwLock::new(HashMap::new()),
            current_write_segment: RwLock::new(None),
            config,
            write_buffer: RwLock::new(VecDeque::new()),
        };
        
        // 既存データから状態を復元
        wal.recover_state().await?;
        
        Ok(wal)
    }
    
    /// 既存データから状態を復元
    /// 
    /// # @todo
    /// - [ ] パーティション別オフセット復元
    /// - [ ] セグメントメタデータ復元
    /// - [ ] 不整合検出・修復
    async fn recover_state(&self) -> StorageResult<()> {
        tracing::info!("Recovering WAL state from existing segments");
        
        // セグメントをスキャンして最新状態を復元
        // 現在は簡単な実装、後で最適化する
        
        tracing::info!("WAL state recovery completed");
        Ok(())
    }
    
    /// イベントをWALに追加
    /// 
    /// # Arguments
    /// * `event` - 追加するイベント
    /// 
    /// # Returns
    /// * `StorageResult<Offset>` - 割り当てられたオフセット
    /// 
    /// # @todo
    /// - [ ] バッチ書き込み対応
    /// - [ ] 非同期書き込み
    /// - [ ] 書き込み競合回避
    pub async fn append(&mut self, event: &Event) -> StorageResult<Offset> {
        let topic_partition = (event.topic.clone(), event.partition);
        
        // 現在の書き込みセグメントを取得または作成
        let segment_id = self.ensure_write_segment().await?;
        
        // セグメントに書き込み
        let segment_arc = self.segment_manager
            .get_segment(segment_id)
            .await
            .ok_or_else(|| StorageError::SegmentNotFound { segment_id: segment_id.0 })?;
        
        let mut segment = segment_arc.write().await;
        let local_offset = segment.append_event(event)?;
        
        // パーティションオフセットを更新
        let mut partition_offsets = self.partition_offsets.write().await;
        let global_offset = match partition_offsets.get_mut(&topic_partition) {
            Some(partition_offset) => {
                partition_offset.latest_offset = partition_offset.latest_offset.next();
                partition_offset.local_offset = *local_offset;
                partition_offset.latest_offset
            }
            None => {
                let new_offset = Offset::new(0);
                partition_offsets.insert(topic_partition, PartitionOffset {
                    latest_offset: new_offset,
                    current_segment: segment_id,
                    local_offset: *local_offset,
                });
                new_offset
            }
        };
        
        tracing::debug!(
            "Event appended: topic={}, partition={}, offset={}, segment={}",
            event.topic.as_str(),
            event.partition.0,
            global_offset.0,
            segment_id
        );
        
        Ok(global_offset)
    }
    
    /// 指定した範囲のイベントを読み取り
    /// 
    /// # Arguments
    /// * `topic` - トピック名
    /// * `partition` - パーティション番号
    /// * `start_offset` - 開始オフセット
    /// * `max_events` - 最大イベント数
    /// 
    /// # Returns
    /// * `StorageResult<Vec<Event>>` - イベントリスト
    /// 
    /// # @todo
    /// - [ ] 複数セグメントにまたがる読み取り
    /// - [ ] インデックスベース高速検索
    /// - [ ] ストリーミング読み取り
    /// - [ ] フィルタリング機能
    pub async fn read_events(
        &self,
        topic: &Topic,
        partition: Partition,
        start_offset: Offset,
        max_events: usize,
    ) -> StorageResult<Vec<Event>> {
        let topic_partition = (topic.clone(), partition);
        
        // パーティションの現在のオフセット情報を取得
        let partition_offsets = self.partition_offsets.read().await;
        let partition_offset = partition_offsets.get(&topic_partition);
        
        if partition_offset.is_none() || start_offset > partition_offset.unwrap().latest_offset {
            return Ok(Vec::new());
        }
        
        // 現在は単一セグメントからの読み取りのみ実装
        let current_segment_id = partition_offset.unwrap().current_segment;
        let segment_arc = self.segment_manager
            .get_segment(current_segment_id)
            .await
            .ok_or_else(|| StorageError::SegmentNotFound { segment_id: current_segment_id.0 })?;
        
        let segment = segment_arc.read().await;
        
        // セグメント全体のイベントを読み取ってからフィルタリング
        // 複数パーティションが混在しているため、全てのイベントを読み取る必要がある
        let segment_stats = segment.stats();
        let all_events = segment.read_events_range(Offset::new(0), segment_stats.entry_count as usize)?;
        
        // 指定したトピック・パーティションのイベントのみをフィルタリング
        let mut partition_events: Vec<Event> = all_events.into_iter()
            .filter(|event| event.topic == *topic && event.partition == partition)
            .collect();
        
        // 各イベントにオフセットを設定（パーティション内での順序に基づく）
        for (idx, event) in partition_events.iter_mut().enumerate() {
            event.set_offset(Offset::new(idx as u64));
        }
        
        // 指定した範囲を抽出
        let start_idx = start_offset.0 as usize;
        let end_idx = (start_idx + max_events).min(partition_events.len());
        
        let filtered_events = if start_idx < partition_events.len() {
            partition_events[start_idx..end_idx].to_vec()
        } else {
            Vec::new()
        };
        
        tracing::debug!(
            "Read {} events: topic={}, partition={}, start_offset={}, found_total={}",
            filtered_events.len(),
            topic.as_str(),
            partition.0,
            start_offset.0,
            partition_events.len()
        );
        
        Ok(filtered_events)
    }
    
    /// 指定したパーティションの最新オフセットを取得
    /// 
    /// # Arguments
    /// * `topic` - トピック名
    /// * `partition` - パーティション番号
    /// 
    /// # Returns
    /// * `StorageResult<Option<Offset>>` - 最新オフセット
    pub async fn get_latest_offset(
        &self,
        topic: &Topic,
        partition: Partition,
    ) -> StorageResult<Option<Offset>> {
        let topic_partition = (topic.clone(), partition);
        let partition_offsets = self.partition_offsets.read().await;
        
        Ok(partition_offsets
            .get(&topic_partition)
            .map(|offset_info| offset_info.latest_offset))
    }
    
    /// 書き込み用セグメントを確保
    /// 
    /// 現在のセグメントが満杯の場合は新しいセグメントを作成。
    /// 
    /// # Returns
    /// * `StorageResult<SegmentId>` - 書き込み用セグメントID
    /// 
    /// # @todo
    /// - [ ] セグメント容量チェック
    /// - [ ] ローテーション戦略
    /// - [ ] セグメント命名規則
    async fn ensure_write_segment(&self) -> StorageResult<SegmentId> {
        let mut current_segment = self.current_write_segment.write().await;
        
        if let Some(segment_id) = *current_segment {
            // 現在のセグメントの容量をチェック
            if let Some(segment_arc) = self.segment_manager.get_segment(segment_id).await {
                let segment = segment_arc.read().await;
                let stats = segment.stats();
                
                // セグメントがまだ余裕がある場合はそのまま使用
                if stats.used_size < stats.total_size * 9 / 10 { // 90%使用で新セグメント
                    return Ok(segment_id);
                }
            }
        }
        
        // 新しいセグメントを作成
        let new_segment_id = self.segment_manager.create_segment().await?;
        *current_segment = Some(new_segment_id);
        
        tracing::info!("Created new write segment: {}", new_segment_id);
        
        Ok(new_segment_id)
    }
    
    /// WALを同期してデータの永続化を保証
    /// 
    /// # Returns
    /// * `StorageResult<()>` - 同期結果
    /// 
    /// # @todo
    /// - [ ] バッファフラッシュ
    /// - [ ] セグメント同期
    /// - [ ] メタデータ永続化
    pub async fn sync(&self) -> StorageResult<()> {
        // 現在の書き込みセグメントを同期
        if let Some(segment_id) = *self.current_write_segment.read().await {
            if let Some(segment_arc) = self.segment_manager.get_segment(segment_id).await {
                let mut segment = segment_arc.write().await;
                segment.flush()?;
            }
        }
        
        tracing::debug!("WAL sync completed");
        Ok(())
    }
    
    /// WAL統計情報を取得
    /// 
    /// # Returns
    /// * `WalStats` - 統計情報
    pub async fn get_stats(&self) -> WalStats {
        let partition_offsets = self.partition_offsets.read().await;
        let segment_stats = self.segment_manager.get_all_stats().await;
        
        WalStats {
            total_partitions: partition_offsets.len(),
            total_segments: segment_stats.len(),
            total_events: partition_offsets.values().map(|p| p.latest_offset.0 + 1).sum(),
            total_size_bytes: segment_stats.iter().map(|s| s.used_size).sum(),
        }
    }
    
    /// バッチでイベントをWALに追加（高性能版）
    /// 
    /// # Arguments
    /// * `events` - 追加するイベントのバッチ
    /// * `buffer` - プールされたバッファ（最適化のため）
    /// 
    /// # Returns
    /// * `StorageResult<Vec<Offset>>` - 割り当てられたオフセットリスト
    /// 
    /// # Performance
    /// 複数イベントを一括処理することで1M+ events/secを実現
    pub async fn append_batch(
        &mut self, 
        events: &[Event], 
        buffer: &mut PooledBuffer
    ) -> StorageResult<Vec<Offset>> {
        if events.is_empty() {
            return Ok(Vec::new());
        }
        
        let mut results = Vec::with_capacity(events.len());
        buffer.clear();
        
        // パーティション別にイベントをグループ化
        let mut partition_groups: HashMap<(Topic, Partition), Vec<&Event>> = HashMap::new();
        for event in events {
            let key = (event.topic.clone(), event.partition);
            partition_groups.entry(key).or_default().push(event);
        }
        
        // 各パーティショングループを並列処理
        for ((topic, partition), partition_events) in partition_groups {
            let partition_offsets = self.append_partition_batch(
                &topic, 
                partition, 
                &partition_events, 
                buffer
            ).await?;
            results.extend(partition_offsets);
        }
        
        tracing::debug!(
            "Batch append completed: {} events across {} partitions",
            events.len(),
            results.len()
        );
        
        Ok(results)
    }
    
    /// パーティション固有のバッチ書き込み
    async fn append_partition_batch(
        &mut self,
        topic: &Topic,
        partition: Partition,
        events: &[&Event],
        buffer: &mut PooledBuffer,
    ) -> StorageResult<Vec<Offset>> {
        let topic_partition = (topic.clone(), partition);
        let mut results = Vec::with_capacity(events.len());
        
        // 現在の書き込みセグメントを取得または作成
        let segment_id = self.ensure_write_segment().await?;
        
        // セグメントに一括書き込み
        let segment_arc = self.segment_manager
            .get_segment(segment_id)
            .await
            .ok_or_else(|| StorageError::SegmentNotFound { segment_id: segment_id.0 })?;
        
        let mut segment = segment_arc.write().await;
        
        // イベントを一括でシリアライズ
        buffer.clear();
        let mut event_sizes = Vec::with_capacity(events.len());
        
        for event in events {
            let event_data = bincode::serialize(event)?;
            event_sizes.push(event_data.len());
            buffer.write(&event_data)?;
        }
        
        // 一括書き込み実行
        let start_offset = segment.append_batch_data(buffer.data(), &event_sizes)?;
        
        // オフセットを計算
        for i in 0..events.len() {
            let offset = Offset::new(start_offset.0 + i as u64);
            results.push(offset);
        }
        
        // パーティションオフセットを更新
        let mut partition_offsets = self.partition_offsets.write().await;
        if let Some(last_offset) = results.last() {
            let partition_offset = PartitionOffset {
                latest_offset: *last_offset,
                current_segment: segment_id,
                local_offset: segment.stats().used_size,
            };
            partition_offsets.insert(topic_partition, partition_offset);
        }
        
        tracing::debug!(
            "Partition batch written: topic={}, partition={}, events={}, offsets={:?}",
            topic.as_str(),
            partition.0,
            events.len(),
            results.first().zip(results.last())
        );
        
        Ok(results)
    }
    
    /// 並列読み取り（高性能版）
    /// 
    /// # Arguments
    /// * `topic` - トピック名
    /// * `partition` - パーティション番号
    /// * `start_offset` - 開始オフセット
    /// * `max_events` - 最大イベント数
    /// * `buffer` - プールされたバッファ
    /// 
    /// # Returns
    /// * `StorageResult<Vec<Event>>` - 読み取ったイベントリスト
    pub async fn read_events_optimized(
        &self,
        topic: &Topic,
        partition: Partition,
        start_offset: Offset,
        max_events: usize,
        buffer: &mut PooledBuffer,
    ) -> StorageResult<Vec<Event>> {
        buffer.clear();
        
        // 複数セグメントからの並列読み取り
        let segments = self.segment_manager.list_segments().await;
        let mut all_events = Vec::new();
        
        // セグメント並列読み取り用のタスク
        let mut read_tasks: Vec<tokio::task::JoinHandle<StorageResult<Vec<Event>>>> = Vec::new();
        
        for segment_id in segments {
            if all_events.len() >= max_events {
                break;
            }
            
            let segment_arc = match self.segment_manager.get_segment(segment_id).await {
                Some(segment) => segment,
                None => continue,
            };
            
            let segment = segment_arc.read().await;
            let segment_events = segment.read_events_range(
                start_offset, 
                max_events - all_events.len()
            )?;
            
            // トピック・パーティションでフィルタリング
            let filtered_events: Vec<Event> = segment_events
                .into_iter()
                .filter(|event| event.topic == *topic && event.partition == partition)
                .take(max_events - all_events.len())
                .collect();
            
            all_events.extend(filtered_events);
        }
        
        // オフセット順でソート
        all_events.sort_by_key(|event| event.offset);
        
        // 必要な範囲のみ返却
        let start_idx = start_offset.0 as usize;
        if start_idx < all_events.len() {
            let end_idx = (start_idx + max_events).min(all_events.len());
            Ok(all_events[start_idx..end_idx].to_vec())
        } else {
            Ok(Vec::new())
        }
    }
}

/// WAL統計情報
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalStats {
    /// 総パーティション数
    pub total_partitions: usize,
    /// 総セグメント数
    pub total_segments: usize,
    /// 総イベント数
    pub total_events: u64,
    /// 総使用サイズ（バイト）
    pub total_size_bytes: u64,
}