write qlogs to a temporary file first, then rename them when done

qlog.go

@@ -0,0 +1,79 @@
+package libp2pquic
+
+import (
+	"bufio"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"os"
+	"time"
+)
+
+var qlogDir string
+
+func init() {
+	qlogDir = os.Getenv("QLOGDIR")
+}
+
+func getLogWriterFor(role string) func([]byte) io.WriteCloser {
+	if len(qlogDir) == 0 {
+		return nil
+	}
+	return func(connID []byte) io.WriteCloser {
+		// create the QLOGDIR, if it doesn't exist
+		if err := os.MkdirAll(qlogDir, 0777); err != nil {
+			log.Errorf("creating the QLOGDIR failed: %s", err)
+			return nil
+		}
+		return newQlogger(role, connID)
+	}
+}
+
+type qlogger struct {
+	filename      string // QLOGDIR/.log_xxx.qlog.gz.swp
+	finalFilename string // QLOGDIR/log_xxx.qlog.gz
+	io.WriteCloser
+}
+
+func newQlogger(role string, connID []byte) io.WriteCloser {
+	t := time.Now().UTC().Format("2006-01-02T15-04-05.999999999UTC")
+	finalFilename := fmt.Sprintf("%s%clog_%s_%s_%x.qlog.gz", qlogDir, os.PathSeparator, t, role, connID)
+	filename := fmt.Sprintf("%s%c.log_%s_%s_%x.qlog.gz.swp", qlogDir, os.PathSeparator, t, role, connID)
+	f, err := os.Create(filename)
+	if err != nil {
+		log.Errorf("unable to create qlog file %s: %s", filename, err)
+		return nil
+	}
+	gz := gzip.NewWriter(f)
+	return &qlogger{
+		filename:      filename,
+		finalFilename: finalFilename,
+		WriteCloser:   newBufferedWriteCloser(bufio.NewWriter(gz), gz),
+	}
+}
+
+func (l *qlogger) Close() error {
+	if err := l.WriteCloser.Close(); err != nil {
+		return err
+	}
+	return os.Rename(l.filename, l.finalFilename)
+}
+
+type bufferedWriteCloser struct {
+	*bufio.Writer
+	io.Closer
+}
+
+func newBufferedWriteCloser(writer *bufio.Writer, closer io.Closer) io.WriteCloser {
+	return &bufferedWriteCloser{
+		Writer: writer,
+		Closer: closer,
+	}
+}
+
+func (h bufferedWriteCloser) Close() error {
+	if err := h.Writer.Flush(); err != nil {
+		return err
+	}
+	return h.Closer.Close()
+}

qlog_test.go

