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

/*
Package packagestest creates temporary projects on disk for testing go tools on.

By changing the exporter used, you can create projects for multiple build
systems from the same description, and run the same tests on them in many
cases.
*/
package packagestest

import (
	"flag"
	"fmt"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"golang.org/x/tools/go/expect"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/internal/span"
)

var (
	skipCleanup = flag.Bool("skip-cleanup", false, "Do not delete the temporary export folders") // for debugging
)

// Module is a representation of a go module.
type Module struct {
	// Name is the base name of the module as it would be in the go.mod file.
	Name string
	// Files is the set of source files for all packages that make up the module.
	// The keys are the file fragment that follows the module name, the value can
	// be a string or byte slice, in which case it is the contents of the
	// file, otherwise it must be a Writer function.
	Files map[string]interface{}

	// Overlay is the set of source file overlays for the module.
	// The keys are the file fragment as in the Files configuration.
	// The values are the in memory overlay content for the file.
	Overlay map[string][]byte
}

// A Writer is a function that writes out a test file.
// It is provided the name of the file to write, and may return an error if it
// cannot write the file.
// These are used as the content of the Files map in a Module.
type Writer func(filename string) error

// Exported is returned by the Export function to report the structure that was produced on disk.
type Exported struct {
	// Config is a correctly configured packages.Config ready to be passed to packages.Load.
	// Exactly what it will contain varies depending on the Exporter being used.
	Config *packages.Config

	// Modules is the module description that was used to produce this exported data set.
	Modules []Module

	ExpectFileSet *token.FileSet // The file set used when parsing expectations

	temp    string                       // the temporary directory that was exported to
	primary string                       // the first non GOROOT module that was exported
	written map[string]map[string]string // the full set of exported files
	notes   []*expect.Note               // The list of expectations extracted from go source files
	markers map[string]span.Range        // The set of markers extracted from go source files
}

// Exporter implementations are responsible for converting from the generic description of some
// test data to a driver specific file layout.
type Exporter interface {
	// Name reports the name of the exporter, used in logging and sub-test generation.
	Name() string
	// Filename reports the system filename for test data source file.
	// It is given the base directory, the module the file is part of and the filename fragment to
	// work from.
	Filename(exported *Exported, module, fragment string) string
	// Finalize is called once all files have been written to write any extra data needed and modify
	// the Config to match. It is handed the full list of modules that were encountered while writing
	// files.
	Finalize(exported *Exported) error
}

// All is the list of known exporters.
// This is used by TestAll to run tests with all the exporters.
var All []Exporter

// TestAll invokes the testing function once for each exporter registered in
// the All global.
// Each exporter will be run as a sub-test named after the exporter being used.
func TestAll(t *testing.T, f func(*testing.T, Exporter)) {
	t.Helper()
	for _, e := range All {
		t.Run(e.Name(), func(t *testing.T) {
			t.Helper()
			f(t, e)
		})
	}
}

// BenchmarkAll invokes the testing function once for each exporter registered in
// the All global.
// Each exporter will be run as a sub-test named after the exporter being used.
func BenchmarkAll(b *testing.B, f func(*testing.B, Exporter)) {
	b.Helper()
	for _, e := range All {
		b.Run(e.Name(), func(b *testing.B) {
			b.Helper()
			f(b, e)
		})
	}
}

// Export is called to write out a test directory from within a test function.
// It takes the exporter and the build system agnostic module descriptions, and
// uses them to build a temporary directory.
// It returns an Exported with the results of the export.
// The Exported.Config is prepared for loading from the exported data.
// You must invoke Exported.Cleanup on the returned value to clean up.
// The file deletion in the cleanup can be skipped by setting the skip-cleanup
// flag when invoking the test, allowing the temporary directory to be left for
// debugging tests.
func Export(t testing.TB, exporter Exporter, modules []Module) *Exported {
	t.Helper()
	dirname := strings.Replace(t.Name(), "/", "_", -1)
	dirname = strings.Replace(dirname, "#", "_", -1) // duplicate subtests get a #NNN suffix.
	temp, err := ioutil.TempDir("", dirname)
	if err != nil {
		t.Fatal(err)
	}
	exported := &Exported{
		Config: &packages.Config{
			Dir:     temp,
			Env:     append(os.Environ(), "GOPACKAGESDRIVER=off"),
			Overlay: make(map[string][]byte),
			Tests:   true,
			Mode:    packages.LoadImports,
		},
		Modules:       modules,
		temp:          temp,
		primary:       modules[0].Name,
		written:       map[string]map[string]string{},
		ExpectFileSet: token.NewFileSet(),
	}
	defer func() {
		if t.Failed() || t.Skipped() {
			exported.Cleanup()
		}
	}()
	for _, module := range modules {
		for fragment, value := range module.Files {
			fullpath := exporter.Filename(exported, module.Name, filepath.FromSlash(fragment))
			written, ok := exported.written[module.Name]
			if !ok {
				written = map[string]string{}
				exported.written[module.Name] = written
			}
			written[fragment] = fullpath
			if err := os.MkdirAll(filepath.Dir(fullpath), 0755); err != nil {
				t.Fatal(err)
			}
			switch value := value.(type) {
			case Writer:
				if err := value(fullpath); err != nil {
					t.Fatal(err)
				}
			case string:
				if err := ioutil.WriteFile(fullpath, []byte(value), 0644); err != nil {
					t.Fatal(err)
				}
			default:
				t.Fatalf("Invalid type %T in files, must be string or Writer", value)
			}
		}
		for fragment, value := range module.Overlay {
			fullpath := exporter.Filename(exported, module.Name, filepath.FromSlash(fragment))
			exported.Config.Overlay[fullpath] = value
		}
	}
	if err := exporter.Finalize(exported); err != nil {
		t.Fatal(err)
	}
	return exported
}

// Script returns a Writer that writes out contents to the file and sets the
// executable bit on the created file.
// It is intended for source files that are shell scripts.
func Script(contents string) Writer {
	return func(filename string) error {
		return ioutil.WriteFile(filename, []byte(contents), 0755)
	}
}

// Link returns a Writer that creates a hard link from the specified source to
// the required file.
// This is used to link testdata files into the generated testing tree.
func Link(source string) Writer {
	return func(filename string) error {
		return os.Link(source, filename)
	}
}

// Symlink returns a Writer that creates a symlink from the specified source to the
// required file.
// This is used to link testdata files into the generated testing tree.
func Symlink(source string) Writer {
	if !strings.HasPrefix(source, ".") {
		if abspath, err := filepath.Abs(source); err == nil {
			if _, err := os.Stat(source); !os.IsNotExist(err) {
				source = abspath
			}
		}
	}
	return func(filename string) error {
		return os.Symlink(source, filename)
	}
}

// Copy returns a Writer that copies a file from the specified source to the
// required file.
// This is used to copy testdata files into the generated testing tree.
func Copy(source string) Writer {
	return func(filename string) error {
		stat, err := os.Stat(source)
		if err != nil {
			return err
		}
		if !stat.Mode().IsRegular() {
			// cannot copy non-regular files (e.g., directories,
			// symlinks, devices, etc.)
			return fmt.Errorf("Cannot copy non regular file %s", source)
		}
		contents, err := ioutil.ReadFile(source)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(filename, contents, stat.Mode())
	}
}

// MustCopyFileTree returns a file set for a module based on a real directory tree.
// It scans the directory tree anchored at root and adds a Copy writer to the
// map for every file found.
// This is to enable the common case in tests where you have a full copy of the
// package in your testdata.
// This will panic if there is any kind of error trying to walk the file tree.
func MustCopyFileTree(root string) map[string]interface{} {
	result := map[string]interface{}{}
	if err := filepath.Walk(filepath.FromSlash(root), func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		fragment, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		result[fragment] = Copy(path)
		return nil
	}); err != nil {
		log.Panic(fmt.Sprintf("MustCopyFileTree failed: %v", err))
	}
	return result
}

// Cleanup removes the temporary directory (unless the --skip-cleanup flag was set)
// It is safe to call cleanup multiple times.
func (e *Exported) Cleanup() {
	if e.temp == "" {
		return
	}
	if *skipCleanup {
		log.Printf("Skipping cleanup of temp dir: %s", e.temp)
		return
	}
	// Make everything read-write so that the Module exporter's module cache can be deleted.
	filepath.Walk(e.temp, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			os.Chmod(path, 0777)
		}
		return nil
	})
	os.RemoveAll(e.temp) // ignore errors
	e.temp = ""
}

// Temp returns the temporary directory that was generated.
func (e *Exported) Temp() string {
	return e.temp
}

// File returns the full path for the given module and file fragment.
func (e *Exported) File(module, fragment string) string {
	if m := e.written[module]; m != nil {
		return m[fragment]
	}
	return ""
}

// FileContents returns the contents of the specified file.
// It will use the overlay if the file is present, otherwise it will read it
// from disk.
func (e *Exported) FileContents(filename string) ([]byte, error) {
	if content, found := e.Config.Overlay[filename]; found {
		return content, nil
	}
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return content, nil
}
<commit_msg>go/packages/packagestest: fix MustCopyFileTree so that file fragments are always slash form<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package packagestest creates temporary projects on disk for testing go tools on.

By changing the exporter used, you can create projects for multiple build
systems from the same description, and run the same tests on them in many
cases.
*/
package packagestest

import (
	"flag"
	"fmt"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"golang.org/x/tools/go/expect"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/internal/span"
)

var (
	skipCleanup = flag.Bool("skip-cleanup", false, "Do not delete the temporary export folders") // for debugging
)

// Module is a representation of a go module.
type Module struct {
	// Name is the base name of the module as it would be in the go.mod file.
	Name string
	// Files is the set of source files for all packages that make up the module.
	// The keys are the file fragment that follows the module name, the value can
	// be a string or byte slice, in which case it is the contents of the
	// file, otherwise it must be a Writer function.
	Files map[string]interface{}

	// Overlay is the set of source file overlays for the module.
	// The keys are the file fragment as in the Files configuration.
	// The values are the in memory overlay content for the file.
	Overlay map[string][]byte
}

// A Writer is a function that writes out a test file.
// It is provided the name of the file to write, and may return an error if it
// cannot write the file.
// These are used as the content of the Files map in a Module.
type Writer func(filename string) error

// Exported is returned by the Export function to report the structure that was produced on disk.
type Exported struct {
	// Config is a correctly configured packages.Config ready to be passed to packages.Load.
	// Exactly what it will contain varies depending on the Exporter being used.
	Config *packages.Config

	// Modules is the module description that was used to produce this exported data set.
	Modules []Module

	ExpectFileSet *token.FileSet // The file set used when parsing expectations

	temp    string                       // the temporary directory that was exported to
	primary string                       // the first non GOROOT module that was exported
	written map[string]map[string]string // the full set of exported files
	notes   []*expect.Note               // The list of expectations extracted from go source files
	markers map[string]span.Range        // The set of markers extracted from go source files
}

// Exporter implementations are responsible for converting from the generic description of some
// test data to a driver specific file layout.
type Exporter interface {
	// Name reports the name of the exporter, used in logging and sub-test generation.
	Name() string
	// Filename reports the system filename for test data source file.
	// It is given the base directory, the module the file is part of and the filename fragment to
	// work from.
	Filename(exported *Exported, module, fragment string) string
	// Finalize is called once all files have been written to write any extra data needed and modify
	// the Config to match. It is handed the full list of modules that were encountered while writing
	// files.
	Finalize(exported *Exported) error
}

// All is the list of known exporters.
// This is used by TestAll to run tests with all the exporters.
var All []Exporter

