<commit_before>// Copyright 2013 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 imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, file := range buildPkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "could not parse %q: %v", file, err)
			continue
		}
		for name := range f.Scope.Objects {
			if ast.IsExported(name) {
				exports[name] = true
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<commit_msg>goimports: print \n at end of errors<commit_after>// Copyright 2013 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 imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v\n", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, file := range buildPkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "could not parse %q: %v\n", file, err)
			continue
		}
		for name := range f.Scope.Objects {
			if ast.IsExported(name) {
				exports[name] = true
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<|endoftext|>


<commit_before>// Contains the node database, storing previously seen nodes and any collected
// metadata about them for QoS purposes.

package discover

import (
	"bytes"
	"encoding/binary"
	"os"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/storage"
)

var (
	nodeDBNilNodeID      = NodeID{}       // Special node ID to use as a nil element.
	nodeDBNodeExpiration = 24 * time.Hour // Time after which an unseen node should be dropped.
	nodeDBCleanupCycle   = time.Hour      // Time period for running the expiration task.
)

// nodeDB stores all nodes we know about.
type nodeDB struct {
	lvl    *leveldb.DB       // Interface to the database itself
	seeder iterator.Iterator // Iterator for fetching possible seed nodes

	runner sync.Once     // Ensures we can start at most one expirer
	quit   chan struct{} // Channel to signal the expiring thread to stop
}

// Schema layout for the node database
var (
	nodeDBVersionKey = []byte("version") // Version of the database to flush if changes
	nodeDBItemPrefix = []byte("n:")      // Identifier to prefix node entries with

	nodeDBDiscoverRoot = ":discover"
	nodeDBDiscoverPing = nodeDBDiscoverRoot + ":lastping"
	nodeDBDiscoverPong = nodeDBDiscoverRoot + ":lastpong"
)

// newNodeDB creates a new node database for storing and retrieving infos about
// known peers in the network. If no path is given, an in-memory, temporary
// database is constructed.
func newNodeDB(path string, version int) (*nodeDB, error) {
	if path == "" {
		return newMemoryNodeDB()
	}
	return newPersistentNodeDB(path, version)
}

// newMemoryNodeDB creates a new in-memory node database without a persistent
// backend.
func newMemoryNodeDB() (*nodeDB, error) {
	db, err := leveldb.Open(storage.NewMemStorage(), nil)
	if err != nil {
		return nil, err
	}
	return &nodeDB{
		lvl:  db,
		quit: make(chan struct{}),
	}, nil
}

// newPersistentNodeDB creates/opens a leveldb backed persistent node database,
// also flushing its contents in case of a version mismatch.
func newPersistentNodeDB(path string, version int) (*nodeDB, error) {
	// Try to open the cache, recovering any corruption
	db, err := leveldb.OpenFile(path, nil)
	if _, iscorrupted := err.(leveldb.ErrCorrupted); iscorrupted {
		db, err = leveldb.RecoverFile(path, nil)
	}
	if err != nil {
		return nil, err
	}
	// The nodes contained in the cache correspond to a certain protocol version.
	// Flush all nodes if the version doesn't match.
	currentVer := make([]byte, binary.MaxVarintLen64)
	currentVer = currentVer[:binary.PutVarint(currentVer, int64(version))]

	blob, err := db.Get(nodeDBVersionKey, nil)
	switch err {
	case leveldb.ErrNotFound:
		// Version not found (i.e. empty cache), insert it
		if err := db.Put(nodeDBVersionKey, currentVer, nil); err != nil {
			db.Close()
			return nil, err
		}

	case nil:
		// Version present, flush if different
		if !bytes.Equal(blob, currentVer) {
			db.Close()
			if err = os.RemoveAll(path); err != nil {
				return nil, err
			}
			return newPersistentNodeDB(path, version)
		}
	}
	return &nodeDB{
		lvl:  db,
		quit: make(chan struct{}),
	}, nil
}

// makeKey generates the leveldb key-blob from a node id and its particular
// field of interest.
func makeKey(id NodeID, field string) []byte {
	if bytes.Equal(id[:], nodeDBNilNodeID[:]) {
		return []byte(field)
	}
	return append(nodeDBItemPrefix, append(id[:], field...)...)
}

// splitKey tries to split a database key into a node id and a field part.
func splitKey(key []byte) (id NodeID, field string) {
	// If the key is not of a node, return it plainly
	if !bytes.HasPrefix(key, nodeDBItemPrefix) {
		return NodeID{}, string(key)
	}
	// Otherwise split the id and field
	item := key[len(nodeDBItemPrefix):]
	copy(id[:], item[:len(id)])
	field = string(item[len(id):])

	return id, field
}

// fetchInt64 retrieves an integer instance associated with a particular
// database key.
func (db *nodeDB) fetchInt64(key []byte) int64 {
	blob, err := db.lvl.Get(key, nil)
	if err != nil {
		return 0
	}
	val, read := binary.Varint(blob)
	if read <= 0 {
		return 0
	}
	return val
}

// storeInt64 update a specific database entry to the current time instance as a
// unix timestamp.
func (db *nodeDB) storeInt64(key []byte, n int64) error {
	blob := make([]byte, binary.MaxVarintLen64)
	blob = blob[:binary.PutVarint(blob, n)]

	return db.lvl.Put(key, blob, nil)
}

// node retrieves a node with a given id from the database.
func (db *nodeDB) node(id NodeID) *Node {
	blob, err := db.lvl.Get(makeKey(id, nodeDBDiscoverRoot), nil)
	if err != nil {
		glog.V(logger.Detail).Infof("failed to retrieve node %v: %v", id, err)
		return nil
	}
	node := new(Node)
	if err := rlp.DecodeBytes(blob, node); err != nil {
		glog.V(logger.Warn).Infof("failed to decode node RLP: %v", err)
		return nil
	}
	return node
}

// updateNode inserts - potentially overwriting - a node into the peer database.
func (db *nodeDB) updateNode(node *Node) error {
	blob, err := rlp.EncodeToBytes(node)
	if err != nil {
		return err
	}
	return db.lvl.Put(makeKey(node.ID, nodeDBDiscoverRoot), blob, nil)
}

// ensureExpirer is a small helper method ensuring that the data expiration
// mechanism is running. If the expiration goroutine is already running, this
// method simply returns.
//
// The goal is to start the data evacuation only after the network successfully
// bootstrapped itself (to prevent dumping potentially useful seed nodes). Since
// it would require significant overhead to exactly trace the first successful
// convergence, it's simpler to "ensure" the correct state when an appropriate
// condition occurs (i.e. a successful bonding), and discard further events.
func (db *nodeDB) ensureExpirer() {
	db.runner.Do(func() { go db.expirer() })
}

// expirer should be started in a go routine, and is responsible for looping ad
// infinitum and dropping stale data from the database.
func (db *nodeDB) expirer() {
	tick := time.Tick(nodeDBCleanupCycle)
	for {
		select {
		case <-tick:
			if err := db.expireNodes(); err != nil {
				glog.V(logger.Error).Infof("Failed to expire nodedb items: %v", err)
			}

		case <-db.quit:
			return
		}
	}
}

// expireNodes iterates over the database and deletes all nodes that have not
// been seen (i.e. received a pong from) for some alloted time.
func (db *nodeDB) expireNodes() error {
	threshold := time.Now().Add(-nodeDBNodeExpiration)

	// Find discovered nodes that are older than the allowance
	it := db.lvl.NewIterator(nil, nil)
	defer it.Release()

	for it.Next() {
		// Skip the item if not a discovery node
		id, field := splitKey(it.Key())
		if field != nodeDBDiscoverRoot {
			continue
		}
		// Skip the node if not expired yet
		if seen := db.lastPong(id); seen.After(threshold) {
			continue
		}
		// Otherwise delete all associated information
		prefix := makeKey(id, "")
		for ok := it.Seek(prefix); ok && bytes.HasPrefix(it.Key(), prefix); ok = it.Next() {
			if err := db.lvl.Delete(it.Key(), nil); err != nil {
				return err
			}
		}
	}
	return nil
}

// lastPing retrieves the time of the last ping packet send to a remote node,
// requesting binding.
func (db *nodeDB) lastPing(id NodeID) time.Time {
	return time.Unix(db.fetchInt64(makeKey(id, nodeDBDiscoverPing)), 0)
}

// updateLastPing updates the last time we tried contacting a remote node.
func (db *nodeDB) updateLastPing(id NodeID, instance time.Time) error {
	return db.storeInt64(makeKey(id, nodeDBDiscoverPing), instance.Unix())
}

// lastPong retrieves the time of the last successful contact from remote node.
func (db *nodeDB) lastPong(id NodeID) time.Time {
	return time.Unix(db.fetchInt64(makeKey(id, nodeDBDiscoverPong)), 0)
}

// updateLastPong updates the last time a remote node successfully contacted.
func (db *nodeDB) updateLastPong(id NodeID, instance time.Time) error {
	return db.storeInt64(makeKey(id, nodeDBDiscoverPong), instance.Unix())
}

// querySeeds retrieves a batch of nodes to be used as potential seed servers
// during bootstrapping the node into the network.
//
// Ideal seeds are the most recently seen nodes (highest probability to be still
// alive), but yet untried. However, since leveldb only supports dumb iteration
// we will instead start pulling in potential seeds that haven't been yet pinged
// since the start of the boot procedure.
//
// If the database runs out of potential seeds, we restart the startup counter
// and start iterating over the peers again.
func (db *nodeDB) querySeeds(n int) []*Node {
	// Create a new seed iterator if none exists
	if db.seeder == nil {
		db.seeder = db.lvl.NewIterator(nil, nil)
	}
	// Iterate over the nodes and find suitable seeds
	nodes := make([]*Node, 0, n)
	for len(nodes) < n && db.seeder.Next() {
		// Iterate until a discovery node is found
		id, field := splitKey(db.seeder.Key())
		if field != nodeDBDiscoverRoot {
			continue
		}
		// Load it as a potential seed
		if node := db.node(id); node != nil {
			nodes = append(nodes, node)
		}
	}
	// Release the iterator if we reached the end
	if len(nodes) == 0 {
		db.seeder.Release()
		db.seeder = nil
	}
	return nodes
}

// close flushes and closes the database files.
func (db *nodeDB) close() {
	if db.seeder != nil {
		db.seeder.Release()
	}
	close(db.quit)
	db.lvl.Close()
}
<commit_msg>p2p/discover: fix api issues caused by leveldb update<commit_after>// Contains the node database, storing previously seen nodes and any collected
// metadata about them for QoS purposes.

package discover

import (
	"bytes"
	"encoding/binary"
	"os"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/storage"
	"github.com/syndtr/goleveldb/leveldb/util"
)

var (
	nodeDBNilNodeID      = NodeID{}       // Special node ID to use as a nil element.
	nodeDBNodeExpiration = 24 * time.Hour // Time after which an unseen node should be dropped.
	nodeDBCleanupCycle   = time.Hour      // Time period for running the expiration task.
)

// nodeDB stores all nodes we know about.
type nodeDB struct {
	lvl    *leveldb.DB       // Interface to the database itself
	seeder iterator.Iterator // Iterator for fetching possible seed nodes

	runner sync.Once     // Ensures we can start at most one expirer
	quit   chan struct{} // Channel to signal the expiring thread to stop
}

// Schema layout for the node database
var (
	nodeDBVersionKey = []byte("version") // Version of the database to flush if changes
	nodeDBItemPrefix = []byte("n:")      // Identifier to prefix node entries with

	nodeDBDiscoverRoot = ":discover"
	nodeDBDiscoverPing = nodeDBDiscoverRoot + ":lastping"
	nodeDBDiscoverPong = nodeDBDiscoverRoot + ":lastpong"
)

// newNodeDB creates a new node database for storing and retrieving infos about
// known peers in the network. If no path is given, an in-memory, temporary
// database is constructed.
func newNodeDB(path string, version int) (*nodeDB, error) {
	if path == "" {
		return newMemoryNodeDB()
	}
	return newPersistentNodeDB(path, version)
}

// newMemoryNodeDB creates a new in-memory node database without a persistent
// backend.
func newMemoryNodeDB() (*nodeDB, error) {
	db, err := leveldb.Open(storage.NewMemStorage(), nil)
	if err != nil {
		return nil, err
	}
	return &nodeDB{
		lvl:  db,
		quit: make(chan struct{}),
	}, nil
}

// newPersistentNodeDB creates/opens a leveldb backed persistent node database,
// also flushing its contents in case of a version mismatch.
func newPersistentNodeDB(path string, version int) (*nodeDB, error) {
	// Try to open the cache, recovering any corruption
	db, err := leveldb.OpenFile(path, nil)
	if _, iscorrupted := err.(*errors.ErrCorrupted); iscorrupted {
		db, err = leveldb.RecoverFile(path, nil)
	}
	if err != nil {
		return nil, err
	}
	// The nodes contained in the cache correspond to a certain protocol version.
	// Flush all nodes if the version doesn't match.
	currentVer := make([]byte, binary.MaxVarintLen64)
	currentVer = currentVer[:binary.PutVarint(currentVer, int64(version))]

	blob, err := db.Get(nodeDBVersionKey, nil)
	switch err {
	case leveldb.ErrNotFound:
		// Version not found (i.e. empty cache), insert it
		if err := db.Put(nodeDBVersionKey, currentVer, nil); err != nil {
			db.Close()
			return nil, err
		}

	case nil:
		// Version present, flush if different
		if !bytes.Equal(blob, currentVer) {
			db.Close()
			if err = os.RemoveAll(path); err != nil {
				return nil, err
			}
			return newPersistentNodeDB(path, version)
		}
	}
	return &nodeDB{
		lvl:  db,
		quit: make(chan struct{}),
	}, nil
}

// makeKey generates the leveldb key-blob from a node id and its particular
// field of interest.
func makeKey(id NodeID, field string) []byte {
	if bytes.Equal(id[:], nodeDBNilNodeID[:]) {
		return []byte(field)
	}
	return append(nodeDBItemPrefix, append(id[:], field...)...)
}

// splitKey tries to split a database key into a node id and a field part.
func splitKey(key []byte) (id NodeID, field string) {
	// If the key is not of a node, return it plainly
	if !bytes.HasPrefix(key, nodeDBItemPrefix) {
		return NodeID{}, string(key)
	}
	// Otherwise split the id and field
	item := key[len(nodeDBItemPrefix):]
	copy(id[:], item[:len(id)])
	field = string(item[len(id):])

	return id, field
}

// fetchInt64 retrieves an integer instance associated with a particular
// database key.
func (db *nodeDB) fetchInt64(key []byte) int64 {
	blob, err := db.lvl.Get(key, nil)
	if err != nil {
		return 0
	}
	val, read := binary.Varint(blob)
	if read <= 0 {
		return 0
	}
	return val
}

// storeInt64 update a specific database entry to the current time instance as a
// unix timestamp.
func (db *nodeDB) storeInt64(key []byte, n int64) error {
	blob := make([]byte, binary.MaxVarintLen64)
	blob = blob[:binary.PutVarint(blob, n)]

	return db.lvl.Put(key, blob, nil)
}

// node retrieves a node with a given id from the database.
func (db *nodeDB) node(id NodeID) *Node {
	blob, err := db.lvl.Get(makeKey(id, nodeDBDiscoverRoot), nil)
	if err != nil {
		glog.V(logger.Detail).Infof("failed to retrieve node %v: %v", id, err)
		return nil
	}
	node := new(Node)
	if err := rlp.DecodeBytes(blob, node); err != nil {
		glog.V(logger.Warn).Infof("failed to decode node RLP: %v", err)
		return nil
	}
	return node
}

// updateNode inserts - potentially overwriting - a node into the peer database.
func (db *nodeDB) updateNode(node *Node) error {
	blob, err := rlp.EncodeToBytes(node)
	if err != nil {
		return err
	}
	return db.lvl.Put(makeKey(node.ID, nodeDBDiscoverRoot), blob, nil)
}

// ensureExpirer is a small helper method ensuring that the data expiration
// mechanism is running. If the expiration goroutine is already running, this
// method simply returns.
//
// The goal is to start the data evacuation only after the network successfully
// bootstrapped itself (to prevent dumping potentially useful seed nodes). Since
// it would require significant overhead to exactly trace the first successful
// convergence, it's simpler to "ensure" the correct state when an appropriate
// condition occurs (i.e. a successful bonding), and discard further events.
func (db *nodeDB) ensureExpirer() {
	db.runner.Do(func() { go db.expirer() })
}

// expirer should be started in a go routine, and is responsible for looping ad
// infinitum and dropping stale data from the database.
func (db *nodeDB) expirer() {
	tick := time.Tick(nodeDBCleanupCycle)
	for {
		select {
		case <-tick:
			if err := db.expireNodes(); err != nil {
				glog.V(logger.Error).Infof("Failed to expire nodedb items: %v", err)
			}

		case <-db.quit:
			return
		}
	}
}

// expireNodes iterates over the database and deletes all nodes that have not
// been seen (i.e. received a pong from) for some alloted time.
func (db *nodeDB) expireNodes() error {
	threshold := time.Now().Add(-nodeDBNodeExpiration)

	// Find discovered nodes that are older than the allowance
	it := db.lvl.NewIterator(nil, nil)
	defer it.Release()

	for it.Next() {
		// Skip the item if not a discovery node
		id, field := splitKey(it.Key())
		if field != nodeDBDiscoverRoot {
			continue
		}
		// Skip the node if not expired yet
		if seen := db.lastPong(id); seen.After(threshold) {
			continue
		}
		// Otherwise delete all associated information
		deleter := db.lvl.NewIterator(util.BytesPrefix(makeKey(id, "")), nil)
		for deleter.Next() {
			if err := db.lvl.Delete(deleter.Key(), nil); err != nil {
				return err
			}
		}
	}
	return nil
}

// lastPing retrieves the time of the last ping packet send to a remote node,
// requesting binding.
func (db *nodeDB) lastPing(id NodeID) time.Time {
	return time.Unix(db.fetchInt64(makeKey(id, nodeDBDiscoverPing)), 0)
}

// updateLastPing updates the last time we tried contacting a remote node.
func (db *nodeDB) updateLastPing(id NodeID, instance time.Time) error {
	return db.storeInt64(makeKey(id, nodeDBDiscoverPing), instance.Unix())
}

// lastPong retrieves the time of the last successful contact from remote node.
func (db *nodeDB) lastPong(id NodeID) time.Time {
	return time.Unix(db.fetchInt64(makeKey(id, nodeDBDiscoverPong)), 0)
}

// updateLastPong updates the last time a remote node successfully contacted.
func (db *nodeDB) updateLastPong(id NodeID, instance time.Time) error {
	return db.storeInt64(makeKey(id, nodeDBDiscoverPong), instance.Unix())
}

// querySeeds retrieves a batch of nodes to be used as potential seed servers
// during bootstrapping the node into the network.
//
// Ideal seeds are the most recently seen nodes (highest probability to be still
// alive), but yet untried. However, since leveldb only supports dumb iteration
// we will instead start pulling in potential seeds that haven't been yet pinged
// since the start of the boot procedure.
//
// If the database runs out of potential seeds, we restart the startup counter
// and start iterating over the peers again.
func (db *nodeDB) querySeeds(n int) []*Node {
	// Create a new seed iterator if none exists
	if db.seeder == nil {
		db.seeder = db.lvl.NewIterator(nil, nil)
	}
	// Iterate over the nodes and find suitable seeds
	nodes := make([]*Node, 0, n)
	for len(nodes) < n && db.seeder.Next() {
		// Iterate until a discovery node is found
		id, field := splitKey(db.seeder.Key())
		if field != nodeDBDiscoverRoot {
			continue
		}
		// Load it as a potential seed
		if node := db.node(id); node != nil {
			nodes = append(nodes, node)
		}
	}
	// Release the iterator if we reached the end
	if len(nodes) == 0 {
		db.seeder.Release()
		db.seeder = nil
	}
	return nodes
}

// close flushes and closes the database files.
func (db *nodeDB) close() {
	if db.seeder != nil {
		db.seeder.Release()
	}
	close(db.quit)
	db.lvl.Close()
}
<|endoftext|>


<commit_before>package couchdb

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/cozy/cozy-stack/couchdb/mango"
)

// Doc is the interface that encapsulate a couchdb document, of any
// serializable type. This interface defines method to set and get the
// ID of the document.
type Doc interface {
	ID() string
	Rev() string
	DocType() string

	SetID(id string)
	SetRev(rev string)
}

// JSONDoc is a map representing a simple json object that implements
// the Doc interface.
type JSONDoc struct {
	M    map[string]interface{}
	Type string
}

// ID returns the identifier field of the document
//   "io.cozy.event/123abc123" == doc.ID()
func (j JSONDoc) ID() string {
	id, ok := j.M["_id"].(string)
	if ok {
		return id
	}
	return ""
}

// Rev returns the revision field of the document
//   "3-1234def1234" == doc.Rev()
func (j JSONDoc) Rev() string {
	rev, ok := j.M["_rev"].(string)
	if ok {
		return rev
	}
	return ""
}

// DocType returns the document type of the document
//   "io.cozy.event" == doc.Doctype()
func (j JSONDoc) DocType() string {
	return j.Type
}

// SetID is used to set the identifier of the document
func (j JSONDoc) SetID(id string) {
	j.M["_id"] = id
}

// SetRev is used to set the revision of the document
func (j JSONDoc) SetRev(rev string) {
	j.M["_rev"] = rev
}

// MarshalJSON implements json.Marshaller by proxying to internal map
func (j JSONDoc) MarshalJSON() ([]byte, error) {
	return json.Marshal(j.M)
}

// UnmarshalJSON implements json.Unmarshaller by proxying to internal map
func (j *JSONDoc) UnmarshalJSON(bytes []byte) error {
	err := json.Unmarshal(bytes, &j.M)
	if err != nil {
		return err
	}
	doctype, ok := j.M["_type"].(string)
	if ok {
		j.Type = doctype
	}
	delete(j.M, "_type")
	return nil
}

// ToMapWithType returns the JSONDoc internal map including its DocType
// its used in request response.
func (j *JSONDoc) ToMapWithType() map[string]interface{} {
	j.M["_type"] = j.DocType()
	return j.M
}

// Get returns the value of one of the db fields
func (j JSONDoc) Get(key string) interface{} {
	return j.M[key]
}

// CouchURL is the URL where to check if CouchDB is up
func CouchURL() string {
	return "http://localhost:5984/"
}

var couchdbClient = &http.Client{}

func makeDBName(dbprefix, doctype string) string {
	// @TODO This should be better analysed
	dbname := dbprefix + doctype
	dbname = strings.Replace(dbname, ".", "-", -1)
	dbname = strings.ToLower(dbname)
	return url.QueryEscape(dbname)
}

func docURL(dbprefix, doctype, id string) string {
	return makeDBName(dbprefix, doctype) + "/" + url.QueryEscape(id)
}

func makeRequest(method, path string, reqbody interface{}, resbody interface{}) error {
	var reqjson []byte
	var err error

	if reqbody != nil {
		reqjson, err = json.Marshal(reqbody)
		if err != nil {
			return err
		}
	}

	fmt.Printf("[couchdb request] %v %v %v\n", method, path, string(reqjson))

	req, err := http.NewRequest(method, CouchURL()+path, bytes.NewReader(reqjson))
	// Possible err = wrong method, unparsable url
	if err != nil {
		return newRequestError(err)
	}
	if reqbody != nil {
		req.Header.Add("Content-Type", "application/json")
	}
	req.Header.Add("Accept", "application/json")
	resp, err := couchdbClient.Do(req)
	// Possible err = mostly connection failure
	if err != nil {
		return newConnectionError(err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	// Possible err = mostly connection failure (hangup)
	if err != nil {
		return newIOReadError(err)
	}

	fmt.Printf("[couchdb response] %v\n", string(body))

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		// Couchdb as returned an error HTTP status code
		return newCouchdbError(resp.StatusCode, body)
	}

	if resbody == nil {
		// dont care about the return value
		return nil
	}
	err = json.Unmarshal(body, &resbody)
	return err
}

func fixErrorNoDatabaseIsWrongDoctype(err error) {
	if IsNoDatabaseError(err) {
		err.(*Error).Reason = "wrong_doctype"
	}
}

// GetDoc fetch a document by its docType and ID, out is filled with
// the document by json.Unmarshal-ing
func GetDoc(dbprefix, doctype, id string, out Doc) error {
	err := makeRequest("GET", docURL(dbprefix, doctype, id), nil, out)
	fixErrorNoDatabaseIsWrongDoctype(err)
	return err
}

// CreateDB creates the necessary database for a doctype
func CreateDB(dbprefix, doctype string) error {
	return makeRequest("PUT", makeDBName(dbprefix, doctype), nil, nil)
}

// DeleteDB destroy the database for a doctype
func DeleteDB(dbprefix, doctype string) error {
	return makeRequest("DELETE", makeDBName(dbprefix, doctype), nil, nil)
}

// ResetDB destroy and recreate the database for a doctype
func ResetDB(dbprefix, doctype string) (err error) {
	err = DeleteDB(dbprefix, doctype)
	if err != nil && !IsNoDatabaseError(err) {
		return err
	}
	return CreateDB(dbprefix, doctype)
}

// Delete destroy a document by its doctype and ID .
// If the document's current rev does not match the one passed,
// a CouchdbError(409 conflict) will be returned.
// This functions returns the tombstone revision as string
func Delete(dbprefix, doctype, id, rev string) (tombrev string, err error) {
	var res updateResponse
	qs := url.Values{"rev": []string{rev}}
	url := docURL(dbprefix, doctype, id) + "?" + qs.Encode()
	err = makeRequest("DELETE", url, nil, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		tombrev = res.Rev
	}
	return
}

// DeleteDoc deletes a struct implementing the couchb.Doc interface
// The document's SetRev will be called with tombstone revision
func DeleteDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()
	rev := doc.Rev()
	tombrev, err := Delete(dbprefix, doctype, id, rev)
	if err == nil {
		doc.SetRev(tombrev)
	}
	return
}

// UpdateDoc update a document. The document ID and Rev should be fillled.
// The doc SetRev function will be called with the new rev.
func UpdateDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()
	rev := doc.Rev()
	if id == "" || rev == "" || doctype == "" {
		return fmt.Errorf("UpdateDoc doc argument should have doctype, id and rev")
	}

	url := docURL(dbprefix, doctype, id)
	var res updateResponse
	err = makeRequest("PUT", url, doc, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		doc.SetRev(res.Rev)
	}
	return err
}

// CreateNamedDoc persist a document with an ID.
// if the document already exist, it will return a 409 error.
// The document ID should be fillled.
// The doc SetRev function will be called with the new rev.
func CreateNamedDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()

	if doc.Rev() != "" || doc.ID() == "" || doctype == "" {
		return fmt.Errorf("CreateNamedDoc should have type and id but no rev")
	}

	url := docURL(dbprefix, doctype, id)
	var res updateResponse
	err = makeRequest("PUT", url, doc, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		doc.SetRev(res.Rev)
	}
	return err
}

