watchlogs/main.go

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */
package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"regexp"
	"time"

	"golang.org/x/term"
	logger "idio.link/go/logger/v3"
)

var (
	Version, Build string
)

var Options = struct {
	patternRE *regexp.Regexp
	pattern   []byte
	command   string

	regex,
	jumpToEnd,
	matchOnce,
	printNonmatching,
	verbose,
	debug,
	quiet bool
}{}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	ctx = context.WithValue(ctx, "log", logger.NewLogger())
	log := ctx.Value("log").(*logger.Logger)

	done := make(chan struct{})
	go func() {
		sig := make(chan os.Signal, 1) // catch signals
		signal.Notify(sig, os.Interrupt, os.Kill)

		for {
			select {
			case <-sig:
				fmt.Println()
				log.Info("Caught SIGINT/SIGTERM.")
				cancel()
				return
			case <-done:
				fmt.Println()
				log.Info("Caught SIGINT/SIGTERM.")
				return
			case <-time.After(100 * time.Millisecond):
				// do nothing; just throttle the cpu
			}
		}
	}()

	handlePatternFlag := func(p string) (err error) {
		Options.pattern = []byte(p)
		Options.patternRE, err = regexp.Compile(p)
		return err
	}
	flags := flag.NewFlagSet("watchlogs", flag.ExitOnError)
	flags.SetOutput(os.Stderr)
	flags.BoolVar(&Options.regex, "regex", false, "interpret pattern as a regular expression")
	flags.BoolVar(&Options.regex, "e", false, "interpret pattern as a regular expression")
	flags.Func("pattern", "process lines matching pattern", handlePatternFlag)
	flags.Func("p", "process lines matching pattern", handlePatternFlag)
	flags.BoolVar(&Options.jumpToEnd, "jump", false, "jump to end of file before processing")
	flags.BoolVar(&Options.jumpToEnd, "j", false, "jump to end of file before processing")
	flags.BoolVar(&Options.matchOnce, "once", false, "stop after first matching line is processed")
	flags.BoolVar(&Options.matchOnce, "1", false, "stop after first matching line is processed")
	flags.BoolVar(&Options.printNonmatching, "n", false, "print all lines, not just matching ones")
	flags.BoolVar(&Options.verbose, "v", false, "increase verbosity")
	flags.BoolVar(&Options.debug, "debug", false, "print debugging information")
	flags.BoolVar(&Options.debug, "d", false, "print debugging information")
	flags.BoolVar(&Options.quiet, "quiet", false, "do not print lines")
	flags.BoolVar(&Options.quiet, "q", false, "do not print lines")
	flags.StringVar(&Options.command, "command", "", "execute command on match, passing the matched line as input")
	flags.Parse(os.Args[1:])

	logPath := flags.Arg(0)

	if len(flags.Args()) == 0 {
		fmt.Fprintf(os.Stderr, "%s %s\n", Version, Build)
		fmt.Fprintf(os.Stderr, "usage: %s <log_file> [<option> ...]\n", os.Args[0])
		flags.PrintDefaults()
		os.Exit(1)
	}
	// TODO option: execute command on match

	switch {
	case Options.debug:
		log.SetLevel(logger.LogLevelDebug)
	case Options.verbose:
		log.SetLevel(logger.LogLevelInfo)
	case Options.quiet:
		log.SetLevel(logger.LogLevelError)
	}

	if Options.pattern == nil {
		handlePatternFlag("")
	}

	ret := mainLoop(ctx, logPath)
	close(done)
	os.Exit(ret)
}

func printStatus(n uint64) {
	if term.IsTerminal(int(os.Stderr.Fd())) &&
		!Options.quiet {
		fmt.Fprintf(os.Stderr, "\x1b[%dG", 1)
		fmt.Printf("Read %d lines.", n)
	}
}

func mainLoop(ctx context.Context, path string) int {
	log := ctx.Value("log").(*logger.Logger)

	var file *os.File
	var lineCnt uint64
	var err error

	if file, err = os.Open(path); err != nil {
		log.Error("Could not open file at %#v: %v", path, err)
		return 1
	}
	if file == nil {
		panic("BUG: Received nil file pointer")
	}
	defer func() {
		if err := file.Close(); err != nil {
			log.Error("Error while closing file: %v", err)
		}
	}()

	if Options.jumpToEnd {
		if _, err := file.Seek(0, 2); err != nil {
			log.Error("Could not jump to end of file: %v", err)
		}
	}

	initBuf := make([]byte, 0, SizeCacheL1)
	buf := initBuf

	for {
		if ctx.Err() != nil {
			return 2
		}

		i := 0
		for i < len(buf) {
			if buf[i] == '\n' {
				break
			}
			i++
		}

		if i != len(buf) {
			line := buf[:i]
			log.Debug("Detected linefeed; processing line %#v", string(line))
			buf = buf[i+1:]
			lineCnt++
			if lineCnt&1023 == 0 {
				printStatus(lineCnt)
			}
			// TODO Decide how to handle commands
			if processLine(line) {
				if !Options.quiet {
					fmt.Println()
					fmt.Println(string(line))
				}
				if Options.matchOnce {
					log.Info("Matched once, as requested; terminating...")
					return 0
				}
			} else if Options.printNonmatching &&
				!Options.quiet {
				fmt.Println(string(line))
			}
			continue
		}

		if len(buf) == cap(initBuf) {
			log.Debug("Line exceeds allocated line buffer of %d bytes", cap(initBuf))
			nextCap := 2 * cap(initBuf)
			if nextCap > MaxSizeBufferLine {
				log.Error("Line exceeds max line buffer allocation of %d bytes", MaxSizeBufferLine)
				return 4
			}
			initBuf = make([]byte, 0, nextCap)
			continue
		}

		if len(buf) == cap(buf) {
			log.Debug("Reached end of buffer; resetting cursor to start of buffer")
			nextBuf := initBuf[0:len(buf)]
			copy(nextBuf, buf)
			buf = nextBuf
			continue
		}

		// TODO Note the difference between extension of
		// existing lines and buffering new lines, and make sure
		// new line slices are cache aligned.
		log.Debug("Attempting to read %d bytes into buffer...", cap(buf)-len(buf))
		n, err := file.Read(buf[len(buf):cap(buf)])
		if errors.Is(err, io.EOF) {
			log.Debug("Reached end of file")
			var rot bool
			rot, err = detectFileRotation(path, file)
			if rot {
				log.Info("File %#v was rotated or truncated; reopening...", path)
				if err := file.Close(); err != nil {
					log.Warn("Error while closing file: %v", err)
				}
				if file, err = os.Open(path); err != nil {
					log.Error("Could not reopen file: %v", err)
					return 1
				}
				continue
			}
			time.Sleep(100 * time.Millisecond)
			continue
		}
		if err != nil {
			log.Error("Unexpected read error: %v", err)
			return 3
		}
		buf = buf[:len(buf)+n]
	}
}

func processLine(line []byte) bool {
	// TODO Allow for more sophisticated processing.
	if Options.regex {
		return Options.patternRE.Match(line)
	}
	return bytes.Contains(line, Options.pattern)
}