<commit_before>// Copyright (c) 2016 Pulcy.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package jobs

import (
	"fmt"
	"strings"

	"github.com/juju/errgo"
)

// LinkName is a name of a link consisting of:
// <job>.<task>[@<instance>] or
// <job>.<taskgroup>.<task>[@instance]
type LinkName string

// NewLinkName assembles a link name from its elements.
func NewLinkName(jn JobName, tgn TaskGroupName, tn TaskName, in InstanceName) LinkName {
	result := fmt.Sprintf("%s.%s.%s", jn, tgn, tn)
	if !in.IsEmpty() {
		result = fmt.Sprintf("%s@%s", result, in)
	}
	return LinkName(result)
}

// String returns a link name in format <job>.<taskgroup>.<task>
func (ln LinkName) String() string {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ""
	}
	return string(NewLinkName(jn, tgn, tn, in))
}

// PrivateDomainName returns the DNS name (in the private namespace) for the given link name.
func (ln LinkName) PrivateDomainName() string {
	return fmt.Sprintf("%s.private", strings.Replace(ln.String(), "@", ".", -1))
}

// WeaveDomainName returns the DNS name (in the weave namespace) for the given link name.
func (ln LinkName) WeaveDomainName() string {
	return fmt.Sprintf("%s.weave.local", ln.EtcdServiceName())
}

// EtcdServiceName returns name of the service as it is used in ETCD.
func (ln LinkName) EtcdServiceName() string {
	return strings.Replace(strings.Replace(ln.String(), ".", "-", -1), "@", "-", -1)
}

// Validate checks if a link name follows a valid format
func (ln LinkName) Validate() error {
	_, _, _, _, err := ln.parse()
	return maskAny(err)
}

// Job returns the job part of the link name.
func (ln LinkName) Job() (JobName, error) {
	jn, _, _, _, err := ln.parse()
	return jn, maskAny(err)
}

// TaskGroup returns the taskgroup part of the link name.
func (ln LinkName) TaskGroup() (TaskGroupName, error) {
	_, tg, _, _, err := ln.parse()
	return tg, maskAny(err)
}

// Task returns the task part of the link name.
func (ln LinkName) Task() (TaskName, error) {
	_, _, tn, _, err := ln.parse()
	return tn, maskAny(err)
}

// Instance returns the instance part of the link name.
func (ln LinkName) Instance() (InstanceName, error) {
	_, _, _, in, err := ln.parse()
	return in, maskAny(err)
}

// HasInstance returns true if there is a specific instance in link name.
func (ln LinkName) HasInstance() bool {
	_, _, _, in, _ := ln.parse()
	return !in.IsEmpty()
}

func (ln LinkName) parse() (JobName, TaskGroupName, TaskName, InstanceName, error) {
	var instanceName InstanceName
	parts := strings.Split(string(ln), "@")
	switch len(parts) {
	case 1:
		// Empty instance name
	case 2:
		instanceName = InstanceName(parts[1])
	default:
		return "", "", "", "", maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}

	parts = strings.Split(parts[0], ".")
	var jobName JobName
	var taskGroupName TaskGroupName
	var taskName TaskName
	switch len(parts) {
	case 2:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[1])
	case 3:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[2])
	default:
		return jobName, taskGroupName, taskName, instanceName,
			maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}
	if err := jobName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskGroupName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := instanceName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	return jobName, taskGroupName, taskName, instanceName, nil
}

func (ln LinkName) normalize() LinkName {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ln
	}
	return NewLinkName(jn, tgn, tn, in)
}
<commit_msg>Fixed weave name for LinkNames with `_` in it<commit_after>// Copyright (c) 2016 Pulcy.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package jobs

import (
	"fmt"
	"strings"

	"github.com/juju/errgo"
)

// LinkName is a name of a link consisting of:
// <job>.<task>[@<instance>] or
// <job>.<taskgroup>.<task>[@instance]
type LinkName string

// NewLinkName assembles a link name from its elements.
func NewLinkName(jn JobName, tgn TaskGroupName, tn TaskName, in InstanceName) LinkName {
	result := fmt.Sprintf("%s.%s.%s", jn, tgn, tn)
	if !in.IsEmpty() {
		result = fmt.Sprintf("%s@%s", result, in)
	}
	return LinkName(result)
}

// String returns a link name in format <job>.<taskgroup>.<task>
func (ln LinkName) String() string {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ""
	}
	return string(NewLinkName(jn, tgn, tn, in))
}

// PrivateDomainName returns the DNS name (in the private namespace) for the given link name.
func (ln LinkName) PrivateDomainName() string {
	return fmt.Sprintf("%s.private", strings.Replace(ln.String(), "@", ".", -1))
}

// WeaveDomainName returns the DNS name (in the weave namespace) for the given link name.
func (ln LinkName) WeaveDomainName() string {
	name := strings.Replace(ln.EtcdServiceName(), "_", "-", -1)
	return fmt.Sprintf("%s.weave.local", name)
}

// EtcdServiceName returns name of the service as it is used in ETCD.
func (ln LinkName) EtcdServiceName() string {
	return strings.Replace(strings.Replace(ln.String(), ".", "-", -1), "@", "-", -1)
}

// Validate checks if a link name follows a valid format
func (ln LinkName) Validate() error {
	_, _, _, _, err := ln.parse()
	return maskAny(err)
}

// Job returns the job part of the link name.
func (ln LinkName) Job() (JobName, error) {
	jn, _, _, _, err := ln.parse()
	return jn, maskAny(err)
}

// TaskGroup returns the taskgroup part of the link name.
func (ln LinkName) TaskGroup() (TaskGroupName, error) {
	_, tg, _, _, err := ln.parse()
	return tg, maskAny(err)
}

// Task returns the task part of the link name.
func (ln LinkName) Task() (TaskName, error) {
	_, _, tn, _, err := ln.parse()
	return tn, maskAny(err)
}

// Instance returns the instance part of the link name.
func (ln LinkName) Instance() (InstanceName, error) {
	_, _, _, in, err := ln.parse()
	return in, maskAny(err)
}

// HasInstance returns true if there is a specific instance in link name.
func (ln LinkName) HasInstance() bool {
	_, _, _, in, _ := ln.parse()
	return !in.IsEmpty()
}

func (ln LinkName) parse() (JobName, TaskGroupName, TaskName, InstanceName, error) {
	var instanceName InstanceName
	parts := strings.Split(string(ln), "@")
	switch len(parts) {
	case 1:
		// Empty instance name
	case 2:
		instanceName = InstanceName(parts[1])
	default:
		return "", "", "", "", maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}

	parts = strings.Split(parts[0], ".")
	var jobName JobName
	var taskGroupName TaskGroupName
	var taskName TaskName
	switch len(parts) {
	case 2:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[1])
	case 3:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[2])
	default:
		return jobName, taskGroupName, taskName, instanceName,
			maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}
	if err := jobName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskGroupName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := instanceName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	return jobName, taskGroupName, taskName, instanceName, nil
}

func (ln LinkName) normalize() LinkName {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ln
	}
	return NewLinkName(jn, tgn, tn, in)
}
<|endoftext|>


<commit_before>package go_config

import (
	"bytes"
	"encoding/json"
	"fmt"
	toml "github.com/pelletier/go-toml"
	yaml "gopkg.in/yaml.v2"
	"io"
)

type ConfigType int

type ConfigParseError struct {
	err error
}

const (
	JSON ConfigType = iota
	YAML
	TOML
)

func (pe ConfigParseError) Error() string {
	return fmt.Sprintf("Config parsing error: %s", pe.err.Error())
}

func ReadConfig(in io.Reader, ct ConfigType, c map[string]interface{}) error {
	buf := new(bytes.Buffer)
	buf.ReadFrom(in)

	switch ct {
	case JSON:
		decoder := json.NewDecoder(buf)
		decoder.UseNumber()
		if err := decoder.Decode(&c); err != nil {
			return ConfigParseError{err}
		}
	case YAML:
		if err := yaml.Unmarshal(buf.Bytes(), &c); err != nil {
			return ConfigParseError{err}
		}
	case TOML:
		tree, err := toml.LoadReader(buf)
		if err != nil {
			return ConfigParseError{err}
		}
		tomlMap := tree.ToMap()
		for k, v := range tomlMap {
			c[k] = v
		}
	}

	return nil
}
<commit_msg>fixed error processing<commit_after>package go_config

import (
	"bytes"
	"encoding/json"
	"fmt"
	toml "github.com/pelletier/go-toml"
	yaml "gopkg.in/yaml.v2"
	"io"
)

type ConfigType int

type ConfigParseError struct {
	err error
}

const (
	JSON ConfigType = iota
	YAML
	TOML
)

func (pe ConfigParseError) Error() string {
	return fmt.Sprintf("Config parsing error: %s", pe.err.Error())
}

func ReadConfig(in io.Reader, ct ConfigType, c map[string]interface{}) error {
	buf := new(bytes.Buffer)
	_, err := buf.ReadFrom(in)
	if err != nil {
		return err
	}

	switch ct {
	case JSON:
		decoder := json.NewDecoder(buf)
		decoder.UseNumber()
		if err := decoder.Decode(&c); err != nil {
			return ConfigParseError{err}
		}
	case YAML:
		if err := yaml.Unmarshal(buf.Bytes(), &c); err != nil {
			return ConfigParseError{err}
		}
	case TOML:
		tree, err := toml.LoadReader(buf)
		if err != nil {
			return ConfigParseError{err}
		}
		tomlMap := tree.ToMap()
		for k, v := range tomlMap {
			c[k] = v
		}
	}

	return nil
}
<|endoftext|>


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

package vm

// mtail programs may be updated while emtail is running, and they will be
// reloaded without having to restart the mtail process. Programs can be
// created and deleted as well, and some configuration systems do an atomic
// rename of the program when it is installed, so mtail is also aware of file
// moves.

import (
	"expvar"
	"html/template"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/spf13/afero"

	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/watcher"
)

var (
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors = expvar.NewMap("prog_load_errors")
)