func createDocOrDb(dbprefix string, doc Doc, response interface{}) (err error) {
	doctype := doc.DocType()
	db := makeDBName(dbprefix, doctype)
	err = makeRequest("POST", db, doc, response)
	if err == nil || !IsNoDatabaseError(err) {
		return
	}

	err = CreateDB(dbprefix, doctype)
	if err == nil {
		err = makeRequest("POST", db, doc, response)
	}
	return
}

// CreateDoc is used to persist the given document in the couchdb
// database. The document's SetRev and SetID function will be called
// with the document's new ID and Rev.
// This function creates a database if this is the first document of its type
func CreateDoc(dbprefix string, doc Doc) (err error) {
	var res *updateResponse

	if doc.ID() != "" {
		err = fmt.Errorf("Can not create document with a defined ID")
		return
	}

	err = createDocOrDb(dbprefix, doc, &res)
	if err != nil {
		return err
	} else if !res.Ok {
		return fmt.Errorf("CouchDB replied with 200 ok=false")
	}

	doc.SetID(res.ID)
	doc.SetRev(res.Rev)
	return nil
}

// DefineIndex define the index on the doctype database
// see query package on how to define an index
func DefineIndex(dbprefix, doctype string, index mango.IndexDefinitionRequest) error {
	url := makeDBName(dbprefix, doctype) + "/_index"
	var response indexCreationResponse
	return makeRequest("POST", url, &index, &response)
}

