<commit_before>package gotak

import (
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"strings"
)

type Diagnostics struct {
	Version      TlsVersion
	CipherSuite  string
	Certificates []*x509.Certificate
	NPN          bool
	NpnStrings   []string
}

func Diagnose(addr string, config *Config) (*Diagnostics, error) {
	if config == nil {
		config = new(Config)
		config.NextProtos = []string{"http/1.1"}
	}

	if !strings.Contains(addr, ":") {
		addr = addr + ":443"
	}

	conn, diag, err := Dial("tcp", addr, config)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	clientConn := httputil.NewClientConn(conn, nil)

	req, err := http.NewRequest("GET", "/favicon.ico", nil)
	if err != nil {
		return nil, err
	}

	_, err = clientConn.Do(req)
	if err != nil {
		return nil, err
	}

	if diag.NpnStrings != nil {
		diag.NPN = true
	}

	return diag, nil
}

func DiagnoseRequest(r *http.Request, config *Config) (*Diagnostics, error) {
	if config == nil {
		config = new(Config)
		config.NextProtos = []string{"http/1.1"}
	}

	if !strings.Contains(addr, ":") {
		addr = addr + ":443"
	}

	conn, diag, err := Dial("tcp", addr, config)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	clientConn := httputil.NewClientConn(conn, nil)

	_, err = clientConn.Do(r)
	if err != nil {
		return nil, err
	}

	if diag.NpnStrings != nil {
		diag.NPN = true
	}

	return diag, nil
}

func (d *Diagnostics) JSON() ([]byte, error) {
	jd := new(jsonDiagnostics)
	jd.Version = d.Version.String()
	jd.CipherSuite = d.CipherSuite
	jd.NPN = d.NpnStrings

	return json.Marshal(jd)
}

func (d *Diagnostics) EncodeJSON(w io.Writer) error {
	jd := new(jsonDiagnostics)
	jd.Version = d.Version.String()
	jd.CipherSuite = d.CipherSuite
	jd.NPN = d.NpnStrings

	return json.NewEncoder(w).Encode(jd)
}

type jsonDiagnostics struct {
	Version     string   `json:"version"`
	CipherSuite string   `json:"cipher_suite"`
	NPN         []string `json:"next_protocol_negotiation,omitempty"`
}

type TlsVersion uint16

const (
	TLS_1_0 TlsVersion = 10
	TLS_1_1            = 11
	TLS_1_2            = 12
)

func (t TlsVersion) String() string {
	switch t {
	case TLS_1_0:
		return "1.0"
	case TLS_1_1:
		return "1.1"
	case TLS_1_2:
		return "1.2"
	default:
		return ""
	}
}

func cryptVersTlsToGotak(vers uint16) (TlsVersion, error) {
	switch vers {
	case VersionSSL30, VersionTLS10:
		return TLS_1_0, nil
	case VersionTLS11:
		return TLS_1_1, nil
	case VersionTLS12:
		return TLS_1_2, nil
	default:
		return 0, fmt.Errorf("Error: Could not parse version %d.", vers)
	}
}
<commit_msg>Minor bug fix<commit_after>package gotak

import (
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"strings"
)

type Diagnostics struct {
	Version      TlsVersion
	CipherSuite  string
	Certificates []*x509.Certificate
	NPN          bool
	NpnStrings   []string
}

func Diagnose(addr string, config *Config) (*Diagnostics, error) {
	if config == nil {
		config = new(Config)
		config.NextProtos = []string{"http/1.1"}
	}

	if !strings.Contains(addr, ":") {
		addr = addr + ":443"
	}

	conn, diag, err := Dial("tcp", addr, config)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	clientConn := httputil.NewClientConn(conn, nil)

	req, err := http.NewRequest("GET", "/favicon.ico", nil)
	if err != nil {
		return nil, err
	}

	_, err = clientConn.Do(req)
	if err != nil {
		return nil, err
	}

	if diag.NpnStrings != nil {
		diag.NPN = true
	}

	return diag, nil
}

func DiagnoseRequest(r *http.Request, config *Config) (*Diagnostics, error) {
	if config == nil {
		config = new(Config)
		config.NextProtos = []string{"http/1.1"}
	}

	addr := r.URL.Host

	if !strings.Contains(addr, ":") {
		addr = addr + ":443"
	}

	conn, diag, err := Dial("tcp", addr, config)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	clientConn := httputil.NewClientConn(conn, nil)

	_, err = clientConn.Do(r)
	if err != nil {
		return nil, err
	}

	if diag.NpnStrings != nil {
		diag.NPN = true
	}

	return diag, nil
}

func (d *Diagnostics) JSON() ([]byte, error) {
	jd := new(jsonDiagnostics)
	jd.Version = d.Version.String()
	jd.CipherSuite = d.CipherSuite
	jd.NPN = d.NpnStrings

	return json.Marshal(jd)
}

func (d *Diagnostics) EncodeJSON(w io.Writer) error {
	jd := new(jsonDiagnostics)
	jd.Version = d.Version.String()
	jd.CipherSuite = d.CipherSuite
	jd.NPN = d.NpnStrings

	return json.NewEncoder(w).Encode(jd)
}

type jsonDiagnostics struct {
	Version     string   `json:"version"`
	CipherSuite string   `json:"cipher_suite"`
	NPN         []string `json:"next_protocol_negotiation,omitempty"`
}

type TlsVersion uint16

const (
	TLS_1_0 TlsVersion = 10
	TLS_1_1            = 11
	TLS_1_2            = 12
)

func (t TlsVersion) String() string {
	switch t {
	case TLS_1_0:
		return "1.0"
	case TLS_1_1:
		return "1.1"
	case TLS_1_2:
		return "1.2"
	default:
		return ""
	}
}

func cryptVersTlsToGotak(vers uint16) (TlsVersion, error) {
	switch vers {
	case VersionSSL30, VersionTLS10:
		return TLS_1_0, nil
	case VersionTLS11:
		return TLS_1_1, nil
	case VersionTLS12:
		return TLS_1_2, nil
	default:
		return 0, fmt.Errorf("Error: Could not parse version %d.", vers)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/pkg/errors"
	"github.com/spf13/afero"
)

// MtailServer contains the state of the main program object.
type MtailServer struct {
	lines chan *tailer.LogLine // Channel of lines from tailer to VM engine.
	store *metrics.Store       // Metrics storage.
	w     watcher.Watcher
	fs    afero.Fs

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	overrideLocation *time.Location // Timezone location to use when parsing timestamps
	bindAddress      string         // address to bind HTTP server
	buildInfo        string         // go build information

	programPath     string   // path to programs to load
	logPathPatterns []string // list of patterns to watch for log files to tail
	logFds          []int    // list of file descriptors to tail

	oneShot      bool // if set, mtail reads log files from the beginning, once, then exits
	compileOnly  bool // if set, mtail compiles programs then exits
	dumpAst      bool // if set, mtail prints the program syntax tree after parse
	dumpAstTypes bool // if set, mtail prints the program syntax tree after type checking
	dumpBytecode bool // if set, mtail prints the program bytecode after code generation

	syslogUseCurrentYear bool // if set, use the current year for timestamps that have no year information
	omitMetricSource     bool // if set, do not link the source program to a metric
	omitProgLabel        bool // if set, do not put the program name in the metric labels
}

// StartTailing constructs a new Tailer and commences sending log lines into
// the lines channel.
func (m *MtailServer) StartTailing() error {
	opts := []func(*tailer.Tailer) error{}
	if m.oneShot {
		opts = append(opts, tailer.OneShot)
	}
	var err error
	m.t, err = tailer.New(m.lines, m.fs, m.w, opts...)
	if err != nil {
		return errors.Wrap(err, "tailer.New")
	}
	for _, pattern := range m.logPathPatterns {
		glog.V(1).Infof("Tail pattern %q", pattern)
		if err = m.t.TailPattern(pattern); err != nil {
			glog.Error(err)
		}
	}
	for _, fd := range m.logFds {
		f := os.NewFile(uintptr(fd), strconv.Itoa(fd))
		if f == nil {
			glog.Errorf("Attempt to reopen fd %q returned nil", fd)
			continue
		}
		if err = m.t.TailHandle(f); err != nil {
			glog.Error(err)
		}
	}
	return nil
}

// InitLoader constructs a new program loader and performs the initial load of program files in the program directory.
func (m *MtailServer) InitLoader() error {
	opts := []func(*vm.Loader) error{}
	if m.compileOnly {
		opts = append(opts, vm.CompileOnly)
		if m.oneShot {
			opts = append(opts, vm.ErrorsAbort)
		}
	}
	if m.dumpAst {
		opts = append(opts, vm.DumpAst)
	}
	if m.dumpAstTypes {
		opts = append(opts, vm.DumpAstTypes)
	}
	if m.dumpBytecode {
		opts = append(opts, vm.DumpBytecode)
	}
	if m.syslogUseCurrentYear {
		opts = append(opts, vm.SyslogUseCurrentYear)
	}
	if m.omitMetricSource {
		opts = append(opts, vm.OmitMetricSource)
	}
	if m.overrideLocation != nil {
		opts = append(opts, vm.OverrideLocation(m.overrideLocation))
	}
	var err error
	m.l, err = vm.NewLoader(m.programPath, m.store, m.lines, m.w, m.fs, opts...)
	if err != nil {
		return err
	}
	if m.programPath == "" {
		return nil
	}
	if errs := m.l.LoadAllPrograms(); errs != nil {
		return errors.Errorf("Compile encountered errors:\n%s", errs)
	}
	return nil
}

// InitExporter sets up an Exporter for this MtailServer.
func (m *MtailServer) InitExporter() (err error) {
	opts := []func(*exporter.Exporter) error{}
	if m.omitProgLabel {
		opts = append(opts, exporter.OmitProgLabel)
	}
	m.e, err = exporter.New(m.store, opts...)
	return
}

const statusTemplate = `

<title>mtail on {{.BindAddress}}</title>

mtail on {{.BindAddress}}

Build: {{.BuildInfo}}

Metrics: json, prometheus, varz

Debug: debug/pprof, debug/vars

`

func (m *MtailServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	t, err := template.New("status").Parse(statusTemplate)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	data := struct {
		BindAddress string
		BuildInfo   string
	}{
		m.bindAddress,
		m.buildInfo,
	}
	w.Header().Add("Content-type", "text/html")
	w.WriteHeader(http.StatusFound)
	if err = t.Execute(w, data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	err = m.l.WriteStatusHTML(w)
	if err != nil {
		glog.Warningf("Error while writing loader status: %s", err)
	}
	err = m.t.WriteStatusHTML(w)
	if err != nil {
		glog.Warningf("Error while writing tailer status: %s", err)
	}
}

// Store sets the metric store in the MtailServer
func Store(store *metrics.Store) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.store = store
		return nil
	}
}

// ProgramPath sets the path to find mtail programs in the MtailServer.
func ProgramPath(path string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.programPath = path
		return nil
	}
}

// LogPathPatterns sets the patterns to find log paths in the MtailServer.
func LogPathPatterns(patterns []string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.logPathPatterns = patterns
		return nil
	}
}

// LogFds sets the file descriptors to read log lines directly in the MtailServer.
func LogFds(fds []int) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.logFds = fds
		return nil
	}
}

// BindAddress sets the HTTP server address in MtailServer.
func BindAddress(address, port string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.bindAddress = net.JoinHostPort(address, port)
		return nil
	}
}

// BuildInfo sets the mtail program build information in the MtailServer.
func BuildInfo(info string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.buildInfo = info
		return nil
	}
}

// OverrideLocation sets the timezone location for log timestamps without any such information.
func OverrideLocation(loc *time.Location) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.overrideLocation = loc
		return nil
	}
}

// OneShot sets one-shot mode in the MtailServer.
func OneShot(m *MtailServer) error {
	m.oneShot = true
	return nil
}

// CompileOnly sets compile-only mode in the MtailServer.
func CompileOnly(m *MtailServer) error {
	m.compileOnly = true
	return nil
}