var (
	// LineCount counts the number of lines read by the virtual machine engine from the input channel.
	LineCount = expvar.NewInt("line_count")
)

const (
	fileExt = ".mtail"
)

// LoadProgs loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (l *Loader) LoadProgs(programPath string) error {
	s, err := os.Stat(programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", programPath)
	}
	if err = l.w.Add(programPath); err != nil {
		glog.Infof("Failed to add watch on %q but continuing: %s", programPath, err)
	}
	switch {
	case s.IsDir():
		fis, err := ioutil.ReadDir(programPath)
		if err != nil {
			return errors.Wrapf(err, "Failed to list programs in %q", programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = l.LoadProg(path.Join(programPath, fi.Name()))
			if err != nil {
				glog.Warning(err)
			}
		}
	default:
		err = l.LoadProg(programPath)
		if err != nil {
			glog.Warning(err)
		}
	}
	return nil
}

// LoadProg loads or reloads a program from the path specified.  The name of
// the program is the basename of the file.
func (l *Loader) LoadProg(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := l.fs.Open(programPath)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer f.Close()
	l.programErrorMu.Lock()
	defer l.programErrorMu.Unlock()
	l.programErrors[name] = l.CompileAndRun(name, f)
	if l.programErrors[name] != nil {
		glog.Infof("Compile errors for %s:\n%s", name, l.programErrors[name])
	}
	return nil
}

const loaderTemplate = `

Program Loader

{{range $name, $errors := $.Errors}}

{{$name}}

{{if $errors}}

{{$errors}}

{{else}}

No compile errors

{{end}}

Total load errors {{index $.Loaderrors $name}}; successes: {{index $.Loadsuccess $name}}

{{end}}
`