// TestAll invokes the testing function once for each exporter registered in
// the All global.
// Each exporter will be run as a sub-test named after the exporter being used.
func TestAll(t *testing.T, f func(*testing.T, Exporter)) {
	t.Helper()
	for _, e := range All {
		t.Run(e.Name(), func(t *testing.T) {
			t.Helper()
			f(t, e)
		})
	}
}

// BenchmarkAll invokes the testing function once for each exporter registered in
// the All global.
// Each exporter will be run as a sub-test named after the exporter being used.
func BenchmarkAll(b *testing.B, f func(*testing.B, Exporter)) {
	b.Helper()
	for _, e := range All {
		b.Run(e.Name(), func(b *testing.B) {
			b.Helper()
			f(b, e)
		})
	}
}

// Export is called to write out a test directory from within a test function.
// It takes the exporter and the build system agnostic module descriptions, and
// uses them to build a temporary directory.
// It returns an Exported with the results of the export.
// The Exported.Config is prepared for loading from the exported data.
// You must invoke Exported.Cleanup on the returned value to clean up.
// The file deletion in the cleanup can be skipped by setting the skip-cleanup
// flag when invoking the test, allowing the temporary directory to be left for
// debugging tests.
func Export(t testing.TB, exporter Exporter, modules []Module) *Exported {
	t.Helper()
	dirname := strings.Replace(t.Name(), "/", "_", -1)
	dirname = strings.Replace(dirname, "#", "_", -1) // duplicate subtests get a #NNN suffix.
	temp, err := ioutil.TempDir("", dirname)
	if err != nil {
		t.Fatal(err)
	}
	exported := &Exported{
		Config: &packages.Config{
			Dir:     temp,
			Env:     append(os.Environ(), "GOPACKAGESDRIVER=off"),
			Overlay: make(map[string][]byte),
			Tests:   true,
			Mode:    packages.LoadImports,
		},
		Modules:       modules,
		temp:          temp,
		primary:       modules[0].Name,
		written:       map[string]map[string]string{},
		ExpectFileSet: token.NewFileSet(),
	}
	defer func() {
		if t.Failed() || t.Skipped() {
			exported.Cleanup()
		}
	}()
	for _, module := range modules {
		for fragment, value := range module.Files {
			fullpath := exporter.Filename(exported, module.Name, filepath.FromSlash(fragment))
			written, ok := exported.written[module.Name]
			if !ok {
				written = map[string]string{}
				exported.written[module.Name] = written
			}
			written[fragment] = fullpath
			if err := os.MkdirAll(filepath.Dir(fullpath), 0755); err != nil {
				t.Fatal(err)
			}
			switch value := value.(type) {
			case Writer:
				if err := value(fullpath); err != nil {
					t.Fatal(err)
				}
			case string:
				if err := ioutil.WriteFile(fullpath, []byte(value), 0644); err != nil {
					t.Fatal(err)
				}
			default:
				t.Fatalf("Invalid type %T in files, must be string or Writer", value)
			}
		}
		for fragment, value := range module.Overlay {
			fullpath := exporter.Filename(exported, module.Name, filepath.FromSlash(fragment))
			exported.Config.Overlay[fullpath] = value
		}
	}
	if err := exporter.Finalize(exported); err != nil {
		t.Fatal(err)
	}
	return exported
}

// Script returns a Writer that writes out contents to the file and sets the
// executable bit on the created file.
// It is intended for source files that are shell scripts.
func Script(contents string) Writer {
	return func(filename string) error {
		return ioutil.WriteFile(filename, []byte(contents), 0755)
	}
}

// Link returns a Writer that creates a hard link from the specified source to
// the required file.
// This is used to link testdata files into the generated testing tree.
func Link(source string) Writer {
	return func(filename string) error {
		return os.Link(source, filename)
	}
}

// Symlink returns a Writer that creates a symlink from the specified source to the
// required file.
// This is used to link testdata files into the generated testing tree.
func Symlink(source string) Writer {
	if !strings.HasPrefix(source, ".") {
		if abspath, err := filepath.Abs(source); err == nil {
			if _, err := os.Stat(source); !os.IsNotExist(err) {
				source = abspath
			}
		}
	}
	return func(filename string) error {
		return os.Symlink(source, filename)
	}
}

// Copy returns a Writer that copies a file from the specified source to the
// required file.
// This is used to copy testdata files into the generated testing tree.
func Copy(source string) Writer {
	return func(filename string) error {
		stat, err := os.Stat(source)
		if err != nil {
			return err
		}
		if !stat.Mode().IsRegular() {
			// cannot copy non-regular files (e.g., directories,
			// symlinks, devices, etc.)
			return fmt.Errorf("Cannot copy non regular file %s", source)
		}
		contents, err := ioutil.ReadFile(source)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(filename, contents, stat.Mode())
	}
}

// MustCopyFileTree returns a file set for a module based on a real directory tree.
// It scans the directory tree anchored at root and adds a Copy writer to the
// map for every file found.
// This is to enable the common case in tests where you have a full copy of the
// package in your testdata.
// This will panic if there is any kind of error trying to walk the file tree.
func MustCopyFileTree(root string) map[string]interface{} {
	result := map[string]interface{}{}
	if err := filepath.Walk(filepath.FromSlash(root), func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		fragment, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		result[filepath.ToSlash(fragment)] = Copy(path)
		return nil
	}); err != nil {
		log.Panic(fmt.Sprintf("MustCopyFileTree failed: %v", err))
	}
	return result
}

// Cleanup removes the temporary directory (unless the --skip-cleanup flag was set)
// It is safe to call cleanup multiple times.
func (e *Exported) Cleanup() {
	if e.temp == "" {
		return
	}
	if *skipCleanup {
		log.Printf("Skipping cleanup of temp dir: %s", e.temp)
		return
	}
	// Make everything read-write so that the Module exporter's module cache can be deleted.
	filepath.Walk(e.temp, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			os.Chmod(path, 0777)
		}
		return nil
	})
	os.RemoveAll(e.temp) // ignore errors
	e.temp = ""
}

// Temp returns the temporary directory that was generated.
func (e *Exported) Temp() string {
	return e.temp
}

// File returns the full path for the given module and file fragment.
func (e *Exported) File(module, fragment string) string {
	if m := e.written[module]; m != nil {
		return m[fragment]
	}
	return ""
}

// FileContents returns the contents of the specified file.
// It will use the overlay if the file is present, otherwise it will read it
// from disk.
func (e *Exported) FileContents(filename string) ([]byte, error) {
	if content, found := e.Config.Overlay[filename]; found {
		return content, nil
	}
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return content, nil
}
<|endoftext|>


<commit_before>package sched

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/tsaf/conf"
	"github.com/StackExchange/tsaf/expr"
)

type Schedule struct {
	sync.Mutex

	Conf          *conf.Conf
	Freq          time.Duration
	Status        map[AlertKey]*State
	Notifications map[AlertKey]map[string]time.Time

	cache     *opentsdb.Cache
	runStates map[AlertKey]Status
	nc        chan interface{}
}

func (s *Schedule) MarshalJSON() ([]byte, error) {
	t := struct {
		Alerts map[string]*conf.Alert
		Freq   time.Duration
		Status map[string]*State
	}{
		s.Conf.Alerts,
		s.Freq,
		make(map[string]*State),
	}
	for k, v := range s.Status {
		if v.Last().Status < ST_WARN {
			continue
		}
		t.Status[k.String()] = v
	}
	return json.Marshal(&t)
}

var DefaultSched = &Schedule{
	Freq: time.Minute * 5,
}

// Loads a configuration into the default schedule
func Load(c *conf.Conf) {
	DefaultSched.Load(c)
}

// Runs the default schedule.
func Run() error {
	return DefaultSched.Run()
}

func (s *Schedule) Load(c *conf.Conf) {
	s.Conf = c
	s.RestoreState()
}

// Restores notification and alert state from the file on disk.
func (s *Schedule) RestoreState() {
	s.Lock()
	defer s.Unlock()
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	s.Notifications = nil
	s.Status = make(map[AlertKey]*State)
	f, err := os.Open(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	dec := gob.NewDecoder(f)
	notifications := make(map[AlertKey]map[string]time.Time)
	if err := dec.Decode(&notifications); err != nil {
		log.Println(err)
		return
	}
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				log.Println("sched: notification not present during restore:", name)
				continue
			}
			_, present = s.Conf.Alerts[ak.Name]
			if !present {
				log.Println("sched: alert not present during restore:", ak.Name)
				continue
			}
			s.AddNotification(ak, n, t)
		}
	}
	for {
		var ak AlertKey
		var st State
		if err := dec.Decode(&ak); err == io.EOF {
			break
		} else if err != nil {
			log.Println(err)
			return
		}
		if err := dec.Decode(&st); err != nil {
			log.Println(err)
			return
		}
		if _, present := s.Conf.Alerts[ak.Name]; !present {
			log.Println("sched: alert no longer present, ignoring:", ak)
			continue
		}
		s.Status[ak] = &st
	}
}

func (s *Schedule) Save() {
	s.Lock()
	defer s.Unlock()
	f, err := os.Create(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	enc := gob.NewEncoder(f)
	if err := enc.Encode(s.Notifications); err != nil {
		log.Println(err)
		return
	}
	for k, v := range s.Status {
		enc.Encode(k)
		enc.Encode(v)
	}
	if err := f.Close(); err != nil {
		log.Println(err)
		return
	}
	log.Println("sched: wrote state to", s.Conf.StateFile)
}

func (s *Schedule) Run() error {
	go func() {
		for {
			time.Sleep(time.Minute)
			s.Save()
		}
	}()
	s.nc = make(chan interface{}, 1)
	go s.Poll()
	for {
		wait := time.After(s.Freq)
		if s.Freq < time.Second {
			return fmt.Errorf("sched: frequency must be > 1 second")
		}
		if s.Conf == nil {
			return fmt.Errorf("sched: nil configuration")
		}
		start := time.Now()
		s.Check()
		fmt.Printf("run at %v took %v\n", start, time.Since(start))
		<-wait
	}
}

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	var timeout time.Duration
	for {
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
		timeout = s.CheckNotifications()
	}
}

// CheckNotifications processes past notification events. It returns the
func (s *Schedule) CheckNotifications() time.Duration {
	s.Lock()
	defer s.Unlock()
	timeout := time.Hour
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				if remaining < timeout {
					timeout = remaining
				}
				s.AddNotification(ak, n, t)
				continue
			}
			st, present := s.Status[ak]
			if !present {
				continue
			}
			a, present := s.Conf.Alerts[ak.Name]
			if !present {
				continue
			}
			s.Notify(st, a, n)
			if n.Timeout < timeout {
				timeout = n.Timeout
			}
		}
	}
	return timeout
}

func (s *Schedule) Check() {
	s.Lock()
	defer s.Unlock()
	s.runStates = make(map[AlertKey]Status)
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	for _, a := range s.Conf.Alerts {
		s.CheckAlert(a)
	}
	changed := false
	for ak, status := range s.runStates {
		state := s.Status[ak]
		change := state.Append(status)
		if change {
			changed = true
		}
		if change && status > ST_NORM {
			var subject = new(bytes.Buffer)
			a := s.Conf.Alerts[ak.Name]
			if err := s.ExecuteSubject(subject, a, state); err != nil {
				log.Println(err)
			}
			state.Subject = subject.String()
			notify := func(notifications map[string]*conf.Notification) {
				for _, n := range notifications {
					s.Notify(state, a, n)
				}
			}
			switch status {
			case ST_CRIT:
				notify(a.CritNotification)
			case ST_WARN:
				notify(a.WarnNotification)
			}
		}
	}
	if changed {
		s.nc <- true
	}
}

func (s *Schedule) CheckAlert(a *conf.Alert) {
	crits := s.CheckExpr(a, a.Crit, true, nil)
	s.CheckExpr(a, a.Warn, false, crits)
}

