fix flush

Author: onursatici

vortex-file/src/segments/ordered.rs

@@ -57,7 +57,7 @@ impl OrderedBuffers {
     pub fn finish_region(&mut self, region: &Region) {
         self.active_regions.remove(&region);
         if let Ok(first) = self.first_region() {
-            if let Some(waker) = self.wakers.get(&first) {
+            if let Some(waker) = self.wakers.remove(&first) {
                 waker.wake_by_ref();
             }
         }
@@ -93,13 +93,15 @@ impl OrderedBuffers {
             .ok_or_else(|| vortex_err!("no active regions"))
     }
 
-    pub fn completed_buffers(&mut self) -> VortexResult<BTreeMap<usize, Vec<ByteBuffer>>> {
-        let completed_until = self.first_region()?;
-        let mut completed = std::mem::take(&mut self.data);
-        // truncate completed to include only completed regions
-        let in_progress = completed.split_off(&completed_until.start);
-        self.data = in_progress;
-        Ok(completed)
+    pub fn take_buffers(&mut self) -> VortexResult<BTreeMap<usize, Vec<ByteBuffer>>> {
+        if self.active_regions.len() > 1 {
+            vortex_bail!("there are more than one active writers");
+        }
+        if !self.wakers.is_empty() {
+            vortex_bail!("there is an inflight write");
+        }
+        self.active_regions = [Region::default()].into();
+        Ok(std::mem::take(&mut self.data))
     }
 
     pub fn next_segment_id(&mut self) -> SegmentId {

vortex-file/src/segments/writer.rs

@@ -35,11 +35,17 @@ impl SegmentWriter for InOrderSegmentWriter {
 }
 
 impl InOrderSegmentWriter {
+    // TODO: split to get &self, should modify self.region to come after all the splits
     pub fn split(self, splits: usize) -> VortexResult<Vec<Self>> {
+        let unwritten_region = Region {
+            start: self.region.start + self.region_offset,
+            end: self.region.end,
+        };
+
         Ok(self
             .buffers
             .lock()
-            .split_region(&self.region, splits)?
+            .split_region(&unwritten_region, splits)?
             .map(|region| Self {
                 buffers: self.buffers.clone(),
                 region,
@@ -61,10 +67,12 @@ impl InOrderSegmentWriter {
         writer: &mut futures::io::Cursor<W>,
         segment_specs: &mut Vec<SegmentSpec>,
     ) -> VortexResult<()> {
-        let completed = {
-            let mut guard = self.buffers.lock();
-            guard.completed_buffers()?
-        };
+        let completed = self.buffers.lock().take_buffers()?;
+        // we are the only writer if here, reclaim the entire region
+        self.region = Region::default();
+        self.region_offset = 0;
+
+        // TODO(os): spawn everything below
         for buffers in completed.into_values() {
             // The API requires us to write these buffers contiguously. Therefore, we can only
             // respect the alignment of the first one.