// WriteStatusHTML writes the current state of the loader as HTML to the given writer w.
func (l *Loader) WriteStatusHTML(w io.Writer) error {
	t, err := template.New("loader").Parse(loaderTemplate)
	if err != nil {
		return err
	}
	l.programErrorMu.RLock()
	defer l.programErrorMu.RUnlock()
	data := struct {
		Errors      map[string]error
		Loaderrors  map[string]string
		Loadsuccess map[string]string
	}{
		l.programErrors,
		make(map[string]string),
		make(map[string]string),
	}
	for name := range l.programErrors {
		if ProgLoadErrors.Get(name) != nil {
			data.Loaderrors[name] = ProgLoadErrors.Get(name).String()
		}
		if ProgLoads.Get(name) != nil {
			data.Loadsuccess[name] = ProgLoads.Get(name).String()
		}
	}
	return t.Execute(w, data)
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (l *Loader) CompileAndRun(name string, input io.Reader) error {
	o := &Options{
		CompileOnly:          l.compileOnly,
		EmitAst:              l.dumpAst,
		EmitAstTypes:         l.dumpAstTypes,
		SyslogUseCurrentYear: l.syslogUseCurrentYear,
		OverrideLocation:     l.overrideLocation,
	}
	v, errs := Compile(name, input, o)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if v == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}

	if l.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode(name))
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.m {
		if !m.Hidden {
			if l.omitMetricSource {
				m.Source = ""
			}
			err := l.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if l.compileOnly {
		return nil
	}

	l.handleMu.Lock()
	defer l.handleMu.Unlock()

	// Stop any previous VM.
	if handle, ok := l.handles[name]; ok {
		glog.Infof("END OF LINE, %s", name)
		close(handle.lines)
		<-handle.done
		glog.Infof("Stopped %s", name)
	}

	l.handles[name] = &vmHandle{make(chan *tailer.LogLine), make(chan struct{})}
	nameCode := nameToCode(name)
	glog.Infof("Program %s has goroutine marker 0x%x", name, nameCode)
	started := make(chan struct{})
	go v.Run(nameCode, l.handles[name].lines, l.handles[name].done, started)
	<-started
	glog.Infof("Started %s", name)

	return nil
}

func nameToCode(name string) uint32 {
	return uint32(uint32(name[0])<<24 | uint32(name[1])<<16 | uint32(name[2])<<8 | uint32(name[3]))
}

// Loader handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the running virtual machines that receive input from the lines
// channel.
type Loader struct {
	w  watcher.Watcher // watches for program changes
	fs afero.Fs        // filesystem interface
	ms *metrics.Store  // pointer to store to pass to compiler

	handles  map[string]*vmHandle // map of program names to virtual machines
	handleMu sync.RWMutex         // guards accesses to handles

	programErrors  map[string]error // errors from the last compile attempt of the program
	programErrorMu sync.RWMutex     // guards access to programErrors

	watcherDone chan struct{} // Synchronise shutdown of the watcher and lines handlers.
	VMsDone     chan struct{} // Notify mtail when all running VMs are shutdown.

	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	dumpAst              bool           // print the AST after parse
	dumpAstTypes         bool           // print the AST after type check
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	omitMetricSource     bool
}

// LoaderOptions contains the required and optional parameters for creating a
// new Loader.
type LoaderOptions struct {
	Store *metrics.Store
	Lines <-chan *tailer.LogLine

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

	CompileOnly          bool           // Compile, don't start execution.
	DumpAst              bool           // print the AST after type check
	DumpAstTypes         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	DumpBytecode         bool           // Instructs the loader to dump the program bytecode after compilation.
	SyslogUseCurrentYear bool           // If true, override empty year with the current in strptime().
	OverrideLocation     *time.Location // if not nil, overrides the timezone in strptime().
	OmitMetricSource     bool           // Don't put the source in the metric when added to the Store.
}

// NewLoader creates a new program loader.  It takes a filesystem watcher
// and a filesystem interface as arguments.  If fs is nil, it will use the
// default filesystem interface.
func NewLoader(o LoaderOptions) (*Loader, error) {
	if o.Store == nil || o.Lines == nil {
		return nil, errors.New("loader needs a store and lines")
	}
	fs := o.FS
	if fs == nil {
		fs = &afero.OsFs{}
	}
	w := o.W
	if w == nil {
		var err error
		w, err = watcher.NewLogWatcher()
		if err != nil {
			return nil, errors.Wrap(err, "Couldn't create a watcher for loader")
		}
	}

	l := &Loader{
		w:                    w,
		ms:                   o.Store,
		fs:                   fs,
		handles:              make(map[string]*vmHandle),
		programErrors:        make(map[string]error),
		watcherDone:          make(chan struct{}),
		VMsDone:              make(chan struct{}),
		compileOnly:          o.CompileOnly,
		dumpAst:              o.DumpAst,
		dumpAstTypes:         o.DumpAstTypes,
		dumpBytecode:         o.DumpBytecode,
		syslogUseCurrentYear: o.SyslogUseCurrentYear,
		overrideLocation:     o.OverrideLocation,
		omitMetricSource:     o.OmitMetricSource,
	}

	eventsChan := l.w.Events()
	go l.processEvents(eventsChan)
	go l.processLines(o.Lines)
	return l, nil
}

type vmHandle struct {
	lines chan *tailer.LogLine
	done  chan struct{}
}

// processEvents manages program lifecycle triggered by events from the
// filesystem watcher.
func (l *Loader) processEvents(events <-chan watcher.Event) {
	defer close(l.watcherDone)
	for event := range events {
		switch event := event.(type) {
		case watcher.DeleteEvent:
			l.UnloadProgram(event.Pathname)
		case watcher.UpdateEvent:
			l.LoadProg(event.Pathname)
		case watcher.CreateEvent:
			l.w.Add(event.Pathname)
		default:
			glog.V(1).Infof("Unexected event type %+#v", event)
		}
	}
}

// processLines provides fanout of the input log lines to each virtual machine
// running.  Upon close of the incoming lines channel, it also communicates
// shutdown to the target VMs via channel close.
func (l *Loader) processLines(lines <-chan *tailer.LogLine) {
	for logline := range lines {
		LineCount.Add(1)
		l.handleMu.RLock()
		for prog := range l.handles {
			l.handles[prog].lines <- logline
		}
		l.handleMu.RUnlock()
	}
	glog.Info("Shutting down loader.")
	l.w.Close()
	<-l.watcherDone
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	for prog := range l.handles {
		close(l.handles[prog].lines)
		<-l.handles[prog].done
		delete(l.handles, prog)
	}
	close(l.VMsDone)
}

// UnloadProgram removes the named program from the watcher to prevent future
// updates, and terminates any currently running VM goroutine.
func (l *Loader) UnloadProgram(pathname string) {
	if err := l.w.Remove(pathname); err != nil {
		glog.Infof("Remove watch on %s failed: %s", pathname, err)
	}
	name := filepath.Base(pathname)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	if handle, ok := l.handles[name]; ok {
		close(handle.lines)
		<-handle.done
		delete(l.handles, name)
	}
}
<commit_msg>Returns the compiler errors when in compile only mode.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

// mtail programs may be updated while emtail is running, and they will be
// reloaded without having to restart the mtail process. Programs can be
// created and deleted as well, and some configuration systems do an atomic
// rename of the program when it is installed, so mtail is also aware of file
// moves.

import (
	"expvar"
	"html/template"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/spf13/afero"

	"github.com/google/mtail/metrics"
	"github.com/google/mtail/tailer"
	"github.com/google/mtail/watcher"
)

var (
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors = expvar.NewMap("prog_load_errors")
)

var (
	// LineCount counts the number of lines read by the virtual machine engine from the input channel.
	LineCount = expvar.NewInt("line_count")
)

const (
	fileExt = ".mtail"
)

// LoadProgs loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (l *Loader) LoadProgs(programPath string) error {
	s, err := os.Stat(programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", programPath)
	}
	if err = l.w.Add(programPath); err != nil {
		glog.Infof("Failed to add watch on %q but continuing: %s", programPath, err)
	}
	switch {
	case s.IsDir():
		fis, err := ioutil.ReadDir(programPath)
		if err != nil {
			return errors.Wrapf(err, "Failed to list programs in %q", programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = l.LoadProg(path.Join(programPath, fi.Name()))
			if err != nil {
				glog.Warning(err)
				if l.compileOnly {
					return err
				}
			}
		}
	default:
		err = l.LoadProg(programPath)
		if err != nil {
			glog.Warning(err)
			if l.compileOnly {
				return err
			}
		}
	}
	return nil
}

// LoadProg loads or reloads a program from the path specified.  The name of
// the program is the basename of the file.
func (l *Loader) LoadProg(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := l.fs.Open(programPath)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer f.Close()
	l.programErrorMu.Lock()
	defer l.programErrorMu.Unlock()
	l.programErrors[name] = l.CompileAndRun(name, f)
	if l.programErrors[name] != nil {
		glog.Infof("Compile errors for %s:\n%s", name, l.programErrors[name])
		if l.compileOnly {
			return l.programErrors[name]
		}
	}
	return nil
}

const loaderTemplate = `

Program Loader

{{range $name, $errors := $.Errors}}

{{$name}}

{{if $errors}}

{{$errors}}

{{else}}

No compile errors

{{end}}

Total load errors {{index $.Loaderrors $name}}; successes: {{index $.Loadsuccess $name}}

{{end}}
`

// WriteStatusHTML writes the current state of the loader as HTML to the given writer w.
func (l *Loader) WriteStatusHTML(w io.Writer) error {
	t, err := template.New("loader").Parse(loaderTemplate)
	if err != nil {
		return err
	}
	l.programErrorMu.RLock()
	defer l.programErrorMu.RUnlock()
	data := struct {
		Errors      map[string]error
		Loaderrors  map[string]string
		Loadsuccess map[string]string
	}{
		l.programErrors,
		make(map[string]string),
		make(map[string]string),
	}
	for name := range l.programErrors {
		if ProgLoadErrors.Get(name) != nil {
			data.Loaderrors[name] = ProgLoadErrors.Get(name).String()
		}
		if ProgLoads.Get(name) != nil {
			data.Loadsuccess[name] = ProgLoads.Get(name).String()
		}
	}
	return t.Execute(w, data)
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (l *Loader) CompileAndRun(name string, input io.Reader) error {
	o := &Options{
		CompileOnly:          l.compileOnly,
		EmitAst:              l.dumpAst,
		EmitAstTypes:         l.dumpAstTypes,
		SyslogUseCurrentYear: l.syslogUseCurrentYear,
		OverrideLocation:     l.overrideLocation,
	}
	v, errs := Compile(name, input, o)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if v == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}

	if l.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode(name))
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.m {
		if !m.Hidden {
			if l.omitMetricSource {
				m.Source = ""
			}
			err := l.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if l.compileOnly {
		return nil
	}

	l.handleMu.Lock()
	defer l.handleMu.Unlock()

	// Stop any previous VM.
	if handle, ok := l.handles[name]; ok {
		glog.Infof("END OF LINE, %s", name)
		close(handle.lines)
		<-handle.done
		glog.Infof("Stopped %s", name)
	}

	l.handles[name] = &vmHandle{make(chan *tailer.LogLine), make(chan struct{})}
	nameCode := nameToCode(name)
	glog.Infof("Program %s has goroutine marker 0x%x", name, nameCode)
	started := make(chan struct{})
	go v.Run(nameCode, l.handles[name].lines, l.handles[name].done, started)
	<-started
	glog.Infof("Started %s", name)

	return nil
}

func nameToCode(name string) uint32 {
	return uint32(uint32(name[0])<<24 | uint32(name[1])<<16 | uint32(name[2])<<8 | uint32(name[3]))
}

// Loader handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the running virtual machines that receive input from the lines
// channel.
type Loader struct {
	w  watcher.Watcher // watches for program changes
	fs afero.Fs        // filesystem interface
	ms *metrics.Store  // pointer to store to pass to compiler

	handles  map[string]*vmHandle // map of program names to virtual machines
	handleMu sync.RWMutex         // guards accesses to handles

	programErrors  map[string]error // errors from the last compile attempt of the program
	programErrorMu sync.RWMutex     // guards access to programErrors

	watcherDone chan struct{} // Synchronise shutdown of the watcher and lines handlers.
	VMsDone     chan struct{} // Notify mtail when all running VMs are shutdown.

	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	dumpAst              bool           // print the AST after parse
	dumpAstTypes         bool           // print the AST after type check
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	omitMetricSource     bool
}

// LoaderOptions contains the required and optional parameters for creating a
// new Loader.
type LoaderOptions struct {
	Store *metrics.Store
	Lines <-chan *tailer.LogLine

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

	CompileOnly          bool           // Compile, don't start execution.
	DumpAst              bool           // print the AST after type check
	DumpAstTypes         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	DumpBytecode         bool           // Instructs the loader to dump the program bytecode after compilation.
	SyslogUseCurrentYear bool           // If true, override empty year with the current in strptime().
	OverrideLocation     *time.Location // if not nil, overrides the timezone in strptime().
	OmitMetricSource     bool           // Don't put the source in the metric when added to the Store.
}

// NewLoader creates a new program loader.  It takes a filesystem watcher
// and a filesystem interface as arguments.  If fs is nil, it will use the
// default filesystem interface.
func NewLoader(o LoaderOptions) (*Loader, error) {
	if o.Store == nil || o.Lines == nil {
		return nil, errors.New("loader needs a store and lines")
	}
	fs := o.FS
	if fs == nil {
		fs = &afero.OsFs{}
	}
	w := o.W
	if w == nil {
		var err error
		w, err = watcher.NewLogWatcher()
		if err != nil {
			return nil, errors.Wrap(err, "Couldn't create a watcher for loader")
		}
	}

	l := &Loader{
		w:                    w,
		ms:                   o.Store,
		fs:                   fs,
		handles:              make(map[string]*vmHandle),
		programErrors:        make(map[string]error),
		watcherDone:          make(chan struct{}),
		VMsDone:              make(chan struct{}),
		compileOnly:          o.CompileOnly,
		dumpAst:              o.DumpAst,
		dumpAstTypes:         o.DumpAstTypes,
		dumpBytecode:         o.DumpBytecode,
		syslogUseCurrentYear: o.SyslogUseCurrentYear,
		overrideLocation:     o.OverrideLocation,
		omitMetricSource:     o.OmitMetricSource,
	}

	eventsChan := l.w.Events()
	go l.processEvents(eventsChan)
	go l.processLines(o.Lines)
	return l, nil
}

type vmHandle struct {
	lines chan *tailer.LogLine
	done  chan struct{}
}

// processEvents manages program lifecycle triggered by events from the
// filesystem watcher.
func (l *Loader) processEvents(events <-chan watcher.Event) {
	defer close(l.watcherDone)
	for event := range events {
		switch event := event.(type) {
		case watcher.DeleteEvent:
			l.UnloadProgram(event.Pathname)
		case watcher.UpdateEvent:
			l.LoadProg(event.Pathname)
		case watcher.CreateEvent:
			l.w.Add(event.Pathname)
		default:
			glog.V(1).Infof("Unexected event type %+#v", event)
		}
	}
}

// processLines provides fanout of the input log lines to each virtual machine
// running.  Upon close of the incoming lines channel, it also communicates
// shutdown to the target VMs via channel close.
func (l *Loader) processLines(lines <-chan *tailer.LogLine) {
	for logline := range lines {
		LineCount.Add(1)
		l.handleMu.RLock()
		for prog := range l.handles {
			l.handles[prog].lines <- logline
		}
		l.handleMu.RUnlock()
	}
	glog.Info("Shutting down loader.")
	l.w.Close()
	<-l.watcherDone
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	for prog := range l.handles {
		close(l.handles[prog].lines)
		<-l.handles[prog].done
		delete(l.handles, prog)
	}
	close(l.VMsDone)
}

// UnloadProgram removes the named program from the watcher to prevent future
// updates, and terminates any currently running VM goroutine.
func (l *Loader) UnloadProgram(pathname string) {
	if err := l.w.Remove(pathname); err != nil {
		glog.Infof("Remove watch on %s failed: %s", pathname, err)
	}
	name := filepath.Base(pathname)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	if handle, ok := l.handles[name]; ok {
		close(handle.lines)
		<-handle.done
		delete(l.handles, name)
	}
}
<|endoftext|>


<commit_before>package torrent

import (
	"errors"
	"io"
	"os"
)

// Accesses torrent data via a client.
type Reader struct {
	t          *Torrent
	pos        int64
	responsive bool
	readahead  int64
}

var _ io.ReadCloser = &Reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *Reader) SetResponsive() {
	r.responsive = true
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *Reader) SetReadahead(readahead int64) {
	r.readahead = readahead
}

func (r *Reader) raisePriorities(off int64, n int) {
	if r.responsive {
		r.t.cl.addUrgentRequests(r.t.torrent, off, n)
	}
	r.t.cl.readRaisePiecePriorities(r.t.torrent, off, int64(n)+r.readahead)
}

func (r *Reader) readable(off int64) (ret bool) {
	// log.Println("readable", off)
	// defer func() {
	// 	log.Println("readable", ret)
	// }()
	req, ok := r.t.offsetRequest(off)
	if !ok {
		panic(off)
	}
	if r.responsive {
		return r.t.haveChunk(req)
	}
	return r.t.pieceComplete(int(req.Index))
}

// How many bytes are available to read. Max is the most we could require.
func (r *Reader) available(off, max int64) (ret int64) {
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *Reader) waitReadable(off int64) {
	r.t.Pieces[off/int64(r.t.usualPieceSize())].Event.Wait()
}

func (r *Reader) ReadAt(b []byte, off int64) (n int, err error) {
	return r.readAt(b, off)
}

func (r *Reader) Read(b []byte) (n int, err error) {
	n, err = r.readAt(b, r.pos)
	r.pos += int64(n)
	return
}

// Must only return EOF at the end of the torrent.
func (r *Reader) readAt(b []byte, pos int64) (n int, err error) {
	// defer func() {
	// 	log.Println(pos, n, err)
	// }()
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	maxLen := r.t.torrent.Info.TotalLength() - pos
	if maxLen <= 0 {
		err = io.EOF
		return
	}
	if int64(len(b)) > maxLen {
		b = b[:maxLen]
	}
again:
	r.raisePriorities(pos, len(b))
	for !r.readable(pos) {
		r.raisePriorities(pos, len(b))
		r.waitReadable(pos)
	}
	avail := r.available(pos, int64(len(b)))
	// log.Println("available", avail)
	b1 := b[:avail]
	n, err = dataReadAt(r.t.data, b1, pos)
	if n != 0 {
		err = nil
		return
	}
	if err == io.ErrUnexpectedEOF {
		goto again
	}
	return
}

func (r *Reader) Close() error {
	r.t = nil
	return nil
}

func (r *Reader) Seek(off int64, whence int) (ret int64, err error) {
	switch whence {
	case os.SEEK_SET:
		r.pos = off
	case os.SEEK_CUR:
		r.pos += off
	case os.SEEK_END:
		r.pos = r.t.torrent.Info.TotalLength() + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos
	return
}
<commit_msg>Add fixes triggered by TestResponsive<commit_after>package torrent

import (
	"errors"
	"io"
	"os"
)

// Accesses torrent data via a client.
type Reader struct {
	t          *Torrent
	pos        int64
	responsive bool
	readahead  int64
}

var _ io.ReadCloser = &Reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *Reader) SetResponsive() {
	r.responsive = true
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *Reader) SetReadahead(readahead int64) {
	r.readahead = readahead
}

func (r *Reader) raisePriorities(off int64, n int) {
	if r.responsive {
		r.t.cl.addUrgentRequests(r.t.torrent, off, n)
	}
	if !r.responsive || r.readahead != 0 {
		r.t.cl.readRaisePiecePriorities(r.t.torrent, off, int64(n)+r.readahead)
	}
}

func (r *Reader) readable(off int64) (ret bool) {
	// log.Println("readable", off)
	// defer func() {
	// 	log.Println("readable", ret)
	// }()
	req, ok := r.t.offsetRequest(off)
	if !ok {
		panic(off)
	}
	if r.responsive {
		return r.t.haveChunk(req)
	}
	return r.t.pieceComplete(int(req.Index))
}

// How many bytes are available to read. Max is the most we could require.
func (r *Reader) available(off, max int64) (ret int64) {
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *Reader) waitReadable(off int64) {
	r.t.Pieces[off/int64(r.t.usualPieceSize())].Event.Wait()
}

func (r *Reader) ReadAt(b []byte, off int64) (n int, err error) {
	for {
		var n1 int
		n1, err = r.readAt(b, off)
		n += n1
		b = b[n1:]
		off += int64(n1)
		if err != nil || len(b) == 0 || n1 == 0 {
			return
		}
	}
}

func (r *Reader) Read(b []byte) (n int, err error) {
	n, err = r.readAt(b, r.pos)
	r.pos += int64(n)
	return
}

// Must only return EOF at the end of the torrent.
func (r *Reader) readAt(b []byte, pos int64) (n int, err error) {
	// defer func() {
	// 	log.Println(pos, n, err)
	// }()
	r.t.cl.mu.Lock()
	defer r.t.cl.mu.Unlock()
	maxLen := r.t.torrent.Info.TotalLength() - pos
	if maxLen <= 0 {
		err = io.EOF
		return
	}
	if int64(len(b)) > maxLen {
		b = b[:maxLen]
	}
again:
	r.raisePriorities(pos, len(b))
	for !r.readable(pos) {
		r.raisePriorities(pos, len(b))
		r.waitReadable(pos)
	}
	avail := r.available(pos, int64(len(b)))
	// log.Println("available", avail)
	b1 := b[:avail]
	n, err = dataReadAt(r.t.data, b1, pos)
	if n != 0 {
		err = nil
		return
	}
	if err == io.ErrUnexpectedEOF {
		goto again
	}
	return
}

func (r *Reader) Close() error {
	r.t = nil
	return nil
}

func (r *Reader) Seek(off int64, whence int) (ret int64, err error) {
	switch whence {
	case os.SEEK_SET:
		r.pos = off
	case os.SEEK_CUR:
		r.pos += off
	case os.SEEK_END:
		r.pos = r.t.torrent.Info.TotalLength() + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 The prometheus-operator Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package framework

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/coreos/prometheus-operator/pkg/client/monitoring/v1alpha1"
	"github.com/coreos/prometheus-operator/pkg/k8sutil"
)

type Framework struct {
	KubeClient  kubernetes.Interface
	MonClient   *v1alpha1.MonitoringV1alpha1Client
	HTTPClient  *http.Client
	MasterHost  string
	Namespace   *v1.Namespace
	OperatorPod *v1.Pod
	ClusterIP   string
}

// Setup setups a test framework and returns it.
func New(ns, kubeconfig, opImage, ip string) (*Framework, error) {
	config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
	if err != nil {
		return nil, err
	}

	cli, err := kubernetes.NewForConfig(config)
	if err != nil {
		return nil, err
	}

	httpc := cli.CoreV1().RESTClient().(*rest.RESTClient).Client
	if err != nil {
		return nil, err
	}

	mclient, err := v1alpha1.NewForConfig(config)
	if err != nil {
		return nil, err
	}

	namespace, err := cli.Core().Namespaces().Create(&v1.Namespace{
		ObjectMeta: metav1.ObjectMeta{
			Name: ns,
		},
	})
	if err != nil {
		return nil, err
	}

	f := &Framework{
		MasterHost: config.Host,
		KubeClient: cli,
		MonClient:  mclient,
		HTTPClient: httpc,
		Namespace:  namespace,
		ClusterIP:  ip,
	}

	err = f.setup(opImage)
	if err != nil {
		return nil, err
	}

	return f, nil
}

func (f *Framework) setup(opImage string) error {
	if err := f.setupPrometheusOperator(opImage); err != nil {
		return err
	}
	return nil
}

func (f *Framework) setupPrometheusOperator(opImage string) error {
	fn, err := filepath.Abs("../../deployment.yaml")
	if err != nil {
		return err
	}

	deployManifest, err := os.Open(fn)
	if err != nil {
		return err
	}

	deploy := v1beta1.Deployment{}
	err = yaml.NewYAMLOrJSONDecoder(deployManifest, 100).Decode(&deploy)
	if err != nil {
		return err
	}
	if opImage != "" {
		// Override operator image used, if specified when running tests.
		deploy.Spec.Template.Spec.Containers[0].Image = opImage
	}

	err = f.createDeployment(&deploy)
	if err != nil {
		return err
	}

	opts := metav1.ListOptions{LabelSelector: fields.SelectorFromSet(fields.Set(deploy.Spec.Template.ObjectMeta.Labels)).String()}
	err = f.WaitForPodsReady(1, opts)
	if err != nil {
		return err
	}

	pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
	if err != nil {
		return err
	}
	f.OperatorPod = &pl.Items[0]

	err = k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRPrometheusName)
	if err != nil {
		return err
	}

	err = k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRServiceMonitorName)
	if err != nil {
		return err
	}

	return k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRAlertmanagerName)
}

