ruda-runtime 0.1.31

Ruda portable host runtime.
use alloc::vec::Vec;
use ruda_core::backtrace::BackTrace;

use super::{ManagedMemoryBinding, ManagedMemoryHandle, PageMapping, Slice, calculate_padding};
use crate::runtime::{
    memory_management::{MemoryLocation, MemoryUsage},
    server::IoError,
    storage::{ComputeStorage, StorageUtilization},
};

pub(crate) struct DedicatedAllocations {
    slices: Vec<Slice>,
    alignment: u64,
    max_alloc_size: u64,
}

impl DedicatedAllocations {
    pub(crate) fn new(max_alloc_size: u64, alignment: u64) -> Self {
        Self { slices: Vec::new(), alignment, max_alloc_size }
    }

    pub(crate) fn alloc<S: ComputeStorage>(
        &mut self,
        storage: &mut S,
        size: u64,
    ) -> Result<ManagedMemoryHandle, IoError> {
        let padding = calculate_padding(size, self.alignment);
        let allocation_size = size.checked_add(padding)
            .filter(|size| *size <= self.max_alloc_size)
            .ok_or_else(|| IoError::BufferTooBig { size, backtrace: BackTrace::capture() })?;
        let index = u32::try_from(self.slices.len()).map_err(|_| IoError::Unknown {
            description: "Dedicated allocation table is full".into(),
            backtrace: BackTrace::capture(),
        })?;
        let mapping = PageMapping::current();
        let storage_handle = mapping.allocate(storage, allocation_size)?;
        let mut slice = Slice::new(storage_handle, padding);
        slice.mapped = mapping.is_mapped();
        slice.storage.utilization = StorageUtilization { offset: 0, size };
        slice.descriptor().update_location(MemoryLocation::dedicated(index));
        let handle = slice.handle.clone();
        self.slices.push(slice);
        Ok(handle)
    }

    pub(crate) fn find(&self, binding: &ManagedMemoryBinding) -> Result<&Slice, IoError> {
        self.slices.get(binding.descriptor().slice())
            .filter(|slice| slice.descriptor() == binding.descriptor())
            .ok_or_else(|| IoError::NotFound {
                reason: "Dedicated allocation no longer exists".into(),
                backtrace: BackTrace::capture(),
            })
    }

    pub(crate) fn bind(
        &mut self,
        reserved: ManagedMemoryHandle,
        assigned: ManagedMemoryHandle,
        cursor: u64,
    ) -> Result<(), IoError> {
        let location = reserved.descriptor().location();
        let slice = self.slices.get_mut(location.slice as usize)
            .filter(|slice| slice.descriptor() == reserved.descriptor())
            .ok_or_else(|| IoError::NotFound {
                reason: "Dedicated reservation no longer exists".into(),
                backtrace: BackTrace::capture(),
            })?;
        assigned.descriptor().update_location(location);
        slice.cursor = cursor;
        slice.handle = assigned;
        Ok(())
    }

    pub(crate) fn materialize<S: ComputeStorage>(
        &mut self, storage: &mut S, binding: &ManagedMemoryBinding,
    ) -> Result<(), IoError> {
        self.find(binding)?;
        let slice = &mut self.slices[binding.descriptor().slice()];
        let size = slice.effective_size();
        slice.materialize(storage, size)
    }

    pub(crate) fn reclaim<S: ComputeStorage>(&mut self, storage: &mut S) {
        let mut next = 0u32;
        let mut reclaimed = false;
        self.slices.retain(|slice| {
            if slice.is_free() {
                if slice.mapped { storage.dealloc(slice.storage.id); }
                reclaimed = true;
                false
            } else {
                slice.descriptor().update_slice(next);
                next += 1;
                true
            }
        });
        if reclaimed { storage.flush(); }
    }

    pub(crate) fn memory_usage(&self) -> MemoryUsage {
        let mut usage = MemoryUsage {
            number_allocs: 0, bytes_in_use: 0, bytes_padding: 0, bytes_reserved: 0,
        };
        for slice in &self.slices {
            usage.bytes_reserved += slice.effective_size();
            if !slice.is_free() {
                usage.number_allocs += 1;
                usage.bytes_in_use += slice.storage.size();
                usage.bytes_padding += slice.padding;
            }
        }
        usage
    }
}