// FindDocs returns all documents matching the passed FindRequest
// documents will be unmarshalled in the provided results slice.
func FindDocs(dbprefix, doctype string, req *FindRequest, results interface{}) error {
	url := makeDBName(dbprefix, doctype) + "/_find"
	// prepare a structure to receive the results
	var response findResponse
	err := makeRequest("POST", url, &req, &response)
	if err != nil {
		return err
	}
	return json.Unmarshal(response.Docs, results)
}

type indexCreationResponse struct {
	Result string `json:"result"`
	Error  string `json:"error"`
	ID     string `json:"id"`
	Name   string `json:"name"`
}

type updateResponse struct {
	ID  string `json:"id"`
	Rev string `json:"rev"`
	Ok  bool   `json:"ok"`
}

type findResponse struct {
	Docs json.RawMessage `json:"docs"`
}

// A FindRequest is a structure containin
type FindRequest struct {
	Selector mango.Filter  `json:"selector"`
	Limit    int           `json:"limit,omitempty"`
	Skip     int           `json:"skip,omitempty"`
	Sort     *mango.SortBy `json:"sort,omitempty"`
	Fields   []string      `json:"fields,omitempty"`
}
<commit_msg>CouchDB: improve unmarshaling performance (streaming in decoder)<commit_after>package couchdb

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/cozy/cozy-stack/couchdb/mango"
)