// DumpAst instructs the MtailServer's compiler to print the AST after parsing.
func DumpAst(m *MtailServer) error {
	m.dumpAst = true
	return nil
}

// DumpAstTypes instructs the MtailServer's copmiler to print the AST after type checking.
func DumpAstTypes(m *MtailServer) error {
	m.dumpAstTypes = true
	return nil
}

// DumpBytecode instructs the MtailServer's compiuler to print the program bytecode after code generation.
func DumpBytecode(m *MtailServer) error {
	m.dumpBytecode = true
	return nil
}

// SyslogUseCurrentYear instructs the MtailServer to use the current year for year-less log timestamp during parsing.
func SyslogUseCurrentYear(m *MtailServer) error {
	m.syslogUseCurrentYear = true
	return nil
}

// OmitProgLabel sets the MtailServer to not put the program name as a label in exported metrics.
func OmitProgLabel(m *MtailServer) error {
	m.omitProgLabel = true
	return nil
}

// OmitMetricSource sets the MtailServer to not link created metrics to their source program.
func OmitMetricSource(m *MtailServer) error {
	m.omitMetricSource = true
	return nil
}

// New creates a MtailServer from the supplied Options.
func New(w watcher.Watcher, fs afero.Fs, options ...func(*MtailServer) error) (*MtailServer, error) {
	m := &MtailServer{
		lines:   make(chan *tailer.LogLine),
		w:       w,
		fs:      fs,
		webquit: make(chan struct{}),
	}
	if err := m.SetOption(options...); err != nil {
		return nil, err
	}
	if err := m.InitExporter(); err != nil {
		return nil, err
	}
	if err := m.InitLoader(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *MtailServer) SetOption(options ...func(*MtailServer) error) error {
	for _, option := range options {
		if err := option(m); err != nil {
			return err
		}
	}
	return nil
}

// WriteMetrics dumps the current state of the metrics store in JSON format to
// the io.Writer.
func (m *MtailServer) WriteMetrics(w io.Writer) error {
	m.store.RLock()
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	m.store.RUnlock()
	if err != nil {
		return errors.Wrap(err, "failed to marshal metrics into json")
	}
	_, err = w.Write(b)
	return err
}

// Serve begins the webserver and awaits a shutdown instruction.
func (m *MtailServer) Serve() error {
	if m.bindAddress == "" {
		return errors.Errorf("No bind address provided.")
	}
	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/varz", http.HandlerFunc(m.e.HandleVarz))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.bindAddress)
		err := http.ListenAndServe(m.bindAddress, nil)
		if err != nil {
			glog.Exit(err)
		}
	}()
	m.WaitForShutdown()
	return nil
}

func (m *MtailServer) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// WaitForShutdown handles shutdown requests from the system or the UI.
func (m *MtailServer) WaitForShutdown() {
	n := make(chan os.Signal, 1)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from HTTP, exiting...")
	}
	if err := m.Close(); err != nil {
		glog.Warning(err)
	}
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *MtailServer) Close() error {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		// If we have a tailer (i.e. not in test) then signal the tailer to
		// shut down, which will cause the watcher to shut down and for the
		// lines channel to close, causing the loader to start shutdown.
		if m.t != nil {
			err := m.t.Close()
			if err != nil {
				glog.Infof("tailer close failed: %s", err)
			}
		} else {
			// Without a tailer, MtailServer has ownership of the lines channel.
			glog.V(2).Info("No tailer, closing lines channel directly.")
			close(m.lines)
		}
		// If we have a loader, wait for it to signal that it has completed shutdown.
		if m.l != nil {
			<-m.l.VMsDone
		} else {
			glog.V(2).Info("No loader, so not waiting for loader shutdown.")
		}
		glog.Info("All done.")
	})
	return nil
}

// Run starts MtailServer's primary function, in which it watches the log
// files for changes and sends any new lines found into the lines channel for
// pick up by the virtual machines. If OneShot mode is enabled, it will exit.
func (m *MtailServer) Run() error {
	if m.compileOnly {
		glog.Info("compile-only is set, exiting")
		return nil
	}
	err := m.StartTailing()
	if err != nil {
		glog.Exitf("tailing failed: %s", err)
	}
	if m.oneShot {
		err := m.Close()
		if err != nil {
			return err
		}
		fmt.Printf("Metrics store:")
		if err := m.WriteMetrics(os.Stdout); err != nil {
			return err
		}
	} else {
		if err := m.Serve(); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Docstring.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/pkg/errors"
	"github.com/spf13/afero"
)

// MtailServer contains the state of the main program object.
type MtailServer struct {
	lines chan *tailer.LogLine // Channel of lines from tailer to VM engine.
	store *metrics.Store       // Metrics storage.
	w     watcher.Watcher
	fs    afero.Fs

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	overrideLocation *time.Location // Timezone location to use when parsing timestamps
	bindAddress      string         // address to bind HTTP server
	buildInfo        string         // go build information

	programPath     string   // path to programs to load
	logPathPatterns []string // list of patterns to watch for log files to tail
	logFds          []int    // list of file descriptors to tail

	oneShot      bool // if set, mtail reads log files from the beginning, once, then exits
	compileOnly  bool // if set, mtail compiles programs then exits
	dumpAst      bool // if set, mtail prints the program syntax tree after parse
	dumpAstTypes bool // if set, mtail prints the program syntax tree after type checking
	dumpBytecode bool // if set, mtail prints the program bytecode after code generation

	syslogUseCurrentYear bool // if set, use the current year for timestamps that have no year information
	omitMetricSource     bool // if set, do not link the source program to a metric
	omitProgLabel        bool // if set, do not put the program name in the metric labels
}

// StartTailing constructs a new Tailer and commences sending log lines into
// the lines channel.
func (m *MtailServer) StartTailing() error {
	opts := []func(*tailer.Tailer) error{}
	if m.oneShot {
		opts = append(opts, tailer.OneShot)
	}
	var err error
	m.t, err = tailer.New(m.lines, m.fs, m.w, opts...)
	if err != nil {
		return errors.Wrap(err, "tailer.New")
	}
	for _, pattern := range m.logPathPatterns {
		glog.V(1).Infof("Tail pattern %q", pattern)
		if err = m.t.TailPattern(pattern); err != nil {
			glog.Error(err)
		}
	}
	for _, fd := range m.logFds {
		f := os.NewFile(uintptr(fd), strconv.Itoa(fd))
		if f == nil {
			glog.Errorf("Attempt to reopen fd %q returned nil", fd)
			continue
		}
		if err = m.t.TailHandle(f); err != nil {
			glog.Error(err)
		}
	}
	return nil
}

// InitLoader constructs a new program loader and performs the initial load of program files in the program directory.
func (m *MtailServer) InitLoader() error {
	opts := []func(*vm.Loader) error{}
	if m.compileOnly {
		opts = append(opts, vm.CompileOnly)
		if m.oneShot {
			opts = append(opts, vm.ErrorsAbort)
		}
	}
	if m.dumpAst {
		opts = append(opts, vm.DumpAst)
	}
	if m.dumpAstTypes {
		opts = append(opts, vm.DumpAstTypes)
	}
	if m.dumpBytecode {
		opts = append(opts, vm.DumpBytecode)
	}
	if m.syslogUseCurrentYear {
		opts = append(opts, vm.SyslogUseCurrentYear)
	}
	if m.omitMetricSource {
		opts = append(opts, vm.OmitMetricSource)
	}
	if m.overrideLocation != nil {
		opts = append(opts, vm.OverrideLocation(m.overrideLocation))
	}
	var err error
	m.l, err = vm.NewLoader(m.programPath, m.store, m.lines, m.w, m.fs, opts...)
	if err != nil {
		return err
	}
	if m.programPath == "" {
		return nil
	}
	if errs := m.l.LoadAllPrograms(); errs != nil {
		return errors.Errorf("Compile encountered errors:\n%s", errs)
	}
	return nil
}

// InitExporter sets up an Exporter for this MtailServer.
func (m *MtailServer) InitExporter() (err error) {
	opts := []func(*exporter.Exporter) error{}
	if m.omitProgLabel {
		opts = append(opts, exporter.OmitProgLabel)
	}
	m.e, err = exporter.New(m.store, opts...)
	return
}

const statusTemplate = `

<title>mtail on {{.BindAddress}}</title>

mtail on {{.BindAddress}}

Build: {{.BuildInfo}}

Metrics: json, prometheus, varz

Debug: debug/pprof, debug/vars

`

func (m *MtailServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	t, err := template.New("status").Parse(statusTemplate)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	data := struct {
		BindAddress string
		BuildInfo   string
	}{
		m.bindAddress,
		m.buildInfo,
	}
	w.Header().Add("Content-type", "text/html")
	w.WriteHeader(http.StatusFound)
	if err = t.Execute(w, data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	err = m.l.WriteStatusHTML(w)
	if err != nil {
		glog.Warningf("Error while writing loader status: %s", err)
	}
	err = m.t.WriteStatusHTML(w)
	if err != nil {
		glog.Warningf("Error while writing tailer status: %s", err)
	}
}

// Store sets the metric store in the MtailServer
func Store(store *metrics.Store) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.store = store
		return nil
	}
}

// ProgramPath sets the path to find mtail programs in the MtailServer.
func ProgramPath(path string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.programPath = path
		return nil
	}
}

// LogPathPatterns sets the patterns to find log paths in the MtailServer.
func LogPathPatterns(patterns []string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.logPathPatterns = patterns
		return nil
	}
}

// LogFds sets the file descriptors to read log lines directly in the MtailServer.
func LogFds(fds []int) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.logFds = fds
		return nil
	}
}

// BindAddress sets the HTTP server address in MtailServer.
func BindAddress(address, port string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.bindAddress = net.JoinHostPort(address, port)
		return nil
	}
}

// BuildInfo sets the mtail program build information in the MtailServer.
func BuildInfo(info string) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.buildInfo = info
		return nil
	}
}

// OverrideLocation sets the timezone location for log timestamps without any such information.
func OverrideLocation(loc *time.Location) func(*MtailServer) error {
	return func(m *MtailServer) error {
		m.overrideLocation = loc
		return nil
	}
}

// OneShot sets one-shot mode in the MtailServer.
func OneShot(m *MtailServer) error {
	m.oneShot = true
	return nil
}

// CompileOnly sets compile-only mode in the MtailServer.
func CompileOnly(m *MtailServer) error {
	m.compileOnly = true
	return nil
}

// DumpAst instructs the MtailServer's compiler to print the AST after parsing.
func DumpAst(m *MtailServer) error {
	m.dumpAst = true
	return nil
}

// DumpAstTypes instructs the MtailServer's copmiler to print the AST after type checking.
func DumpAstTypes(m *MtailServer) error {
	m.dumpAstTypes = true
	return nil
}

// DumpBytecode instructs the MtailServer's compiuler to print the program bytecode after code generation.
func DumpBytecode(m *MtailServer) error {
	m.dumpBytecode = true
	return nil
}

// SyslogUseCurrentYear instructs the MtailServer to use the current year for year-less log timestamp during parsing.
func SyslogUseCurrentYear(m *MtailServer) error {
	m.syslogUseCurrentYear = true
	return nil
}

// OmitProgLabel sets the MtailServer to not put the program name as a label in exported metrics.
func OmitProgLabel(m *MtailServer) error {
	m.omitProgLabel = true
	return nil
}

// OmitMetricSource sets the MtailServer to not link created metrics to their source program.
func OmitMetricSource(m *MtailServer) error {
	m.omitMetricSource = true
	return nil
}

// New creates a MtailServer from the supplied Options.
func New(w watcher.Watcher, fs afero.Fs, options ...func(*MtailServer) error) (*MtailServer, error) {
	m := &MtailServer{
		lines:   make(chan *tailer.LogLine),
		w:       w,
		fs:      fs,
		webquit: make(chan struct{}),
	}
	if err := m.SetOption(options...); err != nil {
		return nil, err
	}
	if err := m.InitExporter(); err != nil {
		return nil, err
	}
	if err := m.InitLoader(); err != nil {
		return nil, err
	}
	return m, nil
}

