log: pull in logrotator, switch to ZSTD compressor

go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/decred/dcrd/lru v1.0.0
 	github.com/gorilla/websocket v1.5.0
 	github.com/jessevdk/go-flags v1.4.0
-	github.com/jrick/logrotate v1.0.0
+	github.com/klauspost/compress v1.17.9
 	github.com/stretchr/testify v1.8.4
 	github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7
 	golang.org/x/crypto v0.22.0

go.sum

@@ -59,10 +59,11 @@ github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpO
 github.com/jessevdk/go-flags v0.0.0-20141203071132-1679536dcc89/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
 github.com/jessevdk/go-flags v1.4.0 h1:4IU2WS7AumrZ/40jfhf4QVDMsQwqA7VEHozFRrGARJA=
 github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
-github.com/jrick/logrotate v1.0.0 h1:lQ1bL/n9mBNeIXoTUoYRlK4dHuNJVofX9oWqBtPnSzI=
 github.com/jrick/logrotate v1.0.0/go.mod h1:LNinyqDIJnpAur+b8yyulnQw/wDuN1+BYKlTRt3OuAQ=
 github.com/kkdai/bstream v0.0.0-20161212061736-f391b8402d23 h1:FOOIBWrEkLgmlgGfMuZT83xIwfPDxEI2OHu6xUmJMFE=
 github.com/kkdai/bstream v0.0.0-20161212061736-f391b8402d23/go.mod h1:J+Gs4SYgM6CZQHDETBtE9HaSEkGmuNXF86RwHhHUvq4=
+github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
+github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=

log.go

@@ -23,7 +23,6 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 
 	"github.com/btcsuite/btclog"
-	"github.com/jrick/logrotate/rotator"
 )
 
 // logWriter implements an io.Writer that outputs to both standard output and