// Doc is the interface that encapsulate a couchdb document, of any
// serializable type. This interface defines method to set and get the
// ID of the document.
type Doc interface {
	ID() string
	Rev() string
	DocType() string

	SetID(id string)
	SetRev(rev string)
}

// JSONDoc is a map representing a simple json object that implements
// the Doc interface.
type JSONDoc struct {
	M    map[string]interface{}
	Type string
}

// ID returns the identifier field of the document
//   "io.cozy.event/123abc123" == doc.ID()
func (j JSONDoc) ID() string {
	id, ok := j.M["_id"].(string)
	if ok {
		return id
	}
	return ""
}

// Rev returns the revision field of the document
//   "3-1234def1234" == doc.Rev()
func (j JSONDoc) Rev() string {
	rev, ok := j.M["_rev"].(string)
	if ok {
		return rev
	}
	return ""
}

// DocType returns the document type of the document
//   "io.cozy.event" == doc.Doctype()
func (j JSONDoc) DocType() string {
	return j.Type
}

// SetID is used to set the identifier of the document
func (j JSONDoc) SetID(id string) {
	j.M["_id"] = id
}

// SetRev is used to set the revision of the document
func (j JSONDoc) SetRev(rev string) {
	j.M["_rev"] = rev
}

// MarshalJSON implements json.Marshaller by proxying to internal map
func (j JSONDoc) MarshalJSON() ([]byte, error) {
	return json.Marshal(j.M)
}

// UnmarshalJSON implements json.Unmarshaller by proxying to internal map
func (j *JSONDoc) UnmarshalJSON(bytes []byte) error {
	err := json.Unmarshal(bytes, &j.M)
	if err != nil {
		return err
	}
	doctype, ok := j.M["_type"].(string)
	if ok {
		j.Type = doctype
	}
	delete(j.M, "_type")
	return nil
}

// ToMapWithType returns the JSONDoc internal map including its DocType
// its used in request response.
func (j *JSONDoc) ToMapWithType() map[string]interface{} {
	j.M["_type"] = j.DocType()
	return j.M
}

// Get returns the value of one of the db fields
func (j JSONDoc) Get(key string) interface{} {
	return j.M[key]
}

// CouchURL is the URL where to check if CouchDB is up
func CouchURL() string {
	return "http://localhost:5984/"
}

var couchdbClient = &http.Client{}

func makeDBName(dbprefix, doctype string) string {
	// @TODO This should be better analysed
	dbname := dbprefix + doctype
	dbname = strings.Replace(dbname, ".", "-", -1)
	dbname = strings.ToLower(dbname)
	return url.QueryEscape(dbname)
}

func docURL(dbprefix, doctype, id string) string {
	return makeDBName(dbprefix, doctype) + "/" + url.QueryEscape(id)
}

func makeRequest(method, path string, reqbody interface{}, resbody interface{}) error {
	var reqjson []byte
	var err error

	if reqbody != nil {
		reqjson, err = json.Marshal(reqbody)
		if err != nil {
			return err
		}
	}

	fmt.Printf("[couchdb request] %v %v %v\n", method, path, string(reqjson))

	req, err := http.NewRequest(method, CouchURL()+path, bytes.NewReader(reqjson))
	// Possible err = wrong method, unparsable url
	if err != nil {
		return newRequestError(err)
	}
	if reqbody != nil {
		req.Header.Add("Content-Type", "application/json")
	}
	req.Header.Add("Accept", "application/json")
	resp, err := couchdbClient.Do(req)
	// Possible err = mostly connection failure
	if err != nil {
		return newConnectionError(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		var body []byte
		body, err = ioutil.ReadAll(resp.Body)
		if err != nil {
			err = newIOReadError(err)
		} else {
			err = newCouchdbError(resp.StatusCode, body)
		}
		fmt.Printf("[couchdb error] %v\n", err.Error())
		return err
	}

	if resbody != nil {
		err = json.NewDecoder(resp.Body).Decode(&resbody)
	}

	return err
}

func fixErrorNoDatabaseIsWrongDoctype(err error) {
	if IsNoDatabaseError(err) {
		err.(*Error).Reason = "wrong_doctype"
	}
}

// GetDoc fetch a document by its docType and ID, out is filled with
// the document by json.Unmarshal-ing
func GetDoc(dbprefix, doctype, id string, out Doc) error {
	err := makeRequest("GET", docURL(dbprefix, doctype, id), nil, out)
	fixErrorNoDatabaseIsWrongDoctype(err)
	return err
}

// CreateDB creates the necessary database for a doctype
func CreateDB(dbprefix, doctype string) error {
	return makeRequest("PUT", makeDBName(dbprefix, doctype), nil, nil)
}

// DeleteDB destroy the database for a doctype
func DeleteDB(dbprefix, doctype string) error {
	return makeRequest("DELETE", makeDBName(dbprefix, doctype), nil, nil)
}

// ResetDB destroy and recreate the database for a doctype
func ResetDB(dbprefix, doctype string) (err error) {
	err = DeleteDB(dbprefix, doctype)
	if err != nil && !IsNoDatabaseError(err) {
		return err
	}
	return CreateDB(dbprefix, doctype)
}

// Delete destroy a document by its doctype and ID .
// If the document's current rev does not match the one passed,
// a CouchdbError(409 conflict) will be returned.
// This functions returns the tombstone revision as string
func Delete(dbprefix, doctype, id, rev string) (tombrev string, err error) {
	var res updateResponse
	qs := url.Values{"rev": []string{rev}}
	url := docURL(dbprefix, doctype, id) + "?" + qs.Encode()
	err = makeRequest("DELETE", url, nil, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		tombrev = res.Rev
	}
	return
}

// DeleteDoc deletes a struct implementing the couchb.Doc interface
// The document's SetRev will be called with tombstone revision
func DeleteDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()
	rev := doc.Rev()
	tombrev, err := Delete(dbprefix, doctype, id, rev)
	if err == nil {
		doc.SetRev(tombrev)
	}
	return
}

// UpdateDoc update a document. The document ID and Rev should be fillled.
// The doc SetRev function will be called with the new rev.
func UpdateDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()
	rev := doc.Rev()
	if id == "" || rev == "" || doctype == "" {
		return fmt.Errorf("UpdateDoc doc argument should have doctype, id and rev")
	}

	url := docURL(dbprefix, doctype, id)
	var res updateResponse
	err = makeRequest("PUT", url, doc, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		doc.SetRev(res.Rev)
	}
	return err
}

// CreateNamedDoc persist a document with an ID.
// if the document already exist, it will return a 409 error.
// The document ID should be fillled.
// The doc SetRev function will be called with the new rev.
func CreateNamedDoc(dbprefix string, doc Doc) (err error) {
	doctype := doc.DocType()
	id := doc.ID()

	if doc.Rev() != "" || doc.ID() == "" || doctype == "" {
		return fmt.Errorf("CreateNamedDoc should have type and id but no rev")
	}

	url := docURL(dbprefix, doctype, id)
	var res updateResponse
	err = makeRequest("PUT", url, doc, &res)
	fixErrorNoDatabaseIsWrongDoctype(err)
	if err == nil {
		doc.SetRev(res.Rev)
	}
	return err
}

func createDocOrDb(dbprefix string, doc Doc, response interface{}) (err error) {
	doctype := doc.DocType()
	db := makeDBName(dbprefix, doctype)
	err = makeRequest("POST", db, doc, response)
	if err == nil || !IsNoDatabaseError(err) {
		return
	}

	err = CreateDB(dbprefix, doctype)
	if err == nil {
		err = makeRequest("POST", db, doc, response)
	}
	return
}

// CreateDoc is used to persist the given document in the couchdb
// database. The document's SetRev and SetID function will be called
// with the document's new ID and Rev.
// This function creates a database if this is the first document of its type
func CreateDoc(dbprefix string, doc Doc) (err error) {
	var res *updateResponse

	if doc.ID() != "" {
		err = fmt.Errorf("Can not create document with a defined ID")
		return
	}

	err = createDocOrDb(dbprefix, doc, &res)
	if err != nil {
		return err
	} else if !res.Ok {
		return fmt.Errorf("CouchDB replied with 200 ok=false")
	}

	doc.SetID(res.ID)
	doc.SetRev(res.Rev)
	return nil
}

// DefineIndex define the index on the doctype database
// see query package on how to define an index
func DefineIndex(dbprefix, doctype string, index mango.IndexDefinitionRequest) error {
	url := makeDBName(dbprefix, doctype) + "/_index"
	var response indexCreationResponse
	return makeRequest("POST", url, &index, &response)
}

// FindDocs returns all documents matching the passed FindRequest
// documents will be unmarshalled in the provided results slice.
func FindDocs(dbprefix, doctype string, req *FindRequest, results interface{}) error {
	url := makeDBName(dbprefix, doctype) + "/_find"
	// prepare a structure to receive the results
	var response findResponse
	err := makeRequest("POST", url, &req, &response)
	if err != nil {
		return err
	}
	return json.Unmarshal(response.Docs, results)
}