// Teardown tears down a previously initialized test environment.
func (f *Framework) Teardown() error {
	if err := f.KubeClient.Core().Services(f.Namespace.Name).Delete("prometheus-operated", nil); err != nil && !k8sutil.IsResourceNotFoundError(err) {
		return err
	}

	if err := f.KubeClient.Core().Services(f.Namespace.Name).Delete("alertmanager-operated", nil); err != nil && !k8sutil.IsResourceNotFoundError(err) {
		return err
	}

	if err := f.KubeClient.Extensions().Deployments(f.Namespace.Name).Delete("prometheus-operator", nil); err != nil {
		return err
	}

	if err := f.KubeClient.Core().Namespaces().Delete(f.Namespace.Name, nil); err != nil {
		return err
	}

	return nil
}

// WaitForPodsReady waits for a selection of Pods to be running and each
// container to pass its readiness check.
func (f *Framework) WaitForPodsReady(expectedReplicas int, opts metav1.ListOptions) error {
	return f.Poll(time.Minute*2, time.Second, func() (bool, error) {
		pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
		if err != nil {
			return false, err
		}

		runningAndReady := 0
		for _, p := range pl.Items {
			isRunningAndReady, err := k8sutil.PodRunningAndReady(p)
			if err != nil {
				return false, err
			}

			if isRunningAndReady {
				runningAndReady++
			}
		}

		if runningAndReady == expectedReplicas {
			return true, nil
		}
		return false, nil
	})
}

func (f *Framework) WaitForPodsRunImage(expectedReplicas int, image string, opts metav1.ListOptions) error {
	return f.Poll(time.Minute*5, time.Second, func() (bool, error) {
		pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
		if err != nil {
			return false, err
		}

		runningImage := 0
		for _, p := range pl.Items {
			if podRunsImage(p, image) {
				runningImage++
			}
		}

		if runningImage == expectedReplicas {
			return true, nil
		}
		return false, nil
	})
}

func (f *Framework) WaitForHTTPSuccessStatusCode(timeout time.Duration, url string) error {
	return f.Poll(time.Minute*5, time.Second, func() (bool, error) {
		resp, err := http.Get(url)
		if err != nil {
			return false, err
		}
		if err == nil && resp.StatusCode == 200 {
			return true, nil
		}
		return false, nil
	})
}

func podRunsImage(p v1.Pod, image string) bool {
	for _, c := range p.Spec.Containers {
		if image == c.Image {
			return true
		}
	}

	return false
}

func (f *Framework) CreateDeployment(kclient kubernetes.Interface, ns string, deploy *v1beta1.Deployment) error {
	if _, err := f.KubeClient.Extensions().Deployments(ns).Create(deploy); err != nil {
		return err
	}

	return nil
}

func (f *Framework) createDeployment(deploy *v1beta1.Deployment) error {
	if _, err := f.KubeClient.Extensions().Deployments(f.Namespace.Name).Create(deploy); err != nil {
		return err
	}

	return nil
}

func (f *Framework) GetLogs(podName, containerName string) (string, error) {
	logs, err := f.KubeClient.Core().RESTClient().Get().
		Resource("pods").
		Namespace(f.Namespace.Name).
		Name(podName).SubResource("log").
		Param("container", containerName).
		Do().
		Raw()
	if err != nil {
		return "", err
	}
	return string(logs), err
}