func (s *Schedule) CheckExpr(a *conf.Alert, e *expr.Expr, isCrit bool, ignore []AlertKey) (alerts []AlertKey) {
	if e == nil {
		return
	}
	results, _, err := e.Execute(s.cache, nil)
	if err != nil {
		// todo: do something here?
		log.Println(err)
		return
	}
Loop:
	for _, r := range results {
		if a.Squelched(r.Group) {
			continue
		}
		ak := AlertKey{a.Name, r.Group.String()}
		for _, v := range ignore {
			if ak == v {
				continue Loop
			}
		}
		state := s.Status[ak]
		if state == nil {
			state = &State{
				Group: r.Group,
			}
			s.Status[ak] = state
		}
		state.Computations = r.Computations
		status := ST_NORM
		if r.Value.(expr.Number) != 0 {
			state.Expr = e.String()
			alerts = append(alerts, ak)
			if isCrit {
				status = ST_CRIT
			} else {
				status = ST_WARN
			}
		}
		if status > s.runStates[ak] {
			s.runStates[ak] = status
		}
	}
	return
}

func (s *Schedule) Notify(st *State, a *conf.Alert, n *conf.Notification) {
	if len(n.Email) > 0 {
		go s.Email(a, n, st)
	}
	if n.Post != nil {
		go s.Post(a, n, st)
	}
	if n.Get != nil {
		go s.Get(a, n, st)
	}
	if n.Print {
		go s.Print(a, n, st)
	}
	if n.Next == nil {
		return
	}
	s.AddNotification(AlertKey{Name: a.Name, Group: st.Group.String()}, n, time.Now().UTC())
}

func (s *Schedule) AddNotification(ak AlertKey, n *conf.Notification, started time.Time) {
	if s.Notifications == nil {
		s.Notifications = make(map[AlertKey]map[string]time.Time)
	}
	if s.Notifications[ak] == nil {
		s.Notifications[ak] = make(map[string]time.Time)
	}
	stn := s.Notifications[ak]
	// Prevent duplicate notifications restarting each other.
	if _, present := stn[n.Name]; !present {
		stn[n.Name] = started
	}
}

type AlertKey struct {
	Name  string
	Group string
}

func (a AlertKey) String() string {
	return a.Name + a.Group
}

type State struct {
	// Most recent event last.
	History      []Event
	Touched      time.Time
	Expr         string
	Group        opentsdb.TagSet
	Computations expr.Computations
	Subject      string
}

func (s *Schedule) Acknowledge(ak AlertKey) {
	s.Lock()
	delete(s.Notifications, ak)
	s.Unlock()
}

func (s *State) Touch() {
	s.Touched = time.Now().UTC()
}

// Appends status to the history if the status is different than the latest
// status. Returns true if state was changed.
func (s *State) Append(status Status) bool {
	s.Touch()
	if len(s.History) == 0 || s.Last().Status != status {
		s.History = append(s.History, Event{status, time.Now().UTC()})
		return true
	}
	return false
}

func (s *State) Last() Event {
	if len(s.History) == 0 {
		return Event{}
	}
	return s.History[len(s.History)-1]
}

type Event struct {
	Status Status
	Time   time.Time // embedding this breaks JSON encoding
}

type Status int

const (
	ST_UNKNOWN Status = iota
	ST_NORM
	ST_WARN
	ST_CRIT
)

func (s Status) String() string {
	switch s {
	case ST_NORM:
		return "normal"
	case ST_WARN:
		return "warning"
	case ST_CRIT:
		return "critical"
	default:
		return "unknown"
	}
}
<commit_msg>Always set subject if displayed<commit_after>package sched

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/tsaf/conf"
	"github.com/StackExchange/tsaf/expr"
)

type Schedule struct {
	sync.Mutex

	Conf          *conf.Conf
	Freq          time.Duration
	Status        map[AlertKey]*State
	Notifications map[AlertKey]map[string]time.Time

	cache     *opentsdb.Cache
	runStates map[AlertKey]Status
	nc        chan interface{}
}

func (s *Schedule) MarshalJSON() ([]byte, error) {
	t := struct {
		Alerts map[string]*conf.Alert
		Freq   time.Duration
		Status map[string]*State
	}{
		s.Conf.Alerts,
		s.Freq,
		make(map[string]*State),
	}
	for k, v := range s.Status {
		if v.Last().Status < ST_WARN {
			continue
		}
		t.Status[k.String()] = v
	}
	return json.Marshal(&t)
}

var DefaultSched = &Schedule{
	Freq: time.Minute * 5,
}

// Loads a configuration into the default schedule
func Load(c *conf.Conf) {
	DefaultSched.Load(c)
}

// Runs the default schedule.
func Run() error {
	return DefaultSched.Run()
}

func (s *Schedule) Load(c *conf.Conf) {
	s.Conf = c
	s.RestoreState()
}

// Restores notification and alert state from the file on disk.
func (s *Schedule) RestoreState() {
	s.Lock()
	defer s.Unlock()
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	s.Notifications = nil
	s.Status = make(map[AlertKey]*State)
	f, err := os.Open(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	dec := gob.NewDecoder(f)
	notifications := make(map[AlertKey]map[string]time.Time)
	if err := dec.Decode(&notifications); err != nil {
		log.Println(err)
		return
	}
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				log.Println("sched: notification not present during restore:", name)
				continue
			}
			_, present = s.Conf.Alerts[ak.Name]
			if !present {
				log.Println("sched: alert not present during restore:", ak.Name)
				continue
			}
			s.AddNotification(ak, n, t)
		}
	}
	for {
		var ak AlertKey
		var st State
		if err := dec.Decode(&ak); err == io.EOF {
			break
		} else if err != nil {
			log.Println(err)
			return
		}
		if err := dec.Decode(&st); err != nil {
			log.Println(err)
			return
		}
		if _, present := s.Conf.Alerts[ak.Name]; !present {
			log.Println("sched: alert no longer present, ignoring:", ak)
			continue
		}
		s.Status[ak] = &st
	}
}

func (s *Schedule) Save() {
	s.Lock()
	defer s.Unlock()
	f, err := os.Create(s.Conf.StateFile)
	if err != nil {
		log.Println(err)
		return
	}
	enc := gob.NewEncoder(f)
	if err := enc.Encode(s.Notifications); err != nil {
		log.Println(err)
		return
	}
	for k, v := range s.Status {
		enc.Encode(k)
		enc.Encode(v)
	}
	if err := f.Close(); err != nil {
		log.Println(err)
		return
	}
	log.Println("sched: wrote state to", s.Conf.StateFile)
}

func (s *Schedule) Run() error {
	go func() {
		for {
			time.Sleep(time.Minute)
			s.Save()
		}
	}()
	s.nc = make(chan interface{}, 1)
	go s.Poll()
	for {
		wait := time.After(s.Freq)
		if s.Freq < time.Second {
			return fmt.Errorf("sched: frequency must be > 1 second")
		}
		if s.Conf == nil {
			return fmt.Errorf("sched: nil configuration")
		}
		start := time.Now()
		s.Check()
		fmt.Printf("run at %v took %v\n", start, time.Since(start))
		<-wait
	}
}

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	var timeout time.Duration
	for {
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
		timeout = s.CheckNotifications()
	}
}

// CheckNotifications processes past notification events. It returns the
func (s *Schedule) CheckNotifications() time.Duration {
	s.Lock()
	defer s.Unlock()
	timeout := time.Hour
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				if remaining < timeout {
					timeout = remaining
				}
				s.AddNotification(ak, n, t)
				continue
			}
			st, present := s.Status[ak]
			if !present {
				continue
			}
			a, present := s.Conf.Alerts[ak.Name]
			if !present {
				continue
			}
			s.Notify(st, a, n)
			if n.Timeout < timeout {
				timeout = n.Timeout
			}
		}
	}
	return timeout
}

func (s *Schedule) Check() {
	s.Lock()
	defer s.Unlock()
	s.runStates = make(map[AlertKey]Status)
	s.cache = opentsdb.NewCache(s.Conf.TsdbHost)
	for _, a := range s.Conf.Alerts {
		s.CheckAlert(a)
	}
	changed := false
	for ak, status := range s.runStates {
		state := s.Status[ak]
		change := state.Append(status)
		if change {
			changed = true
		}
		if status > ST_NORM {
			var subject = new(bytes.Buffer)
			a := s.Conf.Alerts[ak.Name]
			if err := s.ExecuteSubject(subject, a, state); err != nil {
				log.Println(err)
			}
			state.Subject = subject.String()
			if change {
				notify := func(notifications map[string]*conf.Notification) {
					for _, n := range notifications {
						s.Notify(state, a, n)
					}
				}
				switch status {
				case ST_CRIT:
					notify(a.CritNotification)
				case ST_WARN:
					notify(a.WarnNotification)
				}
			}
		}
	}
	if changed {
		s.nc <- true
	}
}

func (s *Schedule) CheckAlert(a *conf.Alert) {
	crits := s.CheckExpr(a, a.Crit, true, nil)
	s.CheckExpr(a, a.Warn, false, crits)
}

func (s *Schedule) CheckExpr(a *conf.Alert, e *expr.Expr, isCrit bool, ignore []AlertKey) (alerts []AlertKey) {
	if e == nil {
		return
	}
	results, _, err := e.Execute(s.cache, nil)
	if err != nil {
		// todo: do something here?
		log.Println(err)
		return
	}
Loop:
	for _, r := range results {
		if a.Squelched(r.Group) {
			continue
		}
		ak := AlertKey{a.Name, r.Group.String()}
		for _, v := range ignore {
			if ak == v {
				continue Loop
			}
		}
		state := s.Status[ak]
		if state == nil {
			state = &State{
				Group: r.Group,
			}
			s.Status[ak] = state
		}
		state.Computations = r.Computations
		status := ST_NORM
		if r.Value.(expr.Number) != 0 {
			state.Expr = e.String()
			alerts = append(alerts, ak)
			if isCrit {
				status = ST_CRIT
			} else {
				status = ST_WARN
			}
		}
		if status > s.runStates[ak] {
			s.runStates[ak] = status
		}
	}
	return
}

func (s *Schedule) Notify(st *State, a *conf.Alert, n *conf.Notification) {
	if len(n.Email) > 0 {
		go s.Email(a, n, st)
	}
	if n.Post != nil {
		go s.Post(a, n, st)
	}
	if n.Get != nil {
		go s.Get(a, n, st)
	}
	if n.Print {
		go s.Print(a, n, st)
	}
	if n.Next == nil {
		return
	}
	s.AddNotification(AlertKey{Name: a.Name, Group: st.Group.String()}, n, time.Now().UTC())
}

func (s *Schedule) AddNotification(ak AlertKey, n *conf.Notification, started time.Time) {
	if s.Notifications == nil {
		s.Notifications = make(map[AlertKey]map[string]time.Time)
	}
	if s.Notifications[ak] == nil {
		s.Notifications[ak] = make(map[string]time.Time)
	}
	stn := s.Notifications[ak]
	// Prevent duplicate notifications restarting each other.
	if _, present := stn[n.Name]; !present {
		stn[n.Name] = started
	}
}

type AlertKey struct {
	Name  string
	Group string
}

func (a AlertKey) String() string {
	return a.Name + a.Group
}

type State struct {
	// Most recent event last.
	History      []Event
	Touched      time.Time
	Expr         string
	Group        opentsdb.TagSet
	Computations expr.Computations
	Subject      string
}

func (s *Schedule) Acknowledge(ak AlertKey) {
	s.Lock()
	delete(s.Notifications, ak)
	s.Unlock()
}

func (s *State) Touch() {
	s.Touched = time.Now().UTC()
}