// SetOption takes one or more option functions and applies them in order to MtailServer.
func (m *MtailServer) SetOption(options ...func(*MtailServer) error) error {
	for _, option := range options {
		if err := option(m); err != nil {
			return err
		}
	}
	return nil
}

// WriteMetrics dumps the current state of the metrics store in JSON format to
// the io.Writer.
func (m *MtailServer) WriteMetrics(w io.Writer) error {
	m.store.RLock()
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	m.store.RUnlock()
	if err != nil {
		return errors.Wrap(err, "failed to marshal metrics into json")
	}
	_, err = w.Write(b)
	return err
}

// Serve begins the webserver and awaits a shutdown instruction.
func (m *MtailServer) Serve() error {
	if m.bindAddress == "" {
		return errors.Errorf("No bind address provided.")
	}
	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/varz", http.HandlerFunc(m.e.HandleVarz))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.bindAddress)
		err := http.ListenAndServe(m.bindAddress, nil)
		if err != nil {
			glog.Exit(err)
		}
	}()
	m.WaitForShutdown()
	return nil
}

func (m *MtailServer) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// WaitForShutdown handles shutdown requests from the system or the UI.
func (m *MtailServer) WaitForShutdown() {
	n := make(chan os.Signal, 1)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from HTTP, exiting...")
	}
	if err := m.Close(); err != nil {
		glog.Warning(err)
	}
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *MtailServer) Close() error {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		// If we have a tailer (i.e. not in test) then signal the tailer to
		// shut down, which will cause the watcher to shut down and for the
		// lines channel to close, causing the loader to start shutdown.
		if m.t != nil {
			err := m.t.Close()
			if err != nil {
				glog.Infof("tailer close failed: %s", err)
			}
		} else {
			// Without a tailer, MtailServer has ownership of the lines channel.
			glog.V(2).Info("No tailer, closing lines channel directly.")
			close(m.lines)
		}
		// If we have a loader, wait for it to signal that it has completed shutdown.
		if m.l != nil {
			<-m.l.VMsDone
		} else {
			glog.V(2).Info("No loader, so not waiting for loader shutdown.")
		}
		glog.Info("All done.")
	})
	return nil
}

// Run starts MtailServer's primary function, in which it watches the log
// files for changes and sends any new lines found into the lines channel for
// pick up by the virtual machines. If OneShot mode is enabled, it will exit.
func (m *MtailServer) Run() error {
	if m.compileOnly {
		glog.Info("compile-only is set, exiting")
		return nil
	}
	err := m.StartTailing()
	if err != nil {
		glog.Exitf("tailing failed: %s", err)
	}
	if m.oneShot {
		err := m.Close()
		if err != nil {
			return err
		}
		fmt.Printf("Metrics store:")
		if err := m.WriteMetrics(os.Stdout); err != nil {
			return err
		}
	} else {
		if err := m.Serve(); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/thoj/go-ircevent"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
	"unicode"
)

const (
	TIME_MESSAGE_FORMAT = "Jan 2 2006 15:04:05"
	TIME_FILE_FORMAT    = "2006-01-02"
	LOG_FILE_DAYS       = 15
)

var commands = make(map[string]func(msg string, from string))
var offlines = make(map[string][]string)

type IRCConfig struct {
	Nick     string
	Username string
	Server   string
	Channels []string
	LogDir   string
	LogFiles map[string]*os.File
}

func main() {
	// FIXME: how much uglier can we get?
	commands[",tell"] = tell

	config, err := getConfig()
	if err != nil {
		fmt.Println(err)
		return
	}

	connection := irc.IRC(config.Nick, config.Username)
	connection.UseTLS = true

	defer closeLogFiles(config)
	addCallbacks(connection, config)

	//Connect to the server
	err = connection.Connect(config.Server)
	if err != nil {
		fmt.Println("Failed to connect.")
		return
	}

	// Connect and wait
	connection.Loop()
}

func parseAndCallCommand(message string, from string) {
	message = strings.TrimSpace(message)
	if !strings.HasPrefix(message, ",") {
		return
	}
	name := strings.FieldsFunc(message, unicode.IsSpace)[0]
	command, ok := commands[name]
	if ok {
		command(strings.TrimSpace(message[len(name):]), from)
	}
}

// Deliver a message to a given user.
// Message is of the form "name <msg>"
// FIXME: allow other seperators!
func tell(message, from string) {
	nick := strings.FieldsFunc(message, unicode.IsSpace)[0]
	message = strings.TrimSpace(message[len(nick):])

	// FIXME: Do we care to check if the user is online/in the channel?
	now := time.Now().UTC()
	message = fmt.Sprintf("At <%s>, %s said -- %s\n", now.Format(TIME_MESSAGE_FORMAT), from, message)

	// JOIN callback looks at the offlines and sends them.
	messages, ok := offlines[nick]
	if ok {
		offlines[nick] = append(messages, message)
	} else {
		messages = []string{message}
		offlines[nick] = messages
	}

}

func sendOfflines(connection *irc.Connection, nick string) {
	for _, message := range offlines[nick] {
		connection.Privmsg(nick, message)
	}
}

func addCallbacks(connection *irc.Connection, config IRCConfig) {

	// Join the channels
	connection.AddCallback("001", func(e *irc.Event) {
		for _, channel := range config.Channels {
			connection.Join(channel)
		}
	})

	connection.AddCallback("JOIN", func(e *irc.Event) {
		channel := e.Arguments[0]
		var message string

		if e.Nick == config.Nick {
			message = fmt.Sprintf("Hello, I'm yet another logbot written in Go.")
			connection.Privmsg(channel, message)
		} else {
			message = fmt.Sprintf("Hello %s, Welcome to %s!", e.Nick, channel)
			connection.Privmsg(channel, message)
			go sendOfflines(connection, e.Nick)
		}

		logMessage(&config, channel, "%s entered %s", e.Nick, channel)
	})

	connection.AddCallback("NICK", func(e *irc.Event) {
		nick := e.Message()
		go sendOfflines(connection, nick)
	})

	connection.AddCallback("PRIVMSG", func(e *irc.Event) {
		channel := e.Arguments[0]
		switch channel {
		case config.Nick:
			connection.Privmsg(e.Nick, "Sorry, I don't accept direct messages!")
		default:
			logMessage(&config, channel, "%s: %s", e.Nick, e.Message())
			go parseAndCallCommand(e.Message(), e.Nick)
		}
	})

	connection.AddCallback("CTCP_ACTION", func(e *irc.Event) {
		channel := e.Arguments[0]
		logMessage(&config, channel, "***%s %s", e.Nick, e.Message())
	})

	connection.AddCallback("PART", func(e *irc.Event) {
		channel := e.Arguments[0]
		logMessage(&config, channel, "%s left %s", e.Nick, channel)
	})

	connection.AddCallback("QUIT", func(e *irc.Event) {
		// FIXME: Quit messages are written to all channels!
		for _, channel := range config.Channels {
			logMessage(&config, channel, "%s quit IRC.", e.Nick)
		}
	})

}

func getConfig() (config IRCConfig, err error) {
	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
	data, err := ioutil.ReadFile(filepath.Join(dir, "config.json"))
	if err != nil {
		return IRCConfig{}, fmt.Errorf("Could not read config file.\n")
	}

	decoder := json.NewDecoder(strings.NewReader(string(data)))
	if err := decoder.Decode(&config); err == io.EOF {
		// Return config ...
	} else if err != nil {
		return IRCConfig{}, fmt.Errorf("Could not read config file.\n")
	}

	config.LogFiles = make(map[string]*os.File)

	return config, nil
}

func closeLogFiles(config IRCConfig) {
	for _, logFile := range config.LogFiles {
		logFile.Close()
	}
}

func logMessage(config *IRCConfig, channel string, format string, args ...interface{}) {
	now := time.Now().UTC()
	logFile := getLogFile(config, channel, now)
	message := fmt.Sprintf("<%s> %s\n", now.Format(TIME_MESSAGE_FORMAT), format)
	message = fmt.Sprintf(message, args...)

	_, err := logFile.WriteString(message)
	if err != nil {
		panic(fmt.Sprintf("Writing is failing %+v\n", err))
	}

}

func getLogFile(config *IRCConfig, channel string, now time.Time) *os.File {
	today := now.Format(TIME_FILE_FORMAT)
	cname := strings.Replace(channel, "#", "", -1)
	logFileName := filepath.Join(config.LogDir, fmt.Sprintf("%s-%s.txt", cname, today))
	var err error

	// FIXME: maps are not safe to use concurrently
	logFile, ok := config.LogFiles[channel]
	switch {

	case ok && logFileName == logFile.Name():
		// Nothing to do.

	default:
		// If existing file open, close it.
		if ok {
			logFile.Close()
		}

		logFile, err = os.OpenFile(logFileName, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			panic(fmt.Sprintf("Log file opening failed: %+v\n", err))
		}

		// FIXME: maps are not safe to use concurrently, but this
		// probably doesn't matter.
		config.LogFiles[channel] = logFile

		// Delete unnecessary files, if any.
		cleanUpLogs(config.LogDir, now, config.Channels)
	}

	return logFile
}

func cleanUpLogs(dir string, now time.Time, channels []string) {
	re_string := fmt.Sprintf("%s-(?P<date>\\d{4}-\\d{2}-\\d{2}).txt", strings.Join(channels, "|"))
	re := regexp.MustCompile(strings.Replace(re_string, "#", "", -1))
	files, _ := ioutil.ReadDir(dir)
	for _, f := range files {
		match := re.FindStringSubmatch(f.Name())
		if match == nil {
			continue
		}

		t, _ := time.Parse(TIME_FILE_FORMAT, match[1])
		if !now.After(t.AddDate(0, 0, LOG_FILE_DAYS)) {
			continue
		}

		fpath := filepath.Join(dir, f.Name())
		fmt.Printf("Deleting log file %s\n", fpath)
		err := os.Remove(fpath)
		if err != nil {
			fmt.Println(err)
		}

	}
}
<commit_msg>Delete messages that were sent.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/thoj/go-ircevent"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
	"unicode"
)

const (
	TIME_MESSAGE_FORMAT = "Jan 2 2006 15:04:05"
	TIME_FILE_FORMAT    = "2006-01-02"
	LOG_FILE_DAYS       = 15
)

var commands = make(map[string]func(msg string, from string))
var offlines = make(map[string][]string)

type IRCConfig struct {
	Nick     string
	Username string
	Server   string
	Channels []string
	LogDir   string
	LogFiles map[string]*os.File
}

func main() {
	// FIXME: how much uglier can we get?
	commands[",tell"] = tell

	config, err := getConfig()
	if err != nil {
		fmt.Println(err)
		return
	}

	connection := irc.IRC(config.Nick, config.Username)
	connection.UseTLS = true

	defer closeLogFiles(config)
	addCallbacks(connection, config)

	//Connect to the server
	err = connection.Connect(config.Server)
	if err != nil {
		fmt.Println("Failed to connect.")
		return
	}

	// Connect and wait
	connection.Loop()
}

func parseAndCallCommand(message string, from string) {
	message = strings.TrimSpace(message)
	if !strings.HasPrefix(message, ",") {
		return
	}
	name := strings.FieldsFunc(message, unicode.IsSpace)[0]
	command, ok := commands[name]
	if ok {
		command(strings.TrimSpace(message[len(name):]), from)
	}
}

// Deliver a message to a given user.
// Message is of the form "name <msg>"
// FIXME: allow other seperators!
func tell(message, from string) {
	nick := strings.FieldsFunc(message, unicode.IsSpace)[0]
	message = strings.TrimSpace(message[len(nick):])

	// FIXME: Do we care to check if the user is online/in the channel?
	now := time.Now().UTC()
	message = fmt.Sprintf("At <%s>, %s said -- %s\n", now.Format(TIME_MESSAGE_FORMAT), from, message)

	// JOIN callback looks at the offlines and sends them.
	messages, ok := offlines[nick]
	if ok {
		offlines[nick] = append(messages, message)
	} else {
		messages = []string{message}
		offlines[nick] = messages
	}

}