type indexCreationResponse struct {
	Result string `json:"result"`
	Error  string `json:"error"`
	ID     string `json:"id"`
	Name   string `json:"name"`
}

type updateResponse struct {
	ID  string `json:"id"`
	Rev string `json:"rev"`
	Ok  bool   `json:"ok"`
}

type findResponse struct {
	Docs json.RawMessage `json:"docs"`
}

// A FindRequest is a structure containin
type FindRequest struct {
	Selector mango.Filter  `json:"selector"`
	Limit    int           `json:"limit,omitempty"`
	Skip     int           `json:"skip,omitempty"`
	Sort     *mango.SortBy `json:"sort,omitempty"`
	Fields   []string      `json:"fields,omitempty"`
}
<|endoftext|>


<commit_before>package jobs

import (
	"context"
	"encoding/json"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/sirupsen/logrus"
)

const (
	// Queued state
	Queued State = "queued"
	// Running state
	Running = "running"
	// Done state
	Done = "done"
	// Errored state
	Errored = "errored"
)

const (
	// WorkerType is the key in JSON for the type of worker
	WorkerType = "worker"
)

type (
	// Broker interface is used to represent a job broker associated to a
	// particular domain. A broker can be used to create jobs that are pushed in
	// the job system.
	Broker interface {
		Start(workersList WorkersList) error
		Shutdown(ctx context.Context) error

		// PushJob will push try to push a new job from the specified job request.
		// This method is asynchronous.
		PushJob(request *JobRequest) (*Job, error)

		// QueueLen returns the total element in the queue of the specified worker
		// type.
		QueueLen(workerType string) (int, error)
	}

	// State represent the state of a job.
	State string

	// Message is a json encoded job message.
	Message json.RawMessage

	// Event is a json encoded value of a realtime.Event
	Event json.RawMessage

	// Job contains all the metadata informations of a Job. It can be
	// marshalled in JSON.
	Job struct {
		JobID      string      `json:"_id,omitempty"`
		JobRev     string      `json:"_rev,omitempty"`
		Domain     string      `json:"domain"`
		WorkerType string      `json:"worker"`
		TriggerID  string      `json:"trigger_id"`
		Message    Message     `json:"message"`
		Event      Event       `json:"event"`
		Debounced  bool        `json:"debounced"`
		Options    *JobOptions `json:"options"`
		State      State       `json:"state"`
		QueuedAt   time.Time   `json:"queued_at"`
		StartedAt  time.Time   `json:"started_at,omitempty"`
		Error      string      `json:"error,omitempty"`
	}

	// JobRequest struct is used to represent a new job request.
	JobRequest struct {
		Domain     string
		WorkerType string
		TriggerID  string
		Message    Message
		Event      Event
		Debounced  bool
		Options    *JobOptions
	}

	// JobOptions struct contains the execution properties of the jobs.
	JobOptions struct {
		MaxExecCount int           `json:"max_exec_count"`
		MaxExecTime  time.Duration `json:"max_exec_time"`
		Timeout      time.Duration `json:"timeout"`
	}
)

var joblog = logger.WithNamespace("jobs")

// ID implements the couchdb.Doc interface
func (j *Job) ID() string { return j.JobID }

// Rev implements the couchdb.Doc interface
func (j *Job) Rev() string { return j.JobRev }

// Clone implements the couchdb.Doc interface
func (j *Job) Clone() couchdb.Doc {
	cloned := *j
	if j.Options != nil {
		tmp := *j.Options
		cloned.Options = &tmp
	}
	if j.Message != nil {
		tmp := j.Message
		j.Message = make([]byte, len(tmp))
		copy(j.Message[:], tmp)
	}
	if j.Event != nil {
		tmp := j.Event
		j.Event = make([]byte, len(tmp))
		copy(j.Event[:], tmp)
	}
	return &cloned
}

// DocType implements the couchdb.Doc interface
func (j *Job) DocType() string { return consts.Jobs }

// SetID implements the couchdb.Doc interface
func (j *Job) SetID(id string) { j.JobID = id }

// SetRev implements the couchdb.Doc interface
func (j *Job) SetRev(rev string) { j.JobRev = rev }

// Valid implements the permissions.Validable interface
func (j *Job) Valid(key, value string) bool {
	switch key {
	case WorkerType:
		return j.WorkerType == value
	}
	return false
}

// ID implements the permissions.Validable interface
func (jr *JobRequest) ID() string { return "" }

// DocType implements the permissions.Validable interface
func (jr *JobRequest) DocType() string { return consts.Jobs }

// Valid implements the permissions.Validable interface
func (jr *JobRequest) Valid(key, value string) bool {
	switch key {
	case WorkerType:
		return jr.WorkerType == value
	}
	return false
}

// Logger returns a logger associated with the job domain
func (j *Job) Logger() *logrus.Entry {
	return logger.WithDomain(j.Domain)
}

// AckConsumed sets the job infos state to Running an sends the new job infos
// on the channel.
func (j *Job) AckConsumed() error {
	j.Logger().Debugf("[jobs] ack_consume %s ", j.ID())
	j.StartedAt = time.Now()
	j.State = Running
	return j.Update()
}

// Ack sets the job infos state to Done an sends the new job infos on the
// channel.
func (j *Job) Ack() error {
	j.Logger().Debugf("[jobs] ack %s ", j.ID())
	j.State = Done
	return j.Update()
}

// Nack sets the job infos state to Errored, set the specified error has the
// error field and sends the new job infos on the channel.
func (j *Job) Nack(err error) error {
	j.Logger().Debugf("[jobs] nack %s ", j.ID())
	j.State = Errored
	j.Error = err.Error()
	return j.Update()
}

// Update updates the job in couchdb
func (j *Job) Update() error {
	return couchdb.UpdateDoc(j.db(), j)
}

// Create creates the job in couchdb
func (j *Job) Create() error {
	return couchdb.CreateDoc(j.db(), j)
}

func (j *Job) db() couchdb.Database {
	return couchdb.SimpleDatabasePrefix(j.Domain)
}

// UnmarshalJSON implements json.Unmarshaler on Message. It should be retro-
// compatible with the old Message representation { Data, Type }.
func (m *Message) UnmarshalJSON(data []byte) error {
	// For retro-compatibility purposes
	var mm struct {
		Data []byte `json:"Data"`
		Type string `json:"Type"`
	}
	if err := json.Unmarshal(data, &mm); err == nil && mm.Type == "json" {
		var v json.RawMessage
		if err = json.Unmarshal(mm.Data, &v); err != nil {
			return err
		}
		*m = Message(v)
		return nil
	}
	var v json.RawMessage
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}
	*m = Message(v)
	return nil
}

// MarshalJSON implements json.Marshaler on Message.
func (m Message) MarshalJSON() ([]byte, error) {
	v := json.RawMessage(m)
	return json.Marshal(v)
}

// NewJob creates a new Job instance from a job request.
func NewJob(req *JobRequest) *Job {
	return &Job{
		Domain:     req.Domain,
		WorkerType: req.WorkerType,
		TriggerID:  req.TriggerID,
		Message:    req.Message,
		Debounced:  req.Debounced,
		Event:      req.Event,
		Options:    req.Options,
		State:      Queued,
		QueuedAt:   time.Now(),
	}
}

// Get returns the informations about a job.
func Get(domain, jobID string) (*Job, error) {
	var job Job
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Jobs, jobID, &job); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundJob
		}
		return nil, err
	}
	return &job, nil
}