// Appends status to the history if the status is different than the latest
// status. Returns true if state was changed.
func (s *State) Append(status Status) bool {
	s.Touch()
	if len(s.History) == 0 || s.Last().Status != status {
		s.History = append(s.History, Event{status, time.Now().UTC()})
		return true
	}
	return false
}

func (s *State) Last() Event {
	if len(s.History) == 0 {
		return Event{}
	}
	return s.History[len(s.History)-1]
}

type Event struct {
	Status Status
	Time   time.Time // embedding this breaks JSON encoding
}

type Status int

const (
	ST_UNKNOWN Status = iota
	ST_NORM
	ST_WARN
	ST_CRIT
)

func (s Status) String() string {
	switch s {
	case ST_NORM:
		return "normal"
	case ST_WARN:
		return "warning"
	case ST_CRIT:
		return "critical"
	default:
		return "unknown"
	}
}
<|endoftext|>


<commit_before>package dmm

import "fmt"

func ExampleExecute() {
    api := New("foobarbazbuzz", "dummy-990")
    api.SetSite(SITE_ALLAGES)
    api.SetService("mono")
    api.SetFloor("dvd")
    api.SetSort("date")
    api.SetLength(1)

    result, err := api.Execute()
    if err != nil {
        fmt.Println(err)
    } else {
        fmt.Println(result)
    }
}

func ExampleExecuteOneLiner() {
    rst, err := New("foobarbazbuzz", "dummy-999").SetSite(SITE_ADULT).SetLength(1).Execute()
    if err != nil {
        fmt.Println(err)
    } else {
        fmt.Println(rst)
    }
}<commit_msg>fix example<commit_after>package dmm

import "fmt"

func ExampleExecute() {
    api := New("foobarbazbuzz", "dummy-990")
    api.SetSite(SITE_ALLAGES)
    api.SetService("mono")
    api.SetFloor("dvd")
    api.SetSort("date")
    api.SetLength(1)

    result, err := api.Execute()
    if err != nil {
        fmt.Println(err)
    } else {
        fmt.Println(result)
    }
}

func ExampleExecute_OneLiner() {
    rst, err := New("foobarbazbuzz", "dummy-999").SetSite(SITE_ADULT).SetLength(1).Execute()
    if err != nil {
        fmt.Println(err)
    } else {
        fmt.Println(rst)
    }
}<|endoftext|>