@@ -52,7 +51,7 @@ var (
 
 	// logRotator is one of the logging outputs.  It should be closed on
 	// application shutdown.
-	logRotator *rotator.Rotator
+	logRotator *Rotator
 
 	adxrLog = backendLog.Logger("ADXR")
 	amgrLog = backendLog.Logger("AMGR")
@@ -115,7 +114,7 @@ func initLogRotator(logFile string) {
 		fmt.Fprintf(os.Stderr, "failed to create log directory: %v\n", err)
 		os.Exit(1)
 	}
-	r, err := rotator.New(logFile, 10*1024, false, 3)
+	r, err := NewRotatorWithCompressor(logFile, 10*1024, false, 3, Zstd)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to create file rotator: %v\n", err)
 		os.Exit(1)

logrotator.go

@@ -0,0 +1,313 @@
+// Copyright (c) 2017, moshee
+// Copyright (c) 2017, Josh Rickmar <[email protected]>
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice, this
+//   list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+
+// Package rotator implements a simple logfile rotator. Logs are read from an
+// io.Reader and are written to a file until they reach a specified size. The
+// log is then gzipped to another file and truncated.
+package main
+
+import (
+	"bufio"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+
+	"github.com/klauspost/compress/zstd"
+)
+
+// nl is a byte slice containing a newline byte.  It is used to avoid creating
+// additional allocations when writing newlines to the log file.
+var nl = []byte{'\n'}
+
+// Compressor represents the supported compression algorithms.
+type Compressor uint8
+
+const (
+	// Gzip is the gzip compression algorithm, implemented in the stdlib.
+	Gzip Compressor = 0
+
+	// Zstd is the zstd compression algorithm.
+	Zstd Compressor = 1
+)
+
+// Validate checks that the Compressor is valid, and returns the appropriate
+// file suffix for the compressed file.
+func (c Compressor) Validate() (string, error) {
+	switch c {
+	case Gzip:
+		return ".gz", nil
+	case Zstd:
+		return ".zst", nil
+
+	default:
+		return "", fmt.Errorf("unknown compression algorithm: %d", c)
+	}
+}
+
+// Init creates a new writer to use for compression.
+func (c Compressor) Init(w io.Writer) (io.WriteCloser, error) {
+	switch c {
+	case Gzip:
+		dst := gzip.NewWriter(w)
+		return dst, nil
+
+	case Zstd:
+		return zstd.NewWriter(w)
+
+	default:
+		return nil, fmt.Errorf("unknown compression algorithm: %d", c)
+	}
+}
+
+// A Rotator writes input to a file, splitting it up into gzipped chunks once
+// the filesize reaches a certain threshold.
+type Rotator struct {
+	size      int64
+	threshold int64
+	maxRolls  int
+	filename  string
+	out       *os.File
+	tee       bool
+	wg        sync.WaitGroup
+	comp      Compressor
+}
+
+// NewRotator returns a new Rotator. The rotator can be used either by reading
+// input from an io.Reader by calling Run, or writing directly to the Rotator
+// with Write. The default compression algorithm is Gzip.
+func NewRotator(filename string, thresholdKB int64, tee bool,
+	maxRolls int) (*Rotator, error) {
+
+	f, err := os.OpenFile(filename, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0644)
+	if err != nil {
+		return nil, err
+	}
+
+	stat, err := f.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	return &Rotator{
+		size:      stat.Size(),
+		threshold: 1000 * thresholdKB,
+		maxRolls:  maxRolls,
+		filename:  filename,
+		out:       f,
+		tee:       tee,
+		comp:      Gzip,
+	}, nil
+}
+
+// NewRotatorWithCompressor returns a new Rotator that will use a specific
+// compression algorithm.
+func NewRotatorWithCompressor(filename string, thresholdKB int64, tee bool,
+	maxRolls int, comp Compressor) (*Rotator, error) {
+
+	_, err := comp.Validate()
+	if err != nil {
+		return nil, err
+	}
+
+	rotator, err := NewRotator(filename, thresholdKB, tee, maxRolls)
+	if err != nil {
+		return nil, err
+	}
+
+	rotator.comp = comp
+
+	return rotator, nil
+}
+
+// Run begins reading lines from the reader and rotating logs as necessary.  Run
+// should not be called concurrently with Write.
+//
+// Prefer to use Rotator as a writer instead to avoid unnecessary scanning of
+// input, as this job is better handled using io.Pipe.
+func (r *Rotator) Run(reader io.Reader) error {
+	in := bufio.NewReader(reader)
+
+	// Rotate file immediately if it is already over the size limit.
+	if r.size >= r.threshold {
+		if err := r.rotate(); err != nil {
+			return err
+		}
+		r.size = 0
+	}
+
+	for {
+		line, isPrefix, err := in.ReadLine()
+		if err != nil {
+			return err
+		}
+
+		n, _ := r.out.Write(line)
+		r.size += int64(n)
+		if r.tee {
+			os.Stdout.Write(line)
+		}
+		if isPrefix {
+			continue
+		}
+
+		m, _ := r.out.Write(nl)
+		if r.tee {
+			os.Stdout.Write(nl)
+		}
+		r.size += int64(m)
+
+		if r.size >= r.threshold {
+			err := r.rotate()
+			if err != nil {
+				return err
+			}
+			r.size = 0
+		}
+	}
+}
+
+// Write implements the io.Writer interface for Rotator.  If p ends in a newline
+// and the file has exceeded the threshold size, the file is rotated.
+func (r *Rotator) Write(p []byte) (n int, err error) {
+	n, _ = r.out.Write(p)
+	r.size += int64(n)
+
+	if r.size >= r.threshold && len(p) > 0 && p[len(p)-1] == '\n' {
+		err := r.rotate()
+		if err != nil {
+			return 0, err
+		}
+		r.size = 0
+	}
+
+	return n, nil
+}
+
+// Close closes the output logfile.
+func (r *Rotator) Close() error {
+	err := r.out.Close()
+	r.wg.Wait()
+	return err
+}
+
+func (r *Rotator) rotate() error {
+	dir := filepath.Dir(r.filename)
+	glob := filepath.Join(dir, filepath.Base(r.filename)+".*")
+	existing, err := filepath.Glob(glob)
+	if err != nil {
+		return err
+	}
+
+	maxNum := 0
+	for _, name := range existing {
+		parts := strings.Split(name, ".")
+		if len(parts) < 2 {
+			continue
+		}
+		numIdx := len(parts) - 1
+		if parts[numIdx] == "gz" {
+			numIdx--
+		}
+		num, err := strconv.Atoi(parts[numIdx])
+		if err != nil {
+			continue
+		}
+		if num > maxNum {
+			maxNum = num
+		}
+	}
+
+	err = r.out.Close()
+	if err != nil {
+		return err
+	}
+	rotname := fmt.Sprintf("%s.%d", r.filename, maxNum+1)
+	err = os.Rename(r.filename, rotname)
+	if err != nil {
+		return err
+	}
+	if r.maxRolls > 0 {
+		for n := maxNum + 1 - r.maxRolls; ; n-- {
+			err := os.Remove(fmt.Sprintf("%s.%d.gz", r.filename, n))
+			if err != nil {
+				break
+			}
+		}
+	}
+	r.out, err = os.OpenFile(r.filename, os.O_CREATE|os.O_RDWR, 0644)
+	if err != nil {
+		return err
+	}
+
+	r.wg.Add(1)
+	go func() {
+		err := compress(rotname, r.comp)
+		if err == nil {
+			os.Remove(rotname)
+		}
+		r.wg.Done()
+	}()
+
+	return nil
+}
+
+func compress(name string, comp Compressor) (err error) {
+	f, err := os.Open(name)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	suffix, err := comp.Validate()
+	if err != nil {
+		return err
+	}
+
+	arc, err := os.OpenFile(
+		name+suffix, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0644,
+	)
+	if err != nil {
+		return err
+	}
+
+	z, err := comp.Init(arc)
+	if err != nil {
+		return err
+	}
+
+	if _, err = io.Copy(z, f); err != nil {
+		return err
+	}
+	if err = z.Close(); err != nil {
+		return err
+	}
+	return arc.Close()
+}