Support native PostgreSQL logging format (#483)

* feat: support native PostgreSQL logging format

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* fix: csv format log output correction

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* refactor: change log_desitation judgment from string to intt

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* chore: put the code under src/logger, remove logger crate

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* feat: log transfer supports large chunk writing

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* chore: remove test code

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* chore: remove no used mod

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* fix: add logger dir declaration

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

* chore: replace PIPE_BUF with sizeof PipeProtoChunk

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

---------

Signed-off-by: jinweios <jinwei.peng@beingthink.com>

Author: JinweiOS

Cargo.lock

@@ -21,6 +21,8 @@ arrayvec.workspace = true
 bincode.workspace = true
 bytemuck.workspace = true
 byteorder.workspace = true
+chrono = "0.4.38"
+csv = "1.3.0"
 env_logger = "0.11.2"
 libc.workspace = true
 log.workspace = true

Cargo.toml

@@ -34,15 +34,14 @@ extern "C" fn _vectors_main(_arg: pgrx::pg_sys::Datum) {
         pub layout: std::alloc::Layout,
     }
     {
-        let mut builder = env_logger::builder();
-        builder.target(env_logger::Target::Stderr);
+        let mut builder = crate::logger::VectorLogger::build();
         #[cfg(not(debug_assertions))]
         {
-            builder.filter(None, log::LevelFilter::Info);
+            builder.filter_level(log::LevelFilter::Info);
         }
         #[cfg(debug_assertions)]
         {
-            builder.filter(None, log::LevelFilter::Trace);
+            builder.filter_level(log::LevelFilter::Trace);
         }
         builder.init();
     }

src/bgworker/mod.rs

@@ -17,6 +17,7 @@ mod error;
 mod gucs;
 mod index;
 mod ipc;
+mod logger;
 mod upgrade;
 mod utils;
 

src/lib.rs

@@ -0,0 +1,111 @@
+use chrono::Local;
+use pgrx::pg_sys::PipeProtoChunk;
+use serde::Serialize;
+use std::io::{stderr, Write};
+use std::{process, vec};
+
+const PIPE_PROTO_DEST_STDERR: u8 = 0x10;
+const PIPE_PROTO_DEST_CSVLOG: u8 = 0x20;
+const PIPE_PROTO_DEST_JSONLOG: u8 = 0x40;
+const PIPE_PROTO_IS_LAST: u8 = 0x01;
+
+const PIPE_CHUNK_SIZE: usize = std::mem::size_of::<PipeProtoChunk>();
+const PIPE_HEADER_SIZE: usize = 9;
+const PIPE_MAX_PAYLOAD: usize = PIPE_CHUNK_SIZE - PIPE_HEADER_SIZE;
+
+struct PipeProtoHeader {
+    nuls: [u8; 2],
+    len: u16,
+    pid: i32,
+    flags: u8,
+}
+
+impl PipeProtoHeader {
+    fn new(flags: u8) -> Self {
+        Self {
+            nuls: [0, 0],
+            len: 0,
+            pid: process::id() as i32,
+            flags,
+        }
+    }
+
+    fn chunk(&self) -> Vec<u8> {
+        let mut chunk = Vec::from(&self.nuls);
+        chunk.extend_from_slice(&self.len.to_le_bytes());
+        chunk.extend_from_slice(&self.pid.to_le_bytes());
+        chunk.push(self.flags);
+        chunk
+    }
+
+    fn set_flag(&mut self, flags: u8) {
+        self.flags = flags
+    }
+
+    fn set_len(&mut self, len: usize) {
+        self.len = len as u16;
+    }
+}
+
+#[derive(Serialize)]
+pub struct Message {
+    timestamp: String,
+    pid: u32,
+    message: String,
+}
+
+impl Message {
+    pub fn new(msg: &str) -> Self {
+        Self {
+            timestamp: Local::now().format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
+            pid: process::id(),
+            message: String::from(msg),
+        }
+    }
+
+    pub fn csv_chunk(&self) {
+        let mut wtr = csv::WriterBuilder::new()
+            .has_headers(false)
+            .from_writer(vec![]);
+        wtr.serialize(self).unwrap();
+        write_pipe_chunks(wtr.into_inner().unwrap(), PIPE_PROTO_DEST_CSVLOG)
+    }
+
+    pub fn stderr_chunk(&self) {
+        let err = format!("{} [{}] LOG:  {}", self.timestamp, self.pid, self.message);
+        let mut data: Vec<u8> = Vec::new();
+        writeln!(&mut data, "{}", &err).unwrap();
+        write_pipe_chunks(data, PIPE_PROTO_DEST_STDERR)
+    }
+
+    pub fn json_chunk(&self) {
+        let json = serde_json::to_string(self).unwrap();
+        let mut data: Vec<u8> = Vec::new();
+        writeln!(&mut data, "{}", &json).unwrap();
+        write_pipe_chunks(data, PIPE_PROTO_DEST_JSONLOG)
+    }
+}
+
+fn write_pipe_chunks(msg_buf: Vec<u8>, flags: u8) {
+    let mut len = msg_buf.len();
+    let mut header = PipeProtoHeader::new(flags);
+    let mut fd = stderr();
+    let mut cursor: usize = 0;
+    while len > PIPE_MAX_PAYLOAD {
+        header.set_len(PIPE_MAX_PAYLOAD);
+        let data = &msg_buf[cursor..(cursor + PIPE_MAX_PAYLOAD)];
+        let mut chunk = Vec::new();
+        chunk.extend_from_slice(&header.chunk());
+        chunk.extend_from_slice(data);
+        let _ = fd.write(&chunk).unwrap();
+        len -= PIPE_MAX_PAYLOAD;
+        cursor += PIPE_MAX_PAYLOAD;
+    }
+    header.set_flag(flags | PIPE_PROTO_IS_LAST);
+    header.set_len(len);
+    let data = &msg_buf[cursor..(cursor + len)];
+    let mut chunk = Vec::new();
+    chunk.extend_from_slice(&header.chunk());
+    chunk.extend_from_slice(data);
+    let _ = fd.write(&chunk).unwrap();
+}

src/logger/message.rs

@@ -0,0 +1,41 @@
+mod message;
+mod protocol;
+
+use log::{set_boxed_logger, set_max_level, LevelFilter, Metadata, Record};
+use protocol::pipe_log;
+
+#[derive(Debug, Clone)]
+pub struct VectorLogger {
+    level: LevelFilter,
+}
+
+impl log::Log for VectorLogger {
+    fn enabled(&self, metadata: &Metadata) -> bool {
+        metadata.level() <= self.level
+    }
+
+    fn log(&self, record: &Record) {
+        if self.enabled(record.metadata()) {
+            pipe_log(&record.args().to_string());
+        }
+    }
+
+    fn flush(&self) {}
+}
+
+impl VectorLogger {
+    pub fn init(&self) {
+        set_boxed_logger(Box::new(self.clone())).unwrap();
+        set_max_level(self.level);
+    }
+
+    pub fn build() -> Self {
+        VectorLogger {
+            level: LevelFilter::Info,
+        }
+    }
+
+    pub fn filter_level(&mut self, level: LevelFilter) {
+        self.level = level;
+    }
+}

src/logger/mod.rs

@@ -0,0 +1,21 @@
+use super::message::Message;
+use pgrx::pg_sys::Log_destination;
+
+const LOG_DESTINATION_STDERR: u32 = 1;
+const LOG_DESTINATION_CSVLOG: u32 = 8;
+const LOG_DESTINATION_JSONLOG: u32 = 16;
+
+pub fn pipe_log(msg: &str) {
+    let message = Message::new(msg);
+    unsafe {
+        if Log_destination as u32 & LOG_DESTINATION_CSVLOG != 0 {
+            message.csv_chunk()
+        }
+        if Log_destination as u32 & LOG_DESTINATION_STDERR != 0 {
+            message.stderr_chunk()
+        }
+        if Log_destination as u32 & LOG_DESTINATION_JSONLOG != 0 {
+            message.json_chunk()
+        }
+    }
+}