feat(wal): Make `wal` WriteOpEntryDecoder an interator

Also, implement drop on `wal_inspect`'s LineProtoWriter and
bubble up flush to the caller.

wal/src/lib.rs

@@ -585,25 +585,31 @@ pub struct WriteOpEntryDecoder {
     reader: ClosedSegmentFileReader,
 }
 
-impl WriteOpEntryDecoder {
+impl From<ClosedSegmentFileReader> for WriteOpEntryDecoder {
     /// Creates a decoder which will use the closed segment file of `reader` to
     /// decode write ops from their on-disk format.
-    pub fn from_closed_segment(reader: ClosedSegmentFileReader) -> Self {
+    fn from(reader: ClosedSegmentFileReader) -> Self {
         Self { reader }
     }
+}
+
+impl Iterator for WriteOpEntryDecoder {
+    type Item = Result<Vec<WriteOpEntry>, DecodeError>;
 
     /// Reads a collection of write op entries in the next WAL entry batch from the
     /// underlying closed segment. A returned Ok(None) indicates that there are no
     /// more entries to be decoded from the underlying segment. A zero-length vector
     /// may be returned if there are no writes in a WAL entry batch, but does not
     /// indicate the decoder is consumed.
-    pub fn next_write_op_entry_batch(&mut self) -> Result<Option<Vec<WriteOpEntry>>, DecodeError> {
+    fn next(&mut self) -> Option<Self::Item> {
-        match self.reader.next_batch().context(FailedToReadWalSnafu)? {
+        match self.reader.next_batch().context(FailedToReadWalSnafu) {
-            Some(batch) => Ok(batch
+            Ok(Some(batch)) => Some(
+                batch
                     .into_iter()
                     .filter_map(|sequenced_op| match sequenced_op.op {
                         WalOp::Write(w) => Some(w),
-                    _ => None,
+                        WalOp::Delete(..) => None,
+                        WalOp::Persist(..) => None,
                     })
                     .map(|w| -> Result<WriteOpEntry, DecodeError> {
                         Ok(WriteOpEntry {
@@ -612,9 +618,10 @@ impl WriteOpEntryDecoder {
                                 .context(UnableToCreateMutableBatchSnafu)?,
                         })
                     })
-                .collect::<Result<Vec<WriteOpEntry>, DecodeError>>()?
+                    .collect::<Self::Item>(),
-                .into()),
+            ),
-            None => Ok(None),
+            Ok(None) => None,
+            Err(e) => Some(Err(e)),
         }
     }
 }
@@ -783,15 +790,15 @@ mod tests {
 
         let (closed, _) = wal.rotate().unwrap();
 
-        let mut decoder = WriteOpEntryDecoder::from_closed_segment(
+        let decoder = WriteOpEntryDecoder::from(
             wal.reader_for_segment(closed.id)
                 .expect("failed to open reader for closed WAL segment"),
         );
 
-        let mut write_op_entries = vec![];
+        let write_op_entries = decoder
-        while let Ok(Some(mut entry_batch)) = decoder.next_write_op_entry_batch() {
+            .into_iter()
-            write_op_entries.append(&mut entry_batch);
+            .flat_map(|r| r.expect("unexpected bad entry"))
-        }
+            .collect::<Vec<_>>();
         // The decoder should find 2 entries, each containing a single table write
         assert_eq!(write_op_entries.len(), 2);
         assert_matches!(write_op_entries.get(0), Some(got_op1) => {

wal_inspect/src/lib.rs

@@ -43,13 +43,38 @@ pub enum WriteError {
 /// Provides namespaced write functionality from table-based mutable batches
 /// to namespaced line protocol output.
 #[derive(Debug)]
-pub struct LineProtoWriter<W, F> {
+pub struct LineProtoWriter<W, F>
+where
+    W: Write,
+{
     namespaced_output: HashMap<NamespaceId, W>,
     new_write_sink: F,
 
     table_name_index: HashMap<TableId, String>,
 }
 
+impl<W, F> LineProtoWriter<W, F>
+where
+    W: Write,
+{
+    /// Flushes all write destinations opened by the [`LineProtoWriter`].
+    pub fn flush(&mut self) -> Result<(), WriteError> {
+        for w in self.namespaced_output.values_mut() {
+            w.flush()?
+        }
+        Ok(())
+    }
+}
+
+impl<W, F> Drop for LineProtoWriter<W, F>
+where
+    W: Write,
+{
+    fn drop(&mut self) {
+        _ = self.flush()
+    }
+}
+
 impl<W, F> LineProtoWriter<W, F>
 where
     W: Write,
@@ -80,13 +105,7 @@ where
             .entry(ns)
             .or_insert((self.new_write_sink)(ns)?);
 
-        match write_batches_as_line_proto(sink, &self.table_name_index, table_batches) {
+        write_batches_as_line_proto(sink, &self.table_name_index, table_batches)
-            Ok(_) => sink.flush().map_err(WriteError::IoError),
-            Err(e) => {
-                _ = sink.flush();
-                Err(e)
-            }
-        }
     }
 }
 
@@ -127,7 +146,7 @@ mod tests {
         iox::wal::v1::sequenced_wal_op::Op, pbdata::v1::DatabaseBatch,
     };
     use mutable_batch_lp::lines_to_batches;
-    use wal::{DecodeError, SequencedWalOp, WriteOpEntryDecoder};
+    use wal::{DecodeError, SequencedWalOp, WriteOpEntry, WriteOpEntryDecoder};
 
     use super::*;
 
@@ -163,32 +182,31 @@ mod tests {
         // Rotate the WAL and create the translator.
         let (closed, _) = wal.rotate().expect("failed to rotate WAL");
 
-        let mut decoder = WriteOpEntryDecoder::from_closed_segment(
+        let decoder = WriteOpEntryDecoder::from(
             wal.reader_for_segment(closed.id())
                 .expect("failed to open reader for closed segment"),
         );
         let mut writer = LineProtoWriter::new(|_| Ok(Vec::<u8>::new()), table_name_index);
 
-        let mut decoded_entries = 0;
+        let decoded_entries = decoder
-        let mut decoded_ops = 0;
+            .into_iter()
-        while let Some(new_entries) = decoder
+            .map(|r| r.expect("unexpected bad entry"))
-            .next_write_op_entry_batch()
+            .collect::<Vec<_>>();
-            .expect("decoder error should not occur")
+        assert_eq!(decoded_entries.len(), 1);
-        {
+        let decoded_ops = decoded_entries
-            decoded_entries += 1;
+            .into_iter()
-            decoded_ops += new_entries.len();
+            .flatten()
+            .collect::<Vec<WriteOpEntry>>();
+        assert_eq!(decoded_ops.len(), 3);
 
-            for entry in new_entries {
+        for entry in decoded_ops {
             writer
                 .write_namespaced_table_batches(entry.namespace, entry.table_batches)
                 .expect("batch write should not fail");
         }
-        }
         // The WAL has been given a single entry containing three write ops
-        assert_eq!(decoded_entries, 1);
-        assert_eq!(decoded_ops, 3);
 
-        let results = writer.namespaced_output;
+        let results = &writer.namespaced_output;
 
         // Assert that the namespaced writes contain ONLY the following:
         //
@@ -269,36 +287,36 @@ mod tests {
         }
 
         // Create the translator and read as much as possible out of the bad segment file
-        let mut decoder = WriteOpEntryDecoder::from_closed_segment(
+        let decoder = WriteOpEntryDecoder::from(
             wal.reader_for_segment(closed.id())
                 .expect("failed to open reader for closed segment"),
         );
         let mut writer = LineProtoWriter::new(|_| Ok(Vec::<u8>::new()), table_name_index);
 
-        let mut decoded_entries = 0;
+        // The translator should be able to read all 2 good entries containing 4 write ops
-        let mut decoded_ops = 0;
+        let decoded_entries = decoder
-        loop {
+            .into_iter()
-            match decoder.next_write_op_entry_batch() {
+            .map_while(|r| {
-                // If the translator returns `None` indicating successful translation
+                r.map_err(|e| match e {
-                // then something is broken.
+                    DecodeError::FailedToReadWal { .. } => None::<()>,
-                Ok(v) => assert_matches!(v, Some(new_entries) => {
+                    _ => panic!("unexpected error"),
-                    decoded_entries += 1;
+                })
-                    decoded_ops += new_entries.len();
+                .ok()
-                    for entry in new_entries {
+            })
-                        writer.write_namespaced_table_batches(entry.namespace, entry.table_batches).expect("batch write should not fail");
+            .collect::<Vec<_>>();
+        assert_eq!(decoded_entries.len(), 2);
+        let decoded_ops = decoded_entries
+            .into_iter()
+            .flatten()
+            .collect::<Vec<WriteOpEntry>>();
+        assert_eq!(decoded_ops.len(), 4);
+        for entry in decoded_ops {
+            writer
+                .write_namespaced_table_batches(entry.namespace, entry.table_batches)
+                .expect("batch write should not fail");
         }
-                }),
-                Err(e) => {
-                    assert_matches!(e, DecodeError::FailedToReadWal { .. });
-                    break;
-                }
-            };
-        }
-        // The translator should have read all 2 good entries containing 4 write ops
-        assert_eq!(decoded_entries, 2);
-        assert_eq!(decoded_ops, 4);
 
-        let results = writer.namespaced_output;
+        let results = &writer.namespaced_output;
 
         // Assert that the namespaced writes contain ONLY the following:
         //