circolog/cmd/circolog-tail/main.go
2019-03-25 16:42:21 +01:00

159 lines
3.6 KiB
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(1 * time.Second):
log.Println("Forced close")
}
return
}
}
}