From 646e0f3581b0cd723345487c9a7ac9073cfb31a5 Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Mon, 12 Sep 2022 14:43:07 -0400
Subject: [PATCH] Write traces to file

---
 pageserver/src/page_service.rs | 54 ++++++++++++++++++++++++++++++----
 pageserver/src/trace.rs        | 19 +++++++-----
 2 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 21d224738c..014ac61580 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -118,6 +118,50 @@ pub struct PagestreamDbSizeResponse {
 }
 
 impl PagestreamFeMessage {
+    pub fn serialize(&self) -> Bytes {
+        let mut bytes = BytesMut::new();
+
+        match self {
+            Self::Exists(req) => {
+                bytes.put_u8(if req.latest { 1 } else { 0 });
+                bytes.put_u64(req.lsn.0);
+                bytes.put_u32(req.rel.spcnode);
+                bytes.put_u32(req.rel.dbnode);
+                bytes.put_u32(req.rel.relnode);
+                bytes.put_u8(req.rel.forknum);
+            }
+
+            Self::Nblocks(req) => {
+                bytes.put_u8(1);
+                bytes.put_u8(if req.latest { 1 } else { 0 });
+                bytes.put_u64(req.lsn.0);
+                bytes.put_u32(req.rel.spcnode);
+                bytes.put_u32(req.rel.dbnode);
+                bytes.put_u32(req.rel.relnode);
+                bytes.put_u8(req.rel.forknum);
+            }
+
+            Self::GetPage(req) => {
+                bytes.put_u8(2);
+                bytes.put_u8(if req.latest { 1 } else { 0 });
+                bytes.put_u64(req.lsn.0);
+                bytes.put_u32(req.rel.spcnode);
+                bytes.put_u32(req.rel.dbnode);
+                bytes.put_u32(req.rel.relnode);
+                bytes.put_u8(req.rel.forknum);
+            }
+
+            Self::DbSize(req) => {
+                bytes.put_u8(3);
+                bytes.put_u8(if req.latest { 1 } else { 0 });
+                bytes.put_u64(req.lsn.0);
+                bytes.put_u32(req.dbnode);
+            }
+        }
+
+        bytes.into()
+    }
+
     fn parse(mut body: Bytes) -> anyhow::Result<PagestreamFeMessage> {
         // TODO these gets can fail
 
@@ -479,15 +523,15 @@ impl PageServerHandler {
                             _ => continue,
                         };
 
+                        // Trace request if needed
+                        if let Some(t) = tracer.as_mut() {
+                            t.trace(&copy_data_bytes)
+                        }
+
                         let zenith_fe_msg = PagestreamFeMessage::parse(copy_data_bytes)?;
                         let tenant_id = tenant_id.to_string();
                         let timeline_id = timeline_id.to_string();
 
-                        // Trace request if needed
-                        if let Some(t) = tracer.as_mut() {
-                            t.trace(&zenith_fe_msg)
-                        }
-
                         let response = match zenith_fe_msg {
                             PagestreamFeMessage::Exists(req) => SMGR_QUERY_TIME
                                 .with_label_values(&["get_rel_exists", &tenant_id, &timeline_id])
diff --git a/pageserver/src/trace.rs b/pageserver/src/trace.rs
index 818fc5f13f..0be3e61b27 100644
--- a/pageserver/src/trace.rs
+++ b/pageserver/src/trace.rs
@@ -1,8 +1,12 @@
-use crate::page_service::PagestreamFeMessage;
-use std::{fs::File, io::Write, path::PathBuf};
+use bytes::Bytes;
+use std::{
+    fs::File,
+    io::{BufWriter, Write},
+    path::PathBuf,
+};
 
 pub struct Tracer {
-    output: File,
+    writer: BufWriter<File>,
 }
 
 impl Drop for Tracer {
@@ -13,16 +17,17 @@ impl Drop for Tracer {
 
 impl Tracer {
     pub fn new(path: PathBuf) -> Self {
+        let file = File::create(path).expect("failed to create trace file");
         Tracer {
-            output: File::create(path).expect("failed to create trace file"),
+            writer: BufWriter::new(file),
         }
     }
 
-    pub fn trace(&mut self, _msg: &PagestreamFeMessage) {
-        // TODO(now) implement
+    pub fn trace(&mut self, msg: &Bytes) {
+        self.writer.write(msg).expect("failed to write trace");
     }
 
     pub fn flush(&mut self) {
-        self.output.flush().expect("failed to flush trace file");
+        self.writer.flush().expect("failed to flush trace file");
     }
 }