func (f *Framework) Poll(timeout, pollInterval time.Duration, pollFunc func() (bool, error)) error {
	t := time.After(timeout)
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()

	for {
		select {
		case <-t:
			return fmt.Errorf("timed out")
		case <-ticker.C:
			b, err := pollFunc()
			if err != nil {
				return err
			}
			if b {
				return nil
			}
		}
	}
}

func (f *Framework) ProxyGetPod(podName string, port string, path string) *rest.Request {
	return f.KubeClient.CoreV1().RESTClient().Get().Prefix("proxy").Namespace(f.Namespace.Name).Resource("pods").Name(podName + ":" + port).Suffix(path)
}
<commit_msg>test: Fix flaky ingress tests<commit_after>// Copyright 2016 The prometheus-operator Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package framework

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/coreos/prometheus-operator/pkg/client/monitoring/v1alpha1"
	"github.com/coreos/prometheus-operator/pkg/k8sutil"
)

type Framework struct {
	KubeClient  kubernetes.Interface
	MonClient   *v1alpha1.MonitoringV1alpha1Client
	HTTPClient  *http.Client
	MasterHost  string
	Namespace   *v1.Namespace
	OperatorPod *v1.Pod
	ClusterIP   string
}

// Setup setups a test framework and returns it.
func New(ns, kubeconfig, opImage, ip string) (*Framework, error) {
	config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
	if err != nil {
		return nil, err
	}

	cli, err := kubernetes.NewForConfig(config)
	if err != nil {
		return nil, err
	}

	httpc := cli.CoreV1().RESTClient().(*rest.RESTClient).Client
	if err != nil {
		return nil, err
	}

	mclient, err := v1alpha1.NewForConfig(config)
	if err != nil {
		return nil, err
	}

	namespace, err := cli.Core().Namespaces().Create(&v1.Namespace{
		ObjectMeta: metav1.ObjectMeta{
			Name: ns,
		},
	})
	if err != nil {
		return nil, err
	}

	f := &Framework{
		MasterHost: config.Host,
		KubeClient: cli,
		MonClient:  mclient,
		HTTPClient: httpc,
		Namespace:  namespace,
		ClusterIP:  ip,
	}

	err = f.setup(opImage)
	if err != nil {
		return nil, err
	}

	return f, nil
}

func (f *Framework) setup(opImage string) error {
	if err := f.setupPrometheusOperator(opImage); err != nil {
		return err
	}
	return nil
}

func (f *Framework) setupPrometheusOperator(opImage string) error {
	fn, err := filepath.Abs("../../deployment.yaml")
	if err != nil {
		return err
	}

	deployManifest, err := os.Open(fn)
	if err != nil {
		return err
	}

	deploy := v1beta1.Deployment{}
	err = yaml.NewYAMLOrJSONDecoder(deployManifest, 100).Decode(&deploy)
	if err != nil {
		return err
	}
	if opImage != "" {
		// Override operator image used, if specified when running tests.
		deploy.Spec.Template.Spec.Containers[0].Image = opImage
	}

	err = f.createDeployment(&deploy)
	if err != nil {
		return err
	}

	opts := metav1.ListOptions{LabelSelector: fields.SelectorFromSet(fields.Set(deploy.Spec.Template.ObjectMeta.Labels)).String()}
	err = f.WaitForPodsReady(1, opts)
	if err != nil {
		return err
	}

	pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
	if err != nil {
		return err
	}
	f.OperatorPod = &pl.Items[0]

	err = k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRPrometheusName)
	if err != nil {
		return err
	}

	err = k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRServiceMonitorName)
	if err != nil {
		return err
	}

	return k8sutil.WaitForTPRReady(f.KubeClient.Core().RESTClient(), v1alpha1.TPRGroup, v1alpha1.TPRVersion, v1alpha1.TPRAlertmanagerName)
}

// Teardown tears down a previously initialized test environment.
func (f *Framework) Teardown() error {
	if err := f.KubeClient.Core().Services(f.Namespace.Name).Delete("prometheus-operated", nil); err != nil && !k8sutil.IsResourceNotFoundError(err) {
		return err
	}

	if err := f.KubeClient.Core().Services(f.Namespace.Name).Delete("alertmanager-operated", nil); err != nil && !k8sutil.IsResourceNotFoundError(err) {
		return err
	}

	if err := f.KubeClient.Extensions().Deployments(f.Namespace.Name).Delete("prometheus-operator", nil); err != nil {
		return err
	}

	if err := f.KubeClient.Core().Namespaces().Delete(f.Namespace.Name, nil); err != nil {
		return err
	}

	return nil
}

// WaitForPodsReady waits for a selection of Pods to be running and each
// container to pass its readiness check.
func (f *Framework) WaitForPodsReady(expectedReplicas int, opts metav1.ListOptions) error {
	return f.Poll(time.Minute*2, time.Second, func() (bool, error) {
		pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
		if err != nil {
			return false, err
		}

		runningAndReady := 0
		for _, p := range pl.Items {
			isRunningAndReady, err := k8sutil.PodRunningAndReady(p)
			if err != nil {
				return false, err
			}

			if isRunningAndReady {
				runningAndReady++
			}
		}

		if runningAndReady == expectedReplicas {
			return true, nil
		}
		return false, nil
	})
}

func (f *Framework) WaitForPodsRunImage(expectedReplicas int, image string, opts metav1.ListOptions) error {
	return f.Poll(time.Minute*5, time.Second, func() (bool, error) {
		pl, err := f.KubeClient.Core().Pods(f.Namespace.Name).List(opts)
		if err != nil {
			return false, err
		}

		runningImage := 0
		for _, p := range pl.Items {
			if podRunsImage(p, image) {
				runningImage++
			}
		}

		if runningImage == expectedReplicas {
			return true, nil
		}
		return false, nil
	})
}

func (f *Framework) WaitForHTTPSuccessStatusCode(timeout time.Duration, url string) error {
	return f.Poll(timeout, time.Second, func() (bool, error) {
		resp, err := http.Get(url)
		if err == nil && resp.StatusCode == 200 {
			return true, nil
		}
		return false, nil
	})
}

func podRunsImage(p v1.Pod, image string) bool {
	for _, c := range p.Spec.Containers {
		if image == c.Image {
			return true
		}
	}

	return false
}

func (f *Framework) CreateDeployment(kclient kubernetes.Interface, ns string, deploy *v1beta1.Deployment) error {
	if _, err := f.KubeClient.Extensions().Deployments(ns).Create(deploy); err != nil {
		return err
	}

	return nil
}

func (f *Framework) createDeployment(deploy *v1beta1.Deployment) error {
	if _, err := f.KubeClient.Extensions().Deployments(f.Namespace.Name).Create(deploy); err != nil {
		return err
	}

	return nil
}

func (f *Framework) GetLogs(podName, containerName string) (string, error) {
	logs, err := f.KubeClient.Core().RESTClient().Get().
		Resource("pods").
		Namespace(f.Namespace.Name).
		Name(podName).SubResource("log").
		Param("container", containerName).
		Do().
		Raw()
	if err != nil {
		return "", err
	}
	return string(logs), err
}

func (f *Framework) Poll(timeout, pollInterval time.Duration, pollFunc func() (bool, error)) error {
	t := time.After(timeout)
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()

	for {
		select {
		case <-t:
			return fmt.Errorf("timed out")
		case <-ticker.C:
			b, err := pollFunc()
			if err != nil {
				return err
			}
			if b {
				return nil
			}
		}
	}
}

func (f *Framework) ProxyGetPod(podName string, port string, path string) *rest.Request {
	return f.KubeClient.CoreV1().RESTClient().Get().Prefix("proxy").Namespace(f.Namespace.Name).Resource("pods").Name(podName + ":" + port).Suffix(path)
}
<|endoftext|>


<commit_before>package alert

import (
	"context"
	"fmt"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
	"github.com/prometheus/common/model"

	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/disruption"
	helper "github.com/openshift/origin/test/extended/util/prometheus"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"
)

// UpgradeTest runs verifies invariants regarding what alerts are allowed to fire
// during the upgrade process.
type UpgradeTest struct {
	url         string
	bearerToken string
	oc          *exutil.CLI
}

func (UpgradeTest) Name() string { return "check-for-alerts" }
func (UpgradeTest) DisplayName() string {
	return "[sig-arch] Check if alerts are firing during or after upgrade success"
}

// Setup creates parameters to query Prometheus
func (t *UpgradeTest) Setup(f *framework.Framework) {
	g.By("Setting up upgrade alert test")

	oc := exutil.NewCLIWithFramework(f)
	url, _, bearerToken, ok := helper.LocatePrometheus(oc)
	if !ok {
		framework.Failf("Prometheus could not be located on this cluster, failing test %s", t.Name())
	}
	t.url = url
	t.bearerToken = bearerToken
	t.oc = oc
	framework.Logf("Post-upgrade alert test setup complete")
}