func sendOfflines(connection *irc.Connection, nick string) {
	for _, message := range offlines[nick] {
		connection.Privmsg(nick, message)
	}
	// FIXME: What if new messages were added while we were iterating?
	delete(offlines, nick)
}

func addCallbacks(connection *irc.Connection, config IRCConfig) {

	// Join the channels
	connection.AddCallback("001", func(e *irc.Event) {
		for _, channel := range config.Channels {
			connection.Join(channel)
		}
	})

	connection.AddCallback("JOIN", func(e *irc.Event) {
		channel := e.Arguments[0]
		var message string

		if e.Nick == config.Nick {
			message = fmt.Sprintf("Hello, I'm yet another logbot written in Go.")
			connection.Privmsg(channel, message)
		} else {
			message = fmt.Sprintf("Hello %s, Welcome to %s!", e.Nick, channel)
			connection.Privmsg(channel, message)
			go sendOfflines(connection, e.Nick)
		}

		logMessage(&config, channel, "%s entered %s", e.Nick, channel)
	})

	connection.AddCallback("NICK", func(e *irc.Event) {
		nick := e.Message()
		go sendOfflines(connection, nick)
	})

	connection.AddCallback("PRIVMSG", func(e *irc.Event) {
		channel := e.Arguments[0]
		switch channel {
		case config.Nick:
			connection.Privmsg(e.Nick, "Sorry, I don't accept direct messages!")
		default:
			logMessage(&config, channel, "%s: %s", e.Nick, e.Message())
			go parseAndCallCommand(e.Message(), e.Nick)
		}
	})

	connection.AddCallback("CTCP_ACTION", func(e *irc.Event) {
		channel := e.Arguments[0]
		logMessage(&config, channel, "***%s %s", e.Nick, e.Message())
	})

	connection.AddCallback("PART", func(e *irc.Event) {
		channel := e.Arguments[0]
		logMessage(&config, channel, "%s left %s", e.Nick, channel)
	})

	connection.AddCallback("QUIT", func(e *irc.Event) {
		// FIXME: Quit messages are written to all channels!
		for _, channel := range config.Channels {
			logMessage(&config, channel, "%s quit IRC.", e.Nick)
		}
	})

}

func getConfig() (config IRCConfig, err error) {
	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
	data, err := ioutil.ReadFile(filepath.Join(dir, "config.json"))
	if err != nil {
		return IRCConfig{}, fmt.Errorf("Could not read config file.\n")
	}

	decoder := json.NewDecoder(strings.NewReader(string(data)))
	if err := decoder.Decode(&config); err == io.EOF {
		// Return config ...
	} else if err != nil {
		return IRCConfig{}, fmt.Errorf("Could not read config file.\n")
	}

	config.LogFiles = make(map[string]*os.File)

	return config, nil
}

func closeLogFiles(config IRCConfig) {
	for _, logFile := range config.LogFiles {
		logFile.Close()
	}
}

func logMessage(config *IRCConfig, channel string, format string, args ...interface{}) {
	now := time.Now().UTC()
	logFile := getLogFile(config, channel, now)
	message := fmt.Sprintf("<%s> %s\n", now.Format(TIME_MESSAGE_FORMAT), format)
	message = fmt.Sprintf(message, args...)

	_, err := logFile.WriteString(message)
	if err != nil {
		panic(fmt.Sprintf("Writing is failing %+v\n", err))
	}

}

func getLogFile(config *IRCConfig, channel string, now time.Time) *os.File {
	today := now.Format(TIME_FILE_FORMAT)
	cname := strings.Replace(channel, "#", "", -1)
	logFileName := filepath.Join(config.LogDir, fmt.Sprintf("%s-%s.txt", cname, today))
	var err error

	// FIXME: maps are not safe to use concurrently
	logFile, ok := config.LogFiles[channel]
	switch {

	case ok && logFileName == logFile.Name():
		// Nothing to do.

	default:
		// If existing file open, close it.
		if ok {
			logFile.Close()
		}

		logFile, err = os.OpenFile(logFileName, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			panic(fmt.Sprintf("Log file opening failed: %+v\n", err))
		}

		// FIXME: maps are not safe to use concurrently, but this
		// probably doesn't matter.
		config.LogFiles[channel] = logFile

		// Delete unnecessary files, if any.
		cleanUpLogs(config.LogDir, now, config.Channels)
	}

	return logFile
}

func cleanUpLogs(dir string, now time.Time, channels []string) {
	re_string := fmt.Sprintf("%s-(?P<date>\\d{4}-\\d{2}-\\d{2}).txt", strings.Join(channels, "|"))
	re := regexp.MustCompile(strings.Replace(re_string, "#", "", -1))
	files, _ := ioutil.ReadDir(dir)
	for _, f := range files {
		match := re.FindStringSubmatch(f.Name())
		if match == nil {
			continue
		}

		t, _ := time.Parse(TIME_FILE_FORMAT, match[1])
		if !now.After(t.AddDate(0, 0, LOG_FILE_DAYS)) {
			continue
		}

		fpath := filepath.Join(dir, f.Name())
		fmt.Printf("Deleting log file %s\n", fpath)
		err := os.Remove(fpath)
		if err != nil {
			fmt.Println(err)
		}

	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/spf13/afero"
)

// Mtail contains the state of the main program object.
type Mtail struct {
	lines chan string    // Channel of lines from tailer to VM engine.
	store *metrics.Store // Metrics storage.

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	o Options // Options passed in at creation time.
}

// OneShot reads the contents of a log file into the lines channel from start to finish, terminating the program at the end.
func (m *Mtail) OneShot(logfile string, print bool) (count int64, err error) {
	glog.Infof("Oneshot %q", logfile)
	l, err := os.Open(logfile)
	if err != nil {
		return 0, fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	defer l.Close()

	r := bufio.NewReader(l)

	if print {
		fmt.Printf("%s: %d MAXPROCS, %d CPUs, ", logfile, runtime.GOMAXPROCS(-1), runtime.NumCPU())
	}

	start := time.Now()

Loop:
	for {
		line, err := r.ReadString('\n')
		line = strings.TrimSuffix(line, "\n")
		switch {
		case err == io.EOF:
			if len(line) > 0 {
				m.lines <- line
			}
			break Loop
		case err != nil:
			return 0, fmt.Errorf("failed to read from %q: %s", logfile, err)
		default:
			m.lines <- line
		}
	}
	duration := time.Since(start)
	count, err = strconv.ParseInt(vm.LineCount.String(), 10, 64)
	if err != nil {
		return
	}
	if print {
		µsPerL := float64(duration.Nanoseconds()) / (float64(count) * 1000)
		fmt.Printf("%d lines, %s total time, %6.3f µs/line\n", count, duration, µsPerL)
	}
	return
}

// StartTailing constructs a new Tailer and commences sending log lines into
// the lines channel.
func (m *Mtail) StartTailing() error {
	o := tailer.Options{Lines: m.lines, W: m.o.W, FS: m.o.FS}
	var err error
	m.t, err = tailer.New(o)
	if err != nil {
		return fmt.Errorf("couldn't create a log tailer: %s", err)
	}

	for _, pathname := range m.o.LogPaths {
		m.t.Tail(pathname)
	}
	for _, fd := range m.o.LogFds {
		f := os.NewFile(uintptr(fd), strconv.Itoa(fd))
		if f == nil {
			glog.Error("Attempt to reopen fd %q returned nil", fd)
			continue
		}
		if e := m.t.TailFile(f); e != nil {
			glog.Error(e)
		}
	}
	return nil
}

// InitLoader constructs a new program loader and performs the inital load of program files in the program directory.
func (m *Mtail) InitLoader() error {
	o := vm.LoaderOptions{Store: m.store, Lines: m.lines, CompileOnly: m.o.CompileOnly, DumpBytecode: m.o.DumpBytecode, SyslogUseCurrentYear: m.o.SyslogUseCurrentYear, W: m.o.W, FS: m.o.FS}
	var err error
	m.l, err = vm.NewLoader(o)
	if err != nil {
		return err
	}
	if m.o.Progs != "" {
		errors := m.l.LoadProgs(m.o.Progs)
		if errors != nil {
			return fmt.Errorf("Compile encountered errors:\n%s", errors)
		}
	}
	return nil
}

func (m *Mtail) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte(`<a href="/json">json</a>, <a href="/metrics">prometheus metrics</a>, <a href="/varz">varz</a>`))
}

// Options contains all the parameters necessary for constructing a new Mtail.
type Options struct {
	Progs                string
	LogPaths             []string
	LogFds               []int
	Port                 string
	OneShot              bool
	OneShotMetrics       bool
	CompileOnly          bool
	DumpBytecode         bool
	SyslogUseCurrentYear bool

	Store *metrics.Store

	W  watcher.Watcher // Not required, will use watcher.LogWatcher if zero.
	FS afero.Fs        // Not required, will use afero.OsFs if zero.
}

// New creates an Mtail from the supplied Options.
func New(o Options) (*Mtail, error) {
	store := o.Store
	if store == nil {
		store = metrics.NewStore()
	}
	m := &Mtail{
		lines:   make(chan string),
		store:   store,
		webquit: make(chan struct{}),
		o:       o}

	err := m.InitLoader()
	if err != nil {
		return nil, err
	}

	m.e, err = exporter.New(exporter.Options{Store: m.store})
	if err != nil {
		return nil, err
	}

	return m, nil
}

// WriteMetrics dumps the current state of the metrics store in JSON format to
// the io.Writer.
func (m *Mtail) WriteMetrics(w io.Writer) error {
	m.store.RLock()
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	m.store.RUnlock()
	if err != nil {
		return fmt.Errorf("failed to marshal metrics into json: %s", err)
	}
	w.Write(b)
	return nil
}

// RunOneShot performs the work of the one_shot commandline flag; after compiling programs mtail will read all of the log files in full, once, dump the metric results at the end, and then exit.
func (m *Mtail) RunOneShot() {
	fmt.Println("Oneshot results:")
	for _, pathname := range m.o.LogPaths {
		_, err := m.OneShot(pathname, true)
		if err != nil {
			glog.Exitf("Failed one shot mode for %q: %s\n", pathname, err)
		}
	}
	if m.o.OneShotMetrics {
		fmt.Printf("Metrics store:")
		if err := m.WriteMetrics(os.Stdout); err != nil {
			glog.Exit(err)
		}
	}
	m.Close()
}

// Serve begins the long-running mode of mtail, in which it watches the log
// files for changes and sends any new lines found into the lines channel for
// pick up by the virtual machines.  It will continue to do so until it is
// signalled to exit.
func (m *Mtail) Serve() {
	err := m.StartTailing()
	if err != nil {
		glog.Exitf("tailing failed: %s", err)
	}

	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/varz", http.HandlerFunc(m.e.HandleVarz))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.o.Port)
		err := http.ListenAndServe(":"+m.o.Port, nil)
		if err != nil {
			glog.Exit(err)
		}
	}()
	m.shutdownHandler()
}

func (m *Mtail) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// shutdownHandler handles external shutdown request events.
func (m *Mtail) shutdownHandler() {
	n := make(chan os.Signal)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from UI, exiting...")
	}
	m.Close()
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *Mtail) Close() {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		if m.t != nil {
			m.t.Close()
		} else {
			glog.Info("No tailer, closing lines channel.")
			close(m.lines)
		}
		if m.l != nil {
			<-m.l.VMsDone
		}
		glog.Info("All done.")
	})
}

// Run starts Mtail in the configuration supplied in Options at creation.
func (m *Mtail) Run() {
	if m.o.OneShot {
		m.RunOneShot()
	} else {
		m.Serve()
	}
}
<commit_msg>Replace glog.Error with glog.Errorf<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/spf13/afero"
)

// Mtail contains the state of the main program object.
type Mtail struct {
	lines chan string    // Channel of lines from tailer to VM engine.
	store *metrics.Store // Metrics storage.

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	o Options // Options passed in at creation time.
}

