circolog/cmd/circolog-tail/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"time"

	"git.lattuga.net/boyska/circolog/data"
	"git.lattuga.net/boyska/circolog/filtering"
	"git.lattuga.net/boyska/circolog/formatter"
	"github.com/gorilla/websocket"
	isatty "github.com/mattn/go-isatty"
	"github.com/mgutz/ansi"
	"gopkg.in/mgo.v2/bson"
)

type BoolAuto uint

const (
	BoolAuto_NO   BoolAuto = iota
	BoolAuto_YES  BoolAuto = iota
	BoolAuto_AUTO BoolAuto = iota
)

func (b *BoolAuto) String() string {
	switch *b {
	case BoolAuto_NO:
		return "no"
	case BoolAuto_YES:
		return "always"
	case BoolAuto_AUTO:
		return "auto"
	}
	return ""
}
func (b *BoolAuto) Set(s string) error {
	switch s {
	case "auto":
		*b = BoolAuto_AUTO
	case "always":
		*b = BoolAuto_YES
	case "no":
		*b = BoolAuto_NO
	default:
		return fmt.Errorf("Invalid value %s", s)
	}
	return nil
}

func main() {
	addr := flag.String("addr", "localhost:9080", "http service address")
	querySocket := flag.String("socket", "", "Path to a unix domain socket for the HTTP server")
	backlogLimit := flag.Int("n", -1, "Limit the backlog length, defaults to no limit (-1)")
	var format formatter.Format
	format = formatter.FormatSyslog
	flag.Var(&format, "fmt", "Output format [syslog|json]")
	var filter filtering.ExprValue
	flag.Var(&filter, "where", "sql-like query to filter logs")
	// TODO: change to color-mode=auto/no/always
	hasColor := BoolAuto_AUTO
	flag.Var(&hasColor, "color", "dis/enable colors; yes/no/auto")
	flag.Parse()

	if hasColor == BoolAuto_NO || (!isatty.IsTerminal(os.Stdout.Fd()) && hasColor != BoolAuto_YES) {
		ansi.DisableColors(true)
	}

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)
	var d *websocket.Dialer
	u := url.URL{Scheme: "ws",
		Host: *addr, // ignored in case of -socket; see the Dialer below
		Path: "/ws",
	}
	q := u.Query()
	q.Set("fmt", "bson")
	if *backlogLimit >= 0 {
		q.Set("l", strconv.Itoa(*backlogLimit))
	}
	u.RawQuery = q.Encode()
	if *querySocket != "" {
		d = &websocket.Dialer{
			NetDial: func(network, addr string) (net.Conn, error) {
				return net.Dial("unix", *querySocket)
			},
			HandshakeTimeout: 45 * time.Second, // same as DefaultDialer
		}

		log.Printf("connecting to %s", *querySocket)
	} else {
		d = websocket.DefaultDialer
		log.Printf("connecting to %s", *addr)
	}

	c, _, err := d.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer c.Close()
	log.Println("connected!", u.String())

	done := make(chan struct{})

	go func() {
		defer close(done)
		for {
			_, serialized, err := c.ReadMessage()
			if err != nil {
				log.Println("close:", err)
				return
			}
			var parsed data.Message
			if err := bson.Unmarshal(serialized, &parsed); err != nil {
				log.Println("invalid BSON", err)
				continue
			}
			if !filter.Validate(parsed) {
				continue
			}
			if err := formatter.WriteFormatted(os.Stdout, format, parsed); err != nil {
				log.Println("error printing", err)
			}
			if format == formatter.FormatSyslog { // oops
				fmt.Println()
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		case <-interrupt:
			log.Println("interrupt")

			// Cleanly close the connection by sending a close message and then waiting (with timeout) for the
			// server to close the connection.
			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Println("write close:", err)
				return
			}
			select {
			case <-done:
				log.Println("Successfully close")
			case <-time.After(5 * time.Second):
				log.Println("Forced close")
			}
			return
		}
	}

}