// Test checks if alerts are firing at various points during upgrade.
// An alert firing during an upgrade is a high severity bug - it either points to a real issue in
// a dependency, or a failure of the component, and therefore must be fixed.
func (t *UpgradeTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	tolerateDuringSkew := exutil.TolerateVersionSkewInTests()
	firingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "openshift-apiserver"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "machine-config"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955300",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubeDaemonSetRolloutStuck"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1943667",
		},
		{
			Selector: map[string]string{"alertname": "KubeAPIErrorBudgetBurn"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1953798",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "AggregatedAPIDown", "namespace": "default", "name": "v1beta1.metrics.k8s.io"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1970624",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "prometheus-k8s"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1949262",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "alertmanager-main"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955489",
		},
	}

	pendingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "ClusterMonitoringOperatorReconciliationErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1932624",
		},
		{
			Selector: map[string]string{"alertname": "KubeClientErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1925698",
		},
	}
	allowedPendingAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "etcdMemberCommunicationSlow"},
			Text:     "Excluded because it triggers during upgrade (detects ~5m of high latency immediately preceeding the end of the test), and we don't want to change the alert because it is correct",
		},
	}

	knownViolations := sets.NewString()
	unexpectedViolations := sets.NewString()
	unexpectedViolationsAsFlakes := sets.NewString()
	debug := sets.NewString()

	g.By("Checking for alerts")

	start := time.Now()

	// Block until upgrade is done
	g.By("Waiting for upgrade to finish before checking for alerts")
	<-done

	// Additonal delay after upgrade completion to allow pending alerts to settle
	g.By("Waiting before checking for alerts")
	time.Sleep(1 * time.Minute)

	ns := t.oc.SetupNamespace()
	execPod := exutil.CreateExecPodOrFail(t.oc.AdminKubeClient(), ns, "execpod")
	defer func() {
		t.oc.AdminKubeClient().CoreV1().Pods(ns).Delete(context.Background(), execPod.Name, *metav1.NewDeleteOptions(1))
	}()

	testDuration := time.Now().Sub(start).Truncate(time.Second)

	// Invariant: The watchdog alert should be firing continuously during the whole upgrade via the thanos
	// querier (which should have no gaps when it queries the individual stores). Allow zero or one changes
	// to the presence of this series (zero if data is preserved over upgrade, one if data is lost on upgrade).
	// This would not catch the alert stopping firing, but we catch that in other places and tests.
	watchdogQuery := fmt.Sprintf(`changes((max((ALERTS{alertstate="firing",alertname="Watchdog",severity="none"}) or (absent(ALERTS{alertstate="firing",alertname="Watchdog",severity="none"})*0)))[%s:1s]) > 1`, testDuration)
	result, err := helper.RunQuery(watchdogQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check watchdog alert over upgrade window")
	if len(result.Data.Result) > 0 {
		if result.Data.Result[0].Value <= 8 {
			unexpectedViolations.Insert(fmt.Sprintf("Watchdog alert had %s changes during the run, which may be a sign of a Prometheus outage in violation of the prometheus query SLO of 100%% uptime during upgrade", result.Data.Result[0].Value))
		} else {
			knownViolations.Insert(fmt.Sprintf("Watchdog alert had %s changes during the run, which may be a sign of a Prometheus outage in violation of the prometheus query SLO of 100%% uptime during upgrade, but is being tracked in https://bugzilla.redhat.com/show_bug.cgi?id=1949262 and is acceptable while that is open", result.Data.Result[0].Value))
		}
	}

	// Invariant: No non-info level alerts should have fired during the upgrade
	firingAlertQuery := fmt.Sprintf(`
sort_desc(
count_over_time(ALERTS{alertstate="firing",severity!="info",alertname!~"Watchdog|AlertmanagerReceiversNotConfigured"}[%[1]s:1s])
) > 0
`, testDuration)
	result, err = helper.RunQuery(firingAlertQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check firing alerts during upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s fired for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := firingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			unexpectedViolations.Insert(violation)
		}
	}

	// Invariant: There should be no pending alerts 1m after the upgrade completes
	pendingAlertQuery := fmt.Sprintf(`
sort_desc(
  time() * ALERTS + 1
  -
  last_over_time((
    time() * ALERTS{alertname!~"Watchdog|AlertmanagerReceiversNotConfigured",alertstate="pending",severity!="info"}
    unless
    ALERTS offset 1s
  )[%[1]s:1s])
)
`, testDuration)
	result, err = helper.RunQuery(pendingAlertQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to retrieve pending alerts after upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s pending for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedPendingAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := pendingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			// treat pending errors as a flake right now because we are still trying to determine the scope
			// TODO: move this to unexpectedViolations later
			unexpectedViolationsAsFlakes.Insert(violation)
		}
	}

	if len(debug) > 0 {
		framework.Logf("Alerts were detected during upgrade which are allowed:\n\n%s", strings.Join(debug.List(), "\n"))
	}
	if len(unexpectedViolations) > 0 {
		if !tolerateDuringSkew {
			framework.Failf("Unexpected alerts fired or pending during the upgrade:\n\n%s", strings.Join(unexpectedViolations.List(), "\n"))
		}
	}
	if flakes := sets.NewString().Union(knownViolations).Union(unexpectedViolations).Union(unexpectedViolationsAsFlakes); len(flakes) > 0 {
		disruption.FrameworkFlakef(f, "Unexpected alert behavior during upgrade:\n\n%s", strings.Join(flakes.List(), "\n"))
	}
	framework.Logf("No alerts fired during upgrade")
}

// Teardown cleans up any remaining resources.
func (t *UpgradeTest) Teardown(f *framework.Framework) {
	// rely on the namespace deletion to clean up everything
}
<commit_msg>test/e2e/upgrade/alert: Temporarily allow HighOverallControlPlaneCPU<commit_after>package alert

import (
	"context"
	"fmt"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
	"github.com/prometheus/common/model"

	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/disruption"
	helper "github.com/openshift/origin/test/extended/util/prometheus"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"
)

// UpgradeTest runs verifies invariants regarding what alerts are allowed to fire
// during the upgrade process.
type UpgradeTest struct {
	url         string
	bearerToken string
	oc          *exutil.CLI
}

func (UpgradeTest) Name() string { return "check-for-alerts" }
func (UpgradeTest) DisplayName() string {
	return "[sig-arch] Check if alerts are firing during or after upgrade success"
}

// Setup creates parameters to query Prometheus
func (t *UpgradeTest) Setup(f *framework.Framework) {
	g.By("Setting up upgrade alert test")

	oc := exutil.NewCLIWithFramework(f)
	url, _, bearerToken, ok := helper.LocatePrometheus(oc)
	if !ok {
		framework.Failf("Prometheus could not be located on this cluster, failing test %s", t.Name())
	}
	t.url = url
	t.bearerToken = bearerToken
	t.oc = oc
	framework.Logf("Post-upgrade alert test setup complete")
}

// Test checks if alerts are firing at various points during upgrade.
// An alert firing during an upgrade is a high severity bug - it either points to a real issue in
// a dependency, or a failure of the component, and therefore must be fixed.
func (t *UpgradeTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	tolerateDuringSkew := exutil.TolerateVersionSkewInTests()
	firingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "openshift-apiserver"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "machine-config"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955300",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubeDaemonSetRolloutStuck"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1943667",
		},
		{
			Selector: map[string]string{"alertname": "KubeAPIErrorBudgetBurn"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1953798",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "AggregatedAPIDown", "namespace": "default", "name": "v1beta1.metrics.k8s.io"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1970624",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "prometheus-k8s"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1949262",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "alertmanager-main"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955489",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed, or after that bug is fixed in a backport to the previous minor.
			Selector: map[string]string{"alertname": "HighOverallControlPlaneCPU"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1985073",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
	}

	pendingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "ClusterMonitoringOperatorReconciliationErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1932624",
		},
		{
			Selector: map[string]string{"alertname": "KubeClientErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1925698",
		},
	}
	allowedPendingAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "etcdMemberCommunicationSlow"},
			Text:     "Excluded because it triggers during upgrade (detects ~5m of high latency immediately preceeding the end of the test), and we don't want to change the alert because it is correct",
		},
	}

	knownViolations := sets.NewString()
	unexpectedViolations := sets.NewString()
	unexpectedViolationsAsFlakes := sets.NewString()
	debug := sets.NewString()

	g.By("Checking for alerts")

	start := time.Now()

	// Block until upgrade is done
	g.By("Waiting for upgrade to finish before checking for alerts")
	<-done

	// Additonal delay after upgrade completion to allow pending alerts to settle
	g.By("Waiting before checking for alerts")
	time.Sleep(1 * time.Minute)

	ns := t.oc.SetupNamespace()
	execPod := exutil.CreateExecPodOrFail(t.oc.AdminKubeClient(), ns, "execpod")
	defer func() {
		t.oc.AdminKubeClient().CoreV1().Pods(ns).Delete(context.Background(), execPod.Name, *metav1.NewDeleteOptions(1))
	}()

	testDuration := time.Now().Sub(start).Truncate(time.Second)

	// Invariant: The watchdog alert should be firing continuously during the whole upgrade via the thanos
	// querier (which should have no gaps when it queries the individual stores). Allow zero or one changes
	// to the presence of this series (zero if data is preserved over upgrade, one if data is lost on upgrade).
	// This would not catch the alert stopping firing, but we catch that in other places and tests.
	watchdogQuery := fmt.Sprintf(`changes((max((ALERTS{alertstate="firing",alertname="Watchdog",severity="none"}) or (absent(ALERTS{alertstate="firing",alertname="Watchdog",severity="none"})*0)))[%s:1s]) > 1`, testDuration)
	result, err := helper.RunQuery(watchdogQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check watchdog alert over upgrade window")
	if len(result.Data.Result) > 0 {
		if result.Data.Result[0].Value <= 8 {
			unexpectedViolations.Insert(fmt.Sprintf("Watchdog alert had %s changes during the run, which may be a sign of a Prometheus outage in violation of the prometheus query SLO of 100%% uptime during upgrade", result.Data.Result[0].Value))
		} else {
			knownViolations.Insert(fmt.Sprintf("Watchdog alert had %s changes during the run, which may be a sign of a Prometheus outage in violation of the prometheus query SLO of 100%% uptime during upgrade, but is being tracked in https://bugzilla.redhat.com/show_bug.cgi?id=1949262 and is acceptable while that is open", result.Data.Result[0].Value))
		}
	}

	// Invariant: No non-info level alerts should have fired during the upgrade
	firingAlertQuery := fmt.Sprintf(`
sort_desc(
count_over_time(ALERTS{alertstate="firing",severity!="info",alertname!~"Watchdog|AlertmanagerReceiversNotConfigured"}[%[1]s:1s])
) > 0
`, testDuration)
	result, err = helper.RunQuery(firingAlertQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check firing alerts during upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s fired for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := firingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			unexpectedViolations.Insert(violation)
		}
	}

	// Invariant: There should be no pending alerts 1m after the upgrade completes
	pendingAlertQuery := fmt.Sprintf(`
sort_desc(
  time() * ALERTS + 1
  -
  last_over_time((
    time() * ALERTS{alertname!~"Watchdog|AlertmanagerReceiversNotConfigured",alertstate="pending",severity!="info"}
    unless
    ALERTS offset 1s
  )[%[1]s:1s])
)
`, testDuration)
	result, err = helper.RunQuery(pendingAlertQuery, ns, execPod.Name, t.url, t.bearerToken)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to retrieve pending alerts after upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s pending for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedPendingAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := pendingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			// treat pending errors as a flake right now because we are still trying to determine the scope
			// TODO: move this to unexpectedViolations later
			unexpectedViolationsAsFlakes.Insert(violation)
		}
	}

	if len(debug) > 0 {
		framework.Logf("Alerts were detected during upgrade which are allowed:\n\n%s", strings.Join(debug.List(), "\n"))
	}
	if len(unexpectedViolations) > 0 {
		if !tolerateDuringSkew {
			framework.Failf("Unexpected alerts fired or pending during the upgrade:\n\n%s", strings.Join(unexpectedViolations.List(), "\n"))
		}
	}
	if flakes := sets.NewString().Union(knownViolations).Union(unexpectedViolations).Union(unexpectedViolationsAsFlakes); len(flakes) > 0 {
		disruption.FrameworkFlakef(f, "Unexpected alert behavior during upgrade:\n\n%s", strings.Join(flakes.List(), "\n"))
	}
	framework.Logf("No alerts fired during upgrade")
}