// OneShot reads the contents of a log file into the lines channel from start to finish, terminating the program at the end.
func (m *Mtail) OneShot(logfile string, print bool) (count int64, err error) {
	glog.Infof("Oneshot %q", logfile)
	l, err := os.Open(logfile)
	if err != nil {
		return 0, fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	defer l.Close()

	r := bufio.NewReader(l)

	if print {
		fmt.Printf("%s: %d MAXPROCS, %d CPUs, ", logfile, runtime.GOMAXPROCS(-1), runtime.NumCPU())
	}

	start := time.Now()

Loop:
	for {
		line, err := r.ReadString('\n')
		line = strings.TrimSuffix(line, "\n")
		switch {
		case err == io.EOF:
			if len(line) > 0 {
				m.lines <- line
			}
			break Loop
		case err != nil:
			return 0, fmt.Errorf("failed to read from %q: %s", logfile, err)
		default:
			m.lines <- line
		}
	}
	duration := time.Since(start)
	count, err = strconv.ParseInt(vm.LineCount.String(), 10, 64)
	if err != nil {
		return
	}
	if print {
		µsPerL := float64(duration.Nanoseconds()) / (float64(count) * 1000)
		fmt.Printf("%d lines, %s total time, %6.3f µs/line\n", count, duration, µsPerL)
	}
	return
}

// StartTailing constructs a new Tailer and commences sending log lines into
// the lines channel.
func (m *Mtail) StartTailing() error {
	o := tailer.Options{Lines: m.lines, W: m.o.W, FS: m.o.FS}
	var err error
	m.t, err = tailer.New(o)
	if err != nil {
		return fmt.Errorf("couldn't create a log tailer: %s", err)
	}

	for _, pathname := range m.o.LogPaths {
		m.t.Tail(pathname)
	}
	for _, fd := range m.o.LogFds {
		f := os.NewFile(uintptr(fd), strconv.Itoa(fd))
		if f == nil {
			glog.Errorf("Attempt to reopen fd %q returned nil", fd)
			continue
		}
		if e := m.t.TailFile(f); e != nil {
			glog.Error(e)
		}
	}
	return nil
}

// InitLoader constructs a new program loader and performs the inital load of program files in the program directory.
func (m *Mtail) InitLoader() error {
	o := vm.LoaderOptions{Store: m.store, Lines: m.lines, CompileOnly: m.o.CompileOnly, DumpBytecode: m.o.DumpBytecode, SyslogUseCurrentYear: m.o.SyslogUseCurrentYear, W: m.o.W, FS: m.o.FS}
	var err error
	m.l, err = vm.NewLoader(o)
	if err != nil {
		return err
	}
	if m.o.Progs != "" {
		errors := m.l.LoadProgs(m.o.Progs)
		if errors != nil {
			return fmt.Errorf("Compile encountered errors:\n%s", errors)
		}
	}
	return nil
}

func (m *Mtail) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte(`<a href="/json">json</a>, <a href="/metrics">prometheus metrics</a>, <a href="/varz">varz</a>`))
}

// Options contains all the parameters necessary for constructing a new Mtail.
type Options struct {
	Progs                string
	LogPaths             []string
	LogFds               []int
	Port                 string
	OneShot              bool
	OneShotMetrics       bool
	CompileOnly          bool
	DumpBytecode         bool
	SyslogUseCurrentYear bool

	Store *metrics.Store

	W  watcher.Watcher // Not required, will use watcher.LogWatcher if zero.
	FS afero.Fs        // Not required, will use afero.OsFs if zero.
}

// New creates an Mtail from the supplied Options.
func New(o Options) (*Mtail, error) {
	store := o.Store
	if store == nil {
		store = metrics.NewStore()
	}
	m := &Mtail{
		lines:   make(chan string),
		store:   store,
		webquit: make(chan struct{}),
		o:       o}

	err := m.InitLoader()
	if err != nil {
		return nil, err
	}

	m.e, err = exporter.New(exporter.Options{Store: m.store})
	if err != nil {
		return nil, err
	}

	return m, nil
}

// WriteMetrics dumps the current state of the metrics store in JSON format to
// the io.Writer.
func (m *Mtail) WriteMetrics(w io.Writer) error {
	m.store.RLock()
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	m.store.RUnlock()
	if err != nil {
		return fmt.Errorf("failed to marshal metrics into json: %s", err)
	}
	w.Write(b)
	return nil
}

// RunOneShot performs the work of the one_shot commandline flag; after compiling programs mtail will read all of the log files in full, once, dump the metric results at the end, and then exit.
func (m *Mtail) RunOneShot() {
	fmt.Println("Oneshot results:")
	for _, pathname := range m.o.LogPaths {
		_, err := m.OneShot(pathname, true)
		if err != nil {
			glog.Exitf("Failed one shot mode for %q: %s\n", pathname, err)
		}
	}
	if m.o.OneShotMetrics {
		fmt.Printf("Metrics store:")
		if err := m.WriteMetrics(os.Stdout); err != nil {
			glog.Exit(err)
		}
	}
	m.Close()
}

// Serve begins the long-running mode of mtail, in which it watches the log
// files for changes and sends any new lines found into the lines channel for
// pick up by the virtual machines.  It will continue to do so until it is
// signalled to exit.
func (m *Mtail) Serve() {
	err := m.StartTailing()
	if err != nil {
		glog.Exitf("tailing failed: %s", err)
	}

	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/varz", http.HandlerFunc(m.e.HandleVarz))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.o.Port)
		err := http.ListenAndServe(":"+m.o.Port, nil)
		if err != nil {
			glog.Exit(err)
		}
	}()
	m.shutdownHandler()
}

func (m *Mtail) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// shutdownHandler handles external shutdown request events.
func (m *Mtail) shutdownHandler() {
	n := make(chan os.Signal)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from UI, exiting...")
	}
	m.Close()
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *Mtail) Close() {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		if m.t != nil {
			m.t.Close()
		} else {
			glog.Info("No tailer, closing lines channel.")
			close(m.lines)
		}
		if m.l != nil {
			<-m.l.VMsDone
		}
		glog.Info("All done.")
	})
}

// Run starts Mtail in the configuration supplied in Options at creation.
func (m *Mtail) Run() {
	if m.o.OneShot {
		m.RunOneShot()
	} else {
		m.Serve()
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"

	_ "net/http/pprof"
)

type mtail struct {
	lines chan string   // Channel of lines from tailer to VM engine.
	store metrics.Store // Metrics storage.

	pathnames []string // pathnames of logs to tail
	progs     string   // directory path containing mital programs to load
	port      string   // port to serve HTTP on

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	oneShot              bool
	compileOnly          bool
	dumpBytecode         bool
	syslogUseCurrentYear bool
}

func (m *mtail) OneShot(logfile string, lines chan string) error {
	defer m.Close()
	l, err := os.Open(logfile)
	if err != nil {
		return fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	defer l.Close()

	r := bufio.NewReader(l)

	for {
		line, err := r.ReadString('\n')
		switch {
		case err == io.EOF:
			return nil
		case err != nil:
			return fmt.Errorf("failed to read from %q: %s", logfile, err)
		default:
			lines <- line
		}
	}
}

func (m *mtail) StartTailing() {
	o := tailer.Options{Lines: m.lines}
	m.t = tailer.New(o)
	if m.t == nil {
		glog.Fatal("Couldn't create a log tailer.")
	}

	for _, pathname := range m.pathnames {
		m.t.Tail(pathname)
	}
}

func (m *mtail) InitLoader() {
	o := vm.LoaderOptions{Store: &m.store, Lines: m.lines, CompileOnly: m.compileOnly, DumpBytecode: m.dumpBytecode, SyslogUseCurrentYear: m.syslogUseCurrentYear}
	m.l = vm.NewLoader(o)
	if m.l == nil {
		glog.Fatal("Couldn't create a program loader.")
	}
	errors := m.l.LoadProgs(m.progs)
	if m.compileOnly || m.dumpBytecode {
		os.Exit(errors)
	}
}

func (m *mtail) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte(`<a href="/json">json</a>, <a href="/metrics">prometheus metrics</a>`))
}

func NewMtail() *mtail {
	return &mtail{
		lines:   make(chan string),
		webquit: make(chan struct{}),
	}
}

type Options struct {
	Progs                string
	Logs                 string
	Port                 string
	OneShot              bool
	CompileOnly          bool
	DumpBytecode         bool
	SyslogUseCurrentYear bool
}

func New(o Options) *mtail {
	if o.Progs == "" {
		glog.Fatalf("No mtail program directory specified; use -progs")
		return nil
	}
	if o.Logs == "" {
		glog.Fatalf("No logs specified to tail; use -logs")
		return nil
	}
	m := NewMtail()
	m.progs = o.Progs
	m.port = o.Port
	m.oneShot = o.OneShot
	m.compileOnly = o.CompileOnly
	m.dumpBytecode = o.DumpBytecode
	m.syslogUseCurrentYear = o.SyslogUseCurrentYear

	for _, pathname := range strings.Split(o.Logs, ",") {
		if pathname != "" {
			m.pathnames = append(m.pathnames, pathname)
		}
	}
	if len(m.pathnames) == 0 {
		glog.Fatal("No logs to tail.")
		return nil
	}

	m.InitLoader()

	m.e = exporter.New(exporter.Options{Store: &m.store})

	return m
}

func (m *mtail) RunOneShot() {
	for _, pathname := range m.pathnames {
		err := m.OneShot(pathname, m.lines)
		if err != nil {
			glog.Fatalf("Failed one shot mode for %q: %s\n", pathname, err)
		}
	}
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	if err != nil {
		glog.Fatalf("Failed to marshal metrics into json: %s", err)
	}
	os.Stdout.Write(b)
	m.e.WriteMetrics()
}

func (m *mtail) Serve() {
	m.StartTailing()

	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.port)
		err := http.ListenAndServe(":"+m.port, nil)
		if err != nil {
			glog.Fatal(err)
		}
	}()
	m.shutdownHandler()
}

func (m *mtail) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// shutdownHandler handles external shutdown request events.
func (m *mtail) shutdownHandler() {
	n := make(chan os.Signal)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from UI, exiting...")
	}
	m.Close()
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *mtail) Close() {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		if m.t != nil {
			m.t.Close()
		} else {
			glog.Info("Closing lines")
			close(m.lines)
		}
		if m.l != nil {
			<-m.l.VMsDone
		}
		glog.Info("All done.")
	})
}

func (m *mtail) Run() {
	if m.oneShot {
		m.RunOneShot()
	} else {
		m.Serve()
	}
}
<commit_msg>Remove pprof from mtail.go's imports.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"

	"github.com/golang/glog"
	"github.com/google/mtail/exporter"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/vm"
)

type mtail struct {
	lines chan string   // Channel of lines from tailer to VM engine.
	store metrics.Store // Metrics storage.

	pathnames []string // pathnames of logs to tail
	progs     string   // directory path containing mital programs to load
	port      string   // port to serve HTTP on

	t *tailer.Tailer     // t tails the watched files and feeds lines to the VMs.
	l *vm.Loader         // l loads programs and manages the VM lifecycle.
	e *exporter.Exporter // e manages the export of metrics from the store.

	webquit   chan struct{} // Channel to signal shutdown from web UI.
	closeOnce sync.Once     // Ensure shutdown happens only once.

	oneShot              bool
	compileOnly          bool
	dumpBytecode         bool
	syslogUseCurrentYear bool
}

func (m *mtail) OneShot(logfile string, lines chan string) error {
	defer m.Close()
	l, err := os.Open(logfile)
	if err != nil {
		return fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	defer l.Close()

	r := bufio.NewReader(l)

	for {
		line, err := r.ReadString('\n')
		switch {
		case err == io.EOF:
			return nil
		case err != nil:
			return fmt.Errorf("failed to read from %q: %s", logfile, err)
		default:
			lines <- line
		}
	}
}

func (m *mtail) StartTailing() {
	o := tailer.Options{Lines: m.lines}
	m.t = tailer.New(o)
	if m.t == nil {
		glog.Fatal("Couldn't create a log tailer.")
	}

	for _, pathname := range m.pathnames {
		m.t.Tail(pathname)
	}
}

func (m *mtail) InitLoader() {
	o := vm.LoaderOptions{Store: &m.store, Lines: m.lines, CompileOnly: m.compileOnly, DumpBytecode: m.dumpBytecode, SyslogUseCurrentYear: m.syslogUseCurrentYear}
	m.l = vm.NewLoader(o)
	if m.l == nil {
		glog.Fatal("Couldn't create a program loader.")
	}
	errors := m.l.LoadProgs(m.progs)
	if m.compileOnly || m.dumpBytecode {
		os.Exit(errors)
	}
}

func (m *mtail) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte(`<a href="/json">json</a>, <a href="/metrics">prometheus metrics</a>`))
}