<commit_before>/*
 * Copyright 2014 Jason Woods.
 *
 * This file is a modification of code from Logstash Forwarder.
 * Copyright 2012-2013 Jordan Sissel and contributors.
 *
 * 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 spooler

import (
  "lc-lib/core"
  "lc-lib/publisher"
  "time"
)

const (
  // Event header is just uint32 at the moment
  event_header_size   = 4
)

type Spooler struct {
  core.PipelineSegment
  core.PipelineConfigReceiver

  config      *core.GeneralConfig
  spool       []*core.EventDescriptor
  spool_size  int
  input       chan *core.EventDescriptor
  output      chan<- []*core.EventDescriptor
  timer_start time.Time
  timer       *time.Timer
}

func NewSpooler(pipeline *core.Pipeline, config *core.GeneralConfig, publisher_imp *publisher.Publisher) *Spooler {
  ret := &Spooler{
    config: config,
    spool:  make([]*core.EventDescriptor, 0, config.SpoolSize),
    input:  make(chan *core.EventDescriptor, 16), // TODO: Make configurable?
    output: publisher_imp.Connect(),
  }

  pipeline.Register(ret)

  return ret
}

func (s *Spooler) Connect() chan<- *core.EventDescriptor {
  return s.input
}

func (s *Spooler) Run() {
  defer func() {
    s.Done()
  }()

  s.timer_start = time.Now()
  s.timer = time.NewTimer(s.config.SpoolTimeout)

SpoolerLoop:
  for {
    select {
    case event := <-s.input:
      if len(s.spool) > 0 && int64(s.spool_size) + int64(len(event.Event)) + event_header_size >= s.config.SpoolMaxBytes {
        log.Debug("Spooler flushing %d events due to spool max bytes (%d/%d - next is %d)", len(s.spool), s.spool_size, s.config.SpoolMaxBytes, len(event.Event) + 4)

        // Can't fit this event in the spool - flush and then queue
        if !s.sendSpool() {
          break SpoolerLoop
        }

        s.resetTimer()
        s.spool_size += len(event.Event) + event_header_size
        s.spool = append(s.spool, event)

        continue
      }

      s.spool_size += len(event.Event) + event_header_size
      s.spool = append(s.spool, event)

      // Flush if full
      if len(s.spool) >= cap(s.spool) {
        log.Debug("Spooler flushing %d events due to spool size reached", len(s.spool))

        if !s.sendSpool() {
          break SpoolerLoop
        }

        s.resetTimer()
      }
    case <-s.timer.C:
      // Flush what we have, if anything
      if len(s.spool) > 0 {
        log.Debug("Spooler flushing %d events due to spool timeout exceeded", len(s.spool))

        if !s.sendSpool() {
          break SpoolerLoop
        }
      }

      s.resetTimer()
    case <-s.OnShutdown():
      break SpoolerLoop
    case config := <-s.OnConfig():
      if !s.reloadConfig(config) {
        break SpoolerLoop
      }
    }
  }

  log.Info("Spooler exiting")
}

func (s *Spooler) sendSpool() bool {
  select {
  case <-s.OnShutdown():
    return false
  case config := <-s.OnConfig():
    if !s.reloadConfig(config) {
      return false
    }
  case s.output <- s.spool:
  }

  s.spool = make([]*core.EventDescriptor, 0, s.config.SpoolSize)
  s.spool_size = 0

  return true
}

func (s *Spooler) resetTimer() {
  s.timer_start = time.Now()
  s.timer.Reset(s.config.SpoolTimeout)
}

func (s *Spooler) reloadConfig(config *core.Config) bool {
  s.config = &config.General

  // Immediate flush?
  passed := time.Now().Sub(s.timer_start)
  if passed >= s.config.SpoolTimeout || len(s.spool) >= int(s.config.SpoolSize) {
    if !s.sendSpool() {
      return false
    }
    s.timer_start = time.Now()
    s.timer.Reset(s.config.SpoolTimeout)
  } else {
    s.timer.Reset(passed - s.config.SpoolTimeout)
  }

  return true
}
<commit_msg>Fix a spooler race condition which could trigger spool timeout flush one or more times immediately after a max spool size flush containing only single figure event counts<commit_after>/*
 * Copyright 2014 Jason Woods.
 *
 * This file is a modification of code from Logstash Forwarder.
 * Copyright 2012-2013 Jordan Sissel and contributors.
 *
 * 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 spooler

import (
  "lc-lib/core"
  "lc-lib/publisher"
  "time"
)

const (
  // Event header is just uint32 at the moment
  event_header_size   = 4
)

type Spooler struct {
  core.PipelineSegment
  core.PipelineConfigReceiver

  config      *core.GeneralConfig
  spool       []*core.EventDescriptor
  spool_size  int
  input       chan *core.EventDescriptor
  output      chan<- []*core.EventDescriptor
  timer_start time.Time
  timer       *time.Timer
}

func NewSpooler(pipeline *core.Pipeline, config *core.GeneralConfig, publisher_imp *publisher.Publisher) *Spooler {
  ret := &Spooler{
    config: config,
    spool:  make([]*core.EventDescriptor, 0, config.SpoolSize),
    input:  make(chan *core.EventDescriptor, 16), // TODO: Make configurable?
    output: publisher_imp.Connect(),
  }

  pipeline.Register(ret)

  return ret
}

func (s *Spooler) Connect() chan<- *core.EventDescriptor {
  return s.input
}

func (s *Spooler) Run() {
  defer func() {
    s.Done()
  }()

  s.timer_start = time.Now()
  s.timer = time.NewTimer(s.config.SpoolTimeout)

SpoolerLoop:
  for {
    select {
    case event := <-s.input:
      if len(s.spool) > 0 && int64(s.spool_size) + int64(len(event.Event)) + event_header_size >= s.config.SpoolMaxBytes {
        log.Debug("Spooler flushing %d events due to spool max bytes (%d/%d - next is %d)", len(s.spool), s.spool_size, s.config.SpoolMaxBytes, len(event.Event) + 4)

        // Can't fit this event in the spool - flush and then queue
        if !s.sendSpool() {
          break SpoolerLoop
        }

        s.resetTimer()
        s.spool_size += len(event.Event) + event_header_size
        s.spool = append(s.spool, event)

        continue
      }

      s.spool_size += len(event.Event) + event_header_size
      s.spool = append(s.spool, event)

      // Flush if full
      if len(s.spool) >= cap(s.spool) {
        log.Debug("Spooler flushing %d events due to spool size reached", len(s.spool))

        if !s.sendSpool() {
          break SpoolerLoop
        }

        s.resetTimer()
      }
    case <-s.timer.C:
      // Flush what we have, if anything
      if len(s.spool) > 0 {
        log.Debug("Spooler flushing %d events due to spool timeout exceeded", len(s.spool))

        if !s.sendSpool() {
          break SpoolerLoop
        }
      }

      s.resetTimer()
    case <-s.OnShutdown():
      break SpoolerLoop
    case config := <-s.OnConfig():
      if !s.reloadConfig(config) {
        break SpoolerLoop
      }
    }
  }

  log.Info("Spooler exiting")
}

func (s *Spooler) sendSpool() bool {
  select {
  case <-s.OnShutdown():
    return false
  case config := <-s.OnConfig():
    if !s.reloadConfig(config) {
      return false
    }
  case s.output <- s.spool:
  }

  s.spool = make([]*core.EventDescriptor, 0, s.config.SpoolSize)
  s.spool_size = 0

  return true
}

func (s *Spooler) resetTimer() {
  s.timer_start = time.Now()

  // Stop the timer, and ensure the channel is empty before restarting it
  s.timer.Stop()
  select {
  case <-s.timer.C:
  default:
  }
  s.timer.Reset(s.config.SpoolTimeout)
}

func (s *Spooler) reloadConfig(config *core.Config) bool {
  s.config = &config.General

  // Immediate flush?
  passed := time.Now().Sub(s.timer_start)
  if passed >= s.config.SpoolTimeout || len(s.spool) >= int(s.config.SpoolSize) {
    if !s.sendSpool() {
      return false
    }
    s.timer_start = time.Now()
    s.timer.Reset(s.config.SpoolTimeout)
  } else {
    s.timer.Reset(passed - s.config.SpoolTimeout)
  }

  return true
}
<|endoftext|>


<commit_before>// Package pca9685 allows interfacing with the pca9685 16-channel, 12-bit PWM Controller through I2C protocol.
package pca9685

import (
	"log"
	"math"
	"sync"
	"time"

	"github.com/kid0m4n/go-rpi/i2c"
)

const (
	clockFreq        = 25000000
	pwmControlPoints = 4096

	mode1RegAddr    = 0x00
	preScaleRegAddr = 0xFE

	pwm0OnLowReg = 0x6
)

// PCA9685 represents a PCA9685 PWM generator.
type PCA9685 struct {
	Bus  i2c.Bus
	Addr byte
	Freq int

	initialized bool
	mu          sync.RWMutex

	Debug bool
}

// New creates a new PCA9685 interface.
func New(bus i2c.Bus, addr byte, freq int) *PCA9685 {
	return &PCA9685{
		Bus:  bus,
		Addr: addr,
		Freq: freq,
	}
}

func (d *PCA9685) mode1Reg() (byte, error) {
	return d.Bus.ReadByteFromReg(d.Addr, mode1RegAddr)
}

func (d *PCA9685) setup() (err error) {
	d.mu.RLock()
	if d.initialized {
		d.mu.RUnlock()
		return
	}
	d.mu.RUnlock()

	d.mu.Lock()
	defer d.mu.Unlock()

	preScaleValue := byte(math.Floor(float64(clockFreq/(pwmControlPoints*d.Freq))+float64(0.5)) - 1)
	if d.Debug {
		log.Printf("pca9685: calculated Prescale value = %#02x", preScaleValue)
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: read MODE1 Reg [regAddr: %#02x] Value: [%v]", mode1RegAddr, mode1Reg)
	}

	if err = d.sleep(); err != nil {
		return
	}

	if err = d.Bus.WriteByteToReg(d.Addr, preScaleRegAddr, byte(preScaleValue)); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: prescale value [%#02x] written to PRE_SCALE Reg [regAddr: %#02x]", preScaleValue, preScaleRegAddr)
	}

	if err = d.wake(); err != nil {
		return
	}

	newmode := ((mode1Reg | 0x01) & 0xDF)
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, newmode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: new mode [%#02x] [disabling register auto increment] written to MODE1 Reg [regAddr: %#02x]", newmode, mode1RegAddr)
	}

	d.initialized = true
	if d.Debug {
		log.Printf("pca9685: driver initialized with pwm freq: %v", d.Freq)
	}

	return
}

// SetPwm sets the ON and OFF time registers for pwm signal shaping.
// channel: 0-15
// onTime/offTime: 0-4095
func (d *PCA9685) SetPwm(channel, onTime, offTime int) (err error) {
	if err = d.setup(); err != nil {
		return
	}

	onTimeLowReg := byte(pwm0OnLowReg + (4 * channel))

	onTimeLow := byte(onTime & 0xFF)
	onTimeHigh := byte(onTime >> 8)
	offTimeLow := byte(offTime & 0xFF)
	offTimeHigh := byte(offTime >> 8)

	if err = d.Bus.WriteByteToReg(d.Addr, onTimeLowReg, onTimeLow); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing On-Time Low [%#02x] to CHAN%v_ON_L reg [RegAddr: %#02x]", onTimeLow, channel, onTimeLowReg)
	}

	onTimeHighReg := onTimeLowReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, onTimeHighReg, onTimeHighReg); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing On-Time High [%#02x] to CHAN%v_ON_H reg [RegAddr: %#02x]", onTimeHigh, channel, onTimeHighReg)
	}

	offTimeLowReg := onTimeHighReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, offTimeLowReg, offTimeLow); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing Off-Time Low [%#02x] to CHAN%v_OFF_L reg [RegAddr: %#02x]", offTimeLow, channel, offTimeLowReg)
	}

	offTimeHighReg := offTimeLowReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, offTimeHighReg, offTimeHigh); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing Off-Time High [%#02x] to CHAN%v_OFF_H reg [RegAddr: %#02x]", offTimeHigh, channel, offTimeHighReg)
	}

	return
}

// Close stops the controller and resets mode and pwm controller registers.
func (d *PCA9685) Close() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	if err = d.sleep(); err != nil {
		return
	}

	if d.Debug {
		log.Println("pca9685: reset request received")
	}

	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, 0x00); err != nil {
		return
	}

	if d.Debug {
		log.Printf("pca9685: cleaning up all PWM control registers")
	}

	for regAddr := 0x0; regAddr <= 0x45; regAddr++ {
		if err = d.Bus.WriteByteToReg(d.Addr, byte(regAddr), 0x00); err != nil {
			return
		}
	}

	if d.Debug {
		log.Printf("pca9685: done Cleaning up all PWM control registers")
	}

	if d.Debug {
		log.Println("pca9685: controller reset")
	}

	return
}

func (d *PCA9685) sleep() (err error) {
	if d.Debug {
		log.Println("pca9685: sleep request received")
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	sleepmode := (mode1Reg & 0x7F) | 0x10
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, sleepmode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: sleep mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", sleepmode, mode1RegAddr)
	}

	if d.Debug {
		log.Println("pca9685: controller set to Sleep mode")
	}

	return
}

// Sleep puts the controller in sleep mode. Does not change the pwm control registers.
func (d *PCA9685) Sleep() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	return d.sleep()
}

func (d *PCA9685) wake() (err error) {
	if d.Debug {
		log.Println("pca9685: wake request received")
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	wakeMode := mode1Reg & 0xEF
	if (mode1Reg & 0x80) == 0x80 {
		if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, wakeMode); err != nil {
			return
		}
		if d.Debug {
			log.Printf("pca9685: wake mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", wakeMode, mode1RegAddr)
		}

		time.Sleep(500 * time.Microsecond)
	}

	restartOpCode := wakeMode | 0x80
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, restartOpCode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: restart mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", restartOpCode, mode1RegAddr)
	}

	return
}

// Wake allows the controller to exit sleep mode and resume with PWM generation.
func (d *PCA9685) Wake() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	return d.wake()
}
<commit_msg>pca9685: lowercase logs<commit_after>// Package pca9685 allows interfacing with the pca9685 16-channel, 12-bit PWM Controller through I2C protocol.
package pca9685

import (
	"log"
	"math"
	"sync"
	"time"

	"github.com/kid0m4n/go-rpi/i2c"
)

const (
	clockFreq        = 25000000
	pwmControlPoints = 4096

	mode1RegAddr    = 0x00
	preScaleRegAddr = 0xFE

	pwm0OnLowReg = 0x6
)

// PCA9685 represents a PCA9685 PWM generator.
type PCA9685 struct {
	Bus  i2c.Bus
	Addr byte
	Freq int

	initialized bool
	mu          sync.RWMutex

	Debug bool
}

// New creates a new PCA9685 interface.
func New(bus i2c.Bus, addr byte, freq int) *PCA9685 {
	return &PCA9685{
		Bus:  bus,
		Addr: addr,
		Freq: freq,
	}
}

func (d *PCA9685) mode1Reg() (byte, error) {
	return d.Bus.ReadByteFromReg(d.Addr, mode1RegAddr)
}

func (d *PCA9685) setup() (err error) {
	d.mu.RLock()
	if d.initialized {
		d.mu.RUnlock()
		return
	}
	d.mu.RUnlock()

	d.mu.Lock()
	defer d.mu.Unlock()

	preScaleValue := byte(math.Floor(float64(clockFreq/(pwmControlPoints*d.Freq))+float64(0.5)) - 1)
	if d.Debug {
		log.Printf("pca9685: calculated prescale value = %#02x", preScaleValue)
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: read MODE1 Reg [regAddr: %#02x] Value: [%v]", mode1RegAddr, mode1Reg)
	}

	if err = d.sleep(); err != nil {
		return
	}

	if err = d.Bus.WriteByteToReg(d.Addr, preScaleRegAddr, byte(preScaleValue)); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: prescale value [%#02x] written to PRE_SCALE Reg [regAddr: %#02x]", preScaleValue, preScaleRegAddr)
	}

	if err = d.wake(); err != nil {
		return
	}

	newmode := ((mode1Reg | 0x01) & 0xDF)
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, newmode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: new mode [%#02x] [disabling register auto increment] written to MODE1 Reg [regAddr: %#02x]", newmode, mode1RegAddr)
	}

	d.initialized = true
	if d.Debug {
		log.Printf("pca9685: driver initialized with pwm freq: %v", d.Freq)
	}

	return
}

// SetPwm sets the ON and OFF time registers for pwm signal shaping.
// channel: 0-15
// onTime/offTime: 0-4095
func (d *PCA9685) SetPwm(channel, onTime, offTime int) (err error) {
	if err = d.setup(); err != nil {
		return
	}

	onTimeLowReg := byte(pwm0OnLowReg + (4 * channel))

	onTimeLow := byte(onTime & 0xFF)
	onTimeHigh := byte(onTime >> 8)
	offTimeLow := byte(offTime & 0xFF)
	offTimeHigh := byte(offTime >> 8)

	if err = d.Bus.WriteByteToReg(d.Addr, onTimeLowReg, onTimeLow); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing on-time low [%#02x] to CHAN%v_ON_L reg [reg: %#02x]", onTimeLow, channel, onTimeLowReg)
	}

	onTimeHighReg := onTimeLowReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, onTimeHighReg, onTimeHighReg); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing on-time high [%#02x] to CHAN%v_ON_H reg [reg: %#02x]", onTimeHigh, channel, onTimeHighReg)
	}

	offTimeLowReg := onTimeHighReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, offTimeLowReg, offTimeLow); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing off-time low [%#02x] to CHAN%v_OFF_L reg [reg: %#02x]", offTimeLow, channel, offTimeLowReg)
	}

	offTimeHighReg := offTimeLowReg + 1
	if err = d.Bus.WriteByteToReg(d.Addr, offTimeHighReg, offTimeHigh); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: writing off-time high [%#02x] to CHAN%v_OFF_H reg [reg: %#02x]", offTimeHigh, channel, offTimeHighReg)
	}

	return
}

// Close stops the controller and resets mode and pwm controller registers.
func (d *PCA9685) Close() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	if err = d.sleep(); err != nil {
		return
	}

	if d.Debug {
		log.Println("pca9685: reset request received")
	}

	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, 0x00); err != nil {
		return
	}

	if d.Debug {
		log.Printf("pca9685: cleaning up all PWM control registers")
	}

	for regAddr := 0x0; regAddr <= 0x45; regAddr++ {
		if err = d.Bus.WriteByteToReg(d.Addr, byte(regAddr), 0x00); err != nil {
			return
		}
	}

	if d.Debug {
		log.Printf("pca9685: done Cleaning up all PWM control registers")
	}

	if d.Debug {
		log.Println("pca9685: controller reset")
	}

	return
}

func (d *PCA9685) sleep() (err error) {
	if d.Debug {
		log.Println("pca9685: sleep request received")
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	sleepmode := (mode1Reg & 0x7F) | 0x10
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, sleepmode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: sleep mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", sleepmode, mode1RegAddr)
	}

	if d.Debug {
		log.Println("pca9685: controller set to Sleep mode")
	}

	return
}

// Sleep puts the controller in sleep mode. Does not change the pwm control registers.
func (d *PCA9685) Sleep() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	return d.sleep()
}

func (d *PCA9685) wake() (err error) {
	if d.Debug {
		log.Println("pca9685: wake request received")
	}

	mode1Reg, err := d.mode1Reg()
	if err != nil {
		return
	}
	wakeMode := mode1Reg & 0xEF
	if (mode1Reg & 0x80) == 0x80 {
		if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, wakeMode); err != nil {
			return
		}
		if d.Debug {
			log.Printf("pca9685: wake mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", wakeMode, mode1RegAddr)
		}

		time.Sleep(500 * time.Microsecond)
	}

	restartOpCode := wakeMode | 0x80
	if err = d.Bus.WriteByteToReg(d.Addr, mode1RegAddr, restartOpCode); err != nil {
		return
	}
	if d.Debug {
		log.Printf("pca9685: restart mode [%#02x] written to MODE1 Reg [regAddr: %#02x]", restartOpCode, mode1RegAddr)
	}

	return
}

// Wake allows the controller to exit sleep mode and resume with PWM generation.
func (d *PCA9685) Wake() (err error) {
	if err = d.setup(); err != nil {
		return
	}

	return d.wake()
}
<|endoftext|>


<commit_before><commit_msg>clear up outstanding TODOs.<commit_after><|endoftext|>


<commit_before>// Copyright 2016-2018 Authors of Cilium
//
// 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 command

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"regexp"

	"github.com/spf13/cobra"
	"k8s.io/client-go/util/jsonpath"
)

var outputOpt string

// OutputJSON returns true if the JSON output option was specified
func OutputJSON() bool {
	return len(outputOpt) > 0
}

//AddJSONOutput adds the -o|--output option to any cmd to export to json
func AddJSONOutput(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&outputOpt, "output", "o", "", "json| jsonpath='{}'")
}

//PrintOutput receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutput(data interface{}) error {
	var re = regexp.MustCompile(`^jsonpath\=(.*)`)

	if outputOpt == "json" {
		return dumpJSON(data, "")
	}

	if re.MatchString(outputOpt) {
		return dumpJSON(data, re.ReplaceAllString(outputOpt, "$1"))
	}

	return fmt.Errorf("Couldn't found output printer")
}

// dumpJSON dump the data variable to the stdout as json.
// If somethings fail, it'll return an error
// If jsonPath is passed, it'll run the json query over data var.
func dumpJSON(data interface{}, jsonPath string) error {

	if len(jsonPath) == 0 {
		result, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't marshal to json: '%s'\n", err)
			return err
		}
		fmt.Println(string(result))
		return nil
	}

	parser := jsonpath.New("").AllowMissingKeys(true)
	if err := parser.Parse(jsonPath); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return err
	}

	buf := new(bytes.Buffer)
	if err := parser.Execute(buf, data); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return err

	}
	fmt.Println(buf.String())
	return nil
}
<commit_msg>command: add new helper functions<commit_after>// Copyright 2016-2018 Authors of Cilium
//
// 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 command

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"regexp"

	"github.com/spf13/cobra"
	"k8s.io/client-go/util/jsonpath"
)

var (
	outputOpt string
	re        = regexp.MustCompile(`^jsonpath\=(.*)`)
)

// OutputJSON returns true if the JSON output option was specified
func OutputJSON() bool {
	return len(outputOpt) > 0
}

//AddJSONOutput adds the -o|--output option to any cmd to export to json
func AddJSONOutput(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&outputOpt, "output", "o", "", "json| jsonpath='{}'")
}

//PrintOutput receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutput(data interface{}) error {
	return PrintOutputWithType(data, outputOpt)
}

//PrintOutputWithType receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutputWithType(data interface{}, outputType string) error {
	if outputType == "json" {
		return dumpJSON(data, "")
	}

	if re.MatchString(outputType) {
		return dumpJSON(data, re.ReplaceAllString(outputType, "$1"))
	}

	return fmt.Errorf("Couldn't found output printer")
}

// DumpJSONToString dumps the contents of data into a string. If jsonpath
// is non-empty, will attempt to do jsonpath filtering using said string.
// Returns the string representation of the JSON in data, or an error if
// any JSON marshaling, parsing operations fail.
func DumpJSONToString(data interface{}, jsonPath string) (string, error) {
	if len(jsonPath) == 0 {
		result, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't marshal to json: '%s'\n", err)
			return "", err
		}
		fmt.Println(string(result))
		return "", nil
	}

	parser := jsonpath.New("").AllowMissingKeys(true)
	if err := parser.Parse(jsonPath); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err
	}

	buf := new(bytes.Buffer)
	if err := parser.Execute(buf, data); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err

	}
	return buf.String(), nil
}

// dumpJSON dump the data variable to the stdout as json.
// If somethings fail, it'll return an error
// If jsonPath is passed, it'll run the json query over data var.
func dumpJSON(data interface{}, jsonPath string) error {
	strOut, err := DumpJSONToString(data, jsonPath)
	if err != nil {
		return err
	}
	fmt.Println(strOut)
	return nil
}
<|endoftext|>


<commit_before><commit_msg>fixing link<commit_after><|endoftext|>


<commit_before>package models_test

import (
	"fmt"
	"github.com/APTrust/exchange/dpn"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestNewSyncResult(t *testing.T) {
	result := dpn.NewSyncResult("aptrust")
	require.NotNil(t, result)
	assert.Equal(t, "aptrust", result.NodeName)
	assert.NotNil(t, result.FetchCounts)
	assert.NotNil(t, result.SyncCounts)
	assert.NotNil(t, result.Errors)
}

func TestAddToFetchCount(t *testing.T) {
	result := dpn.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddToFetchCount(dpn.DPNTypeBag, 1)
	assert.Equal(t, 1, result.FetchCounts[dpn.DPNTypeBag])
	result.AddToFetchCount(dpn.DPNTypeBag, 4)
	assert.Equal(t, 5, result.FetchCounts[dpn.DPNTypeBag])
}

func TestAddToSyncCount(t *testing.T) {
	result := dpn.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddToSyncCount(dpn.DPNTypeBag, 1)
	assert.Equal(t, 1, result.SyncCounts[dpn.DPNTypeBag])
	result.AddToSyncCount(dpn.DPNTypeBag, 4)
	assert.Equal(t, 5, result.SyncCounts[dpn.DPNTypeBag])
}

func TestAddError(t *testing.T) {
	result := dpn.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 1"))
	assert.Equal(t, 1, len(result.Errors[dpn.DPNTypeBag]))
	assert.Equal(t, "Error 1", result.Errors[dpn.DPNTypeBag][0].Error())
	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 2"))
	assert.Equal(t, 2, len(result.Errors[dpn.DPNTypeBag]))
	assert.Equal(t, "Error 2", result.Errors[dpn.DPNTypeBag][1].Error())
}

func TestHasErrors(t *testing.T) {
	result := dpn.NewSyncResult("aptrust")
	require.NotNil(t, result)
	assert.False(t, result.HasErrors(dpn.DPNTypeBag))
	assert.False(t, result.HasErrors(""))

	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 1"))
	assert.Equal(t, 1, len(result.Errors[dpn.DPNTypeBag]))
	assert.True(t, result.HasErrors(dpn.DPNTypeBag))
	assert.True(t, result.HasErrors(""))
}
<commit_msg>More namespace changes<commit_after>package models_test

import (
	"fmt"
	"github.com/APTrust/exchange/dpn"
	"github.com/APTrust/exchange/dpn/models"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestNewSyncResult(t *testing.T) {
	result := models.NewSyncResult("aptrust")
	require.NotNil(t, result)
	assert.Equal(t, "aptrust", result.NodeName)
	assert.NotNil(t, result.FetchCounts)
	assert.NotNil(t, result.SyncCounts)
	assert.NotNil(t, result.Errors)
}

func TestAddToFetchCount(t *testing.T) {
	result := models.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddToFetchCount(dpn.DPNTypeBag, 1)
	assert.Equal(t, 1, result.FetchCounts[dpn.DPNTypeBag])
	result.AddToFetchCount(dpn.DPNTypeBag, 4)
	assert.Equal(t, 5, result.FetchCounts[dpn.DPNTypeBag])
}

func TestAddToSyncCount(t *testing.T) {
	result := models.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddToSyncCount(dpn.DPNTypeBag, 1)
	assert.Equal(t, 1, result.SyncCounts[dpn.DPNTypeBag])
	result.AddToSyncCount(dpn.DPNTypeBag, 4)
	assert.Equal(t, 5, result.SyncCounts[dpn.DPNTypeBag])
}

func TestAddError(t *testing.T) {
	result := models.NewSyncResult("aptrust")
	require.NotNil(t, result)
	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 1"))
	assert.Equal(t, 1, len(result.Errors[dpn.DPNTypeBag]))
	assert.Equal(t, "Error 1", result.Errors[dpn.DPNTypeBag][0].Error())
	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 2"))
	assert.Equal(t, 2, len(result.Errors[dpn.DPNTypeBag]))
	assert.Equal(t, "Error 2", result.Errors[dpn.DPNTypeBag][1].Error())
}

func TestHasErrors(t *testing.T) {
	result := models.NewSyncResult("aptrust")
	require.NotNil(t, result)
	assert.False(t, result.HasErrors(dpn.DPNTypeBag))
	assert.False(t, result.HasErrors(""))

	result.AddError(dpn.DPNTypeBag, fmt.Errorf("Error 1"))
	assert.Equal(t, 1, len(result.Errors[dpn.DPNTypeBag]))
	assert.True(t, result.HasErrors(dpn.DPNTypeBag))
	assert.True(t, result.HasErrors(""))
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 wait

import (
	"errors"
	"math/rand"
	"time"

	"k8s.io/kubernetes/pkg/util/runtime"
)

// For any test of the style:
//   ...
//   <- time.After(timeout):
//      t.Errorf("Timed out")
// The value for timeout should effectively be "forever." Obviously we don't want our tests to truly lock up forever, but 30s
// is long enough that it is effectively forever for the things that can slow down a run on a heavily contended machine
// (GC, seeks, etc), but not so long as to make a developer ctrl-c a test run if they do happen to break that test.
var ForeverTestTimeout = time.Second * 30

// NeverStop may be passed to Until to make it never stop.
var NeverStop <-chan struct{} = make(chan struct{})

// Forever is syntactic sugar on top of Until
func Forever(f func(), period time.Duration) {
	Until(f, period, NeverStop)
}

// Until loops until stop channel is closed, running f every period.
// Until is syntactic sugar on top of JitterUntil with zero jitter
// factor, with sliding = true (which means the timer for period
// starts after the f completes).
func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
	JitterUntil(f, period, 0.0, true, stopCh)
}

// NonSlidingUntil loops until stop channel is closed, running f every
// period. NonSlidingUntil is syntactic sugar on top of JitterUntil
// with zero jitter factor, with sliding = false (meaning the timer for
// period starts at the same time as the function starts).
func NonSlidingUntil(f func(), period time.Duration, stopCh <-chan struct{}) {
	JitterUntil(f, period, 0.0, false, stopCh)
}

// JitterUntil loops until stop channel is closed, running f every period.
// If jitterFactor is positive, the period is jittered before every run of f.
// If jitterFactor is not positive, the period is unchanged.
// Catches any panics, and keeps going. f may not be invoked if
// stop channel is already closed. Pass NeverStop to Until if you
// don't want it stop.
func JitterUntil(f func(), period time.Duration, jitterFactor float64, sliding bool, stopCh <-chan struct{}) {
	for {

		select {
		case <-stopCh:
			return
		default:
		}

		jitteredPeriod := period
		if jitterFactor > 0.0 {
			jitteredPeriod = Jitter(period, jitterFactor)
		}

		var t *time.Timer
		if !sliding {
			t = time.NewTimer(jitteredPeriod)
		}

		func() {
			defer runtime.HandleCrash()
			f()
		}()

		if sliding {
			t = time.NewTimer(jitteredPeriod)
		}

		// NOTE: b/c there is no priority selection in golang
		// it is possible for this to race, meaning we could
		// trigger t.C and stopCh, and t.C select falls through.
		// In order to mitigate we re-check stopCh at the beginning
		// of every loop to prevent extra executions of f().
		select {
		case <-stopCh:
			return
		case <-t.C:
		}
	}
}

// Jitter returns a time.Duration between duration and duration + maxFactor * duration,
// to allow clients to avoid converging on periodic behavior.  If maxFactor is 0.0, a
// suggested default value will be chosen.
func Jitter(duration time.Duration, maxFactor float64) time.Duration {
	if maxFactor <= 0.0 {
		maxFactor = 1.0
	}
	wait := duration + time.Duration(rand.Float64()*maxFactor*float64(duration))
	return wait
}

// ErrWaitTimeout is returned when the condition exited without success
var ErrWaitTimeout = errors.New("timed out waiting for the condition")

// ConditionFunc returns true if the condition is satisfied, or an error
// if the loop should be aborted.
type ConditionFunc func() (done bool, err error)

// Backoff is parameters applied to a Backoff function.
type Backoff struct {
	Duration time.Duration
	Factor   float64
	Jitter   float64
	Steps    int
}

// ExponentialBackoff repeats a condition check up to steps times, increasing the wait
// by multipling the previous duration by factor. If jitter is greater than zero,
// a random amount of each duration is added (between duration and duration*(1+jitter)).
// If the condition never returns true, ErrWaitTimeout is returned. All other errors
// terminate immediately.
func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
	duration := backoff.Duration
	for i := 0; i < backoff.Steps; i++ {
		if i != 0 {
			adjusted := duration
			if backoff.Jitter > 0.0 {
				adjusted = Jitter(duration, backoff.Jitter)
			}
			time.Sleep(adjusted)
			duration = time.Duration(float64(duration) * backoff.Factor)
		}
		if ok, err := condition(); err != nil || ok {
			return err
		}
	}
	return ErrWaitTimeout
}

// Poll tries a condition func until it returns true, an error, or the timeout
// is reached. condition will always be invoked at least once but some intervals
// may be missed if the condition takes too long or the time window is too short.
// If you want to Poll something forever, see PollInfinite.
// Poll always waits the interval before the first check of the condition.
func Poll(interval, timeout time.Duration, condition ConditionFunc) error {
	return pollInternal(poller(interval, timeout), condition)
}

func pollInternal(wait WaitFunc, condition ConditionFunc) error {
	done := make(chan struct{})
	defer close(done)
	return WaitFor(wait, condition, done)
}

// PollImmediate is identical to Poll, except that it performs the first check
// immediately, not waiting interval beforehand.
func PollImmediate(interval, timeout time.Duration, condition ConditionFunc) error {
	return pollImmediateInternal(poller(interval, timeout), condition)
}

func pollImmediateInternal(wait WaitFunc, condition ConditionFunc) error {
	done, err := condition()
	if err != nil {
		return err
	}
	if done {
		return nil
	}
	return pollInternal(wait, condition)
}

// PollInfinite polls forever.
func PollInfinite(interval time.Duration, condition ConditionFunc) error {
	done := make(chan struct{})
	defer close(done)
	return PollUntil(interval, condition, done)
}

// PollImmediateInfinite is identical to PollInfinite, except that it
// performs the first check immediately, not waiting interval
// beforehand.
func PollImmediateInfinite(interval time.Duration, condition ConditionFunc) error {
	done, err := condition()
	if err != nil {
		return err
	}
	if done {
		return nil
	}
	return PollInfinite(interval, condition)
}

// PollUntil is like Poll, but it takes a stop change instead of total duration
func PollUntil(interval time.Duration, condition ConditionFunc, stopCh <-chan struct{}) error {
	return WaitFor(poller(interval, 0), condition, stopCh)
}

// WaitFunc creates a channel that receives an item every time a test
// should be executed and is closed when the last test should be invoked.
type WaitFunc func(done <-chan struct{}) <-chan struct{}

// WaitFor gets a channel from wait(), and then invokes fn once for every value
// placed on the channel and once more when the channel is closed.  If fn
// returns an error the loop ends and that error is returned, and if fn returns
// true the loop ends and nil is returned. ErrWaitTimeout will be returned if
// the channel is closed without fn ever returning true.
func WaitFor(wait WaitFunc, fn ConditionFunc, done <-chan struct{}) error {
	c := wait(done)
	for {
		_, open := <-c
		ok, err := fn()
		if err != nil {
			return err
		}
		if ok {
			return nil
		}
		if !open {
			break
		}
	}
	return ErrWaitTimeout
}

// poller returns a WaitFunc that will send to the channel every
// interval until timeout has elapsed and then close the channel.
// Over very short intervals you may receive no ticks before
// the channel is closed.  If timeout is 0, the channel
// will never be closed.
func poller(interval, timeout time.Duration) WaitFunc {
	return WaitFunc(func(done <-chan struct{}) <-chan struct{} {
		ch := make(chan struct{})

		go func() {
			defer close(ch)

			tick := time.NewTicker(interval)
			defer tick.Stop()

			var after <-chan time.Time
			if timeout != 0 {
				// time.After is more convenient, but it
				// potentially leaves timers around much longer
				// than necessary if we exit early.
				timer := time.NewTimer(timeout)
				after = timer.C
				defer timer.Stop()
			}

			for {
				select {
				case <-tick.C:
					// If the consumer isn't ready for this signal drop it and
					// check the other channels.
					select {
					case ch <- struct{}{}:
					default:
					}
				case <-after:
					return
				case <-done:
					return
				}
			}
		}()

		return ch
	})
}
<commit_msg>Expand documention for wait package<commit_after>/*
Copyright 2014 The Kubernetes 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 wait

import (
	"errors"
	"math/rand"
	"time"

	"k8s.io/kubernetes/pkg/util/runtime"
)

// For any test of the style:
//   ...
//   <- time.After(timeout):
//      t.Errorf("Timed out")
// The value for timeout should effectively be "forever." Obviously we don't want our tests to truly lock up forever, but 30s
// is long enough that it is effectively forever for the things that can slow down a run on a heavily contended machine
// (GC, seeks, etc), but not so long as to make a developer ctrl-c a test run if they do happen to break that test.
var ForeverTestTimeout = time.Second * 30

// NeverStop may be passed to Until to make it never stop.
var NeverStop <-chan struct{} = make(chan struct{})

// Forever calls f every period for ever.
//
// Forever is syntactic sugar on top of Until.
func Forever(f func(), period time.Duration) {
	Until(f, period, NeverStop)
}

// Until loops until stop channel is closed, running f every period.
//
// Until is syntactic sugar on top of JitterUntil with zero jitter factor and
// with sliding = true (which means the timer for period starts after the f
// completes).
func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
	JitterUntil(f, period, 0.0, true, stopCh)
}

// NonSlidingUntil loops until stop channel is closed, running f every
// period.
//
// NonSlidingUntil is syntactic sugar on top of JitterUntil with zero jitter
// factor, with sliding = false (meaning the timer for period starts at the same
// time as the function starts).
func NonSlidingUntil(f func(), period time.Duration, stopCh <-chan struct{}) {
	JitterUntil(f, period, 0.0, false, stopCh)
}

// JitterUntil loops until stop channel is closed, running f every period.
//
// If jitterFactor is positive, the period is jittered before every run of f.
// If jitterFactor is not positive, the period is unchanged and not jitterd.
//
// If slidingis true, the period is computed after f runs. If it is false then
// period includes the runtime for f.
//
// Close stopCh to stop. f may not be invoked if stop channel is already
// closed. Pass NeverStop to if you don't want it stop.
func JitterUntil(f func(), period time.Duration, jitterFactor float64, sliding bool, stopCh <-chan struct{}) {
	for {

		select {
		case <-stopCh:
			return
		default:
		}

		jitteredPeriod := period
		if jitterFactor > 0.0 {
			jitteredPeriod = Jitter(period, jitterFactor)
		}

		var t *time.Timer
		if !sliding {
			t = time.NewTimer(jitteredPeriod)
		}

		func() {
			defer runtime.HandleCrash()
			f()
		}()

		if sliding {
			t = time.NewTimer(jitteredPeriod)
		}

		// NOTE: b/c there is no priority selection in golang
		// it is possible for this to race, meaning we could
		// trigger t.C and stopCh, and t.C select falls through.
		// In order to mitigate we re-check stopCh at the beginning
		// of every loop to prevent extra executions of f().
		select {
		case <-stopCh:
			return
		case <-t.C:
		}
	}
}

// Jitter returns a time.Duration between duration and duration + maxFactor *
// duration.
//
// This allows clients to avoid converging on periodic behavior. If maxFactor
// is 0.0, a suggested default value will be chosen.
func Jitter(duration time.Duration, maxFactor float64) time.Duration {
	if maxFactor <= 0.0 {
		maxFactor = 1.0
	}
	wait := duration + time.Duration(rand.Float64()*maxFactor*float64(duration))
	return wait
}

// ErrWaitTimeout is returned when the condition exited without success.
var ErrWaitTimeout = errors.New("timed out waiting for the condition")

// ConditionFunc returns true if the condition is satisfied, or an error
// if the loop should be aborted.
type ConditionFunc func() (done bool, err error)

// Backoff holds parameters applied to a Backoff function.
type Backoff struct {
	Duration time.Duration // the base duration
	Factor   float64       // Duration is multipled by factor each iteration
	Jitter   float64       // The amount of jitter applied each iteration
	Steps    int           // Exit with error after this many steps
}

// ExponentialBackoff repeats a condition check with exponential backoff.
//
// It checks the condition up to Steps times, increasing the wait by multipling
// the previous duration by Factor.
//
// If Jitter is greater than zero, a random amount of each duration is added
// (between duration and duration*(1+jitter)).
//
// If the condition never returns true, ErrWaitTimeout is returned. All other
// errors terminate immediately.
func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
	duration := backoff.Duration
	for i := 0; i < backoff.Steps; i++ {
		if i != 0 {
			adjusted := duration
			if backoff.Jitter > 0.0 {
				adjusted = Jitter(duration, backoff.Jitter)
			}
			time.Sleep(adjusted)
			duration = time.Duration(float64(duration) * backoff.Factor)
		}
		if ok, err := condition(); err != nil || ok {
			return err
		}
	}
	return ErrWaitTimeout
}

// Poll tries a condition func until it returns true, an error, or the timeout
// is reached.
//
// Poll always waits the interval before the run of 'condition'.
// 'condition' will always be invoked at least once.
//
// Some intervals may be missed if the condition takes too long or the time
// window is too short.
//
// If you want to Poll something forever, see PollInfinite.
func Poll(interval, timeout time.Duration, condition ConditionFunc) error {
	return pollInternal(poller(interval, timeout), condition)
}

func pollInternal(wait WaitFunc, condition ConditionFunc) error {
	return WaitFor(wait, condition, NeverStop)
}

// PollImmediate tries a condition func until it returns true, an error, or the timeout
// is reached.
//
// Poll always checks 'condition' before waiting for the interval. 'condition'
// will always be invoked at least once.
//
// Some intervals may be missed if the condition takes too long or the time
// window is too short.
//
// If you want to Poll something forever, see PollInfinite.
func PollImmediate(interval, timeout time.Duration, condition ConditionFunc) error {
	return pollImmediateInternal(poller(interval, timeout), condition)
}

func pollImmediateInternal(wait WaitFunc, condition ConditionFunc) error {
	done, err := condition()
	if err != nil {
		return err
	}
	if done {
		return nil
	}
	return pollInternal(wait, condition)
}

// PollInfinite tries a condition func until it returns true or an error
//
// PollInfinite always waits the interval before the run of 'condition'.
//
// Some intervals may be missed if the condition takes too long or the time
// window is too short.
func PollInfinite(interval time.Duration, condition ConditionFunc) error {
	done := make(chan struct{})
	defer close(done)
	return PollUntil(interval, condition, done)
}

// PollImmediateInfinite tries a condition func until it returns true or an error
//
// PollImmediateInfinite runs the 'condition' before waiting for the interval.
//
// Some intervals may be missed if the condition takes too long or the time
// window is too short.
func PollImmediateInfinite(interval time.Duration, condition ConditionFunc) error {
	done, err := condition()
	if err != nil {
		return err
	}
	if done {
		return nil
	}
	return PollInfinite(interval, condition)
}

// PollUntil tries a condition func until it returns true, an error or stopCh is
// closed.
//
// PolUntil always waits interval before the first run of 'condition'.
// 'condition' will always be invoked at least once.
func PollUntil(interval time.Duration, condition ConditionFunc, stopCh <-chan struct{}) error {
	return WaitFor(poller(interval, 0), condition, stopCh)
}

// WaitFunc creates a channel that receives an item every time a test
// should be executed and is closed when the last test should be invoked.
type WaitFunc func(done <-chan struct{}) <-chan struct{}

// WaitFor continually checks 'fn' as driven by 'wait'.
//
// WaitFor gets a channel from 'wait()'', and then invokes 'fn' once for every value
// placed on the channel and once more when the channel is closed.
//
// If 'fn' returns an error the loop ends and that error is returned, and if
// 'fn' returns true the loop ends and nil is returned.
//
// ErrWaitTimeout will be returned if the channel is closed without fn ever
// returning true.
func WaitFor(wait WaitFunc, fn ConditionFunc, done <-chan struct{}) error {
	c := wait(done)
	for {
		_, open := <-c
		ok, err := fn()
		if err != nil {
			return err
		}
		if ok {
			return nil
		}
		if !open {
			break
		}
	}
	return ErrWaitTimeout
}

// poller returns a WaitFunc that will send to the channel every interval until
// timeout has elapsed and then closes the channel.
//
// Over very short intervals you may receive no ticks before the channel is
// closed. A timeout of 0 is interpreted as an infinity.
//
// Output ticks are not buffered. If the channel is not ready to receive an
// item, the tick is skipped.
func poller(interval, timeout time.Duration) WaitFunc {
	return WaitFunc(func(done <-chan struct{}) <-chan struct{} {
		ch := make(chan struct{})

		go func() {
			defer close(ch)

			tick := time.NewTicker(interval)
			defer tick.Stop()

			var after <-chan time.Time
			if timeout != 0 {
				// time.After is more convenient, but it
				// potentially leaves timers around much longer
				// than necessary if we exit early.
				timer := time.NewTimer(timeout)
				after = timer.C
				defer timer.Stop()
			}

			for {
				select {
				case <-tick.C:
					// If the consumer isn't ready for this signal drop it and
					// check the other channels.
					select {
					case ch <- struct{}{}:
					default:
					}
				case <-after:
					return
				case <-done:
					return
				}
			}
		}()

		return ch
	})
}
<|endoftext|>


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

package visitors

import (
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/linuxboot/fiano/pkg/uefi"
)

// Table prints the GUIDS, types and sizes as a compact table.
type Table struct {
	W         *tabwriter.Writer
	Layout    bool
	Depth     int
	indent    int
	offset    uint64
	curOffset uint64
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Table) Run(f uefi.Firmware) error {
	return f.Apply(v)
}

// Visit applies the Table visitor to any Firmware type.
func (v *Table) Visit(f uefi.Firmware) error {
	var offset uint64
	switch f := f.(type) {
	case *uefi.FlashImage:
		if v.Depth > 0 { // Depth <= 0 means all
			v.Depth++
		}
		return v.printFirmware(f, "Image", "", "", 0, 0)
	case *uefi.FirmwareVolume:
		return v.printFirmware(f, "FV", f.FileSystemGUID.String(), "", v.offset+f.FVOffset, v.offset+f.FVOffset+f.DataOffset)
	case *uefi.File:
		// TODO: make name part of the file node
		return v.printFirmware(f, "File", f.Header.GUID.String(), f.Header.Type, v.curOffset, v.curOffset+f.DataOffset)
	case *uefi.Section:
		// Reset offset to O for (compressed) section content
		return v.printFirmware(f, "Sec", f.String(), f.Type, v.curOffset, 0)
	case *uefi.FlashDescriptor:
		return v.printFirmware(f, "IFD", "", "", 0, 0)
	case *uefi.BIOSRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
		}
		return v.printFirmware(f, "BIOS", "", "", offset, offset)
	case *uefi.BIOSPadding:
		return v.printFirmware(f, "BIOS Pad", "", "", v.offset+f.Offset, 0)
	case *uefi.NVarStore:
		return v.printFirmware(f, "NVAR Store", "", "", v.curOffset, v.curOffset)
	case *uefi.NVar:
		return v.printFirmware(f, "NVAR", f.GUID.String(), f, v.curOffset, v.curOffset+uint64(f.DataOffset))
	case *uefi.RawRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
		}
		return v.printFirmware(f, f.Type().String(), "", "", offset, offset)
	default:
		return v.printFirmware(f, fmt.Sprintf("%T", f), "", "", 0, 0)
	}
}

func indent(n int) string {
	return strings.Repeat(" ", n)
}

func (v *Table) printFirmware(f uefi.Firmware, node, name, typez interface{}, offset, dataOffset uint64) error {
	if v.W == nil {
		v.W = tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
		defer func() { v.W.Flush() }()
		if v.Layout {
			fmt.Fprintf(v.W, "%sNode\tGUID/Name\tOffset\tSize\n", indent(v.indent))
		} else {
			fmt.Fprintf(v.W, "%sNode\tGUID/Name\tType\tSize\n", indent(v.indent))
		}
	}
	length := uint64(len(f.Buf()))
	if typez == "" {
		if uefi.IsErased(f.Buf(), uefi.Attributes.ErasePolarity) {
			typez = "(empty)"
		}
	}
	v.printRow(node, name, typez, offset, length)

	v2 := *v
	v2.indent++
	v2.offset = dataOffset
	v2.curOffset = v2.offset
	if v.Depth <= 0 || v.indent < v.Depth {

		if err := f.ApplyChildren(&v2); err != nil {
			return err
		}
	}
	v.curOffset += length
	switch f := f.(type) {
	case *uefi.FirmwareVolume:
		// Print free space at the end of the volume
		v2.printRow("Free", "", "", offset+length-f.FreeSpace, f.FreeSpace)
	case *uefi.NVarStore:
		// Print free space and GUID store
		v2.printRow("Free", "", "", offset+f.FreeSpaceOffset, f.GUIDStoreOffset-f.FreeSpaceOffset)
		v2.printRow("GUIDStore", "", fmt.Sprintf("%d GUID", len(f.GUIDStore)), offset+f.GUIDStoreOffset, f.Length-f.GUIDStoreOffset)
	case *uefi.File:
		// Align
		v.curOffset = uefi.Align8(v.curOffset)
	}
	return nil
}

func (v *Table) printRow(node, name, typez interface{}, offset, length uint64) {
	if v.Layout {
		if name == "" {
			name = typez
		}
		fmt.Fprintf(v.W, "%s%v\t%v\t%#08x\t%#08x\n", indent(v.indent), node, name, offset, length)
	} else {
		fmt.Fprintf(v.W, "%s%v\t%v\t%v\t%#8x\n", indent(v.indent), node, name, typez, length)
	}
}

func init() {
	RegisterCLI("table", "print out important information in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{}, nil
	})
	RegisterCLI("layout-table", "print out offset and size information of top level firmware volumes in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{Layout: true, Depth: 1}, nil
	})
	RegisterCLI("layout-table-full", "print out offset and size information in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{Layout: true}, nil
	})
}
<commit_msg>visitors/table: refactored to use a printRow function pointer<commit_after>// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/linuxboot/fiano/pkg/uefi"
)

// Table prints the GUIDS, types and sizes as a compact table.
type Table struct {
	W         *tabwriter.Writer
	Layout    bool
	Depth     int
	indent    int
	offset    uint64
	curOffset uint64
	printRow  func(v *Table, node, name, typez interface{}, offset, length uint64)
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Table) Run(f uefi.Firmware) error {
	return f.Apply(v)
}

// Visit applies the Table visitor to any Firmware type.
func (v *Table) Visit(f uefi.Firmware) error {
	var offset uint64
	switch f := f.(type) {
	case *uefi.FlashImage:
		if v.Depth > 0 { // Depth <= 0 means all
			v.Depth++
		}
		return v.printFirmware(f, "Image", "", "", 0, 0)
	case *uefi.FirmwareVolume:
		return v.printFirmware(f, "FV", f.FileSystemGUID.String(), "", v.offset+f.FVOffset, v.offset+f.FVOffset+f.DataOffset)
	case *uefi.File:
		// TODO: make name part of the file node
		return v.printFirmware(f, "File", f.Header.GUID.String(), f.Header.Type, v.curOffset, v.curOffset+f.DataOffset)
	case *uefi.Section:
		// Reset offset to O for (compressed) section content
		return v.printFirmware(f, "Sec", f.String(), f.Type, v.curOffset, 0)
	case *uefi.FlashDescriptor:
		return v.printFirmware(f, "IFD", "", "", 0, 0)
	case *uefi.BIOSRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
		}
		return v.printFirmware(f, "BIOS", "", "", offset, offset)
	case *uefi.BIOSPadding:
		return v.printFirmware(f, "BIOS Pad", "", "", v.offset+f.Offset, 0)
	case *uefi.NVarStore:
		return v.printFirmware(f, "NVAR Store", "", "", v.curOffset, v.curOffset)
	case *uefi.NVar:
		return v.printFirmware(f, "NVAR", f.GUID.String(), f, v.curOffset, v.curOffset+uint64(f.DataOffset))
	case *uefi.RawRegion:
		if f.FRegion != nil {
			offset = uint64(f.FRegion.BaseOffset())
		}
		return v.printFirmware(f, f.Type().String(), "", "", offset, offset)
	default:
		return v.printFirmware(f, fmt.Sprintf("%T", f), "", "", 0, 0)
	}
}

func indent(n int) string {
	return strings.Repeat(" ", n)
}

func (v *Table) printFirmware(f uefi.Firmware, node, name, typez interface{}, offset, dataOffset uint64) error {
	// Init: Print title and select printRow func
	if v.W == nil {
		v.W = tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
		defer func() { v.W.Flush() }()
		if v.Layout {
			fmt.Fprintf(v.W, "%sNode\tGUID/Name\tOffset\tSize\n", indent(v.indent))
			v.printRow = printRowLayout
		} else {
			fmt.Fprintf(v.W, "%sNode\tGUID/Name\tType\tSize\n", indent(v.indent))
			v.printRow = printRowStd
		}
	}

	// Prepare data and print
	length := uint64(len(f.Buf()))
	if typez == "" {
		if uefi.IsErased(f.Buf(), uefi.Attributes.ErasePolarity) {
			typez = "(empty)"
		}
	}
	v.printRow(v, node, name, typez, offset, length)

	// Compute offset and visit children
	v2 := *v
	v2.indent++
	v2.offset = dataOffset
	v2.curOffset = v2.offset
	if v.Depth <= 0 || v.indent < v.Depth {
		if err := f.ApplyChildren(&v2); err != nil {
			return err
		}
	}
	v.curOffset += length

	// Print footer
	switch f := f.(type) {
	case *uefi.FirmwareVolume:
		// Print free space at the end of the volume
		v2.printRow(&v2, "Free", "", "", offset+length-f.FreeSpace, f.FreeSpace)
	case *uefi.NVarStore:
		// Print free space and GUID store
		v2.printRow(&v2, "Free", "", "", offset+f.FreeSpaceOffset, f.GUIDStoreOffset-f.FreeSpaceOffset)
		v2.printRow(&v2, "GUIDStore", "", fmt.Sprintf("%d GUID", len(f.GUIDStore)), offset+f.GUIDStoreOffset, f.Length-f.GUIDStoreOffset)
	case *uefi.File:
		// Align
		v.curOffset = uefi.Align8(v.curOffset)
	}
	return nil
}

func printRowLayout(v *Table, node, name, typez interface{}, offset, length uint64) {
	if name == "" {
		name = typez
	}
	fmt.Fprintf(v.W, "%s%v\t%v\t%#08x\t%#08x\n", indent(v.indent), node, name, offset, length)
}

func printRowStd(v *Table, node, name, typez interface{}, offset, length uint64) {
	fmt.Fprintf(v.W, "%s%v\t%v\t%v\t%#8x\n", indent(v.indent), node, name, typez, length)
}

func init() {
	RegisterCLI("table", "print out important information in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{}, nil
	})
	RegisterCLI("layout-table", "print out offset and size information of top level firmware volumes in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{Layout: true, Depth: 1}, nil
	})
	RegisterCLI("layout-table-full", "print out offset and size information in a pretty table", 0, func(args []string) (uefi.Visitor, error) {
		return &Table{Layout: true}, nil
	})
}
<|endoftext|>