Feature/new registry (#2612)

* wip

* overhaul boot process

* wip: new registry

* wip

* wip

* wip

* wip

* wip

* wip

* os registry complete

* ui fixes

* fixes

* fixes

* more fixes

* fix merkle archive

core/startos/src/util/io.rs

@@ -1,11 +1,15 @@
+use std::collections::VecDeque;
 use std::future::Future;
 use std::io::Cursor;
 use std::os::unix::prelude::MetadataExt;
 use std::path::{Path, PathBuf};
+use std::pin::Pin;
 use std::sync::atomic::AtomicU64;
-use std::task::Poll;
+use std::sync::Arc;
+use std::task::{Poll, Waker};
 use std::time::Duration;
 
+use bytes::{Buf, BytesMut};
 use futures::future::{BoxFuture, Fuse};
 use futures::{AsyncSeek, FutureExt, TryStreamExt};
 use helpers::NonDetachingJoinHandle;
@@ -15,6 +19,7 @@ use tokio::io::{
     duplex, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream, ReadBuf, WriteHalf,
 };
 use tokio::net::TcpStream;
+use tokio::sync::Notify;
 use tokio::time::{Instant, Sleep};
 
 use crate::prelude::*;
@@ -312,6 +317,7 @@ impl AsyncRead for BufferedWriteReader {
 
 pub trait CursorExt {
     fn pure_read(&mut self, buf: &mut ReadBuf<'_>);
+    fn remaining_slice(&self) -> &[u8];
 }
 
 impl<T: AsRef<[u8]>> CursorExt for Cursor<T> {
@@ -324,6 +330,10 @@ impl<T: AsRef<[u8]>> CursorExt for Cursor<T> {
         buf.put_slice(&self.get_ref().as_ref()[self.position() as usize..end]);
         self.set_position(end as u64);
     }
+    fn remaining_slice(&self) -> &[u8] {
+        let len = self.position().min(self.get_ref().as_ref().len() as u64);
+        &self.get_ref().as_ref()[(len as usize)..]
+    }
 }
 
 #[pin_project::pin_project]
@@ -744,3 +754,256 @@ pub async fn rename(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> Result<(),
         .await
         .with_ctx(|_| (ErrorKind::Filesystem, lazy_format!("mv {src:?} -> {dst:?}")))
 }
+
+fn poll_flush_prefix<W: AsyncWrite>(
+    mut writer: Pin<&mut W>,
+    cx: &mut std::task::Context<'_>,
+    prefix: &mut VecDeque<Cursor<Vec<u8>>>,
+    flush_writer: bool,
+) -> Poll<Result<(), std::io::Error>> {
+    while let Some(mut cur) = prefix.pop_front() {
+        let buf = cur.remaining_slice();
+        if !buf.is_empty() {
+            match writer.as_mut().poll_write(cx, buf)? {
+                Poll::Ready(n) if n == buf.len() => (),
+                Poll::Ready(n) => {
+                    cur.advance(n);
+                    prefix.push_front(cur);
+                }
+                Poll::Pending => {
+                    prefix.push_front(cur);
+                    return Poll::Pending;
+                }
+            }
+        }
+    }
+    if flush_writer {
+        writer.poll_flush(cx)
+    } else {
+        Poll::Ready(Ok(()))
+    }
+}
+
+fn poll_write_prefix_buf<W: AsyncWrite>(
+    mut writer: Pin<&mut W>,
+    cx: &mut std::task::Context<'_>,
+    prefix: &mut VecDeque<Cursor<Vec<u8>>>,
+    buf: &[u8],
+) -> Poll<Result<usize, std::io::Error>> {
+    futures::ready!(poll_flush_prefix(writer.as_mut(), cx, prefix, false)?);
+    writer.poll_write(cx, buf)
+}
+
+#[pin_project::pin_project]
+pub struct TeeWriter<W1, W2> {
+    capacity: usize,
+    buffer1: VecDeque<Cursor<Vec<u8>>>,
+    buffer2: VecDeque<Cursor<Vec<u8>>>,
+    #[pin]
+    writer1: W1,
+    #[pin]
+    writer2: W2,
+}
+impl<W1: AsyncWrite, W2: AsyncWrite> TeeWriter<W1, W2> {
+    pub fn new(writer1: W1, writer2: W2, capacity: usize) -> Self {
+        Self {
+            capacity,
+            buffer1: VecDeque::new(),
+            buffer2: VecDeque::new(),
+            writer1,
+            writer2,
+        }
+    }
+}
+
+impl<W1: AsyncWrite + Unpin, W2: AsyncWrite + Unpin> TeeWriter<W1, W2> {
+    pub async fn into_inner(mut self) -> Result<(W1, W2), Error> {
+        self.flush().await?;
+
+        Ok((self.writer1, self.writer2))
+    }
+}
+impl<W1: AsyncWrite, W2: AsyncWrite> AsyncWrite for TeeWriter<W1, W2> {
+    fn poll_write(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        mut buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let mut this = self.project();
+        let buffer_size = this
+            .buffer1
+            .iter()
+            .chain(this.buffer2.iter())
+            .map(|b| b.get_ref().len())
+            .sum::<usize>();
+        if buffer_size < *this.capacity {
+            let to_write = std::cmp::min(*this.capacity - buffer_size, buf.len());
+            buf = &buf[0..to_write];
+        } else {
+            match (
+                poll_flush_prefix(this.writer1.as_mut(), cx, &mut this.buffer1, false)?,
+                poll_flush_prefix(this.writer2.as_mut(), cx, &mut this.buffer2, false)?,
+            ) {
+                (Poll::Ready(()), Poll::Ready(())) => (),
+                _ => return Poll::Pending,
+            }
+        }
+        let (w1, w2) = match (
+            poll_write_prefix_buf(this.writer1.as_mut(), cx, &mut this.buffer1, buf)?,
+            poll_write_prefix_buf(this.writer2.as_mut(), cx, &mut this.buffer2, buf)?,
+        ) {
+            (Poll::Pending, Poll::Pending) => return Poll::Pending,
+            (Poll::Ready(n), Poll::Pending) => (n, 0),
+            (Poll::Pending, Poll::Ready(n)) => (0, n),
+            (Poll::Ready(n1), Poll::Ready(n2)) => (n1, n2),
+        };
+        if w1 > w2 {
+            this.buffer2.push_back(Cursor::new(buf[w2..w1].to_vec()));
+        } else if w1 < w2 {
+            this.buffer1.push_back(Cursor::new(buf[w1..w2].to_vec()));
+        }
+        Poll::Ready(Ok(std::cmp::max(w1, w2)))
+    }
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let mut this = self.project();
+        match (
+            poll_flush_prefix(this.writer1, cx, &mut this.buffer1, true)?,
+            poll_flush_prefix(this.writer2, cx, &mut this.buffer2, true)?,
+        ) {
+            (Poll::Ready(()), Poll::Ready(())) => Poll::Ready(Ok(())),
+            _ => Poll::Pending,
+        }
+    }
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        self.poll_flush(cx)
+    }
+}
+
+#[pin_project::pin_project]
+pub struct ParallelBlake3Writer {
+    #[pin]
+    hasher: NonDetachingJoinHandle<blake3::Hash>,
+    buffer: Arc<(std::sync::Mutex<(BytesMut, Vec<Waker>, bool)>, Notify)>,
+    capacity: usize,
+}
+impl ParallelBlake3Writer {
+    /// memory usage can be as much as 2x capacity
+    pub fn new(capacity: usize) -> Self {
+        let buffer = Arc::new((
+            std::sync::Mutex::new((BytesMut::new(), Vec::<Waker>::new(), false)),
+            Notify::new(),
+        ));
+        let hasher_buffer = buffer.clone();
+        Self {
+            hasher: tokio::spawn(async move {
+                let mut hasher = blake3::Hasher::new();
+                let mut to_hash = BytesMut::new();
+                let mut notified;
+                while {
+                    let mut guard = hasher_buffer.0.lock().unwrap();
+                    let (buffer, wakers, shutdown) = &mut *guard;
+                    std::mem::swap(buffer, &mut to_hash);
+                    let wakers = std::mem::take(wakers);
+                    let shutdown = *shutdown;
+                    notified = hasher_buffer.1.notified();
+                    drop(guard);
+                    if to_hash.len() > 128 * 1024
+                    /* 128 KiB */
+                    {
+                        hasher.update_rayon(&to_hash);
+                    } else {
+                        hasher.update(&to_hash);
+                    }
+                    to_hash.truncate(0);
+                    for waker in wakers {
+                        waker.wake();
+                    }
+                    !shutdown && to_hash.len() == 0
+                } {
+                    notified.await;
+                }
+                hasher.finalize()
+            })
+            .into(),
+            buffer,
+            capacity,
+        }
+    }
+
+    pub async fn finalize(mut self) -> Result<blake3::Hash, Error> {
+        self.shutdown().await?;
+        self.hasher.await.with_kind(ErrorKind::Unknown)
+    }
+}
+impl AsyncWrite for ParallelBlake3Writer {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let this = self.project();
+        let mut guard = this.buffer.0.lock().map_err(|_| {
+            std::io::Error::new(std::io::ErrorKind::Other, eyre!("hashing thread panicked"))
+        })?;
+        let (buffer, wakers, shutdown) = &mut *guard;
+        if !*shutdown {
+            if buffer.len() < *this.capacity {
+                let to_write = std::cmp::min(*this.capacity - buffer.len(), buf.len());
+                buffer.extend_from_slice(&buf[0..to_write]);
+                if buffer.len() >= *this.capacity / 2 {
+                    this.buffer.1.notify_waiters();
+                }
+                Poll::Ready(Ok(to_write))
+            } else {
+                wakers.push(cx.waker().clone());
+                Poll::Pending
+            }
+        } else {
+            Poll::Ready(Err(std::io::Error::new(
+                std::io::ErrorKind::BrokenPipe,
+                eyre!("write after shutdown"),
+            )))
+        }
+    }
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let this = self.project();
+        let mut guard = this.buffer.0.lock().map_err(|_| {
+            std::io::Error::new(std::io::ErrorKind::Other, eyre!("hashing thread panicked"))
+        })?;
+        let (buffer, wakers, _) = &mut *guard;
+        if buffer.is_empty() {
+            Poll::Ready(Ok(()))
+        } else {
+            wakers.push(cx.waker().clone());
+            this.buffer.1.notify_waiters();
+            Poll::Pending
+        }
+    }
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        futures::ready!(self.as_mut().poll_flush(cx)?);
+        let this = self.project();
+        let mut guard = this.buffer.0.lock().map_err(|_| {
+            std::io::Error::new(std::io::ErrorKind::Other, eyre!("hashing thread panicked"))
+        })?;
+        let (buffer, wakers, shutdown) = &mut *guard;
+        if *shutdown && buffer.len() == 0 {
+            return Poll::Ready(Ok(()));
+        }
+        wakers.push(cx.waker().clone());
+        *shutdown = true;
+        this.buffer.1.notify_waiters();
+        Poll::Pending
+    }
+}