func NewMtail() *mtail {
	return &mtail{
		lines:   make(chan string),
		webquit: make(chan struct{}),
	}
}

type Options struct {
	Progs                string
	Logs                 string
	Port                 string
	OneShot              bool
	CompileOnly          bool
	DumpBytecode         bool
	SyslogUseCurrentYear bool
}

func New(o Options) *mtail {
	if o.Progs == "" {
		glog.Fatalf("No mtail program directory specified; use -progs")
		return nil
	}
	if o.Logs == "" {
		glog.Fatalf("No logs specified to tail; use -logs")
		return nil
	}
	m := NewMtail()
	m.progs = o.Progs
	m.port = o.Port
	m.oneShot = o.OneShot
	m.compileOnly = o.CompileOnly
	m.dumpBytecode = o.DumpBytecode
	m.syslogUseCurrentYear = o.SyslogUseCurrentYear

	for _, pathname := range strings.Split(o.Logs, ",") {
		if pathname != "" {
			m.pathnames = append(m.pathnames, pathname)
		}
	}
	if len(m.pathnames) == 0 {
		glog.Fatal("No logs to tail.")
		return nil
	}

	m.InitLoader()

	m.e = exporter.New(exporter.Options{Store: &m.store})

	return m
}

func (m *mtail) RunOneShot() {
	for _, pathname := range m.pathnames {
		err := m.OneShot(pathname, m.lines)
		if err != nil {
			glog.Fatalf("Failed one shot mode for %q: %s\n", pathname, err)
		}
	}
	b, err := json.MarshalIndent(m.store.Metrics, "", "  ")
	if err != nil {
		glog.Fatalf("Failed to marshal metrics into json: %s", err)
	}
	os.Stdout.Write(b)
	m.e.WriteMetrics()
}

func (m *mtail) Serve() {
	m.StartTailing()

	http.Handle("/", m)
	http.HandleFunc("/json", http.HandlerFunc(m.e.HandleJSON))
	http.HandleFunc("/metrics", http.HandlerFunc(m.e.HandlePrometheusMetrics))
	http.HandleFunc("/quitquitquit", http.HandlerFunc(m.handleQuit))
	m.e.StartMetricPush()

	go func() {
		glog.Infof("Listening on port %s", m.port)
		err := http.ListenAndServe(":"+m.port, nil)
		if err != nil {
			glog.Fatal(err)
		}
	}()
	m.shutdownHandler()
}

func (m *mtail) handleQuit(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Header().Add("Allow", "POST")
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	fmt.Fprintf(w, "Exiting...")
	close(m.webquit)
}

// shutdownHandler handles external shutdown request events.
func (m *mtail) shutdownHandler() {
	n := make(chan os.Signal)
	signal.Notify(n, os.Interrupt, syscall.SIGTERM)
	select {
	case <-n:
		glog.Info("Received SIGTERM, exiting...")
	case <-m.webquit:
		glog.Info("Received Quit from UI, exiting...")
	}
	m.Close()
}

// Close handles the graceful shutdown of this mtail instance, ensuring that it only occurs once.
func (m *mtail) Close() {
	m.closeOnce.Do(func() {
		glog.Info("Shutdown requested.")
		if m.t != nil {
			m.t.Close()
		} else {
			glog.Info("Closing lines channel.")
			close(m.lines)
		}
		if m.l != nil {
			<-m.l.VMsDone
		}
		glog.Info("All done.")
	})
}

func (m *mtail) Run() {
	if m.oneShot {
		m.RunOneShot()
	} else {
		m.Serve()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	//"log"
	"sort"
	"strconv"
	"strings"

	"github.com/gonum/stat"
)

var (
	colors = []string{
		"rgba(31,119,180,1)",
		"rgba(255,127,14,1)",
		"rgba(44,160,44,1)",
		"rgba(214,39,40,1)",
		"rgba(148,103,189,1)",
		"rgba(140,86,75,1)",
		"rgba(227,119,194,1)",
	}
)

const (
	maxBallNum       = 59
	balls            = 7
	avgMarkerSize    = 8
	regressionLinear = "linear"
	regressionPoly   = "polynomial"
	graphTypeScatter = "scatter"
	graphTypeBar     = "bar"
	graphTypeLine    = "line"
)

type dataset struct {
	X           []string  `json:"x"`
	Y           []float64 `json:"y"`
	Z           []int     `json:"z"` // number
	Name        string    `json:"name"`
	Mode        string    `json:"mode"`
	Type        string    `json:"type"`
	Line        line      `json:"line"`
	Marker      marker    `json:"marker"`
	ConnectGaps bool      `json:"connectgaps"`
}

type datasetB struct {
	X           []string  `json:"x"`
	Y           []float64 `json:"y"`
	Z           []float64 `json:"z"` // number
	Name        string    `json:"name"`
	Mode        string    `json:"mode"`
	Type        string    `json:"type"`
	Line        line      `json:"line"`
	Marker      markerB   `json:"marker"`
	ConnectGaps bool      `json:"connectgaps"`
}

type marker struct {
	Colour  string  `json:"color"`
	Size    float64 `json:"size"`
	Line    line    `json:"line"`
	Opacity float64 `json:"opacity"`
	Symbol  string  `json:"symbol"`
}

type markerB struct {
	Colour   string    `json:"color"`
	Size     []float64 `json:"size"`
	SizeMode string    `json:"sizemode"`
	SizeRef  float64   `json:"sizeref"`
	Line     line      `json:"line"`
	Opacity  float64   `json:"opacity"`
	Symbol   string    `json:"symbol"`
}

type line struct {
	Width   float64 `json:"width"`
	Colour  string  `json:"color"`
	Shape   string  `json:"shape"`
	Dash    string  `json:"dash"`
	Opacity float64 `json:"opacity"`
}

func graphResultsTimeSeries(records <-chan dbRow, bestFit bool, t string) []dataset {
	data := make([]dataset, balls)

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {

				switch t {
				case graphTypeScatter:
					data[ball] = dataset{
						Mode: "markers",
						Marker: marker{
							Colour:  colors[ball],
							Size:    avgMarkerSize,
							Opacity: 1,
						},
					}

				case graphTypeLine:
					data[ball] = dataset{
						Mode: "lines",
						Line: line{
							Width: 1,
						},
					}
				} // END SWITCH

				data[ball].Name = label(ball)
			}
			data[ball].X = append(data[ball].X, fmt.Sprintf("%s:%d:%s", row.Date.Format(formatYYYYMMDD), row.Set, row.Machine))
			data[ball].Y = append(data[ball].Y, float64(row.Num[ball]))
		}

		i++
	}

	if bestFit && t == graphTypeScatter {
		data = append(data, regressionSets(data, regressionLinear)...)
	}

	return data
}

func graphResultsFreqDist(records <-chan dbRow, bestFit bool, t string) []dataset {
	data := make([]dataset, balls)

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {

				switch t {
				case graphTypeScatter:
					data[ball] = dataset{
						Mode: "markers",
						Marker: marker{
							Colour:  colors[ball],
							Size:    avgMarkerSize,
							Opacity: 1,
						},
						X: freqDistXLabels(),
						Y: make([]float64, maxBallNum),
					}

				case graphTypeBar:
					data[ball] = dataset{
						Type: graphTypeBar,
						X:    freqDistXLabels(),
						Y:    make([]float64, maxBallNum),
					}

				} // END SWITCH

				data[ball].Name = label(ball)
			}

			data[ball].Y[row.Num[ball]-1]++
		}
		i++
	}

	/*
		if bestFit && t != graphTypeBar {
			data = append(data, regressionSets(data, regressionPoly)...)
		}
	*/

	return data
}

func graphResultsRawScatter3D(records <-chan dbRow) []dataset {
	data := make([]dataset, balls)

	// x: machine
	// y: set
	// z: ball result

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {
				data[ball] = dataset{
					Mode: "markers",
					Type: "scatter3d",
					Marker: marker{
						Size:    avgMarkerSize / 2,
						Opacity: 0.9,
						Line:    line{Width: 0.1},
					},
				}

				data[ball].Name = label(ball)
			}
			data[ball].X = append(data[ball].X, row.Machine)
			data[ball].Y = append(data[ball].Y, float64(row.Set))
			data[ball].Z = append(data[ball].Z, row.Num[ball])
		}
		i++
	}

	return data
}

func graphMSFreqDistScatter3D(m map[string]int) []dataset {
	data := dataset{
		Type: "scatter3d",
		Mode: "markers",
		Marker: marker{
			Size:    avgMarkerSize,
			Opacity: 0.9,
			Line: line{
				Width: 0.1,
			},
		},
	}

	l := []string{}
	for k := range m {
		l = append(l, k)
	}
	sort.Strings(l)

	for _, k := range l {
		s := strings.Split(k, ":")
		y, _ := strconv.Atoi(s[1])

		data.X = append(data.X, s[0])       // MAchine
		data.Y = append(data.Y, float64(y)) // Set
		data.Z = append(data.Z, m[k])       // Frequency
	}

	return []dataset{data}
}

func graphMSFreqDistBubble(m map[string]int) []datasetB {
	data := datasetB{
		Type: "scatter",
		Mode: "markers",
		Marker: markerB{
			Opacity: 0.8,
			Line:    line{Width: 0.1},
		},
	}

	l := []string{}
	for k := range m {
		l = append(l, k)
	}
	sort.Strings(l)

	for _, k := range l {
		s := strings.Split(k, ":")
		z, _ := strconv.Atoi(s[1])

		data.X = append(data.X, s[0])
		data.Y = append(data.Y, float64(z))
		data.Marker.Size = append(data.Marker.Size, float64(m[k])*2)
	}

	return []datasetB{data}
}

func regressionSets(data []dataset, t string) []dataset {
	// Calculate and append regressionLinear regressions for each set
	r, rX := make([]dataset, balls), make([]float64, len(data[0].Y))

	// Generate numerical X axis data
	for i := range rX {
		rX[i] = float64(i) + 1
	}

	// Iterate existing sets and create new regression sets
	for i, set := range data {
		r[i] = dataset{
			Name: set.Name,
			Mode: "lines",
			Line: line{
				Dash:   "dot",
				Width:  2,
				Colour: colors[i],
			},
			ConnectGaps: true,
			X:           set.X,
		}

		switch t {
		case regressionLinear:
			r[i].Y = linRegYData(rX, set.Y, false)
		case regressionPoly:
			r[i].Y = polRegYData(rX, set.Y, false) // Not currently working because I suck at maths
		}
	}

	return r
}

func polRegYData(prX, prY []float64, subzero bool) []float64 {
	// @TODO: polynomial regression sets for non-linear best fits
	y := make([]float64, len(prX))
	//a, b := stat.LinearRegression(prX, prY, nil, false)
	//r2 := stat.RSquared(prX, prY, nil, a, b)
	//log.Println(r2)

	return y
}

func linRegYData(lrX, lrY []float64, subzero bool) []float64 {
	a, b := stat.LinearRegression(lrX, lrY, nil, false)

	y := make([]float64, len(lrX))
	for idx, x := range lrX {
		n := a + (b * x)
		if subzero {
			y[idx] = n
		} else if n < 0 {
			y[idx] = 0
		} else {
			y[idx] = n
		}

	}
	return y
}

func freqDistXLabels() []string {
	var x []string
	for i := 0; i < maxBallNum; i++ {
		x = append(x, fmt.Sprintf("%d", i+1))
	}
	return x
}