@@ -0,0 +1,83 @@
+package libp2pquic
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io/ioutil"
+	"os"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+type nopCloser struct{}
+
+func (nopCloser) Close() error { return nil }
+
+var _ = Describe("qlogger", func() {
+	var origQlogDir string
+
+	BeforeEach(func() {
+		origQlogDir = qlogDir
+		d, err := ioutil.TempDir("", "libp2p-quic-transport-test")
+		Expect(err).ToNot(HaveOccurred())
+		fmt.Fprintf(GinkgoWriter, "Creating temporary directory: %s\n", d)
+		qlogDir = d
+	})
+
+	AfterEach(func() {
+		qlogDir = origQlogDir
+	})
+
+	getFile := func() os.FileInfo {
+		files, err := ioutil.ReadDir(qlogDir)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(files).To(HaveLen(1))
+		return files[0]
+	}
+
+	It("saves a qlog", func() {
+		logger := newQlogger("server", []byte{0xde, 0xad, 0xbe, 0xef})
+		file := getFile()
+		Expect(string(file.Name()[0])).To(Equal("."))
+		Expect(file.Name()).To(HaveSuffix(".qlog.gz.swp"))
+		// close the logger. This should move the file.
+		Expect(logger.Close()).To(Succeed())
+		file = getFile()
+		Expect(string(file.Name()[0])).ToNot(Equal("."))
+		Expect(file.Name()).To(HaveSuffix(".qlog.gz"))
+		Expect(file.Name()).To(And(
+			ContainSubstring("server"),
+			ContainSubstring("deadbeef"),
+		))
+	})
+
+	It("buffers", func() {
+		logger := newQlogger("server", []byte("connid"))
+		initialSize := getFile().Size()
+		// Do a small write.
+		// Since the writter is buffered, this should not be written to disk yet.
+		logger.Write([]byte("foobar"))
+		Expect(getFile().Size()).To(Equal(initialSize))
+		// Close the logger. This should flush the buffer to disk.
+		Expect(logger.Close()).To(Succeed())
+		finalSize := getFile().Size()
+		fmt.Fprintf(GinkgoWriter, "initial log file size: %d, final log file size: %d\n", initialSize, finalSize)
+		Expect(finalSize).To(BeNumerically(">", initialSize))
+	})
+
+	It("compresses", func() {
+		logger := newQlogger("server", []byte("connid"))
+		logger.Write([]byte("foobar"))
+		Expect(logger.Close()).To(Succeed())
+		compressed, err := ioutil.ReadFile(qlogDir + "/" + getFile().Name())
+		Expect(err).ToNot(HaveOccurred())
+		Expect(compressed).ToNot(Equal("foobar"))
+		gz, err := gzip.NewReader(bytes.NewReader(compressed))
+		Expect(err).ToNot(HaveOccurred())
+		data, err := ioutil.ReadAll(gz)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(data).To(Equal([]byte("foobar")))
+	})
+})

transport.go

@@ -1,15 +1,10 @@
 package libp2pquic
 
 import (
-	"bufio"
-	"compress/gzip"
 	"context"
 	"errors"
-	"fmt"
 	"io"
 	"net"
-	"os"
-	"time"
 
 	"github.com/minio/sha256-simd"
 	"golang.org/x/crypto/hkdf"
@@ -134,8 +129,8 @@ func NewTransport(key ic.PrivKey, psk pnet.PSK, filters *filter.Filters) (tpt.Tr
 		serverConfig: config,
 		clientConfig: config.Clone(),
 	}
-	t.serverConfig.GetLogWriter = t.GetLogWriterFor("server")
-	t.clientConfig.GetLogWriter = t.GetLogWriterFor("client")
+	t.serverConfig.GetLogWriter = getLogWriterFor("server")
+	t.clientConfig.GetLogWriter = getLogWriterFor("client")
 	return t, nil
 }
 
@@ -195,29 +190,6 @@ func (t *transport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (tp
 	}, nil
 }
 
-func (t *transport) GetLogWriterFor(role string) func([]byte) io.WriteCloser {
-	qlogDir := os.Getenv("QLOGDIR")
-	if len(qlogDir) == 0 {
-		return nil
-	}
-	return func(connID []byte) io.WriteCloser {
-		// create the QLOGDIR, if it doesn't exist
-		if err := os.MkdirAll(qlogDir, 0777); err != nil {
-			log.Errorf("creating the QLOGDIR failed: %s", err)
-			return nil
-		}
-		t := time.Now().Format(time.RFC3339Nano)
-		filename := fmt.Sprintf("%s/log_%s_%s_%x.qlog.gz", qlogDir, t, role, connID)
-		f, err := os.Create(filename)
-		if err != nil {
-			log.Errorf("unable to create qlog file %s: %s", filename, err)
-			return nil
-		}
-		gz := gzip.NewWriter(f)
-		return newBufferedWriteCloser(bufio.NewWriter(gz), gz)
-	}
-}
-
 // Don't use mafmt.QUIC as we don't want to dial DNS addresses. Just /ip{4,6}/udp/quic
 var dialMatcher = mafmt.And(mafmt.IP, mafmt.Base(ma.P_UDP), mafmt.Base(ma.P_QUIC))