// GetQueuedJobs returns the list of jobs which states is "queued" or "running"
func GetQueuedJobs(domain, workerType string) ([]*Job, error) {
	var results []*Job
	db := couchdb.SimpleDatabasePrefix(domain)
	req := &couchdb.FindRequest{
		UseIndex: "by-worker-and-state",
		Selector: mango.And(
			mango.Equal("worker", workerType),
			mango.Exists("state"), // XXX it is needed by couchdb to use the index
			mango.Or(
				mango.Equal("state", Queued),
				mango.Equal("state", Running),
			),
		),
		Limit: 200,
	}
	err := couchdb.FindDocs(db, consts.Jobs, req, &results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

// NewMessage returns a json encoded data
func NewMessage(data interface{}) (Message, error) {
	b, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	return Message(b), nil
}

// NewEvent return a json encoded realtime.Event
func NewEvent(data *realtime.Event) (Event, error) {
	b, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	return Event(b), nil
}

// Unmarshal can be used to unmarshal the encoded message value in the
// specified interface's type.
func (m Message) Unmarshal(msg interface{}) error {
	if m == nil {
		return ErrMessageNil
	}
	return json.Unmarshal(m, &msg)
}

// Unmarshal can be used to unmarshal the encoded message value in the
// specified interface's type.
func (e Event) Unmarshal(evt interface{}) error {
	if e == nil {
		return ErrMessageNil
	}
	return json.Unmarshal(e, &evt)
}

// Clone clones the worker config
func (w *WorkerConfig) Clone() *WorkerConfig {
	return &WorkerConfig{
		WorkerInit:   w.WorkerInit,
		WorkerFunc:   w.WorkerFunc,
		WorkerCommit: w.WorkerCommit,
		Concurrency:  w.Concurrency,
		MaxExecCount: w.MaxExecCount,
		MaxExecTime:  w.MaxExecTime,
		Timeout:      w.Timeout,
		RetryDelay:   w.RetryDelay,
	}
}

var (
	_ permissions.Validable = (*JobRequest)(nil)
	_ permissions.Validable = (*Job)(nil)
)
<commit_msg>Add a finished_at field to jobs<commit_after>package jobs

import (
	"context"
	"encoding/json"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/sirupsen/logrus"
)

const (
	// Queued state
	Queued State = "queued"
	// Running state
	Running = "running"
	// Done state
	Done = "done"
	// Errored state
	Errored = "errored"
)

const (
	// WorkerType is the key in JSON for the type of worker
	WorkerType = "worker"
)

type (
	// Broker interface is used to represent a job broker associated to a
	// particular domain. A broker can be used to create jobs that are pushed in
	// the job system.
	Broker interface {
		Start(workersList WorkersList) error
		Shutdown(ctx context.Context) error

		// PushJob will push try to push a new job from the specified job request.
		// This method is asynchronous.
		PushJob(request *JobRequest) (*Job, error)

		// QueueLen returns the total element in the queue of the specified worker
		// type.
		QueueLen(workerType string) (int, error)
	}

	// State represent the state of a job.
	State string

	// Message is a json encoded job message.
	Message json.RawMessage

	// Event is a json encoded value of a realtime.Event
	Event json.RawMessage

	// Job contains all the metadata informations of a Job. It can be
	// marshalled in JSON.
	Job struct {
		JobID      string      `json:"_id,omitempty"`
		JobRev     string      `json:"_rev,omitempty"`
		Domain     string      `json:"domain"`
		WorkerType string      `json:"worker"`
		TriggerID  string      `json:"trigger_id"`
		Message    Message     `json:"message"`
		Event      Event       `json:"event"`
		Debounced  bool        `json:"debounced"`
		Options    *JobOptions `json:"options"`
		State      State       `json:"state"`
		QueuedAt   time.Time   `json:"queued_at"`
		StartedAt  time.Time   `json:"started_at"`
		FinishedAt time.Time   `json:"finished_at"`
		Error      string      `json:"error,omitempty"`
	}

	// JobRequest struct is used to represent a new job request.
	JobRequest struct {
		Domain     string
		WorkerType string
		TriggerID  string
		Message    Message
		Event      Event
		Debounced  bool
		Options    *JobOptions
	}

	// JobOptions struct contains the execution properties of the jobs.
	JobOptions struct {
		MaxExecCount int           `json:"max_exec_count"`
		MaxExecTime  time.Duration `json:"max_exec_time"`
		Timeout      time.Duration `json:"timeout"`
	}
)

var joblog = logger.WithNamespace("jobs")

// ID implements the couchdb.Doc interface
func (j *Job) ID() string { return j.JobID }

// Rev implements the couchdb.Doc interface
func (j *Job) Rev() string { return j.JobRev }

// Clone implements the couchdb.Doc interface
func (j *Job) Clone() couchdb.Doc {
	cloned := *j
	if j.Options != nil {
		tmp := *j.Options
		cloned.Options = &tmp
	}
	if j.Message != nil {
		tmp := j.Message
		j.Message = make([]byte, len(tmp))
		copy(j.Message[:], tmp)
	}
	if j.Event != nil {
		tmp := j.Event
		j.Event = make([]byte, len(tmp))
		copy(j.Event[:], tmp)
	}
	return &cloned
}

// DocType implements the couchdb.Doc interface
func (j *Job) DocType() string { return consts.Jobs }

// SetID implements the couchdb.Doc interface
func (j *Job) SetID(id string) { j.JobID = id }

// SetRev implements the couchdb.Doc interface
func (j *Job) SetRev(rev string) { j.JobRev = rev }

// Valid implements the permissions.Validable interface
func (j *Job) Valid(key, value string) bool {
	switch key {
	case WorkerType:
		return j.WorkerType == value
	}
	return false
}

// ID implements the permissions.Validable interface
func (jr *JobRequest) ID() string { return "" }

// DocType implements the permissions.Validable interface
func (jr *JobRequest) DocType() string { return consts.Jobs }

// Valid implements the permissions.Validable interface
func (jr *JobRequest) Valid(key, value string) bool {
	switch key {
	case WorkerType:
		return jr.WorkerType == value
	}
	return false
}

// Logger returns a logger associated with the job domain
func (j *Job) Logger() *logrus.Entry {
	return logger.WithDomain(j.Domain)
}

// AckConsumed sets the job infos state to Running an sends the new job infos
// on the channel.
func (j *Job) AckConsumed() error {
	j.Logger().Debugf("[jobs] ack_consume %s ", j.ID())
	j.StartedAt = time.Now()
	j.State = Running
	return j.Update()
}

// Ack sets the job infos state to Done an sends the new job infos on the
// channel.
func (j *Job) Ack() error {
	j.Logger().Debugf("[jobs] ack %s ", j.ID())
	j.FinishedAt = time.Now()
	j.State = Done
	return j.Update()
}

// Nack sets the job infos state to Errored, set the specified error has the
// error field and sends the new job infos on the channel.
func (j *Job) Nack(err error) error {
	j.Logger().Debugf("[jobs] nack %s ", j.ID())
	j.FinishedAt = time.Now()
	j.State = Errored
	j.Error = err.Error()
	return j.Update()
}

// Update updates the job in couchdb
func (j *Job) Update() error {
	return couchdb.UpdateDoc(j.db(), j)
}

// Create creates the job in couchdb
func (j *Job) Create() error {
	return couchdb.CreateDoc(j.db(), j)
}

func (j *Job) db() couchdb.Database {
	return couchdb.SimpleDatabasePrefix(j.Domain)
}

// UnmarshalJSON implements json.Unmarshaler on Message. It should be retro-
// compatible with the old Message representation { Data, Type }.
func (m *Message) UnmarshalJSON(data []byte) error {
	// For retro-compatibility purposes
	var mm struct {
		Data []byte `json:"Data"`
		Type string `json:"Type"`
	}
	if err := json.Unmarshal(data, &mm); err == nil && mm.Type == "json" {
		var v json.RawMessage
		if err = json.Unmarshal(mm.Data, &v); err != nil {
			return err
		}
		*m = Message(v)
		return nil
	}
	var v json.RawMessage
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}
	*m = Message(v)
	return nil
}

// MarshalJSON implements json.Marshaler on Message.
func (m Message) MarshalJSON() ([]byte, error) {
	v := json.RawMessage(m)
	return json.Marshal(v)
}

// NewJob creates a new Job instance from a job request.
func NewJob(req *JobRequest) *Job {
	return &Job{
		Domain:     req.Domain,
		WorkerType: req.WorkerType,
		TriggerID:  req.TriggerID,
		Message:    req.Message,
		Debounced:  req.Debounced,
		Event:      req.Event,
		Options:    req.Options,
		State:      Queued,
		QueuedAt:   time.Now(),
	}
}

// Get returns the informations about a job.
func Get(domain, jobID string) (*Job, error) {
	var job Job
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Jobs, jobID, &job); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundJob
		}
		return nil, err
	}
	return &job, nil
}

// GetQueuedJobs returns the list of jobs which states is "queued" or "running"
func GetQueuedJobs(domain, workerType string) ([]*Job, error) {
	var results []*Job
	db := couchdb.SimpleDatabasePrefix(domain)
	req := &couchdb.FindRequest{
		UseIndex: "by-worker-and-state",
		Selector: mango.And(
			mango.Equal("worker", workerType),
			mango.Exists("state"), // XXX it is needed by couchdb to use the index
			mango.Or(
				mango.Equal("state", Queued),
				mango.Equal("state", Running),
			),
		),
		Limit: 200,
	}
	err := couchdb.FindDocs(db, consts.Jobs, req, &results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

// NewMessage returns a json encoded data
func NewMessage(data interface{}) (Message, error) {
	b, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	return Message(b), nil
}

// NewEvent return a json encoded realtime.Event
func NewEvent(data *realtime.Event) (Event, error) {
	b, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	return Event(b), nil
}

// Unmarshal can be used to unmarshal the encoded message value in the
// specified interface's type.
func (m Message) Unmarshal(msg interface{}) error {
	if m == nil {
		return ErrMessageNil
	}
	return json.Unmarshal(m, &msg)
}

// Unmarshal can be used to unmarshal the encoded message value in the
// specified interface's type.
func (e Event) Unmarshal(evt interface{}) error {
	if e == nil {
		return ErrMessageNil
	}
	return json.Unmarshal(e, &evt)
}

// Clone clones the worker config
func (w *WorkerConfig) Clone() *WorkerConfig {
	return &WorkerConfig{
		WorkerInit:   w.WorkerInit,
		WorkerFunc:   w.WorkerFunc,
		WorkerCommit: w.WorkerCommit,
		Concurrency:  w.Concurrency,
		MaxExecCount: w.MaxExecCount,
		MaxExecTime:  w.MaxExecTime,
		Timeout:      w.Timeout,
		RetryDelay:   w.RetryDelay,
	}
}

var (
	_ permissions.Validable = (*JobRequest)(nil)
	_ permissions.Validable = (*Job)(nil)
)
<|endoftext|>


<commit_before>package task

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/golang/glog"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/zk"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
)

var (
	ErrBadConfig = errors.New("bad-config")
	ErrStopped   = errors.New("stopped")
	ErrTimeout   = errors.New("timeout")
)

type Runtime struct {
	Task

	zk zk.ZK

	status chan []byte
	stdout chan []byte
	stderr chan []byte
	stdin  chan []byte

	options interface{}
	done    bool
	ready   bool
	lock    sync.Mutex
	error   error

	stdoutBuff *bytes.Buffer
}

func (this *Task) Copy() (*Task, error) {
	var buff bytes.Buffer
	enc := gob.NewEncoder(&buff)
	dec := gob.NewDecoder(&buff)
	err := enc.Encode(this)
	if err != nil {
		return nil, err
	}
	copy := new(Task)
	err = dec.Decode(copy)
	if err != nil {
		return nil, err
	}
	return copy, nil
}

func (this *Task) Validate() error {
	switch {
	case !this.Info.Valid():
		return ErrBadConfig
	case !this.Status.Valid():
		return ErrBadConfig
	case len(this.Success) > 0 && !this.Success.Valid():
		return ErrBadConfig
	case len(this.Error) > 0 && !this.Error.Valid():
		return ErrBadConfig
	case this.Cmd != nil:
		_, err := exec.LookPath(this.Cmd.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Task) Init(zkc zk.ZK, options ...interface{}) (*Runtime, error) {
	if err := this.Validate(); err != nil {
		return nil, err
	}

	task := Runtime{
		Task: *this,
		zk:   zkc,
	}
	if len(options) > 0 {
		task.options = options[0]
	}
	task.status = make(chan []byte)

	if task.Task.Stdout != nil {
		task.stdout = make(chan []byte)
	}

	if task.Task.Stderr != nil {
		task.stderr = make(chan []byte)
	}

	now := time.Now()
	task.Stat.Started = &now

	if task.zk != nil && task.Info != "" {
		err := zk.CreateOrSet(task.zk, task.Info, task.Stat)
		glog.Infoln("Info=", task.Info)
		if err != nil {
			return nil, err
		}
	}
	return &task, nil
}

func (this *Runtime) Stop() {
	this.lock.Lock()
	defer this.lock.Unlock()

	if this.done {
		return
	}

	if this.stdout != nil {
		this.stdout <- nil
	}
	if this.stderr != nil {
		this.stderr <- nil
	}
	this.Log("Stop")
	this.status <- nil

	this.done = true
}

func (this *Runtime) Stdin() io.Reader {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetReader(*this.Task.Stdin, c)
	} else {
		return nil
	}
}

func (this *Runtime) PublishStdin() io.Writer {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stdin, c)
	} else {
		glog.Warningln("Error getting stdin.", "Topic=", *this.Task.Stdin, "Err=", err)
		return nil
	}
}