func label(ball int) string {
	if ball < 6 {
		return fmt.Sprintf("Ball %d", ball+1)
	}

	return "Bonus"
}
<commit_msg>spliney mcsplineface<commit_after>package main

import (
	"fmt"
	//"log"
	"sort"
	"strconv"
	"strings"

	"github.com/gonum/stat"
)

var (
	colors = []string{
		"rgba(31,119,180,1)",
		"rgba(255,127,14,1)",
		"rgba(44,160,44,1)",
		"rgba(214,39,40,1)",
		"rgba(148,103,189,1)",
		"rgba(140,86,75,1)",
		"rgba(227,119,194,1)",
	}
)

const (
	maxBallNum       = 59
	balls            = 7
	avgMarkerSize    = 8
	regressionLinear = "linear"
	regressionPoly   = "polynomial"
	graphTypeScatter = "scatter"
	graphTypeBar     = "bar"
	graphTypeLine    = "line"
)

type dataset struct {
	X           []string  `json:"x"`
	Y           []float64 `json:"y"`
	Z           []int     `json:"z"` // number
	Name        string    `json:"name"`
	Mode        string    `json:"mode"`
	Type        string    `json:"type"`
	Line        line      `json:"line"`
	Marker      marker    `json:"marker"`
	ConnectGaps bool      `json:"connectgaps"`
}

type datasetB struct {
	X           []string  `json:"x"`
	Y           []float64 `json:"y"`
	Z           []float64 `json:"z"` // number
	Name        string    `json:"name"`
	Mode        string    `json:"mode"`
	Type        string    `json:"type"`
	Line        line      `json:"line"`
	Marker      markerB   `json:"marker"`
	ConnectGaps bool      `json:"connectgaps"`
}

type marker struct {
	Colour  string  `json:"color"`
	Size    float64 `json:"size"`
	Line    line    `json:"line"`
	Opacity float64 `json:"opacity"`
	Symbol  string  `json:"symbol"`
}

type markerB struct {
	Colour   string    `json:"color"`
	Size     []float64 `json:"size"`
	SizeMode string    `json:"sizemode"`
	SizeRef  float64   `json:"sizeref"`
	Line     line      `json:"line"`
	Opacity  float64   `json:"opacity"`
	Symbol   string    `json:"symbol"`
}

type line struct {
	Width   float64 `json:"width"`
	Colour  string  `json:"color"`
	Shape   string  `json:"shape"`
	Dash    string  `json:"dash"`
	Opacity float64 `json:"opacity"`
}

func graphResultsTimeSeries(records <-chan dbRow, bestFit bool, t string) []dataset {
	data := make([]dataset, balls)

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {

				switch t {
				case graphTypeScatter:
					data[ball] = dataset{
						Mode: "markers",
						Marker: marker{
							Colour:  colors[ball],
							Size:    avgMarkerSize,
							Opacity: 1,
						},
					}

				case graphTypeLine:
					data[ball] = dataset{
						Mode: "lines",
						Line: line{
							Width: 1,
							Shape: "spline",
						},
					}
				} // END SWITCH

				data[ball].Name = label(ball)
			}
			data[ball].X = append(data[ball].X, fmt.Sprintf("%s:%d:%s", row.Date.Format(formatYYYYMMDD), row.Set, row.Machine))
			data[ball].Y = append(data[ball].Y, float64(row.Num[ball]))
		}

		i++
	}

	if bestFit && t == graphTypeScatter {
		data = append(data, regressionSets(data, regressionLinear)...)
	}

	return data
}

func graphResultsFreqDist(records <-chan dbRow, bestFit bool, t string) []dataset {
	data := make([]dataset, balls)

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {

				switch t {
				case graphTypeScatter:
					data[ball] = dataset{
						Mode: "markers",
						Marker: marker{
							Colour:  colors[ball],
							Size:    avgMarkerSize,
							Opacity: 1,
						},
						X: freqDistXLabels(),
						Y: make([]float64, maxBallNum),
					}

				case graphTypeBar:
					data[ball] = dataset{
						Type: graphTypeBar,
						X:    freqDistXLabels(),
						Y:    make([]float64, maxBallNum),
					}

				} // END SWITCH

				data[ball].Name = label(ball)
			}

			data[ball].Y[row.Num[ball]-1]++
		}
		i++
	}

	/*
		if bestFit && t != graphTypeBar {
			data = append(data, regressionSets(data, regressionPoly)...)
		}
	*/

	return data
}

func graphResultsRawScatter3D(records <-chan dbRow) []dataset {
	data := make([]dataset, balls)

	// x: machine
	// y: set
	// z: ball result

	i := 0
	for row := range records {
		for ball := 0; ball < balls; ball++ {
			if i == 0 {
				data[ball] = dataset{
					Mode: "markers",
					Type: "scatter3d",
					Marker: marker{
						Size:    avgMarkerSize / 2,
						Opacity: 0.9,
						Line:    line{Width: 0.1},
					},
				}

				data[ball].Name = label(ball)
			}
			data[ball].X = append(data[ball].X, row.Machine)
			data[ball].Y = append(data[ball].Y, float64(row.Set))
			data[ball].Z = append(data[ball].Z, row.Num[ball])
		}
		i++
	}

	return data
}

func graphMSFreqDistScatter3D(m map[string]int) []dataset {
	data := dataset{
		Type: "scatter3d",
		Mode: "markers",
		Marker: marker{
			Size:    avgMarkerSize,
			Opacity: 0.9,
			Line: line{
				Width: 0.1,
			},
		},
	}

	l := []string{}
	for k := range m {
		l = append(l, k)
	}
	sort.Strings(l)

	for _, k := range l {
		s := strings.Split(k, ":")
		y, _ := strconv.Atoi(s[1])

		data.X = append(data.X, s[0])       // MAchine
		data.Y = append(data.Y, float64(y)) // Set
		data.Z = append(data.Z, m[k])       // Frequency
	}

	return []dataset{data}
}

func graphMSFreqDistBubble(m map[string]int) []datasetB {
	data := datasetB{
		Type: "scatter",
		Mode: "markers",
		Marker: markerB{
			Opacity: 0.8,
			Line:    line{Width: 0.1},
		},
	}

	l := []string{}
	for k := range m {
		l = append(l, k)
	}
	sort.Strings(l)

	for _, k := range l {
		s := strings.Split(k, ":")
		z, _ := strconv.Atoi(s[1])

		data.X = append(data.X, s[0])
		data.Y = append(data.Y, float64(z))
		data.Marker.Size = append(data.Marker.Size, float64(m[k])*2)
	}

	return []datasetB{data}
}

func regressionSets(data []dataset, t string) []dataset {
	// Calculate and append regressionLinear regressions for each set
	r, rX := make([]dataset, balls), make([]float64, len(data[0].Y))

	// Generate numerical X axis data
	for i := range rX {
		rX[i] = float64(i) + 1
	}

	// Iterate existing sets and create new regression sets
	for i, set := range data {
		r[i] = dataset{
			Name: set.Name,
			Mode: "lines",
			Line: line{
				Dash:   "dot",
				Width:  2,
				Colour: colors[i],
			},
			ConnectGaps: true,
			X:           set.X,
		}

		switch t {
		case regressionLinear:
			r[i].Y = linRegYData(rX, set.Y, false)
		case regressionPoly:
			r[i].Y = polRegYData(rX, set.Y, false) // Not currently working because I suck at maths
		}
	}

	return r
}

func polRegYData(prX, prY []float64, subzero bool) []float64 {
	// @TODO: polynomial regression sets for non-linear best fits
	y := make([]float64, len(prX))
	//a, b := stat.LinearRegression(prX, prY, nil, false)
	//r2 := stat.RSquared(prX, prY, nil, a, b)
	//log.Println(r2)

	return y
}

func linRegYData(lrX, lrY []float64, subzero bool) []float64 {
	a, b := stat.LinearRegression(lrX, lrY, nil, false)

	y := make([]float64, len(lrX))
	for idx, x := range lrX {
		n := a + (b * x)
		if subzero {
			y[idx] = n
		} else if n < 0 {
			y[idx] = 0
		} else {
			y[idx] = n
		}

	}
	return y
}

func freqDistXLabels() []string {
	var x []string
	for i := 0; i < maxBallNum; i++ {
		x = append(x, fmt.Sprintf("%d", i+1))
	}
	return x
}

func label(ball int) string {
	if ball < 6 {
		return fmt.Sprintf("Ball %d", ball+1)
	}

	return "Bonus"
}
<|endoftext|>


<commit_before>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package graph

// All a node needs to do is identify itself. This allows the user to pass in nodes more
// interesting than an int, but also allow us to reap the benefits of having a map-storable,
// comparable type.
type Node interface {
	ID() int
}

// Allows edges to do something more interesting that just be a group of nodes. While the methods
// are called Head and Tail, they are not considered directed unless the given interface specifies
// otherwise.
type Edge interface {
	Head() Node
	Tail() Node
}

// A Graph implements the behavior of an undirected graph.
//
// All methods in Graph are implicitly undirected. Graph algorithms that care about directionality
// will intelligently choose the DirectedGraph behavior if that interface is also implemented,
// even if the function itself only takes in a Graph (or a super-interface of graph).
type Graph interface {
	// NodeExists returns true when node is currently in the graph.
	NodeExists(Node) bool

	// NodeList returns a list of all nodes in no particular order, useful for
	// determining things like if a graph is fully connected. The caller is
	// free to modify this list. Implementations should construct a new list
	// and not return internal representation.
	NodeList() []Node

	// Neighbors returns all nodes connected by any edge to this node.
	Neighbors(Node) []Node

	// EdgeBetween returns an edge between node and neighbor such that
	// Head is one argument and Tail is the other. If no
	// such edge exists, this function returns nil.
	EdgeBetween(node, neighbor Node) Edge
}

// Directed graphs are characterized by having seperable Heads and Tails in their edges.
// That is, if node1 goes to node2, that does not necessarily imply that node2 goes to node1.
//
// While it's possible for a directed graph to have fully reciprocal edges (i.e. the graph is
// symmetric) -- it is not required to be. The graph is also required to implement Graph
// because in many cases it can be useful to know all neighbors regardless of direction.
type DirectedGraph interface {
	Graph
	// Successors gives the nodes connected by OUTBOUND edges.
	// If the graph is an undirected graph, this set is equal to Predecessors.
	Successors(Node) []Node

	// EdgeTo returns an edge between node and successor such that
	// Head returns node and Tail returns successor, if no
	// such edge exists, this function returns nil.
	EdgeTo(node, successor Node) Edge

	// Predecessors gives the nodes connected by INBOUND edges.
	// If the graph is an undirected graph, this set is equal to Successors.
	Predecessors(Node) []Node
}

// Returns all undirected edges in the graph
type EdgeLister interface {
	EdgeList() []Edge
}

type EdgeListGraph interface {
	Graph
	EdgeLister
}

// Returns all directed edges in the graph.
type DirectedEdgeLister interface {
	DirectedEdgeList() []Edge
}

type DirectedEdgeListGraph interface {
	Graph
	DirectedEdgeLister
}

// A crunch graph forces a sparse graph to become a dense graph. That is, if the node IDs are
// [1,4,9,7] it would "crunch" the ids into the contiguous block [0,1,2,3]. Order is not
// required to be preserved between the non-cruched and crunched instances (that means in
// the example above 0 may correspond to 4 or 7 or 9, not necessarily 1).
//
// All dense graphs must have the first ID as 0.
type CrunchGraph interface {
	Graph
	Crunch()
}

// A Graph that implements Coster has an actual cost between adjacent nodes, also known as a
// weighted graph. If a graph implements coster and a function needs to read cost (e.g. A*),
// this function will take precedence over the Uniform Cost function (all weights are 1) if "nil"
// is passed in for the function argument.
//
// If the argument is nil, or the edge is invalid for some reason, this should return math.Inf(1)
type Coster interface {
	Cost(edge Edge) float64
}

type CostGraph interface {
	Coster
	Graph
}

type CostDirectedGraph interface {
	Coster
	DirectedGraph
}