// Teardown cleans up any remaining resources.
func (t *UpgradeTest) Teardown(f *framework.Framework) {
	// rely on the namespace deletion to clean up everything
}
<|endoftext|>


<commit_before>// VNC client implementation.

package vnc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"reflect"

	"github.com/kward/go-vnc/go/metrics"
	"golang.org/x/net/context"
)

// Connect negotiates a connection to a VNC server.
func Connect(ctx context.Context, c net.Conn, cfg *ClientConfig) (*ClientConn, error) {
	conn := NewClientConn(c, cfg)

	if err := conn.processContext(ctx); err != nil {
		log.Fatal("invalid context: %v", err)
	}

	if err := conn.protocolVersionHandshake(ctx); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.securityHandshake(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.securityResultHandshake(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.clientInit(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.serverInit(); err != nil {
		conn.Close()
		return nil, err
	}

	return conn, nil
}

// A ClientConfig structure is used to configure a ClientConn. After
// one has been passed to initialize a connection, it must not be modified.
type ClientConfig struct {
	secType uint8 // The negotiated security type.

	// A slice of ClientAuth methods. Only the first instance that is
	// suitable by the server will be used to authenticate.
	Auth []ClientAuth

	// Password for servers that require authentication.
	Password string

	// Exclusive determines whether the connection is shared with other
	// clients. If true, then all other clients connected will be
	// disconnected when a connection is established to the VNC server.
	Exclusive bool

	// The channel that all messages received from the server will be
	// sent on. If the channel blocks, then the goroutine reading data
	// from the VNC server may block indefinitely. It is up to the user
	// of the library to ensure that this channel is properly read.
	// If this is not set, then all messages will be discarded.
	ServerMessageCh chan ServerMessage

	// A slice of supported messages that can be read from the server.
	// This only needs to contain NEW server messages, and doesn't
	// need to explicitly contain the RFC-required messages.
	ServerMessages []ServerMessage
}

// NewClientConfig returns a populated ClientConfig.
func NewClientConfig(p string) *ClientConfig {
	return &ClientConfig{
		Auth: []ClientAuth{
			&ClientAuthNone{},
			&ClientAuthVNC{p},
		},
		Password: p,
		ServerMessages: []ServerMessage{
			&FramebufferUpdate{},
			&SetColorMapEntries{},
			&Bell{},
			&ServerCutText{},
		},
	}
}

// The ClientConn type holds client connection information.
type ClientConn struct {
	c               net.Conn
	config          *ClientConfig
	protocolVersion string

	// If the pixel format uses a color map, then this is the color
	// map that is used. This should not be modified directly, since
	// the data comes from the server.
	// Definition in §5 - Representation of Pixel Data.
	colorMap ColorMap

	// Name associated with the desktop, sent from the server.
	desktopName string

	// Encodings supported by the client. This should not be modified
	// directly. Instead, SetEncodings should be used.
	encodings Encodings

	// Height of the frame buffer in pixels, sent from the server.
	fbHeight uint16

	// Width of the frame buffer in pixels, sent from the server.
	fbWidth uint16

	// The pixel format associated with the connection. This shouldn't
	// be modified. If you wish to set a new pixel format, use the
	// SetPixelFormat method.
	pixelFormat PixelFormat

	// Track metrics on system performance.
	metrics map[string]metrics.Metric

	// Flag to determine whether debug output should be provided.
	debug bool
}

func NewClientConn(c net.Conn, cfg *ClientConfig) *ClientConn {
	return &ClientConn{
		c:           c,
		config:      cfg,
		encodings:   Encodings{&RawEncoding{}},
		pixelFormat: PixelFormat32bit,
		metrics: map[string]metrics.Metric{
			"bytes-received": &metrics.Gauge{},
			"bytes-sent":     &metrics.Gauge{},
		},
	}
}

// Close a connection to a VNC server.
func (c *ClientConn) Close() error {
	log.Print("VNC Client connection closed.")
	return c.c.Close()
}

// DesktopName returns the server provided desktop name.
func (c *ClientConn) DesktopName() string {
	return c.desktopName
}

// setDesktopName stores the server provided desktop name.
func (c *ClientConn) setDesktopName(name string) {
	if c.debug {
		log.Printf("desktopName: %v", name)
	}
	c.desktopName = name
}

// Encodings returns the server provided encodings.
func (c *ClientConn) Encodings() Encodings {
	return c.encodings
}

// FramebufferHeight returns the server provided framebuffer height.
func (c *ClientConn) FramebufferHeight() uint16 {
	return c.fbHeight
}

// setFramebufferHeight stores the server provided framebuffer height.
func (c *ClientConn) setFramebufferHeight(height uint16) {
	if c.debug {
		log.Printf("framebufferHeight: %v", height)
	}
	c.fbHeight = height
}

// FramebufferWidth returns the server provided framebuffer width.
func (c *ClientConn) FramebufferWidth() uint16 {
	return c.fbWidth
}

// setFramebufferWidth stores the server provided framebuffer width.
func (c *ClientConn) setFramebufferWidth(width uint16) {
	if c.debug {
		log.Printf("framebufferWidth: %v", width)
	}
	c.fbWidth = width
}

// SetDebugging [dis-]enables debugging.
func (c *ClientConn) SetDebug(debug bool) {
	log.Printf("debug: %v", debug)
	c.debug = debug
}

// ListenAndHandle listens to a VNC server and handles server messages.
func (c *ClientConn) ListenAndHandle() error {
	if c.debug {
		log.Println("ListenAndHanndle()")
	}
	defer c.Close()

	if c.config.ServerMessages == nil {
		return NewVNCError("Client config error: ServerMessages undefined")
	}
	serverMessages := make(map[uint8]ServerMessage)
	for _, m := range c.config.ServerMessages {
		serverMessages[m.Type()] = m
	}

	for {
		var messageType uint8
		if err := c.receive(&messageType); err != nil {
			log.Print("error: reading from server")
			break
		}
		if c.debug {
			log.Printf("message-type: %v", messageType)
		}

		msg, ok := serverMessages[messageType]
		if !ok {
			// Unsupported message type! Bad!
			log.Printf("error unsupported message-type: %v", messageType)
			break
		}

		parsedMsg, err := msg.Read(c)
		if err != nil {
			log.Printf("error parsing message; %v", err)
			break
		}

		if c.config.ServerMessageCh == nil {
			log.Print("ignoring message; no server message channel")
			continue
		}

		c.config.ServerMessageCh <- parsedMsg
	}

	return nil
}

// receive a packet from the network.
func (c *ClientConn) receive(data interface{}) error {
	if err := binary.Read(c.c, binary.BigEndian, data); err != nil {
		return err
	}
	c.metrics["bytes-received"].Adjust(int64(binary.Size(data)))
	return nil
}

// receiveN receives N packets from the network.
func (c *ClientConn) receiveN(data interface{}, n int) error {
	if n == 0 {
		log.Println("warn: receiveN requested no data")
		return nil
	}

	switch data.(type) {
	case *[]uint8:
		var v uint8
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			slice := data.(*[]uint8)
			*slice = append(*slice, v)
		}
	case *[]int32:
		var v int32
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			slice := data.(*[]int32)
			*slice = append(*slice, v)
		}
	case *bytes.Buffer:
		var v byte
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			buf := data.(*bytes.Buffer)
			buf.WriteByte(v)
		}
	default:
		return NewVNCError(fmt.Sprintf("unrecognized data type %v", reflect.TypeOf(data)))
	}
	c.metrics["bytes-received"].Adjust(int64(binary.Size(data)))
	return nil
}

// send a packet to the network.
func (c *ClientConn) send(data interface{}) error {
	if err := binary.Write(c.c, binary.BigEndian, data); err != nil {
		return err
	}
	c.metrics["bytes-sent"].Adjust(int64(binary.Size(data)))
	return nil
}

// sendN sends N packets to the network.
// func (c *ClientConn) sendN(data interface{}, n int) error {
// 	var buf bytes.Buffer
// 	switch data := data.(type) {
// 	case []uint8:
// 		for _, d := range data {
// 			if err := binary.Write(&buf, binary.BigEndian, &d); err != nil {
// 				return err
// 			}
// 		}
// 	case []int32:
// 		for _, d := range data {
// 			if err := binary.Write(&buf, binary.BigEndian, &d); err != nil {
// 				return err
// 			}
// 		}
// 	default:
// 		return NewVNCError(fmt.Sprintf("unable to send data; unrecognized data type %v", reflect.TypeOf(data)))
// 	}
// 	if err := binary.Write(c.c, binary.BigEndian, buf.Bytes()); err != nil {
// 		return err
// 	}
// 	c.metrics["bytes-sent"].Adjust(int64(binary.Size(data)))
// 	return nil
// }