func (this *Runtime) CaptureStdout() {
	this.stdoutBuff = new(bytes.Buffer)
}

func (this *Runtime) GetCapturedStdout() []byte {
	if this.stdoutBuff != nil {
		return this.stdoutBuff.Bytes()
	}
	return nil
}

func (this *Runtime) Stdout() io.Writer {
	var stdout io.Writer = os.Stdout
	if this.Task.Stdout != nil {
		if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
			stdout = pubsub.GetWriter(*this.Task.Stdout, c)
		} else {
			glog.Fatalln("Error getting stdout.", "Topic=", *this.Task.Stdout, "Err=", err)
			return nil
		}
	}
	if this.stdoutBuff != nil {
		stdout = io.MultiWriter(stdout, this.stdoutBuff)
	}
	return stdout
}

func (this *Runtime) Stderr() io.Writer {
	if this.Task.Stderr == nil {
		return os.Stderr
	}
	if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stderr, c)
	} else {
		glog.Fatalln("Error getting stderr.", "Topic=", *this.Task.Stderr, "Err=", err)
		return nil
	}
}

func (this *Runtime) Log(m ...string) {
	if this.done {
		return
	}
	source := ""
	_, file, line, ok := runtime.Caller(1)
	if ok {
		source = fmt.Sprintf("%s:%d", file, line)
	}

	s := strings.Join(m, " ")
	this.status <- []byte(s)
	glog.Infoln(source, m)
}

func (this *Runtime) Running() bool {
	return !this.done
}

func (this *Runtime) ApplyEnvAndFuncs(env map[string]interface{}, funcs map[string]interface{}) error {
	if this.Task.Cmd == nil {
		return nil
	}

	this.lock.Lock()
	defer this.lock.Unlock()

	applied, err := this.Task.Cmd.ApplySubstitutions(env, funcs)
	if err != nil {
		return err
	}
	this.Task.Cmd = applied
	return nil
}

func (this *Runtime) set_defaults() {
	if len(this.Task.Success) == 0 {
		this.Task.Success = this.Info.Member("success")
	}

	if len(this.Task.Error) == 0 {
		this.Task.Error = this.Info.Member("error")
	}

	if len(this.Status) > 0 {
		if this.Task.Stdout == nil {
			t := this.Status.Sub("stdout")
			this.Task.Stdout = &t
		}
		if this.Task.Stderr == nil {
			t := this.Status.Sub("stderr")
			this.Task.Stderr = &t
		}
	}
}

func (this *Runtime) Start() (chan error, error) {

	this.set_defaults()

	if _, _, err := this.start_streams(); err != nil {
		return nil, err
	}

	if err := this.block_on_triggers(); err == zk.ErrTimeout {
		return nil, ErrTimeout
	}

	// Run the actual task
	if this.Task.Cmd != nil {
		return this.exec()
	}
	return nil, nil
}

func (this *Runtime) block_on_triggers() error {
	if this.Cmd == nil {
		return nil
	}

	if this.Trigger == nil {
		return nil
	}

	// TODO - take into account ordering of cron vs registry.

	if this.Trigger.Registry != nil {
		trigger := zk.NewConditions(*this.Trigger.Registry, this.zk)
		// So now just block until the condition is true
		return trigger.Wait()
	}

	return nil
}

func (this *Runtime) exec() (chan error, error) {
	cmd := exec.Command(this.Cmd.Path, this.Cmd.Args...)
	cmd.Dir = this.Cmd.Dir
	cmd.Env = this.Cmd.Env

	if this.Task.Stdin != nil {
		sub, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options)
		if err != nil {
			return nil, err
		}
		stdin, err := sub.Subscribe(*this.Task.Stdin)
		if err != nil {
			return nil, err
		}
		wr, err := cmd.StdinPipe()
		if err != nil {
			return nil, err
		}
		go func() {
			// We need to do some special processing of input so that we can
			// terminate a session. Otherwise, this will just loop forever
			// because the pubsub topic will not go away -- even if it's a unique topic.
			for {
				m := <-stdin
				fmt.Printf(">> %s", string(m))
				switch {
				case strings.Index(string(m), "#bye") == 0:
					wr.Close()
					return
				default:
					wr.Write(m) // Need newline for shell to interpret
				}
			}
		}()
	}
	cmd.Stdout = this.Stdout()
	cmd.Stderr = this.Stderr()

	process_done := make(chan error)
	go func() {
		cmd.Start()

		// Wait for cmd to complete even if we have no more stdout/stderr
		if err := cmd.Wait(); err != nil {
			this.Error(err.Error())
			process_done <- err
			return
		}

		ps := cmd.ProcessState
		if ps == nil {
			this.Error(ErrCommandUnknown.Error())
			process_done <- ErrCommandUnknown
			return
		}

		glog.Infoln("Process pid=", ps.Pid(), "Exited=", ps.Exited(), "Success=", ps.Success())

		if !ps.Success() {
			this.Error(ErrExecFailed.Error())
			process_done <- ErrExecFailed
			return
		} else {
			this.Success(this.GetCapturedStdout())
			process_done <- nil
			return
		}
	}()

	return process_done, nil
}

func (this *Runtime) start_streams() (stdout, stderr chan<- []byte, err error) {
	this.lock.Lock()
	defer func() {
		this.error = err
		this.lock.Unlock()
	}()

	if this.error != nil {
		return nil, nil, this.error
	}

	if this.ready {
		return this.stdout, this.stderr, nil
	}

	if this.done {
		return nil, nil, ErrStopped
	}
	go func() {
		for {
			m := <-this.status
			if m == nil {
				break
			}
			if c, err := this.Task.Status.Broker().PubSub(this.Id, this.options); err == nil {
				c.Publish(this.Task.Status, m)
			} else {
				glog.Warningln("Cannot publish:", this.Task.Status.String(), "Err=", err)
			}
		}
	}()
	if this.stdout != nil {
		glog.Infoln("Starting stream for stdout:", this.Task.Stdout.String())
		go func() {
			for {
				m := <-this.stdout
				if m == nil {
					break
				}
				if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stdout, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stdout.String(), "Err=", err)
				}

			}
		}()
		this.Log("Sending stdout to", this.Task.Stdout.Path())
	}
	if this.stderr != nil {
		glog.Infoln("Starting stream for stderr:", this.Task.Stderr.String())
		go func() {
			for {
				m := <-this.stderr
				if m == nil {
					break
				}
				if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stderr, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stderr.String(), "Err=", err)
				}
			}
		}()
		this.Log("Sending stderr to", this.Task.Stderr.Path())
	}
	this.ready = true
	return this.stdout, this.stderr, nil
}

func (this *Runtime) Success(output interface{}) error {

	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		return nil
	}

	if this.done {
		return ErrStopped
	}

	switch output.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Success, output.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Success, output.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(output)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Success, value)
		if err != nil {
			return err
		}
	}
	this.Log("Success", "Result written to", this.Task.Success.Path())

	now := time.Now()
	this.Stat.Success = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Success", "Completed")
	return nil
}

func (this *Runtime) Error(error interface{}) error {
	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		return nil
	}

	if this.done {
		return ErrStopped
	}
	switch error.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Error, error.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Error, error.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(error)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Error, value)
		if err != nil {
			return err
		}
	}
	this.Log("Error", "Error written to", this.Task.Error.Path())

	now := time.Now()
	this.Stat.Error = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Error", "Stop")
	return nil
}
<commit_msg>Adding stdin interceptor to support filtering and exiting of session<commit_after>package task

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/golang/glog"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/zk"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
)

var (
	ErrBadConfig = errors.New("bad-config")
	ErrStopped   = errors.New("stopped")
	ErrTimeout   = errors.New("timeout")
)

type Runtime struct {
	Task

	zk zk.ZK

	status chan []byte
	stdout chan []byte
	stderr chan []byte
	stdin  chan []byte

	options interface{}
	done    bool
	ready   bool
	lock    sync.Mutex
	error   error

	stdoutBuff       *bytes.Buffer
	stdinInterceptor func(string) (string, bool)
}