// A graph that implements HeuristicCoster implements a heuristic between any two given nodes.
// Like Coster, if a graph implements this and a function needs a heuristic cost (e.g. A*), this
// function will take precedence over the Null Heuristic (always returns 0) if "nil" is passed in
// for the function argument. If HeuristicCost is not intended to be used, it can be implemented as
// the null heuristic (always returns 0)
type HeuristicCoster interface {
	// HeuristicCost returns a heuristic cost between any two nodes.
	HeuristicCost(n1, n2 Node) float64
}

// A Mutable is a graph that can have arbitrary nodes and edges added or removed.
//
// Anything implementing Mutable is required to store the actual argument. So if AddNode(myNode) is
// called and later a user calls on the graph graph.NodeList(), the node added by AddNode must be
// an the exact node, not a new node with the same ID.
//
// In any case where conflict is possible (e.g. adding two nodes with the same ID), the later
// call always supercedes the earlier one.
//
// Functions will generally expect one of MutableGraph or MutableDirectedGraph and not Mutable
// itself. That said, any function that takes Mutable[x], the destination mutable should
// always be a different graph than the source.
type Mutable interface {
	// NewNode adds a node with an arbitrary ID and returns the new, unique ID
	// used.
	NewNode() Node

	// Adds a node to the graph. If this is called multiple times for the same ID, the newer node
	// overwrites the old one.
	AddNode(Node)

	// RemoveNode removes a node from the graph, as well as any edges
	// attached to it. If no such node exists, this is a no-op, not an error.
	RemoveNode(Node)
}

// MutableGraph is an interface ensuring the implementation of the ability to construct
// an arbitrary undirected graph. It is very important to note that any implementation
// of MutableGraph absolutely cannot safely implement the DirectedGraph interface.
//
// A MutableGraph is required to store any Edge argument in the same way Mutable must
// store a Node argument -- any retrieval call is required to return the exact supplied edge.
// This is what makes it incompatible with DirectedGraph.
//
// The reasoning is this: if you call AddUndirectedEdge(Edge{head,tail}); you are required
// to return the exact edge passed in when a retrieval method (EdgeTo/EdgeBetween) is called.
// If I call EdgeTo(tail,head), this means that since the edge exists, and was added as
// Edge{head,tail} this function MUST return Edge{head,tail}. However, EdgeTo requires this
// be returned as Edge{tail,head}. Thus there's a conflict that cannot be resolved between the
// two interface requirements.
type MutableGraph interface {
	CostGraph
	Mutable

	// Like EdgeBetween in Graph, AddUndirectedEdge adds an edge between two nodes.
	// If one or both nodes do not exist, the graph is expected to add them. However,
	// if the nodes already exist it should NOT replace existing nodes with e.Head() or
	// e.Tail(). Overwriting nodes should explicitly be done with another call to AddNode()
	AddUndirectedEdge(e Edge, cost float64)

	// RemoveEdge clears the stored edge between two nodes. Calling this will never
	// remove a node. If the edge does not exist this is a no-op, not an error.
	RemoveUndirectedEdge(Edge)
}

// MutableDirectedGraph is an interface that ensures one can construct an arbitrary directed
// graph. Naturally, a MutableDirectedGraph works for both undirected and directed cases,
// but simply using a MutableGraph may be cleaner. As the documentation for MutableGraph
// notes, however, a graph cannot safely implement MutableGraph and MutableDirectedGraph
// at the same time, because of the functionality of a EdgeTo in DirectedGraph.
type MutableDirectedGraph interface {
	CostDirectedGraph
	Mutable

	// Like EdgeTo in DirectedGraph, AddDirectedEdge adds an edge FROM head TO tail.
	// If one or both nodes do not exist, the graph is expected to add them. However,
	// if the nodes already exist it should NOT replace existing nodes with e.Head() or
	// e.Tail(). Overwriting nodes should explicitly be done with another call to AddNode()
	AddDirectedEdge(e Edge, cost float64)

	// Removes an edge FROM e.Head TO e.Tail. If no such edge exists, this is a no-op,
	// not an error.
	RemoveDirectedEdge(Edge)
}

// A function that returns the cost of following an edge
type CostFunc func(Edge) float64

// Estimates the cost of travelling between two nodes
type HeuristicCostFunc func(Node, Node) float64
<commit_msg>Elide unused name<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package graph

// All a node needs to do is identify itself. This allows the user to pass in nodes more
// interesting than an int, but also allow us to reap the benefits of having a map-storable,
// comparable type.
type Node interface {
	ID() int
}

// Allows edges to do something more interesting that just be a group of nodes. While the methods
// are called Head and Tail, they are not considered directed unless the given interface specifies
// otherwise.
type Edge interface {
	Head() Node
	Tail() Node
}

// A Graph implements the behavior of an undirected graph.
//
// All methods in Graph are implicitly undirected. Graph algorithms that care about directionality
// will intelligently choose the DirectedGraph behavior if that interface is also implemented,
// even if the function itself only takes in a Graph (or a super-interface of graph).
type Graph interface {
	// NodeExists returns true when node is currently in the graph.
	NodeExists(Node) bool

	// NodeList returns a list of all nodes in no particular order, useful for
	// determining things like if a graph is fully connected. The caller is
	// free to modify this list. Implementations should construct a new list
	// and not return internal representation.
	NodeList() []Node

	// Neighbors returns all nodes connected by any edge to this node.
	Neighbors(Node) []Node

	// EdgeBetween returns an edge between node and neighbor such that
	// Head is one argument and Tail is the other. If no
	// such edge exists, this function returns nil.
	EdgeBetween(node, neighbor Node) Edge
}

// Directed graphs are characterized by having seperable Heads and Tails in their edges.
// That is, if node1 goes to node2, that does not necessarily imply that node2 goes to node1.
//
// While it's possible for a directed graph to have fully reciprocal edges (i.e. the graph is
// symmetric) -- it is not required to be. The graph is also required to implement Graph
// because in many cases it can be useful to know all neighbors regardless of direction.
type DirectedGraph interface {
	Graph
	// Successors gives the nodes connected by OUTBOUND edges.
	// If the graph is an undirected graph, this set is equal to Predecessors.
	Successors(Node) []Node

	// EdgeTo returns an edge between node and successor such that
	// Head returns node and Tail returns successor, if no
	// such edge exists, this function returns nil.
	EdgeTo(node, successor Node) Edge

	// Predecessors gives the nodes connected by INBOUND edges.
	// If the graph is an undirected graph, this set is equal to Successors.
	Predecessors(Node) []Node
}

// Returns all undirected edges in the graph
type EdgeLister interface {
	EdgeList() []Edge
}

type EdgeListGraph interface {
	Graph
	EdgeLister
}

// Returns all directed edges in the graph.
type DirectedEdgeLister interface {
	DirectedEdgeList() []Edge
}

type DirectedEdgeListGraph interface {
	Graph
	DirectedEdgeLister
}

// A crunch graph forces a sparse graph to become a dense graph. That is, if the node IDs are
// [1,4,9,7] it would "crunch" the ids into the contiguous block [0,1,2,3]. Order is not
// required to be preserved between the non-cruched and crunched instances (that means in
// the example above 0 may correspond to 4 or 7 or 9, not necessarily 1).
//
// All dense graphs must have the first ID as 0.
type CrunchGraph interface {
	Graph
	Crunch()
}

// A Graph that implements Coster has an actual cost between adjacent nodes, also known as a
// weighted graph. If a graph implements coster and a function needs to read cost (e.g. A*),
// this function will take precedence over the Uniform Cost function (all weights are 1) if "nil"
// is passed in for the function argument.
//
// If the argument is nil, or the edge is invalid for some reason, this should return math.Inf(1)
type Coster interface {
	Cost(Edge) float64
}

type CostGraph interface {
	Coster
	Graph
}

type CostDirectedGraph interface {
	Coster
	DirectedGraph
}

// A graph that implements HeuristicCoster implements a heuristic between any two given nodes.
// Like Coster, if a graph implements this and a function needs a heuristic cost (e.g. A*), this
// function will take precedence over the Null Heuristic (always returns 0) if "nil" is passed in
// for the function argument. If HeuristicCost is not intended to be used, it can be implemented as
// the null heuristic (always returns 0)
type HeuristicCoster interface {
	// HeuristicCost returns a heuristic cost between any two nodes.
	HeuristicCost(n1, n2 Node) float64
}

// A Mutable is a graph that can have arbitrary nodes and edges added or removed.
//
// Anything implementing Mutable is required to store the actual argument. So if AddNode(myNode) is
// called and later a user calls on the graph graph.NodeList(), the node added by AddNode must be
// an the exact node, not a new node with the same ID.
//
// In any case where conflict is possible (e.g. adding two nodes with the same ID), the later
// call always supercedes the earlier one.
//
// Functions will generally expect one of MutableGraph or MutableDirectedGraph and not Mutable
// itself. That said, any function that takes Mutable[x], the destination mutable should
// always be a different graph than the source.
type Mutable interface {
	// NewNode adds a node with an arbitrary ID and returns the new, unique ID
	// used.
	NewNode() Node

	// Adds a node to the graph. If this is called multiple times for the same ID, the newer node
	// overwrites the old one.
	AddNode(Node)

	// RemoveNode removes a node from the graph, as well as any edges
	// attached to it. If no such node exists, this is a no-op, not an error.
	RemoveNode(Node)
}

// MutableGraph is an interface ensuring the implementation of the ability to construct
// an arbitrary undirected graph. It is very important to note that any implementation
// of MutableGraph absolutely cannot safely implement the DirectedGraph interface.
//
// A MutableGraph is required to store any Edge argument in the same way Mutable must
// store a Node argument -- any retrieval call is required to return the exact supplied edge.
// This is what makes it incompatible with DirectedGraph.
//
// The reasoning is this: if you call AddUndirectedEdge(Edge{head,tail}); you are required
// to return the exact edge passed in when a retrieval method (EdgeTo/EdgeBetween) is called.
// If I call EdgeTo(tail,head), this means that since the edge exists, and was added as
// Edge{head,tail} this function MUST return Edge{head,tail}. However, EdgeTo requires this
// be returned as Edge{tail,head}. Thus there's a conflict that cannot be resolved between the
// two interface requirements.
type MutableGraph interface {
	CostGraph
	Mutable

	// Like EdgeBetween in Graph, AddUndirectedEdge adds an edge between two nodes.
	// If one or both nodes do not exist, the graph is expected to add them. However,
	// if the nodes already exist it should NOT replace existing nodes with e.Head() or
	// e.Tail(). Overwriting nodes should explicitly be done with another call to AddNode()
	AddUndirectedEdge(e Edge, cost float64)

	// RemoveEdge clears the stored edge between two nodes. Calling this will never
	// remove a node. If the edge does not exist this is a no-op, not an error.
	RemoveUndirectedEdge(Edge)
}

// MutableDirectedGraph is an interface that ensures one can construct an arbitrary directed
// graph. Naturally, a MutableDirectedGraph works for both undirected and directed cases,
// but simply using a MutableGraph may be cleaner. As the documentation for MutableGraph
// notes, however, a graph cannot safely implement MutableGraph and MutableDirectedGraph
// at the same time, because of the functionality of a EdgeTo in DirectedGraph.
type MutableDirectedGraph interface {
	CostDirectedGraph
	Mutable

	// Like EdgeTo in DirectedGraph, AddDirectedEdge adds an edge FROM head TO tail.
	// If one or both nodes do not exist, the graph is expected to add them. However,
	// if the nodes already exist it should NOT replace existing nodes with e.Head() or
	// e.Tail(). Overwriting nodes should explicitly be done with another call to AddNode()
	AddDirectedEdge(e Edge, cost float64)

	// Removes an edge FROM e.Head TO e.Tail. If no such edge exists, this is a no-op,
	// not an error.
	RemoveDirectedEdge(Edge)
}

// A function that returns the cost of following an edge
type CostFunc func(Edge) float64

// Estimates the cost of travelling between two nodes
type HeuristicCostFunc func(Node, Node) float64
<|endoftext|>