func (c *ClientConn) processContext(ctx context.Context) error {
	if debug := ctx.Value("debug"); debug != nil {
		c.SetDebug(debug.(bool))
	}

	if mpv := ctx.Value("vnc_max_proto_version"); mpv != nil && mpv != "" {
		log.Printf("vnc_max_proto_version: %v", mpv)
		vers := []string{"3.3", "3.8"}
		valid := false
		for _, v := range vers {
			if mpv == v {
				valid = true
				break
			}
		}
		if !valid {
			return fmt.Errorf("Invalid max protocol version %v; supported versions are %v", mpv, vers)
		}
	}

	return nil
}
<commit_msg>added DebugMetrics()<commit_after>// VNC client implementation.

package vnc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"reflect"

	"github.com/kward/go-vnc/go/metrics"
	"golang.org/x/net/context"
)

// Connect negotiates a connection to a VNC server.
func Connect(ctx context.Context, c net.Conn, cfg *ClientConfig) (*ClientConn, error) {
	conn := NewClientConn(c, cfg)

	if err := conn.processContext(ctx); err != nil {
		log.Fatal("invalid context: %v", err)
	}

	if err := conn.protocolVersionHandshake(ctx); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.securityHandshake(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.securityResultHandshake(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.clientInit(); err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.serverInit(); err != nil {
		conn.Close()
		return nil, err
	}

	return conn, nil
}

// A ClientConfig structure is used to configure a ClientConn. After
// one has been passed to initialize a connection, it must not be modified.
type ClientConfig struct {
	secType uint8 // The negotiated security type.

	// A slice of ClientAuth methods. Only the first instance that is
	// suitable by the server will be used to authenticate.
	Auth []ClientAuth

	// Password for servers that require authentication.
	Password string

	// Exclusive determines whether the connection is shared with other
	// clients. If true, then all other clients connected will be
	// disconnected when a connection is established to the VNC server.
	Exclusive bool

	// The channel that all messages received from the server will be
	// sent on. If the channel blocks, then the goroutine reading data
	// from the VNC server may block indefinitely. It is up to the user
	// of the library to ensure that this channel is properly read.
	// If this is not set, then all messages will be discarded.
	ServerMessageCh chan ServerMessage

	// A slice of supported messages that can be read from the server.
	// This only needs to contain NEW server messages, and doesn't
	// need to explicitly contain the RFC-required messages.
	ServerMessages []ServerMessage
}

// NewClientConfig returns a populated ClientConfig.
func NewClientConfig(p string) *ClientConfig {
	return &ClientConfig{
		Auth: []ClientAuth{
			&ClientAuthNone{},
			&ClientAuthVNC{p},
		},
		Password: p,
		ServerMessages: []ServerMessage{
			&FramebufferUpdate{},
			&SetColorMapEntries{},
			&Bell{},
			&ServerCutText{},
		},
	}
}

// The ClientConn type holds client connection information.
type ClientConn struct {
	c               net.Conn
	config          *ClientConfig
	protocolVersion string

	// If the pixel format uses a color map, then this is the color
	// map that is used. This should not be modified directly, since
	// the data comes from the server.
	// Definition in §5 - Representation of Pixel Data.
	colorMap ColorMap

	// Name associated with the desktop, sent from the server.
	desktopName string

	// Encodings supported by the client. This should not be modified
	// directly. Instead, SetEncodings should be used.
	encodings Encodings

	// Height of the frame buffer in pixels, sent from the server.
	fbHeight uint16

	// Width of the frame buffer in pixels, sent from the server.
	fbWidth uint16

	// The pixel format associated with the connection. This shouldn't
	// be modified. If you wish to set a new pixel format, use the
	// SetPixelFormat method.
	pixelFormat PixelFormat

	// Track metrics on system performance.
	metrics map[string]metrics.Metric

	// Flag to determine whether debug output should be provided.
	debug bool
}

func NewClientConn(c net.Conn, cfg *ClientConfig) *ClientConn {
	return &ClientConn{
		c:           c,
		config:      cfg,
		encodings:   Encodings{&RawEncoding{}},
		pixelFormat: PixelFormat32bit,
		metrics: map[string]metrics.Metric{
			"bytes-received": &metrics.Gauge{},
			"bytes-sent":     &metrics.Gauge{},
		},
	}
}

// Close a connection to a VNC server.
func (c *ClientConn) Close() error {
	log.Print("VNC Client connection closed.")
	return c.c.Close()
}

// DesktopName returns the server provided desktop name.
func (c *ClientConn) DesktopName() string {
	return c.desktopName
}

// setDesktopName stores the server provided desktop name.
func (c *ClientConn) setDesktopName(name string) {
	if c.debug {
		log.Printf("desktopName: %v", name)
	}
	c.desktopName = name
}

// Encodings returns the server provided encodings.
func (c *ClientConn) Encodings() Encodings {
	return c.encodings
}

// FramebufferHeight returns the server provided framebuffer height.
func (c *ClientConn) FramebufferHeight() uint16 {
	return c.fbHeight
}

// setFramebufferHeight stores the server provided framebuffer height.
func (c *ClientConn) setFramebufferHeight(height uint16) {
	if c.debug {
		log.Printf("framebufferHeight: %v", height)
	}
	c.fbHeight = height
}

// FramebufferWidth returns the server provided framebuffer width.
func (c *ClientConn) FramebufferWidth() uint16 {
	return c.fbWidth
}

// setFramebufferWidth stores the server provided framebuffer width.
func (c *ClientConn) setFramebufferWidth(width uint16) {
	if c.debug {
		log.Printf("framebufferWidth: %v", width)
	}
	c.fbWidth = width
}

// SetDebugging [dis-]enables debugging.
func (c *ClientConn) SetDebug(debug bool) {
	log.Printf("debug: %v", debug)
	c.debug = debug
}

// ListenAndHandle listens to a VNC server and handles server messages.
func (c *ClientConn) ListenAndHandle() error {
	if c.debug {
		log.Println("ListenAndHanndle()")
	}
	defer c.Close()

	if c.config.ServerMessages == nil {
		return NewVNCError("Client config error: ServerMessages undefined")
	}
	serverMessages := make(map[uint8]ServerMessage)
	for _, m := range c.config.ServerMessages {
		serverMessages[m.Type()] = m
	}

	for {
		var messageType uint8
		if err := c.receive(&messageType); err != nil {
			log.Print("error: reading from server")
			break
		}
		if c.debug {
			log.Printf("message-type: %v", messageType)
		}

		msg, ok := serverMessages[messageType]
		if !ok {
			// Unsupported message type! Bad!
			log.Printf("error unsupported message-type: %v", messageType)
			break
		}

		parsedMsg, err := msg.Read(c)
		if err != nil {
			log.Printf("error parsing message; %v", err)
			break
		}

		if c.config.ServerMessageCh == nil {
			log.Print("ignoring message; no server message channel")
			continue
		}

		c.config.ServerMessageCh <- parsedMsg
	}

	return nil
}

// receive a packet from the network.
func (c *ClientConn) receive(data interface{}) error {
	if err := binary.Read(c.c, binary.BigEndian, data); err != nil {
		return err
	}
	c.metrics["bytes-received"].Adjust(int64(binary.Size(data)))
	return nil
}

// receiveN receives N packets from the network.
func (c *ClientConn) receiveN(data interface{}, n int) error {
	if n == 0 {
		log.Println("warn: receiveN requested no data")
		return nil
	}

	switch data.(type) {
	case *[]uint8:
		var v uint8
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			slice := data.(*[]uint8)
			*slice = append(*slice, v)
		}
	case *[]int32:
		var v int32
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			slice := data.(*[]int32)
			*slice = append(*slice, v)
		}
	case *bytes.Buffer:
		var v byte
		for i := 0; i < n; i++ {
			if err := binary.Read(c.c, binary.BigEndian, &v); err != nil {
				return err
			}
			buf := data.(*bytes.Buffer)
			buf.WriteByte(v)
		}
	default:
		return NewVNCError(fmt.Sprintf("unrecognized data type %v", reflect.TypeOf(data)))
	}
	c.metrics["bytes-received"].Adjust(int64(binary.Size(data)))
	return nil
}

// send a packet to the network.
func (c *ClientConn) send(data interface{}) error {
	if err := binary.Write(c.c, binary.BigEndian, data); err != nil {
		return err
	}
	c.metrics["bytes-sent"].Adjust(int64(binary.Size(data)))
	return nil
}

// sendN sends N packets to the network.
// func (c *ClientConn) sendN(data interface{}, n int) error {
// 	var buf bytes.Buffer
// 	switch data := data.(type) {
// 	case []uint8:
// 		for _, d := range data {
// 			if err := binary.Write(&buf, binary.BigEndian, &d); err != nil {
// 				return err
// 			}
// 		}
// 	case []int32:
// 		for _, d := range data {
// 			if err := binary.Write(&buf, binary.BigEndian, &d); err != nil {
// 				return err
// 			}
// 		}
// 	default:
// 		return NewVNCError(fmt.Sprintf("unable to send data; unrecognized data type %v", reflect.TypeOf(data)))
// 	}
// 	if err := binary.Write(c.c, binary.BigEndian, buf.Bytes()); err != nil {
// 		return err
// 	}
// 	c.metrics["bytes-sent"].Adjust(int64(binary.Size(data)))
// 	return nil
// }

func (c *ClientConn) processContext(ctx context.Context) error {
	if debug := ctx.Value("debug"); debug != nil {
		c.SetDebug(debug.(bool))
	}

	if mpv := ctx.Value("vnc_max_proto_version"); mpv != nil && mpv != "" {
		log.Printf("vnc_max_proto_version: %v", mpv)
		vers := []string{"3.3", "3.8"}
		valid := false
		for _, v := range vers {
			if mpv == v {
				valid = true
				break
			}
		}
		if !valid {
			return fmt.Errorf("Invalid max protocol version %v; supported versions are %v", mpv, vers)
		}
	}

	return nil
}

func (c *ClientConn) DebugMetrics() {
	log.Println("Metrics:")
	for name, metric := range c.metrics {
		log.Printf("  %v: %v", name, metric.Value())
	}
}
<|endoftext|>