func (this *Task) Copy() (*Task, error) {
	var buff bytes.Buffer
	enc := gob.NewEncoder(&buff)
	dec := gob.NewDecoder(&buff)
	err := enc.Encode(this)
	if err != nil {
		return nil, err
	}
	copy := new(Task)
	err = dec.Decode(copy)
	if err != nil {
		return nil, err
	}
	return copy, nil
}

func (this *Task) Validate() error {
	switch {
	case !this.Info.Valid():
		return ErrBadConfig
	case !this.Status.Valid():
		return ErrBadConfig
	case len(this.Success) > 0 && !this.Success.Valid():
		return ErrBadConfig
	case len(this.Error) > 0 && !this.Error.Valid():
		return ErrBadConfig
	case this.Cmd != nil:
		_, err := exec.LookPath(this.Cmd.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Task) Init(zkc zk.ZK, options ...interface{}) (*Runtime, error) {
	if err := this.Validate(); err != nil {
		return nil, err
	}

	task := Runtime{
		Task: *this,
		zk:   zkc,
	}
	if len(options) > 0 {
		task.options = options[0]
	}
	task.status = make(chan []byte)

	if task.Task.Stdout != nil {
		task.stdout = make(chan []byte)
	}

	if task.Task.Stderr != nil {
		task.stderr = make(chan []byte)
	}

	// Default interceptor
	task.stdinInterceptor = func(in string) (string, bool) {
		return in, strings.Index(in, "#bye") != 0
	}

	now := time.Now()
	task.Stat.Started = &now

	if task.zk != nil && task.Info != "" {
		err := zk.CreateOrSet(task.zk, task.Info, task.Stat)
		glog.Infoln("Info=", task.Info)
		if err != nil {
			return nil, err
		}
	}
	return &task, nil
}

func (this *Runtime) Stop() {
	this.lock.Lock()
	defer this.lock.Unlock()

	if this.done {
		return
	}

	if this.stdout != nil {
		this.stdout <- nil
	}
	if this.stderr != nil {
		this.stderr <- nil
	}
	this.Log("Stop")
	this.status <- nil

	this.done = true
}

func (this *Runtime) StdinInterceptor(f func(string) (string, bool)) {
	this.stdinInterceptor = f
}

func (this *Runtime) Stdin() io.Reader {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetReader(*this.Task.Stdin, c)
	} else {
		return nil
	}
}

func (this *Runtime) PublishStdin() io.Writer {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stdin, c)
	} else {
		glog.Warningln("Error getting stdin.", "Topic=", *this.Task.Stdin, "Err=", err)
		return nil
	}
}

func (this *Runtime) CaptureStdout() {
	this.stdoutBuff = new(bytes.Buffer)
}

func (this *Runtime) GetCapturedStdout() []byte {
	if this.stdoutBuff != nil {
		return this.stdoutBuff.Bytes()
	}
	return nil
}

func (this *Runtime) Stdout() io.Writer {
	var stdout io.Writer = os.Stdout
	if this.Task.Stdout != nil {
		if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
			stdout = pubsub.GetWriter(*this.Task.Stdout, c)
		} else {
			glog.Fatalln("Error getting stdout.", "Topic=", *this.Task.Stdout, "Err=", err)
			return nil
		}
	}
	if this.stdoutBuff != nil {
		stdout = io.MultiWriter(stdout, this.stdoutBuff)
	}
	return stdout
}

func (this *Runtime) Stderr() io.Writer {
	if this.Task.Stderr == nil {
		return os.Stderr
	}
	if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stderr, c)
	} else {
		glog.Fatalln("Error getting stderr.", "Topic=", *this.Task.Stderr, "Err=", err)
		return nil
	}
}

func (this *Runtime) Log(m ...string) {
	if this.done {
		return
	}
	source := ""
	_, file, line, ok := runtime.Caller(1)
	if ok {
		source = fmt.Sprintf("%s:%d", file, line)
	}

	s := strings.Join(m, " ")
	this.status <- []byte(s)
	glog.Infoln(source, m)
}

func (this *Runtime) Running() bool {
	return !this.done
}

func (this *Runtime) ApplyEnvAndFuncs(env map[string]interface{}, funcs map[string]interface{}) error {
	if this.Task.Cmd == nil {
		return nil
	}

	this.lock.Lock()
	defer this.lock.Unlock()

	applied, err := this.Task.Cmd.ApplySubstitutions(env, funcs)
	if err != nil {
		return err
	}
	this.Task.Cmd = applied
	return nil
}

func (this *Runtime) set_defaults() {
	if len(this.Task.Success) == 0 {
		this.Task.Success = this.Info.Member("success")
	}

	if len(this.Task.Error) == 0 {
		this.Task.Error = this.Info.Member("error")
	}

	if len(this.Status) > 0 {
		if this.Task.Stdout == nil {
			t := this.Status.Sub("stdout")
			this.Task.Stdout = &t
		}
		if this.Task.Stderr == nil {
			t := this.Status.Sub("stderr")
			this.Task.Stderr = &t
		}
	}
}

func (this *Runtime) Start() (chan error, error) {

	this.set_defaults()

	if _, _, err := this.start_streams(); err != nil {
		return nil, err
	}

	if err := this.block_on_triggers(); err == zk.ErrTimeout {
		return nil, ErrTimeout
	}

	// Run the actual task
	if this.Task.Cmd != nil {
		return this.exec()
	}
	return nil, nil
}

func (this *Runtime) block_on_triggers() error {
	if this.Cmd == nil {
		return nil
	}

	if this.Trigger == nil {
		return nil
	}

	// TODO - take into account ordering of cron vs registry.

	if this.Trigger.Registry != nil {
		trigger := zk.NewConditions(*this.Trigger.Registry, this.zk)
		// So now just block until the condition is true
		return trigger.Wait()
	}

	return nil
}

func (this *Runtime) exec() (chan error, error) {
	cmd := exec.Command(this.Cmd.Path, this.Cmd.Args...)
	cmd.Dir = this.Cmd.Dir
	cmd.Env = this.Cmd.Env

	if this.Task.Stdin != nil {
		sub, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options)
		if err != nil {
			return nil, err
		}
		stdin, err := sub.Subscribe(*this.Task.Stdin)
		if err != nil {
			return nil, err
		}
		wr, err := cmd.StdinPipe()
		if err != nil {
			return nil, err
		}

		go func() {
			// We need to do some special processing of input so that we can
			// terminate a session. Otherwise, this will just loop forever
			// because the pubsub topic will not go away -- even if it's a unique topic.
			for {
				m := <-stdin
				if l, ok := this.stdinInterceptor(string(m)); ok {
					fmt.Printf(">> %s", l)
					wr.Write([]byte(l))
				} else {
					wr.Close()
					return
				}
			}
		}()
	}
	cmd.Stdout = this.Stdout()
	cmd.Stderr = this.Stderr()

	process_done := make(chan error)
	go func() {
		cmd.Start()

		// Wait for cmd to complete even if we have no more stdout/stderr
		if err := cmd.Wait(); err != nil {
			this.Error(err.Error())
			process_done <- err
			return
		}

		ps := cmd.ProcessState
		if ps == nil {
			this.Error(ErrCommandUnknown.Error())
			process_done <- ErrCommandUnknown
			return
		}

		glog.Infoln("Process pid=", ps.Pid(), "Exited=", ps.Exited(), "Success=", ps.Success())

		if !ps.Success() {
			this.Error(ErrExecFailed.Error())
			process_done <- ErrExecFailed
			return
		} else {
			this.Success(this.GetCapturedStdout())
			process_done <- nil
			return
		}
	}()

	return process_done, nil
}

func (this *Runtime) start_streams() (stdout, stderr chan<- []byte, err error) {
	this.lock.Lock()
	defer func() {
		this.error = err
		this.lock.Unlock()
	}()

	if this.error != nil {
		return nil, nil, this.error
	}

	if this.ready {
		return this.stdout, this.stderr, nil
	}

	if this.done {
		return nil, nil, ErrStopped
	}
	go func() {
		for {
			m := <-this.status
			if m == nil {
				break
			}
			if c, err := this.Task.Status.Broker().PubSub(this.Id, this.options); err == nil {
				c.Publish(this.Task.Status, m)
			} else {
				glog.Warningln("Cannot publish:", this.Task.Status.String(), "Err=", err)
			}
		}
	}()
	if this.stdout != nil {
		glog.Infoln("Starting stream for stdout:", this.Task.Stdout.String())
		go func() {
			for {
				m := <-this.stdout
				if m == nil {
					break
				}
				if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stdout, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stdout.String(), "Err=", err)
				}

			}
		}()
		this.Log("Sending stdout to", this.Task.Stdout.Path())
	}
	if this.stderr != nil {
		glog.Infoln("Starting stream for stderr:", this.Task.Stderr.String())
		go func() {
			for {
				m := <-this.stderr
				if m == nil {
					break
				}
				if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stderr, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stderr.String(), "Err=", err)
				}
			}
		}()
		this.Log("Sending stderr to", this.Task.Stderr.Path())
	}
	this.ready = true
	return this.stdout, this.stderr, nil
}

func (this *Runtime) Success(output interface{}) error {

	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		return nil
	}

	if this.done {
		return ErrStopped
	}

	switch output.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Success, output.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Success, output.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(output)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Success, value)
		if err != nil {
			return err
		}
	}
	this.Log("Success", "Result written to", this.Task.Success.Path())

	now := time.Now()
	this.Stat.Success = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Success", "Completed")
	return nil
}

func (this *Runtime) Error(error interface{}) error {
	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		return nil
	}

	if this.done {
		return ErrStopped
	}
	switch error.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Error, error.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Error, error.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(error)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Error, value)
		if err != nil {
			return err
		}
	}
	this.Log("Error", "Error written to", this.Task.Error.Path())

	now := time.Now()
	this.Stat.Error = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Error", "Stop")
	return nil
}
<|endoftext|>