1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

//! A naive LRU cache layer for `BlockDevice`
use super::*;
use alloc::{vec, vec::Vec};
use spin::{Mutex, MutexGuard};

pub struct BlockCache<T: BlockDevice> {
    device: T,
    bufs: Vec<Mutex<Buf>>,
    lru: Mutex<LRU>,
}

struct Buf {
    status: BufStatus,
    data: Vec<u8>,
}

enum BufStatus {
    /// buffer is unused
    Unused,
    /// buffer has been read from disk
    Valid(BlockId),
    /// buffer needs to be written to disk
    Dirty(BlockId),
}

impl<T: BlockDevice> BlockCache<T> {
    pub fn new(device: T, capacity: usize) -> Self {
        let mut bufs = Vec::new();
        bufs.resize_with(capacity, || {
            Mutex::new(Buf {
                status: BufStatus::Unused,
                data: vec![0; 1 << T::BLOCK_SIZE_LOG2 as usize],
            })
        });
        let lru = Mutex::new(LRU::new(capacity));
        BlockCache { device, bufs, lru }
    }

    /// Get a buffer for `block_id` with any status
    fn get_buf(&self, block_id: BlockId) -> MutexGuard<Buf> {
        let (i, buf) = self._get_buf(block_id);
        self.lru.lock().visit(i);
        buf
    }

    fn _get_buf(&self, block_id: BlockId) -> (usize, MutexGuard<Buf>) {
        for (i, buf) in self.bufs.iter().enumerate() {
            if let Some(lock) = buf.try_lock() {
                match lock.status {
                    BufStatus::Valid(id) if id == block_id => return (i, lock),
                    BufStatus::Dirty(id) if id == block_id => return (i, lock),
                    _ => {}
                }
            }
        }
        self.get_unused()
    }

    /// Get an unused buffer
    fn get_unused(&self) -> (usize, MutexGuard<Buf>) {
        for (i, buf) in self.bufs.iter().enumerate() {
            if let Some(lock) = buf.try_lock() {
                if let BufStatus::Unused = lock.status {
                    return (i, lock);
                }
            }
        }
        let victim_id = self.lru.lock().victim();
        let mut victim = self.bufs[victim_id].lock();
        self.write_back(&mut victim).expect("failed to write back");
        victim.status = BufStatus::Unused;
        (victim_id, victim)
    }

    /// Write back data if buffer is dirty
    fn write_back(&self, buf: &mut Buf) -> Result<()> {
        if let BufStatus::Dirty(block_id) = buf.status {
            self.device.write_at(block_id, &buf.data)?;
            buf.status = BufStatus::Valid(block_id);
        }
        Ok(())
    }
}

impl<T: BlockDevice> Drop for BlockCache<T> {
    fn drop(&mut self) {
        BlockDevice::sync(self).expect("failed to sync");
    }
}

impl<T: BlockDevice> BlockDevice for BlockCache<T> {
    const BLOCK_SIZE_LOG2: u8 = T::BLOCK_SIZE_LOG2;

    fn read_at(&self, block_id: BlockId, buffer: &mut [u8]) -> Result<()> {
        let mut buf = self.get_buf(block_id);
        match buf.status {
            BufStatus::Unused => {
                // read from device
                self.device.read_at(block_id, &mut buf.data)?;
                buf.status = BufStatus::Valid(block_id);
            }
            _ => {}
        }
        let len = 1 << Self::BLOCK_SIZE_LOG2 as usize;
        buffer[..len].copy_from_slice(&buf.data);
        Ok(())
    }

    fn write_at(&self, block_id: BlockId, buffer: &[u8]) -> Result<()> {
        let mut buf = self.get_buf(block_id);
        buf.status = BufStatus::Dirty(block_id);
        let len = 1 << Self::BLOCK_SIZE_LOG2 as usize;
        buf.data.copy_from_slice(&buffer[..len]);
        Ok(())
    }

    fn sync(&self) -> Result<()> {
        for buf in self.bufs.iter() {
            self.write_back(&mut buf.lock())?;
        }
        self.device.sync()?;
        Ok(())
    }
}

/// Doubly circular linked list LRU manager
struct LRU {
    prev: Vec<usize>,
    next: Vec<usize>,
}

impl LRU {
    fn new(size: usize) -> Self {
        LRU {
            prev: (size - 1..size).chain(0..size - 1).collect(),
            next: (1..size).chain(0..1).collect(),
        }
    }
    /// Visit element `id`, move it to head.
    fn visit(&mut self, id: usize) {
        if id == 0 || id >= self.prev.len() {
            return;
        }
        self._list_remove(id);
        self._list_insert_head(id);
    }
    /// Get a victim at tail.
    fn victim(&self) -> usize {
        self.prev[0]
    }
    fn _list_remove(&mut self, id: usize) {
        let prev = self.prev[id];
        let next = self.next[id];
        self.prev[next] = prev;
        self.next[prev] = next;
    }
    fn _list_insert_head(&mut self, id: usize) {
        let head = self.next[0];
        self.prev[id] = 0;
        self.next[id] = head;
        self.next[0] = id;
        self.prev[head] = id;
    }
}