<commit_before>// Copyright 2016 The clang-server 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 pathutil

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"

	"github.com/pkgutil/osutil"
	"github.com/zchee/clang-server/internal/hashutil"
	xdgbasedir "github.com/zchee/go-xdgbasedir"
)

const dirname = "clang-server"

// CacheDir return the clang-server cache directory path.
func CacheDir() string {
	dir := filepath.Join(xdgbasedir.CacheHome(), dirname)
	if !osutil.IsExist(dir) {
		if err := os.MkdirAll(dir, 0755); err != nil {
			log.Fatal(err)
		}
	}
	return dir
}

func ProjectCacheDir(root string) string {
	// assume the absolude directory path
	id := hashutil.NewHashString(root)
	dir := filepath.Join(CacheDir(), filepath.Base(root)+"."+hashutil.EncodeToString(id))
	return dir
}

func ProjectASTDir(root string) string {
	cacheDir := ProjectCacheDir(root)
	astDir := filepath.Join(cacheDir, "ast")
	if osutil.IsNotExist(astDir) {
		if err := os.MkdirAll(astDir, 0755); err != nil {
			log.Fatal(err)
		}
	}
	return astDir
}

var rootMarks = []string{
	".git",                      // git
	"LICENSE",                   // project license
	".gitignore",                // git
	".travis.yml", "circle.yml", // CI service config files
	"CMakeLists.txt",                                                   // CMake
	"autogen.sh", "configure", "Makefile.am", "Makefile.in", "INSTALL", // GNU Autotools
	".hg", ".svn", ".bzr", "_darcs", ".tup", // typical vcs directories
}

// FindProjectRoot finds the project root directory path from path.
func FindProjectRoot(path string) (string, error) {
	if path == "" {
		return "", errors.New("project root is blank")
	}

	if !filepath.IsAbs(path) {
		abs, err := filepath.Abs(path)
		if err == nil {
			path = abs
		}
	}

	first := path
	for path != "/" {
		for _, c := range rootMarks {
			if p := filepath.Join(path, c); osutil.IsExist(p) {
				return path, nil
			}
		}
		path = filepath.Dir(path)
	}

	return "", fmt.Errorf("couldn't find project root in %s", first)
}
<commit_msg>internal/pathutil: add godoc comments<commit_after>// Copyright 2016 The clang-server 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 pathutil

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"

	"github.com/pkgutil/osutil"
	"github.com/zchee/clang-server/internal/hashutil"
	xdgbasedir "github.com/zchee/go-xdgbasedir"
)

const dirname = "clang-server"

// CacheDir return the clang-server cache directory path.
func CacheDir() string {
	dir := filepath.Join(xdgbasedir.CacheHome(), dirname)
	if !osutil.IsExist(dir) {
		if err := os.MkdirAll(dir, 0755); err != nil {
			log.Fatal(err)
		}
	}
	return dir
}

// ProjectCacheDir return the project path based cache directory path.
func ProjectCacheDir(root string) string {
	// assume the absolude directory path
	id := hashutil.NewHashString(root)
	dir := filepath.Join(CacheDir(), filepath.Base(root)+"."+hashutil.EncodeToString(id))
	return dir
}

// ProjectASTDir return the project path based AST cache directory path.
func ProjectASTDir(root string) string {
	cacheDir := ProjectCacheDir(root)
	astDir := filepath.Join(cacheDir, "ast")
	if osutil.IsNotExist(astDir) {
		if err := os.MkdirAll(astDir, 0755); err != nil {
			log.Fatal(err)
		}
	}
	return astDir
}

var rootMarks = []string{
	".git",                      // git
	"LICENSE",                   // project license
	".gitignore",                // git
	".travis.yml", "circle.yml", // CI service config files
	"CMakeLists.txt",                                                   // CMake
	"autogen.sh", "configure", "Makefile.am", "Makefile.in", "INSTALL", // GNU Autotools
	".hg", ".svn", ".bzr", "_darcs", ".tup", // typical vcs directories
}

// FindProjectRoot finds the project root directory path from path.
func FindProjectRoot(path string) (string, error) {
	if path == "" {
		return "", errors.New("project root is blank")
	}

	if !filepath.IsAbs(path) {
		abs, err := filepath.Abs(path)
		if err == nil {
			path = abs
		}
	}

	first := path
	for path != "/" {
		for _, c := range rootMarks {
			if p := filepath.Join(path, c); osutil.IsExist(p) {
				return path, nil
			}
		}
		path = filepath.Dir(path)
	}

	return "", fmt.Errorf("couldn't find project root in %s", first)
}
<|endoftext|>


<commit_before>package redirect

import (
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Debian/debiman/internal/tag"
	"golang.org/x/text/language"
)

type IndexEntry struct {
	Suite     string // TODO: enum to save space
	Binarypkg string // TODO: sort by popcon, TODO: use a string pool
	Section   string // TODO: use a string pool
	Language  string // TODO: type: would it make sense to use language.Tag?
}

type Index struct {
	Entries  map[string][]IndexEntry
	Suites   map[string]bool
	Langs    map[string]bool
	Sections map[string]bool
}

// TODO(later): the default suite should be the latest stable release
const defaultSuite = "jessie"
const defaultLanguage = "en"

// bestLanguageMatch is like bestLanguageMatch in render.go, but for the redirector index. TODO: can we de-duplicate the code?
func bestLanguageMatch(t []language.Tag, options []IndexEntry) IndexEntry {
	sort.SliceStable(options, func(i, j int) bool {
		// ensure that en comes first, so that language.Matcher treats it as default
		if options[i].Language == "en" && options[j].Language != "en" {
			return true
		}
		return options[i].Language < options[j].Language
	})

	tags := make([]language.Tag, len(options))
	for idx, m := range options {
		tag, err := tag.FromLocale(m.Language)
		if err != nil {
			panic(fmt.Sprintf("Cannot get language.Tag from locale %q: %v", m.Language, err))
		}
		tags[idx] = tag
	}

	matcher := language.NewMatcher(tags)
	tag, _, _ := matcher.Match(t...)
	for idx, t := range tags {
		if t == tag {
			return options[idx]
		}
	}
	return options[0]
}

func (i Index) splitDir(path string) (suite string, binarypkg string) {
	dir := strings.TrimPrefix(filepath.Dir(path), "/")
	parts := strings.Split(dir, "/")
	if len(parts) == 0 {
		return "", ""
	}
	if len(parts) == 1 {
		if i.Suites[parts[0]] {
			return parts[0], ""
		} else {
			return "", parts[0]
		}
	}
	return parts[0], parts[1]
}

func (i Index) splitBase(path string) (name string, section string, lang string) {
	base := filepath.Base(path)
	// the first part can contain dots, so we need to “split from the right”
	parts := strings.Split(base, ".")
	if len(parts) == 1 {
		return base, "", ""
	}

	// The last part can either be a language or a section
	consumed := 0
	if l := parts[len(parts)-1]; i.Langs[l] {
		lang = l
		consumed++
	}
	// The second to last part (if enough parts are present) can
	// be a section (because the language was already specified).
	if len(parts) > 1+consumed {
		if s := parts[len(parts)-1-consumed]; i.Sections[s] {
			section = s
			consumed++
		}
	}

	return strings.Join(parts[:len(parts)-consumed], "."),
		section,
		lang
}

func (i Index) Redirect(r *http.Request) (string, error) {
	path := r.URL.Path
	for strings.HasSuffix(path, ".html") || strings.HasSuffix(path, ".gz") {
		path = strings.TrimSuffix(path, ".html")
		path = strings.TrimSuffix(path, ".gz")
	}

	suite, binarypkg := i.splitDir(path)
	name, section, lang := i.splitBase(path)

	fullyQualified := func() bool {
		return suite != "" && binarypkg != "" && section != "" && lang != ""
	}
	concat := func() string {
		return "/" + suite + "/" + binarypkg + "/" + name + "." + section + "." + lang + ".html"
	}

	log.Printf("path %q -> suite = %q, binarypkg = %q, name = %q, section = %q, lang = %q", path, suite, binarypkg, name, section, lang)

	if fullyQualified() {
		return concat(), nil
	}

	if suite == "" {
		suite = defaultSuite
	}

	if fullyQualified() {
		return concat(), nil
	}

	entries, ok := i.Entries[name]
	if !ok {
		// TODO: this should result in a good 404 page.
		return "", fmt.Errorf("No such man page: name=%q", name)
	}

	if binarypkg == "" {
		binarypkg = entries[0].Binarypkg
	}

	if fullyQualified() {
		return concat(), nil
	}

	// TODO: use pointers
	filtered := make([]IndexEntry, 0, len(entries))
	for _, e := range entries {
		if e.Binarypkg != binarypkg {
			continue
		}
		if e.Suite != suite {
			continue
		}
		if section != "" && e.Section != section {
			continue
		}
		filtered = append(filtered, e)
	}

	if len(filtered) == 0 {
		return "", fmt.Errorf("No such manpage found")
	}

	if section == "" {
		// TODO(later): respect the section preference cookie (+test)
		sort.SliceStable(filtered, func(i, j int) bool {
			return filtered[i].Section < filtered[j].Section
		})
		section = filtered[0].Section
	}

	if fullyQualified() {
		return concat(), nil
	}

	if lang == "" {
		lfiltered := make([]IndexEntry, 0, len(filtered))
		for _, f := range filtered {
			if f.Section != section {
				continue
			}
			lfiltered = append(lfiltered, f)
		}

		t, _, _ := language.ParseAcceptLanguage(r.Header.Get("Accept-Language"))
		// ignore err: t == nil results in the default language
		best := bestLanguageMatch(t, lfiltered)
		lang = best.Language
	}

	return concat(), nil
}
<commit_msg>redirect: don’t settle on binarypkg too early<commit_after>package redirect

import (
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"sort"
	"strings"

	"github.com/Debian/debiman/internal/tag"
	"golang.org/x/text/language"
)

type IndexEntry struct {
	Suite     string // TODO: enum to save space
	Binarypkg string // TODO: sort by popcon, TODO: use a string pool
	Section   string // TODO: use a string pool
	Language  string // TODO: type: would it make sense to use language.Tag?
}

type Index struct {
	Entries  map[string][]IndexEntry
	Suites   map[string]bool
	Langs    map[string]bool
	Sections map[string]bool
}

// TODO(later): the default suite should be the latest stable release
const defaultSuite = "jessie"
const defaultLanguage = "en"

// bestLanguageMatch is like bestLanguageMatch in render.go, but for the redirector index. TODO: can we de-duplicate the code?
func bestLanguageMatch(t []language.Tag, options []IndexEntry) IndexEntry {
	sort.SliceStable(options, func(i, j int) bool {
		// ensure that en comes first, so that language.Matcher treats it as default
		if options[i].Language == "en" && options[j].Language != "en" {
			return true
		}
		return options[i].Language < options[j].Language
	})

	tags := make([]language.Tag, len(options))
	for idx, m := range options {
		tag, err := tag.FromLocale(m.Language)
		if err != nil {
			panic(fmt.Sprintf("Cannot get language.Tag from locale %q: %v", m.Language, err))
		}
		tags[idx] = tag
	}

	matcher := language.NewMatcher(tags)
	tag, _, _ := matcher.Match(t...)
	for idx, t := range tags {
		if t == tag {
			return options[idx]
		}
	}
	return options[0]
}

func (i Index) splitDir(path string) (suite string, binarypkg string) {
	dir := strings.TrimPrefix(filepath.Dir(path), "/")
	parts := strings.Split(dir, "/")
	if len(parts) == 0 {
		return "", ""
	}
	if len(parts) == 1 {
		if i.Suites[parts[0]] {
			return parts[0], ""
		} else {
			return "", parts[0]
		}
	}
	return parts[0], parts[1]
}

func (i Index) splitBase(path string) (name string, section string, lang string) {
	base := filepath.Base(path)
	// the first part can contain dots, so we need to “split from the right”
	parts := strings.Split(base, ".")
	if len(parts) == 1 {
		return base, "", ""
	}

	// The last part can either be a language or a section
	consumed := 0
	if l := parts[len(parts)-1]; i.Langs[l] {
		lang = l
		consumed++
	}
	// The second to last part (if enough parts are present) can
	// be a section (because the language was already specified).
	if len(parts) > 1+consumed {
		if s := parts[len(parts)-1-consumed]; i.Sections[s] {
			section = s
			consumed++
		}
	}

	return strings.Join(parts[:len(parts)-consumed], "."),
		section,
		lang
}

func (i Index) Redirect(r *http.Request) (string, error) {
	path := r.URL.Path
	for strings.HasSuffix(path, ".html") || strings.HasSuffix(path, ".gz") {
		path = strings.TrimSuffix(path, ".html")
		path = strings.TrimSuffix(path, ".gz")
	}

	suite, binarypkg := i.splitDir(path)
	name, section, lang := i.splitBase(path)

	fullyQualified := func() bool {
		return suite != "" && binarypkg != "" && section != "" && lang != ""
	}
	concat := func() string {
		return "/" + suite + "/" + binarypkg + "/" + name + "." + section + "." + lang + ".html"
	}

	log.Printf("path %q -> suite = %q, binarypkg = %q, name = %q, section = %q, lang = %q", path, suite, binarypkg, name, section, lang)

	if fullyQualified() {
		return concat(), nil
	}

	if suite == "" {
		suite = defaultSuite
	}

	if fullyQualified() {
		return concat(), nil
	}

	entries, ok := i.Entries[name]
	if !ok {
		// TODO: this should result in a good 404 page.
		return "", fmt.Errorf("No such man page: name=%q", name)
	}

	if fullyQualified() {
		return concat(), nil
	}

	// TODO: use pointers
	filtered := make([]IndexEntry, 0, len(entries))
	for _, e := range entries {
		if binarypkg != "" && e.Binarypkg != binarypkg {
			continue
		}
		if e.Suite != suite {
			continue
		}
		if section != "" && e.Section != section {
			continue
		}
		filtered = append(filtered, e)
	}

	if len(filtered) == 0 {
		return "", fmt.Errorf("No such manpage found")
	}

	if section == "" {
		// TODO(later): respect the section preference cookie (+test)
		sort.SliceStable(filtered, func(i, j int) bool {
			return filtered[i].Section < filtered[j].Section
		})
		section = filtered[0].Section
	}

	if fullyQualified() {
		return concat(), nil
	}

	if lang == "" {
		lfiltered := make([]IndexEntry, 0, len(filtered))
		for _, f := range filtered {
			if f.Section != section {
				continue
			}
			lfiltered = append(lfiltered, f)
		}

		t, _, _ := language.ParseAcceptLanguage(r.Header.Get("Accept-Language"))
		// ignore err: t == nil results in the default language
		best := bestLanguageMatch(t, lfiltered)
		lang = best.Language
		if binarypkg == "" {
			binarypkg = best.Binarypkg
		}
	}

	if binarypkg == "" {
		binarypkg = filtered[0].Binarypkg
	}

	return concat(), nil
}
<|endoftext|>


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

package gossip

import (
	"bytes"
	"encoding/gob"
	"net"
	"time"

	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
)

const (
	// maxWaitForNewGossip is maximum wait for new gossip before a
	// peer is considered a poor source of good gossip and is GC'd.
	maxWaitForNewGossip = 1 * time.Minute
)

// init pre-registers net.UnixAddr and net.TCPAddr concrete types with
// gob. If other implementations of net.Addr are passed, they must be
// added here as well.
func init() {
	gob.Register(&net.TCPAddr{})
	gob.Register(&net.UnixAddr{})
	gob.Register(&util.RawAddr{})
}

// client is a client-side RPC connection to a gossip peer node.
type client struct {
	peerID      proto.NodeID  // Peer node ID; 0 until first gossip response
	addr        net.Addr      // Peer node network address
	rpcClient   *rpc.Client   // RPC client
	forwardAddr net.Addr      // Set if disconnected with an alternate addr
	lastFresh   int64         // Last wall time client received fresh info
	err         error         // Set if client experienced an error
	closer      chan struct{} // Client shutdown channel
}

// newClient creates and returns a client struct.
func newClient(addr net.Addr) *client {
	return &client{
		addr:   addr,
		closer: make(chan struct{}),
	}
}

// start dials the remote addr and commences gossip once connected.
// Upon exit, signals client is done by pushing it onto the done
// channel. If the client experienced an error, its err field will
// be set. This method starts client processing in a goroutine and
// returns immediately.
func (c *client) start(g *Gossip, done chan *client, context *rpc.Context, stopper *util.Stopper) {
	stopper.RunWorker(func() {
		c.rpcClient = rpc.NewClient(c.addr, nil, context)
		select {
		case <-c.rpcClient.Ready:
			// Success!
		case <-c.rpcClient.Closed:
			c.err = util.Errorf("gossip client failed to connect")
			done <- c
			return
		}

		// Start gossipping and wait for disconnect or error.
		c.lastFresh = time.Now().UnixNano()
		c.err = c.gossip(g, stopper)
		if c.err != nil {
			c.rpcClient.Close()
		}
		done <- c
	})
}

// close stops the client gossip loop and returns immediately.
func (c *client) close() {
	close(c.closer)
}

// gossip loops, sending deltas of the infostore and receiving deltas
// in turn. If an alternate is proposed on response, the client addr
// is modified and method returns for forwarding by caller.
func (c *client) gossip(g *Gossip, stopper *util.Stopper) error {
	localMaxSeq := int64(0)
	remoteMaxSeq := int64(-1)
	for {
		// Compute the delta of local node's infostore to send with request.
		g.mu.Lock()
		delta := g.is.delta(c.peerID, localMaxSeq)
		g.mu.Unlock()
		var deltaBytes []byte
		if delta != nil {
			localMaxSeq = delta.MaxSeq
			var buf bytes.Buffer
			if err := gob.NewEncoder(&buf).Encode(delta); err != nil {
				return util.Errorf("infostore could not be encoded: %s", err)
			}
			deltaBytes = buf.Bytes()
		}

		// Send gossip with timeout.
		args := &proto.GossipRequest{
			NodeID: g.is.NodeID,
			Addr:   *proto.FromNetAddr(g.is.NodeAddr),
			LAddr:  *proto.FromNetAddr(c.rpcClient.LocalAddr()),
			MaxSeq: remoteMaxSeq,
			Delta:  deltaBytes,
		}
		reply := &proto.GossipResponse{}
		gossipCall := c.rpcClient.Go("Gossip.Gossip", args, reply, nil)
		select {
		case <-gossipCall.Done:
			if gossipCall.Error != nil {
				return gossipCall.Error
			}
		case <-c.rpcClient.Closed:
			return util.Error("client closed")
		case <-c.closer:
			return nil
		case <-stopper.ShouldStop():
			return nil
		case <-time.After(g.interval * 10):
			return util.Errorf("timeout after: %s", g.interval*10)
		}

		// Handle remote forwarding.
		if reply.Alternate != nil {
			var err error
			if c.forwardAddr, err = reply.Alternate.NetAddr(); err != nil {
				return util.Errorf("unable to resolve alternate address: %s: %s", reply.Alternate, err)
			}
			return util.Errorf("received forward from %s to %s", c.addr, reply.Alternate)
		}

		// Combine remote node's infostore delta with ours.
		now := time.Now().UnixNano()
		if reply.Delta != nil {
			delta := &infoStore{}
			if err := gob.NewDecoder(bytes.NewBuffer(reply.Delta)).Decode(delta); err != nil {
				return util.Errorf("infostore could not be decoded: %s", err)
			}
			log.V(1).Infof("received gossip reply delta from %s: %s", c.addr, delta)
			g.mu.Lock()
			c.peerID = delta.NodeID
			g.outgoing.addNode(c.peerID)
			freshCount := g.is.combine(delta)
			if freshCount > 0 {
				c.lastFresh = now
			}
			remoteMaxSeq = delta.MaxSeq

			// If we have the sentinel gossip, we're considered connected.
			g.checkHasConnected()
			g.mu.Unlock()
		}

		// Check whether this outgoing client is duplicating work already
		// being done by an incoming client. To avoid mutual shutdown, we
		// only shutdown our client if our node ID is less than the peer's.
		nodeID := g.GetNodeID()
		if g.hasIncoming(c.peerID) && nodeID < c.peerID {
			return util.Errorf("stopping outgoing client %d @ %s; already have incoming", c.peerID, c.addr)
		}
		// Check whether peer node is too boring--disconnect if yes.
		if nodeID != c.peerID && (now-c.lastFresh) > int64(maxWaitForNewGossip) {
			return util.Errorf("peer is too boring")
		}
	}
}
<commit_msg>fix #727<commit_after>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package gossip

import (
	"bytes"
	"encoding/gob"
	"net"
	"time"

	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
)

const (
	// maxWaitForNewGossip is maximum wait for new gossip before a
	// peer is considered a poor source of good gossip and is GC'd.
	maxWaitForNewGossip = 1 * time.Minute
)

// init pre-registers net.UnixAddr and net.TCPAddr concrete types with
// gob. If other implementations of net.Addr are passed, they must be
// added here as well.
func init() {
	gob.Register(&net.TCPAddr{})
	gob.Register(&net.UnixAddr{})
	gob.Register(&util.RawAddr{})
}

// client is a client-side RPC connection to a gossip peer node.
type client struct {
	peerID      proto.NodeID  // Peer node ID; 0 until first gossip response
	addr        net.Addr      // Peer node network address
	rpcClient   *rpc.Client   // RPC client
	forwardAddr net.Addr      // Set if disconnected with an alternate addr
	lastFresh   int64         // Last wall time client received fresh info
	err         error         // Set if client experienced an error
	closer      chan struct{} // Client shutdown channel
}

// newClient creates and returns a client struct.
func newClient(addr net.Addr) *client {
	return &client{
		addr:   addr,
		closer: make(chan struct{}),
	}
}

// start dials the remote addr and commences gossip once connected.
// Upon exit, signals client is done by pushing it onto the done
// channel. If the client experienced an error, its err field will
// be set. This method starts client processing in a goroutine and
// returns immediately.
func (c *client) start(g *Gossip, done chan *client, context *rpc.Context, stopper *util.Stopper) {
	stopper.RunWorker(func() {
		c.rpcClient = rpc.NewClient(c.addr, nil, context)
		select {
		case <-c.rpcClient.Ready:
			// Success!
		case <-c.rpcClient.Closed:
			c.err = util.Errorf("gossip client failed to connect")
			done <- c
			return
		}

		// Start gossipping and wait for disconnect or error.
		c.lastFresh = time.Now().UnixNano()
		c.err = c.gossip(g, stopper)
		if c.err != nil {
			c.rpcClient.Close()
		}
		done <- c
	})
}

// close stops the client gossip loop and returns immediately.
func (c *client) close() {
	close(c.closer)
}

// gossip loops, sending deltas of the infostore and receiving deltas
// in turn. If an alternate is proposed on response, the client addr
// is modified and method returns for forwarding by caller.
func (c *client) gossip(g *Gossip, stopper *util.Stopper) error {
	localMaxSeq := int64(0)
	remoteMaxSeq := int64(-1)
	for {
		// Compute the delta of local node's infostore to send with request.
		g.mu.Lock()
		delta := g.is.delta(c.peerID, localMaxSeq)
		nodeID := g.is.NodeID // needs to be accessed with the lock held
		g.mu.Unlock()
		var deltaBytes []byte
		if delta != nil {
			localMaxSeq = delta.MaxSeq
			var buf bytes.Buffer
			if err := gob.NewEncoder(&buf).Encode(delta); err != nil {
				return util.Errorf("infostore could not be encoded: %s", err)
			}
			deltaBytes = buf.Bytes()
		}

		// Send gossip with timeout.
		args := &proto.GossipRequest{
			NodeID: nodeID,
			Addr:   *proto.FromNetAddr(g.is.NodeAddr),
			LAddr:  *proto.FromNetAddr(c.rpcClient.LocalAddr()),
			MaxSeq: remoteMaxSeq,
			Delta:  deltaBytes,
		}
		reply := &proto.GossipResponse{}
		gossipCall := c.rpcClient.Go("Gossip.Gossip", args, reply, nil)
		select {
		case <-gossipCall.Done:
			if gossipCall.Error != nil {
				return gossipCall.Error
			}
		case <-c.rpcClient.Closed:
			return util.Error("client closed")
		case <-c.closer:
			return nil
		case <-stopper.ShouldStop():
			return nil
		case <-time.After(g.interval * 10):
			return util.Errorf("timeout after: %s", g.interval*10)
		}

		// Handle remote forwarding.
		if reply.Alternate != nil {
			var err error
			if c.forwardAddr, err = reply.Alternate.NetAddr(); err != nil {
				return util.Errorf("unable to resolve alternate address: %s: %s", reply.Alternate, err)
			}
			return util.Errorf("received forward from %s to %s", c.addr, reply.Alternate)
		}

		// Combine remote node's infostore delta with ours.
		now := time.Now().UnixNano()
		if reply.Delta != nil {
			delta := &infoStore{}
			if err := gob.NewDecoder(bytes.NewBuffer(reply.Delta)).Decode(delta); err != nil {
				return util.Errorf("infostore could not be decoded: %s", err)
			}
			log.V(1).Infof("received gossip reply delta from %s: %s", c.addr, delta)
			g.mu.Lock()
			c.peerID = delta.NodeID
			g.outgoing.addNode(c.peerID)
			freshCount := g.is.combine(delta)
			if freshCount > 0 {
				c.lastFresh = now
			}
			remoteMaxSeq = delta.MaxSeq

			// If we have the sentinel gossip, we're considered connected.
			g.checkHasConnected()
			g.mu.Unlock()
		}

		// Check whether this outgoing client is duplicating work already
		// being done by an incoming client. To avoid mutual shutdown, we
		// only shutdown our client if our node ID is less than the peer's.
		if g.hasIncoming(c.peerID) && nodeID < c.peerID {
			return util.Errorf("stopping outgoing client %d @ %s; already have incoming", c.peerID, c.addr)
		}
		// Check whether peer node is too boring--disconnect if yes.
		if nodeID != c.peerID && (now-c.lastFresh) > int64(maxWaitForNewGossip) {
			return util.Errorf("peer is too boring")
		}
	}
}
<|endoftext|>


<commit_before>package graph

import (
	"encoding/json"
	"fmt"

	"github.com/cayleygraph/cayley/clog"
	"github.com/cayleygraph/cayley/quad"
	"golang.org/x/net/context"
)

// IterateChain is a chain-enabled helper to setup iterator execution.
type IterateChain struct {
	ctx context.Context
	it  Iterator
	qs  QuadStore

	paths    bool
	optimize bool

	limit int
	n     int
}

// Iterate is a set of helpers for iteration. Context may be used to cancel execution.
// Iterator will be optimized and closed after execution.
//
// By default, iteration has no limit and includes sub-paths.
func Iterate(ctx context.Context, it Iterator) *IterateChain {
	if ctx == nil {
		ctx = context.Background()
	}
	return &IterateChain{
		ctx: ctx, it: it,
		limit: -1, paths: true,
		optimize: true,
	}
}
func (c *IterateChain) next() bool {
	select {
	case <-c.ctx.Done():
		return false
	default:
	}
	ok := (c.limit < 0 || c.n < c.limit) && c.it.Next()
	if ok {
		c.n++
	}
	return ok
}
func (c *IterateChain) nextPath() bool {
	select {
	case <-c.ctx.Done():
		return false
	default:
	}
	ok := c.paths && (c.limit < 0 || c.n < c.limit) && c.it.NextPath()
	if ok {
		c.n++
	}
	return ok
}
func (c *IterateChain) start() {
	if c.optimize {
		c.it, _ = c.it.Optimize()
		if c.qs != nil {
			c.it, _ = c.qs.OptimizeIterator(c.it)
		}
	}
	if !clog.V(2) {
		return
	}
	if b, err := json.MarshalIndent(c.it.Describe(), "", "  "); err != nil {
		clog.Infof("failed to format description: %v", err)
	} else {
		clog.Infof("%s", b)
	}
}
func (c *IterateChain) end() {
	c.it.Close()
	if !clog.V(2) {
		return
	}
	if b, err := json.MarshalIndent(DumpStats(c.it), "", "  "); err != nil {
		clog.Infof("failed to format stats: %v", err)
	} else {
		clog.Infof("%s", b)
	}
}

// Limit limits a total number of results returned.
func (c *IterateChain) Limit(n int) *IterateChain {
	c.limit = n
	return c
}

// Paths switches iteration over sub-paths (with it.NextPath).
// Defaults to true.
func (c *IterateChain) Paths(enable bool) *IterateChain {
	c.paths = enable
	return c
}

// On sets a default quad store for iteration. If qs was set, it may be omitted in other functions.
func (c *IterateChain) On(qs QuadStore) *IterateChain {
	c.qs = qs
	return c
}

// UnOptimized disables iterator optimization.
func (c *IterateChain) UnOptimized() *IterateChain {
	c.optimize = false
	return c
}

// Each will run a provided callback for each result of the iterator.
func (c *IterateChain) Each(fnc func(Value)) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()

	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		default:
		}
		fnc(c.it.Result())
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			default:
			}
			fnc(c.it.Result())
		}
	}
	return c.it.Err()
}

// All will return all results of an iterator.
func (c *IterateChain) Count() (int64, error) {
	c.start()
	defer c.end()
	if err := c.it.Err(); err != nil {
		return 0, err
	}
	if size, exact := c.it.Size(); exact {
		return size, nil
	}
	done := c.ctx.Done()
	var cnt int64
iteration:
	for c.next() {
		select {
		case <-done:
			break iteration
		default:
		}
		cnt++
		for c.nextPath() {
			select {
			case <-done:
				break iteration
			default:
			}
			cnt++
		}
	}
	return cnt, c.it.Err()
}

// All will return all results of an iterator.
func (c *IterateChain) All() ([]Value, error) {
	c.start()
	defer c.end()
	done := c.ctx.Done()
	var out []Value
iteration:
	for c.next() {
		select {
		case <-done:
			break iteration
		default:
		}
		out = append(out, c.it.Result())
		for c.nextPath() {
			select {
			case <-done:
				break iteration
			default:
			}
			out = append(out, c.it.Result())
		}
	}
	return out, c.it.Err()
}

// Send will send each result of the iterator to the provided channel.
//
// Channel will NOT be closed when function returns.
func (c *IterateChain) Send(out chan<- Value) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()
	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		case out <- c.it.Result():
		}
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			case out <- c.it.Result():
			}
		}
	}
	return c.it.Err()
}

// TagEach will run a provided tag map callback for each result of the iterator.
func (c *IterateChain) TagEach(fnc func(map[string]Value)) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()

	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		default:
		}
		tags := make(map[string]Value)
		c.it.TagResults(tags)
		fnc(tags)
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			default:
			}
			tags := make(map[string]Value)
			c.it.TagResults(tags)
			fnc(tags)
		}
	}
	return c.it.Err()
}

var errNoQuadStore = fmt.Errorf("no quad store in Iterate")

// EachValue is an analog of Each, but it will additionally call NameOf
// for each graph.Value before passing it to a callback.
func (c *IterateChain) EachValue(qs QuadStore, fnc func(quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	// TODO(dennwc): batch NameOf?
	return c.Each(func(v Value) {
		if nv := c.qs.NameOf(v); nv != nil {
			fnc(nv)
		}
	})
}

// EachValuePair is an analog of Each, but it will additionally call NameOf
// for each graph.Value before passing it to a callback. Original value will be passed as well.
func (c *IterateChain) EachValuePair(qs QuadStore, fnc func(Value, quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	// TODO(dennwc): batch NameOf?
	return c.Each(func(v Value) {
		if nv := c.qs.NameOf(v); nv != nil {
			fnc(v, nv)
		}
	})
}

// AllValues is an analog of All, but it will additionally call NameOf
// for each graph.Value before returning the results slice.
func (c *IterateChain) AllValues(qs QuadStore) ([]quad.Value, error) {
	var out []quad.Value
	err := c.EachValue(qs, func(v quad.Value) {
		out = append(out, v)
	})
	return out, err
}

// SendValues is an analog of Send, but it will additionally call NameOf
// for each graph.Value before sending it to a channel.
func (c *IterateChain) SendValues(qs QuadStore, out chan<- quad.Value) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	c.start()
	defer c.end()
	done := c.ctx.Done()
	send := func(v Value) error {
		nv := c.qs.NameOf(c.it.Result())
		if nv == nil {
			return nil
		}
		select {
		case <-done:
			return c.ctx.Err()
		case out <- c.qs.NameOf(c.it.Result()):
		}
		return nil
	}
	for c.next() {
		if err := send(c.it.Result()); err != nil {
			return err
		}
		for c.nextPath() {
			if err := send(c.it.Result()); err != nil {
				return err
			}
		}
	}
	return c.it.Err()
}

// TagValues is an analog of TagEach, but it will additionally call NameOf
// for each graph.Value before passing the map to a callback.
func (c *IterateChain) TagValues(qs QuadStore, fnc func(map[string]quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	return c.TagEach(func(m map[string]Value) {
		vm := make(map[string]quad.Value, len(m))
		for k, v := range m {
			vm[k] = c.qs.NameOf(v) // TODO(dennwc): batch NameOf?
		}
		fnc(vm)
	})
}
<commit_msg>iterate: add First helper<commit_after>package graph

import (
	"encoding/json"
	"fmt"

	"github.com/cayleygraph/cayley/clog"
	"github.com/cayleygraph/cayley/quad"
	"golang.org/x/net/context"
)

// IterateChain is a chain-enabled helper to setup iterator execution.
type IterateChain struct {
	ctx context.Context
	it  Iterator
	qs  QuadStore

	paths    bool
	optimize bool

	limit int
	n     int
}

// Iterate is a set of helpers for iteration. Context may be used to cancel execution.
// Iterator will be optimized and closed after execution.
//
// By default, iteration has no limit and includes sub-paths.
func Iterate(ctx context.Context, it Iterator) *IterateChain {
	if ctx == nil {
		ctx = context.Background()
	}
	return &IterateChain{
		ctx: ctx, it: it,
		limit: -1, paths: true,
		optimize: true,
	}
}
func (c *IterateChain) next() bool {
	select {
	case <-c.ctx.Done():
		return false
	default:
	}
	ok := (c.limit < 0 || c.n < c.limit) && c.it.Next()
	if ok {
		c.n++
	}
	return ok
}
func (c *IterateChain) nextPath() bool {
	select {
	case <-c.ctx.Done():
		return false
	default:
	}
	ok := c.paths && (c.limit < 0 || c.n < c.limit) && c.it.NextPath()
	if ok {
		c.n++
	}
	return ok
}
func (c *IterateChain) start() {
	if c.optimize {
		c.it, _ = c.it.Optimize()
		if c.qs != nil {
			c.it, _ = c.qs.OptimizeIterator(c.it)
		}
	}
	if !clog.V(2) {
		return
	}
	if b, err := json.MarshalIndent(c.it.Describe(), "", "  "); err != nil {
		clog.Infof("failed to format description: %v", err)
	} else {
		clog.Infof("%s", b)
	}
}
func (c *IterateChain) end() {
	c.it.Close()
	if !clog.V(2) {
		return
	}
	if b, err := json.MarshalIndent(DumpStats(c.it), "", "  "); err != nil {
		clog.Infof("failed to format stats: %v", err)
	} else {
		clog.Infof("%s", b)
	}
}

// Limit limits a total number of results returned.
func (c *IterateChain) Limit(n int) *IterateChain {
	c.limit = n
	return c
}

// Paths switches iteration over sub-paths (with it.NextPath).
// Defaults to true.
func (c *IterateChain) Paths(enable bool) *IterateChain {
	c.paths = enable
	return c
}

// On sets a default quad store for iteration. If qs was set, it may be omitted in other functions.
func (c *IterateChain) On(qs QuadStore) *IterateChain {
	c.qs = qs
	return c
}

// UnOptimized disables iterator optimization.
func (c *IterateChain) UnOptimized() *IterateChain {
	c.optimize = false
	return c
}

// Each will run a provided callback for each result of the iterator.
func (c *IterateChain) Each(fnc func(Value)) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()

	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		default:
		}
		fnc(c.it.Result())
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			default:
			}
			fnc(c.it.Result())
		}
	}
	return c.it.Err()
}

// All will return all results of an iterator.
func (c *IterateChain) Count() (int64, error) {
	c.start()
	defer c.end()
	if err := c.it.Err(); err != nil {
		return 0, err
	}
	if size, exact := c.it.Size(); exact {
		return size, nil
	}
	done := c.ctx.Done()
	var cnt int64
iteration:
	for c.next() {
		select {
		case <-done:
			break iteration
		default:
		}
		cnt++
		for c.nextPath() {
			select {
			case <-done:
				break iteration
			default:
			}
			cnt++
		}
	}
	return cnt, c.it.Err()
}

// All will return all results of an iterator.
func (c *IterateChain) All() ([]Value, error) {
	c.start()
	defer c.end()
	done := c.ctx.Done()
	var out []Value
iteration:
	for c.next() {
		select {
		case <-done:
			break iteration
		default:
		}
		out = append(out, c.it.Result())
		for c.nextPath() {
			select {
			case <-done:
				break iteration
			default:
			}
			out = append(out, c.it.Result())
		}
	}
	return out, c.it.Err()
}

// First will return a first result of an iterator. It returns nil if iterator is empty.
func (c *IterateChain) First() (Value, error) {
	c.start()
	defer c.end()
	if !c.next() {
		return nil, c.it.Err()
	}
	return c.it.Result(), nil
}

// Send will send each result of the iterator to the provided channel.
//
// Channel will NOT be closed when function returns.
func (c *IterateChain) Send(out chan<- Value) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()
	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		case out <- c.it.Result():
		}
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			case out <- c.it.Result():
			}
		}
	}
	return c.it.Err()
}

// TagEach will run a provided tag map callback for each result of the iterator.
func (c *IterateChain) TagEach(fnc func(map[string]Value)) error {
	c.start()
	defer c.end()
	done := c.ctx.Done()

	for c.next() {
		select {
		case <-done:
			return c.ctx.Err()
		default:
		}
		tags := make(map[string]Value)
		c.it.TagResults(tags)
		fnc(tags)
		for c.nextPath() {
			select {
			case <-done:
				return c.ctx.Err()
			default:
			}
			tags := make(map[string]Value)
			c.it.TagResults(tags)
			fnc(tags)
		}
	}
	return c.it.Err()
}

var errNoQuadStore = fmt.Errorf("no quad store in Iterate")

// EachValue is an analog of Each, but it will additionally call NameOf
// for each graph.Value before passing it to a callback.
func (c *IterateChain) EachValue(qs QuadStore, fnc func(quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	// TODO(dennwc): batch NameOf?
	return c.Each(func(v Value) {
		if nv := c.qs.NameOf(v); nv != nil {
			fnc(nv)
		}
	})
}

// EachValuePair is an analog of Each, but it will additionally call NameOf
// for each graph.Value before passing it to a callback. Original value will be passed as well.
func (c *IterateChain) EachValuePair(qs QuadStore, fnc func(Value, quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	// TODO(dennwc): batch NameOf?
	return c.Each(func(v Value) {
		if nv := c.qs.NameOf(v); nv != nil {
			fnc(v, nv)
		}
	})
}

// AllValues is an analog of All, but it will additionally call NameOf
// for each graph.Value before returning the results slice.
func (c *IterateChain) AllValues(qs QuadStore) ([]quad.Value, error) {
	var out []quad.Value
	err := c.EachValue(qs, func(v quad.Value) {
		out = append(out, v)
	})
	return out, err
}

// FirstValue is an analog of First, but it does lookup of a value in QuadStore.
func (c *IterateChain) FirstValue(qs QuadStore) (quad.Value, error) {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return nil, errNoQuadStore
	}
	v, err := c.First()
	if err != nil || v == nil {
		return nil, err
	}
	// TODO: return an error from NameOf once we have it exposed
	return c.qs.NameOf(v), nil
}

// SendValues is an analog of Send, but it will additionally call NameOf
// for each graph.Value before sending it to a channel.
func (c *IterateChain) SendValues(qs QuadStore, out chan<- quad.Value) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	c.start()
	defer c.end()
	done := c.ctx.Done()
	send := func(v Value) error {
		nv := c.qs.NameOf(c.it.Result())
		if nv == nil {
			return nil
		}
		select {
		case <-done:
			return c.ctx.Err()
		case out <- c.qs.NameOf(c.it.Result()):
		}
		return nil
	}
	for c.next() {
		if err := send(c.it.Result()); err != nil {
			return err
		}
		for c.nextPath() {
			if err := send(c.it.Result()); err != nil {
				return err
			}
		}
	}
	return c.it.Err()
}

// TagValues is an analog of TagEach, but it will additionally call NameOf
// for each graph.Value before passing the map to a callback.
func (c *IterateChain) TagValues(qs QuadStore, fnc func(map[string]quad.Value)) error {
	if qs != nil {
		c.qs = qs
	}
	if c.qs == nil {
		return errNoQuadStore
	}
	return c.TagEach(func(m map[string]Value) {
		vm := make(map[string]quad.Value, len(m))
		for k, v := range m {
			vm[k] = c.qs.NameOf(v) // TODO(dennwc): batch NameOf?
		}
		fnc(vm)
	})
}
<|endoftext|>


<commit_before>package driver

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log"
        "os"
        "os/exec"
        "regexp"
        "runtime"
        "strings"
        "syscall"
        "time"

        "github.com/hashicorp/nomad/client/config"
        "github.com/hashicorp/nomad/nomad/structs"
)

var (
        reRktVersion  = regexp.MustCompile("rkt version ([\\d\\.]+).+")
        reAppcVersion = regexp.MustCompile("appc version ([\\d\\.]+).+")
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
        DriverContext
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
        proc   *os.Process
        name   string
        waitCh chan error
        doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
        Pid  int
        Name string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
        return &RktDriver{*ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
        // Only enable if we are root when running on non-windows systems.
        if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
                d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
                return false, nil
        }

        outBytes, err := exec.Command("rkt", "version").Output()
        if err != nil {
                return false, nil
        }
        out := strings.TrimSpace(string(outBytes))

        rktMatches := reRktVersion.FindStringSubmatch(out)
        appcMatches := reRktVersion.FindStringSubmatch(out)
        if len(rktMatches) != 2 || len(appcMatches) != 2 {
                return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
        }

        node.Attributes["driver.rkt"] = "1"
        node.Attributes["driver.rkt.version"] = rktMatches[0]
        node.Attributes["driver.appc.version"] = appcMatches[1]

        return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
        trust_prefix, ok := task.Config["trust_prefix"]
        if !ok || trust_prefix == "" {
                return nil, fmt.Errorf("Missing trust prefix for rkt")
        }
        trust_args := []string{
                "sudo",
                "rkt",
                "trust",
                "--prefix=" + trust_prefix,
        }
        // Add the given trust prefix
        var outBuf, errBuf bytes.Buffer
        cmd := exec.Command(trust_args[0], trust_args[1:]...)
        cmd.Stdout = &outBuf
        cmd.Stderr = &errBuf
        d.logger.Printf("[DEBUG] Starting rkt command: %q", strings.Join(trust_args, " "))
        if err := cmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, outBuf.String(), errBuf.String())
        }
        d.logger.Printf("[DEBUG] Added trust prefix: %q", trust_prefix)

        name, ok := task.Config["name"]
        if !ok || name == "" {
                return nil, fmt.Errorf("Missing ACI name for rkt")
        }
        run_args := []string{
                "sudo",
                "rkt",
                "run", "--interactive", name,
        }
        // Run the ACI
        var aoutBuf, aerrBuf bytes.Buffer
        acmd := exec.Command(run_args[0], run_args[1:]...)
        acmd.Stdout = &aoutBuf
        acmd.Stderr = &aerrBuf
        d.logger.Printf("[DEBUG] Starting rkt command: %q", strings.Join(run_args, " "))
        if err := acmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, aoutBuf.String(), aerrBuf.String())
        }
        d.logger.Printf("[DEBUG] Started ACI: %q", name)
        h := &rktHandle{
                proc:   cmd.Process,
                name:   name,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }
        go h.run()
        return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
        // Parse the handle
        pidBytes := []byte(strings.TrimPrefix(handleID, "RKT:"))
        qpid := &rktPID{}
        if err := json.Unmarshal(pidBytes, qpid); err != nil {
                return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
        }

        // Find the process
        proc, err := os.FindProcess(qpid.Pid)
        if proc == nil || err != nil {
                return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
        }

        // Return a driver handle
        h := &rktHandle{
                proc:   proc,
                name:   qpid.Name,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }

        go h.run()
        return h, nil
}

func (h *rktHandle) ID() string {
        // Return a handle to the PID
        pid := &rktPID{
                Pid:  h.proc.Pid,
                Name: h.name,
        }
        data, err := json.Marshal(pid)
        if err != nil {
                log.Printf("[ERR] failed to marshal rkt PID to JSON: %s", err)
        }
        return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan error {
        return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
        // Update is not possible
        return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
//
// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *rktHandle) Kill() error {
        h.proc.Signal(os.Interrupt)
        select {
        case <-h.doneCh:
                return nil
        case <-time.After(5 * time.Second):
                return h.proc.Kill()
        }
}

func (h *rktHandle) run() {
        ps, err := h.proc.Wait()
        close(h.doneCh)
        if err != nil {
                h.waitCh <- err
        } else if !ps.Success() {
                h.waitCh <- fmt.Errorf("task exited with error")
        }
        close(h.waitCh)
}
<commit_msg>Fix name of prefix<commit_after>package driver

import (
        "bytes"
        "encoding/json"
        "fmt"
        "log"
        "os"
        "os/exec"
        "regexp"
        "runtime"
        "strings"
        "syscall"
        "time"

        "github.com/hashicorp/nomad/client/config"
        "github.com/hashicorp/nomad/nomad/structs"
)

var (
        reRktVersion  = regexp.MustCompile("rkt version ([\\d\\.]+).+")
        reAppcVersion = regexp.MustCompile("appc version ([\\d\\.]+).+")
)

// RktDriver is a driver for running images via Rkt
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type RktDriver struct {
        DriverContext
}

// rktHandle is returned from Start/Open as a handle to the PID
type rktHandle struct {
        proc   *os.Process
        name   string
        waitCh chan error
        doneCh chan struct{}
}

// rktPID is a struct to map the pid running the process to the vm image on
// disk
type rktPID struct {
        Pid  int
        Name string
}

// NewRktDriver is used to create a new exec driver
func NewRktDriver(ctx *DriverContext) Driver {
        return &RktDriver{*ctx}
}

func (d *RktDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
        // Only enable if we are root when running on non-windows systems.
        if runtime.GOOS != "windows" && syscall.Geteuid() != 0 {
                d.logger.Printf("[DEBUG] driver.rkt: must run as root user, disabling")
                return false, nil
        }

        outBytes, err := exec.Command("rkt", "version").Output()
        if err != nil {
                return false, nil
        }
        out := strings.TrimSpace(string(outBytes))

        rktMatches := reRktVersion.FindStringSubmatch(out)
        appcMatches := reRktVersion.FindStringSubmatch(out)
        if len(rktMatches) != 2 || len(appcMatches) != 2 {
                return false, fmt.Errorf("Unable to parse Rkt version string: %#v", rktMatches)
        }

        node.Attributes["driver.rkt"] = "1"
        node.Attributes["driver.rkt.version"] = rktMatches[0]
        node.Attributes["driver.appc.version"] = appcMatches[1]

        return true, nil
}

// Run an existing Rkt image.
func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
        trust_prefix, ok := task.Config["trust_prefix"]
        if !ok || trust_prefix == "" {
                return nil, fmt.Errorf("Missing trust prefix for rkt")
        }
        trust_args := []string{
                "sudo",
                "rkt",
                "trust",
                "--prefix=" + trust_prefix,
        }
        // Add the given trust prefix
        var outBuf, errBuf bytes.Buffer
        cmd := exec.Command(trust_args[0], trust_args[1:]...)
        cmd.Stdout = &outBuf
        cmd.Stderr = &errBuf
        d.logger.Printf("[DEBUG] Starting rkt command: %q", strings.Join(trust_args, " "))
        if err := cmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, outBuf.String(), errBuf.String())
        }
        d.logger.Printf("[DEBUG] Added trust prefix: %q", trust_prefix)

        name, ok := task.Config["name"]
        if !ok || name == "" {
                return nil, fmt.Errorf("Missing ACI name for rkt")
        }
        run_args := []string{
                "sudo",
                "rkt",
                "run", "--interactive", name,
        }
        // Run the ACI
        var aoutBuf, aerrBuf bytes.Buffer
        acmd := exec.Command(run_args[0], run_args[1:]...)
        acmd.Stdout = &aoutBuf
        acmd.Stderr = &aerrBuf
        d.logger.Printf("[DEBUG] Starting rkt command: %q", strings.Join(run_args, " "))
        if err := acmd.Start(); err != nil {
                return nil, fmt.Errorf(
                        "Error running rkt: %s\n\nOutput: %s\n\nError: %s",
                        err, aoutBuf.String(), aerrBuf.String())
        }
        d.logger.Printf("[DEBUG] Started ACI: %q", name)
        h := &rktHandle{
                proc:   cmd.Process,
                name:   name,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }
        go h.run()
        return h, nil
}

func (d *RktDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
        // Parse the handle
        pidBytes := []byte(strings.TrimPrefix(handleID, "Rkt:"))
        qpid := &rktPID{}
        if err := json.Unmarshal(pidBytes, qpid); err != nil {
                return nil, fmt.Errorf("failed to parse Rkt handle '%s': %v", handleID, err)
        }

        // Find the process
        proc, err := os.FindProcess(qpid.Pid)
        if proc == nil || err != nil {
                return nil, fmt.Errorf("failed to find Rkt PID %d: %v", qpid.Pid, err)
        }

        // Return a driver handle
        h := &rktHandle{
                proc:   proc,
                name:   qpid.Name,
                doneCh: make(chan struct{}),
                waitCh: make(chan error, 1),
        }

        go h.run()
        return h, nil
}

func (h *rktHandle) ID() string {
        // Return a handle to the PID
        pid := &rktPID{
                Pid:  h.proc.Pid,
                Name: h.name,
        }
        data, err := json.Marshal(pid)
        if err != nil {
                log.Printf("[ERR] failed to marshal rkt PID to JSON: %s", err)
        }
        return fmt.Sprintf("Rkt:%s", string(data))
}

func (h *rktHandle) WaitCh() chan error {
        return h.waitCh
}

func (h *rktHandle) Update(task *structs.Task) error {
        // Update is not possible
        return nil
}

// Kill is used to terminate the task. We send an Interrupt
// and then provide a 5 second grace period before doing a Kill.
//
// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *rktHandle) Kill() error {
        h.proc.Signal(os.Interrupt)
        select {
        case <-h.doneCh:
                return nil
        case <-time.After(5 * time.Second):
                return h.proc.Kill()
        }
}

func (h *rktHandle) run() {
        ps, err := h.proc.Wait()
        close(h.doneCh)
        if err != nil {
                h.waitCh <- err
        } else if !ps.Success() {
                h.waitCh <- fmt.Errorf("task exited with error")
        }
        close(h.waitCh)
}
<|endoftext|>


<commit_before>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/nsqio/go-nsq"
	"net/http"
	"os"
	"strconv"
	"sync"
//	"time"
)

type APTFetcher struct {
	Context        *context.Context
	FetchChannel   chan *FetchData
	RecordChannel  chan *FetchData
	CleanupChannel chan *FetchData
	WaitGroup      sync.WaitGroup
}

type FetchData struct {
	WorkItem        *models.WorkItem
	WorkItemState   *models.WorkItemState
	IngestManifest  *models.IngestManifest
}

func NewATPFetcher(_context *context.Context) (*APTFetcher) {
	fetcher := &APTFetcher{
		Context: _context,
	}
	// Set up buffered channels
	fetcherBufferSize := _context.Config.FetchWorker.NetworkConnections * 4
	workerBufferSize := _context.Config.FetchWorker.Workers * 10
	fetcher.FetchChannel = make(chan *FetchData, fetcherBufferSize)
	fetcher.RecordChannel = make(chan *FetchData, workerBufferSize)
	fetcher.CleanupChannel = make(chan *FetchData, workerBufferSize)
	// Set up a limited number of go routines
	for i := 0; i < _context.Config.FetchWorker.NetworkConnections; i++ {
		go fetcher.fetch()
	}
	for i := 0; i < _context.Config.FetchWorker.Workers; i++ {
		go fetcher.cleanup()
		go fetcher.record()
	}
	return fetcher
}

func (fetcher *APTFetcher) HandleMessage(message *nsq.Message) (error) {
	message.DisableAutoResponse()
	fetchData, err := fetcher.initFetchData(message)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return err
	}
	resp := fetcher.Context.PharosClient.WorkItemStateSave(fetchData.WorkItemState)
	if resp.Error != nil {
		return resp.Error
	}
	fetchData.WorkItem, err = fetcher.markWorkItemAsStarted(fetchData.WorkItem)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return err
	}
	fetcher.FetchChannel <- fetchData
	return nil
}

func (fetcher *APTFetcher) fetch() {
//	for manifest := range fetcher.FetchChannel {
//
//	}
}

func (fetcher *APTFetcher) cleanup() {
//	for manifest := range fetcher.FetchChannel {
//
//	}
}

func (fetcher *APTFetcher) record() {
	// for manifest := range fetcher.FetchChannel {

	// }
}

func (fetcher *APTFetcher) initFetchData (message *nsq.Message) (*FetchData, error) {
	workItem, err := fetcher.getWorkItem(message)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	workItemState, err := fetcher.getWorkItemState(workItem)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	ingestManifest, err := workItemState.IngestManifest()
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	fetchData := &FetchData{
		WorkItem: workItem,
		WorkItemState: workItemState,
		IngestManifest: ingestManifest,
	}
	return fetchData, err
}

// Returns the WorkItem record from Pharos that has the WorkItemId
// specified in the NSQ message.
func (fetcher *APTFetcher) getWorkItem(message *nsq.Message) (*models.WorkItem, error) {
	workItemId, err := strconv.Atoi(string(message.Body))
	if err != nil {
		return nil, fmt.Errorf("Could not get WorkItemId from NSQ message body: %v", err)
	}
	resp := fetcher.Context.PharosClient.WorkItemGet(workItemId)
	if resp.Error != nil {
		return nil, fmt.Errorf("Error getting WorkItem %d from Pharos: %v", err)
	}
	workItem := resp.WorkItem()
	if workItem == nil {
		return nil, fmt.Errorf("Pharos returned nil for WorkItem %d", workItemId)
	}
	return workItem, nil
}

// Returns the WorkItemState record from Pharos with the specified workItem.Id,
// or creates a new WorkItemState (if necessary) and returns that. If this is
// the first time we've attempted to ingest this item, we'll have to crate a
// new WorkItemState.
func (fetcher *APTFetcher) getWorkItemState(workItem *models.WorkItem) (*models.WorkItemState, error) {
	var workItemState *models.WorkItemState
	var err error
	resp := fetcher.Context.PharosClient.WorkItemStateGet(workItem.Id)
	if resp.Response.StatusCode == http.StatusNotFound {
		// Record has not been created yet, so build a new one now.
		workItemState, err = fetcher.initWorkItemState(workItem)
		if err != nil {
			return nil, err
		}
	} else if resp.Error != nil {
		// We got some other 4xx/5xx error from the Pharos REST service.
		return nil, fmt.Errorf("Error getting WorkItemState for WorkItem %d from Pharos: %v", resp.Error)
	} else {
		// We didn't get a 404 or any other error. The WorkItemState should be in
		// the response.
		workItemState = resp.WorkItemState()
		if workItemState == nil {
			return nil, fmt.Errorf("Pharos returned nil for WorkItemState with WorkItem id %d", workItem.Id)
		}
	}
	return workItemState, nil
}

func (fetcher *APTFetcher) initWorkItemState (workItem *models.WorkItem) (*models.WorkItemState, error) {
	ingestManifest := models.NewIngestManifest()
	ingestManifest.WorkItemId = workItem.Id
	ingestManifest.S3Bucket = workItem.Bucket
	ingestManifest.S3Key = workItem.Name
	ingestManifest.ETag = workItem.ETag
	workItemState := models.NewWorkItemState(workItem.Id, constants.ActionIngest, "")
	err := workItemState.SetStateFromIngestManifest(ingestManifest)
	if err != nil {
		return nil, err
	}
	return workItemState, nil
}

func (fetcher *APTFetcher) markWorkItemAsStarted (workItem *models.WorkItem) (*models.WorkItem, error) {
	hostname, _ := os.Hostname()
	if hostname == "" { hostname = "apt_fetcher_host" }
	workItem.Node = hostname
	workItem.Status = constants.StatusStarted
	workItem.Pid = os.Getpid()
	workItem.Note = "Fetching bag from receiving bucket."
	resp := fetcher.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		return nil, resp.Error
	}
	return resp.WorkItem(), nil
}

// This is for direct testing without NSQ.
func (fetcher *APTFetcher) RunWithoutNsq(fetchData *FetchData) {
	fetcher.WaitGroup.Add(1)
	fetcher.FetchChannel <- fetchData
	fetcher.Context.MessageLog.Debug("Put %s into Fluctus channel", fetchData.IngestManifest.S3Key)
	fetcher.WaitGroup.Wait()
}
<commit_msg>Working on fetcher<commit_after>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/network"
	"github.com/APTrust/exchange/util"
	"github.com/APTrust/exchange/util/fileutil"
	"github.com/nsqio/go-nsq"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
//	"time"
)

type APTFetcher struct {
	Context        *context.Context
	FetchChannel   chan *FetchData
	RecordChannel  chan *FetchData
	CleanupChannel chan *FetchData
	WaitGroup      sync.WaitGroup
}

type FetchData struct {
	WorkItem        *models.WorkItem
	WorkItemState   *models.WorkItemState
	IngestManifest  *models.IngestManifest
}

func NewATPFetcher(_context *context.Context) (*APTFetcher) {
	fetcher := &APTFetcher{
		Context: _context,
	}
	// Set up buffered channels
	fetcherBufferSize := _context.Config.FetchWorker.NetworkConnections * 4
	workerBufferSize := _context.Config.FetchWorker.Workers * 10
	fetcher.FetchChannel = make(chan *FetchData, fetcherBufferSize)
	fetcher.RecordChannel = make(chan *FetchData, workerBufferSize)
	fetcher.CleanupChannel = make(chan *FetchData, workerBufferSize)
	// Set up a limited number of go routines
	for i := 0; i < _context.Config.FetchWorker.NetworkConnections; i++ {
		go fetcher.fetch()
	}
	for i := 0; i < _context.Config.FetchWorker.Workers; i++ {
		go fetcher.cleanup()
		go fetcher.record()
	}
	return fetcher
}

func (fetcher *APTFetcher) HandleMessage(message *nsq.Message) (error) {

	// Set up our fetch data. Most of this comes from Pharos;
	// some of it we have to build fresh.
	fetchData, err := fetcher.initFetchData(message)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return err
	}
	// Save the state of this item in Pharos.
	resp := fetcher.Context.PharosClient.WorkItemStateSave(fetchData.WorkItemState)
	if resp.Error != nil {
		return resp.Error
	}
	// Tell Pharos that we've started work on the item.
	fetchData.WorkItem, err = fetcher.recordFetchStarted(fetchData.WorkItem)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return err
	}

	// NSQ message autoresponse periodically tells the queue
	// that the message is still being processed. This doesn't
	// work for us in cases where we're fetching a file that's
	// 100GB+ in size. We need to manually Touch() NSQ periodically
	// to let the queue know that we're still actively working on
	// the message. Otherwise, NSQ thinks it timed out and sends
	// the message to a new worker.
	message.DisableAutoResponse()

	// Now get to work.
	fetcher.FetchChannel <- fetchData
	return nil
}

func (fetcher *APTFetcher) fetch() {
	for fetchData := range fetcher.FetchChannel {
		// Tell NSQ we're working on this
		fetchData.IngestManifest.NSQMessage.Touch()

		fetchData.IngestManifest.Fetch.Start()
		fetchData.IngestManifest.Fetch.Attempted = true

		err := fetcher.downloadFile(fetchData)

		// Download may have taken 1 second or 3 hours.
		// Remind NSQ that we're still on this.
		fetchData.IngestManifest.NSQMessage.Touch()

		if err != nil {
			fetchData.IngestManifest.Fetch.AddError(err.Error())
		}
		fetcher.CleanupChannel <- fetchData
	}
}

func (fetcher *APTFetcher) cleanup() {
	for fetchData := range fetcher.CleanupChannel {
		tarFile := fetchData.IngestManifest.Object.IngestTarFilePath
		if fetchData.IngestManifest.Fetch.HasErrors() && fileutil.FileExists(tarFile) {
			// Most likely bad md5 digest, but perhaps also a partial download.
			fetcher.Context.MessageLog.Info("Deleting due to download error: %s",
				tarFile)
			os.Remove(tarFile)
		}
		fetcher.RecordChannel <- fetchData
	}
}

func (fetcher *APTFetcher) record() {
//	for fetchData := range fetcher.RecordChannel {
		// Call fetchData.IngestManifest.Fetch.Finish()

		// Log WorkItemState
		// Save WorkItemState to Pharos

		// If no errors:
		// Set WorkItem stage to StageValidate, status to StatusPending, node=nil, pid=0
		// Finish the NSQ message

		// If transient errors:
		// Set WorkItem node=nil, pid=0
		// Requeue the NSQ message

		// If fatal errors:
		// Set WorkItem node=nil, pid=0, retry=false, needs_admin_review=true
		// Finish the NSQ message
//	}
}

// Set up the basic pieces of data we'll need to process a fetch request.
func (fetcher *APTFetcher) initFetchData (message *nsq.Message) (*FetchData, error) {
	workItem, err := fetcher.getWorkItem(message)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	workItemState, err := fetcher.getWorkItemState(workItem)
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	ingestManifest, err := workItemState.IngestManifest()
	if err != nil {
		fetcher.Context.MessageLog.Error(err.Error())
		return nil, err
	}
	fetchData := &FetchData{
		WorkItem: workItem,
		WorkItemState: workItemState,
		IngestManifest: ingestManifest,
	}

	// instIdentifier is, e.g., virginia.edu, ncsu.edu, etc.
	// We'll download the tar file from the receiving bucket to
	// something like /mnt/apt/data/virginia.edu/name_of_bag.tar
	// See IngestTarFilePath below.
	instIdentifier := util.OwnerOf(fetchData.IngestManifest.S3Bucket)

	// Set some basic info on our IntellectualObject
	fetchData.IngestManifest.Object.BagName = util.CleanBagName(fetchData.IngestManifest.S3Key)
	fetchData.IngestManifest.Object.Institution = instIdentifier
	//fetchData.IngestManifest.Object.InstitutionId =
	fetchData.IngestManifest.Object.IngestS3Bucket = fetchData.IngestManifest.S3Bucket
	fetchData.IngestManifest.Object.IngestS3Key = fetchData.IngestManifest.S3Key
	fetchData.IngestManifest.Object.IngestTarFilePath = filepath.Join(
		fetcher.Context.Config.TarDirectory,
		instIdentifier, fetchData.IngestManifest.S3Key)

	return fetchData, err
}

// Returns the WorkItem record from Pharos that has the WorkItemId
// specified in the NSQ message.
func (fetcher *APTFetcher) getWorkItem(message *nsq.Message) (*models.WorkItem, error) {
	workItemId, err := strconv.Atoi(string(message.Body))
	if err != nil {
		return nil, fmt.Errorf("Could not get WorkItemId from NSQ message body: %v", err)
	}
	resp := fetcher.Context.PharosClient.WorkItemGet(workItemId)
	if resp.Error != nil {
		return nil, fmt.Errorf("Error getting WorkItem %d from Pharos: %v", err)
	}
	workItem := resp.WorkItem()
	if workItem == nil {
		return nil, fmt.Errorf("Pharos returned nil for WorkItem %d", workItemId)
	}
	return workItem, nil
}

// Returns the WorkItemState record from Pharos with the specified workItem.Id,
// or creates a new WorkItemState (if necessary) and returns that. If this is
// the first time we've attempted to ingest this item, we'll have to crate a
// new WorkItemState.
func (fetcher *APTFetcher) getWorkItemState(workItem *models.WorkItem) (*models.WorkItemState, error) {
	var workItemState *models.WorkItemState
	var err error
	resp := fetcher.Context.PharosClient.WorkItemStateGet(workItem.Id)
	if resp.Response.StatusCode == http.StatusNotFound {
		// Record has not been created yet, so build a new one now.
		workItemState, err = fetcher.initWorkItemState(workItem)
		if err != nil {
			return nil, err
		}
	} else if resp.Error != nil {
		// We got some other 4xx/5xx error from the Pharos REST service.
		return nil, fmt.Errorf("Error getting WorkItemState for WorkItem %d from Pharos: %v", resp.Error)
	} else {
		// We didn't get a 404 or any other error. The WorkItemState should be in
		// the response.
		workItemState = resp.WorkItemState()
		if workItemState == nil {
			return nil, fmt.Errorf("Pharos returned nil for WorkItemState with WorkItem id %d", workItem.Id)
		}
	}
	return workItemState, nil
}

// Create a new WorkItemState object for this WorkItem.
// We do this only when Pharos doesn't already have a WorkItemState
// object, which is often the case when ingesting new bags.
func (fetcher *APTFetcher) initWorkItemState (workItem *models.WorkItem) (*models.WorkItemState, error) {
	ingestManifest := models.NewIngestManifest()
	ingestManifest.WorkItemId = workItem.Id
	ingestManifest.S3Bucket = workItem.Bucket
	ingestManifest.S3Key = workItem.Name
	ingestManifest.ETag = workItem.ETag
	workItemState := models.NewWorkItemState(workItem.Id, constants.ActionIngest, "")
	err := workItemState.SetStateFromIngestManifest(ingestManifest)
	if err != nil {
		return nil, err
	}
	return workItemState, nil
}

// Tell Pharos we've started work on this item.
func (fetcher *APTFetcher) recordFetchStarted (workItem *models.WorkItem) (*models.WorkItem, error) {
	hostname, _ := os.Hostname()
	if hostname == "" { hostname = "apt_fetcher_host" }
	workItem.Node = hostname
	workItem.Stage = constants.StageFetch
	workItem.Status = constants.StatusStarted
	workItem.Pid = os.Getpid()
	workItem.Note = "Fetching bag from receiving bucket."
	resp := fetcher.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		return nil, resp.Error
	}
	return resp.WorkItem(), nil
}

// Download the file, and update the IngestManifest while we're at it.
func (fetcher *APTFetcher) downloadFile (fetchData *FetchData) (error) {
	downloader := network.NewS3Download(
		constants.AWSVirginia,
		fetchData.IngestManifest.S3Bucket,
		fetchData.IngestManifest.S3Key,
		fetchData.IngestManifest.Object.IngestTarFilePath,
		true,    // calculate md5 checksum on the entire tar file
		false,   // calculate sha256 checksum on the entire tar file
	)

	// It's fairly common for very large bags to fail more than
	// once on transient network errors (e.g. "Connection reset by peer")
	// So we give this several tries.
	for i := 0; i < 10; i++ {
		downloader.Fetch()
		if downloader.ErrorMessage == "" {
			fetcher.Context.MessageLog.Info("Fetched %s/%s after %d attempts",
				fetchData.IngestManifest.S3Bucket,
				fetchData.IngestManifest.S3Key,
				i + 1)
			break
		}
	}

	// Return now if we failed.
	if downloader.ErrorMessage != "" {
		return fmt.Errorf("Error fetching %s/%s: %v",
			fetchData.IngestManifest.S3Bucket,
			fetchData.IngestManifest.S3Key,
			downloader.ErrorMessage)
	}

	fetchData.IngestManifest.Object.IngestSize = downloader.BytesCopied
	fetchData.IngestManifest.Object.IngestRemoteMd5 = *downloader.Response.ETag
	fetchData.IngestManifest.Object.IngestLocalMd5 = downloader.Md5Digest

	// The ETag for S3 object uploaded via single-part upload is
	// the file's md5 digest. For objects uploaded via multi-part
	// upload, the ETag is calculated differently and includes a
	// dash near the end, followed by the number of parts in the
	// multipart upload. We can't use that kind of ETag to verify
	// the md5 checksum that we calculated.
	fetchData.IngestManifest.Object.IngestMd5Verifiable = strings.Contains(downloader.Md5Digest, "-")
	if fetchData.IngestManifest.Object.IngestMd5Verifiable {
		fetchData.IngestManifest.Object.IngestMd5Verified = (
			fetchData.IngestManifest.Object.IngestRemoteMd5 == fetchData.IngestManifest.Object.IngestLocalMd5)
	}

	return nil
}

// This is for direct testing without NSQ.
func (fetcher *APTFetcher) RunWithoutNsq(fetchData *FetchData) {
	fetcher.WaitGroup.Add(1)
	fetcher.FetchChannel <- fetchData
	fetcher.Context.MessageLog.Debug("Put %s into Fluctus channel", fetchData.IngestManifest.S3Key)
	fetcher.WaitGroup.Wait()
}
<|endoftext|>


<commit_before>package store

import (
	"github.com/Sirupsen/logrus"
	"utils"
)

func (f File) MMeta() utils.M {
	return utils.M(f.Meta())
}

func (f *File) SetContentType(v string) *File {
	utils.M(f.Meta()).Set(ContentTypeKey, v)
	return f
}

func (f File) ContentType() string {

	return utils.Map(f.Meta()).String(ContentTypeKey)
}

func (f *File) MMapData() utils.M {
	return utils.M(f.MapData())
}

func (f *File) SetMapData(v map[string]interface{}) *File {
	f.File.SetMapData(v)
	return f
}

func (f File) TextData() string {

	rawData := f.RawData().Bytes()

	if len(rawData) > 1024*1024*10 {
		logrus.WithField("length", len(rawData)).Warning("file raw data as text: to long")
		return ""
	}

	return string(rawData)
}

func (f *File) SetTextData(src string) *File {
	f.RawData().Write([]byte(src))
	return f
}

func (f File) IsImage() bool {
	return getTypeNameFromContentType(f.ContentType()) == "image"
}

func (f File) IsText() bool {
	return getTypeNameFromContentType(f.ContentType()) == "text"
}

func (f File) IsRaw() bool {
	return getTypeNameFromContentType(f.ContentType()) == "raw"
}
<commit_msg>add helpful function for file<commit_after>package store

import (
	"github.com/Sirupsen/logrus"
	"utils"
)

func (f File) MMeta() utils.M {
	return utils.M(f.Meta())
}

func (f *File) SetContentType(v string) *File {
	utils.M(f.Meta()).Set(ContentTypeKey, v)
	return f
}

func (f File) ContentType() string {

	return utils.Map(f.Meta()).String(ContentTypeKey)
}

func (f *File) MMapData() utils.M {
	return utils.M(f.MapData())
}

func (f *File) SetMapData(v map[string]interface{}) *File {
	f.File.SetMapData(v)
	return f
}

func (f File) TextData() string {

	rawData := f.RawData().Bytes()

	if len(rawData) > 1024*1024*10 {
		logrus.WithField("length", len(rawData)).Warning("file raw data as text: to long")
		return ""
	}

	return string(rawData)
}

func (f *File) SetTextData(src string) *File {
	f.RawData().Write([]byte(src))
	return f
}

func (f *File) SetRawData(src []byte) *File {
	f.RawData().Write(src)
	return f
}

func (f File) IsImage() bool {
	return getTypeNameFromContentType(f.ContentType()) == "image"
}

func (f File) IsText() bool {
	return getTypeNameFromContentType(f.ContentType()) == "text"
}

func (f File) IsRaw() bool {
	return getTypeNameFromContentType(f.ContentType()) == "raw"
}
<|endoftext|>


<commit_before>package handler

import (
//	"github.com/materials-commons/config"
)

const (
	// Use the default handler
	Default = "Default"

	// Use the Override handler
	Override = "Override"

	// Use the environment handler
	Environment = "Environment"

	// Use the ini handler
	Ini = "Ini"
)
<commit_msg>Create handler that acts like github.com/spf13/viper.<commit_after>package handler

import (
	"github.com/materials-commons/config"
)

const (
	// Use the default handler
	Default = "Default"

	// Use the Override handler
	Override = "Override"

	// Use the environment handler
	Environment = "Environment"

	// Use the ini handler
	Ini = "Ini"
)

// Viper implements github.com/spf13/viper
func Viper(loader config.Loader) config.Handler {
	return LowercaseKey(
		Prioritized(
			NameHandler(Default, Map()),
			NameHandler(Override, Map()),
			NameHandler(Ini, Loader(loader))))
}

// ViperCaseSensitive implements github.com/spf13/viper except that keys
// are case sensitive.
func ViperCaseSensitive(loader config.Loader) config.Handler {
	return Prioritized(
		NameHandler(Default, Map()),
		NameHandler(Override, Map()),
		NameHandler(Ini, Loader(loader)))
}

// ViperEx implements github.com/spf13/viper with the addition of environment
// variables checked before checking for values in the ini file(s).
func ViperEx(loader config.Loader) config.Handler {
	return LowercaseKey(
		Prioritized(
			NameHandler(Default, Map()),
			NameHandler(Override, Map()),
			NameHandler(Environment, Env()),
			NameHandler(Ini, Loader(loader))))
}

// ViperExCaseSensitive implements ViperEx except that keys are case sensitive.
func ViperExCaseSensitive(loader config.Loader) config.Handler {
	return Prioritized(
		NameHandler(Default, Map()),
		NameHandler(Override, Map()),
		NameHandler(Environment, Env()),
		NameHandler(Ini, Loader(loader)))
}
<|endoftext|>


<commit_before>package libvirt

import (
	"log"
	"strings"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/mitchellh/go-vnc"
	"github.com/mitchellh/multistep"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	/*
		config := state.Get("config").(*Config)
		httpPort := state.Get("http_port").(uint)
		hostIp := state.Get("host_ip").(string)
		ui := state.Get("ui").(packer.Ui)

		// Get the VNC IP / Port
		virshOut, _, err := virsh("vncdisplay", config.VMName)
		vncAddr := strings.TrimSpace(virshOut)
		colon := strings.LastIndex(vncAddr, ":")
		if colon < 0 || colon > len(vncAddr)-2 {
			err := fmt.Errorf("Error parsing VNC address: %s", vncAddr)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
		port, err := strconv.ParseUint(vncAddr[colon+1:], 10, 16)
		if err != nil {
			err := fmt.Errorf("Error parsing VNC port: %s", err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
		vncAddr = fmt.Sprintf("%s:%d", vncAddr[:colon], 5900+port)

		log.Printf("VNC Address for VM: %s", vncAddr)

		// Connect to VNC
		ui.Say("Connecting to VM via VNC")
		nc, err := net.Dial("tcp", vncAddr)
		if err != nil {
			err := fmt.Errorf("Error connecting to VNC: %s", err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
		defer nc.Close()

		c, err := vnc.Client(nc, &vnc.ClientConfig{Exclusive: true})
		if err != nil {
			err := fmt.Errorf("Error handshaking with VNC: %s", err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}
		defer c.Close()

		log.Printf("Connected to VNC desktop: %s", c.DesktopName)

		log.Printf("Host IP for the VM: %s", hostIp)

		tplData := &bootCommandTemplateData{
			hostIp,
			httpPort,
			config.VMName,
		}

		ui.Say("Typing the boot command over VNC...")
		for _, command := range config.BootCommand {
			command, err := config.tpl.Process(command, tplData)
			if err != nil {
				err := fmt.Errorf("Error preparing boot command: %s", err)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

			// Check for interrupts between typing things so we can cancel
			// since this isn't the fastest thing.
			if _, ok := state.GetOk(multistep.StateCancelled); ok {
				return multistep.ActionHalt
			}

			vncSendString(c, command)
		}
	*/
	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func vncSendString(c *vnc.ClientConn, original string) {
	special := make(map[string]uint32)
	special["<bs>"] = 0xFF08
	special["<del>"] = 0xFFFF
	special["<enter>"] = 0xFF0D
	special["<esc>"] = 0xFF1B
	special["<f1>"] = 0xFFBE
	special["<f2>"] = 0xFFBF
	special["<f3>"] = 0xFFC0
	special["<f4>"] = 0xFFC1
	special["<f5>"] = 0xFFC2
	special["<f6>"] = 0xFFC3
	special["<f7>"] = 0xFFC4
	special["<f8>"] = 0xFFC5
	special["<f9>"] = 0xFFC6
	special["<f10>"] = 0xFFC7
	special["<f11>"] = 0xFFC8
	special["<f12>"] = 0xFFC9
	special["<return>"] = 0xFF0D
	special["<tab>"] = 0xFF09

	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"

	// TODO(mitchellh): Ripe for optimizations of some point, perhaps.
	for len(original) > 0 {
		var keyCode uint32
		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		for specialCode, specialValue := range special {
			if strings.HasPrefix(original, specialCode) {
				log.Printf("Special code '%s' found, replacing with: %d", specialCode, specialValue)
				keyCode = specialValue
				original = original[len(specialCode):]
				break
			}
		}

		if keyCode == 0 {
			r, size := utf8.DecodeRuneInString(original)
			original = original[size:]
			keyCode = uint32(r)
			keyShift = unicode.IsUpper(r) || strings.ContainsRune(shiftedChars, r)

			log.Printf("Sending char '%c', code %d, shift %v", r, keyCode, keyShift)
		}

		if keyShift {
			c.KeyEvent(KeyLeftShift, true)
		}

		time.Sleep(5 * time.Millisecond)
		c.KeyEvent(keyCode, true)
		time.Sleep(5 * time.Millisecond)
		c.KeyEvent(keyCode, false)

		if keyShift {
			c.KeyEvent(KeyLeftShift, false)
		}
	}
}
<commit_msg>fix<commit_after>package libvirt

import (
	"fmt"
	"log"
	"strings"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/alexzorin/libvirt-go"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	//	httpPort := state.Get("http_port").(uint)
	//	hostIp := state.Get("host_ip").(string)
	ui := state.Get("ui").(packer.Ui)

	var lvd libvirt.VirDomain
	lv, err := libvirt.NewVirConnection(config.LibvirtUrl)
	if err != nil {
		err := fmt.Errorf("Error connecting to libvirt: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lv.CloseConnection()
	if lvd, err = lv.LookupDomainByName(config.VMName); err != nil {
		err := fmt.Errorf("Error lookup domain: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lvd.Free()

	//	tplData := &bootCommandTemplateData{
	//		hostIp,
	//		httpPort,
	//		config.VMName,
	//	}

	ui.Say("Typing the boot command...")
	for _, command := range config.BootCommand {
		//		command, err := config.tpl.Process(command, tplData)
		//		if err != nil {
		//			err := fmt.Errorf("Error preparing boot command: %s", err)
		//			state.Put("error", err)
		//			ui.Error(err.Error())
		//			return multistep.ActionHalt
		//		}

		// Check for interrupts between typing things so we can cancel
		// since this isn't the fastest thing.
		if _, ok := state.GetOk(multistep.StateCancelled); ok {
			return multistep.ActionHalt
		}

		vncSendString(lvd, command)
	}

	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func vncSendString(d libvirt.VirDomain, original string) {
	special := make(map[string]uint32)
	special["<bs>"] = 0xFF08
	special["<del>"] = 0xFFFF
	special["<enter>"] = 0xFF0D
	special["<esc>"] = 0xFF1B
	special["<f1>"] = 0xFFBE
	special["<f2>"] = 0xFFBF
	special["<f3>"] = 0xFFC0
	special["<f4>"] = 0xFFC1
	special["<f5>"] = 0xFFC2
	special["<f6>"] = 0xFFC3
	special["<f7>"] = 0xFFC4
	special["<f8>"] = 0xFFC5
	special["<f9>"] = 0xFFC6
	special["<f10>"] = 0xFFC7
	special["<f11>"] = 0xFFC8
	special["<f12>"] = 0xFFC9
	special["<return>"] = 0xFF0D
	special["<tab>"] = 0xFF09

	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"

	// TODO(mitchellh): Ripe for optimizations of some point, perhaps.
	for len(original) > 0 {
		var keyCode uint32
		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		for specialCode, specialValue := range special {
			if strings.HasPrefix(original, specialCode) {
				log.Printf("Special code '%s' found, replacing with: %d", specialCode, specialValue)
				keyCode = specialValue
				original = original[len(specialCode):]
				break
			}
		}

		if keyCode == 0 {
			r, size := utf8.DecodeRuneInString(original)
			original = original[size:]
			keyCode = uint32(r)
			keyShift = unicode.IsUpper(r) || strings.ContainsRune(shiftedChars, r)

			log.Printf("Sending char '%c', code %d, shift %v", r, keyCode, keyShift)
		}

		//		if keyShift {
		//			c.KeyEvent(KeyLeftShift, true)
		//		}

		//		time.Sleep(5 * time.Millisecond)
		//VIR_KEYCODE_SET_LINUX, VIR_KEYCODE_SET_USB, VIR_KEYCODE_SET_RFB, VIR_KEYCODE_SET_WIN32, VIR_KEYCODE_SET_XT_KBD
		d.SendKey(libvirt.VIR_KEYCODE_SET_XT_KBD, 50, []uint{uint(keyCode)}, 0)
		//		c.KeyEvent(keyCode, true)
		//		time.Sleep(5 * time.Millisecond)
		//		c.KeyEvent(keyCode, false)

		//		if keyShift {
		//			c.KeyEvent(KeyLeftShift, false)
		//		}
	}
}
<|endoftext|>


<commit_before>package libvirt

import (
	"fmt"
	"log"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"gopkg.in/alexzorin/libvirt-go.v2"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	//	httpPort := state.Get("http_port").(uint)
	//	hostIp := state.Get("host_ip").(string)
	ui := state.Get("ui").(packer.Ui)

	var lvd libvirt.VirDomain
	lv, err := libvirt.NewVirConnection(config.LibvirtUrl)
	if err != nil {
		err := fmt.Errorf("Error connecting to libvirt: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lv.CloseConnection()
	if lvd, err = lv.LookupDomainByName(config.VMName); err != nil {
		err := fmt.Errorf("Error lookup domain: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lvd.Free()

	//	tplData := &bootCommandTemplateData{
	//		hostIp,
	//		httpPort,
	//		config.VMName,
	//	}

	ui.Say("Typing the boot command...")
	for _, command := range config.BootCommand {
		//		command, err := config.tpl.Process(command, tplData)
		//		if err != nil {
		//			err := fmt.Errorf("Error preparing boot command: %s", err)
		//			state.Put("error", err)
		//			ui.Error(err.Error())
		//			return multistep.ActionHalt
		//		}

		// Check for interrupts between typing things so we can cancel
		// since this isn't the fastest thing.
		if _, ok := state.GetOk(multistep.StateCancelled); ok {
			return multistep.ActionHalt
		}

		sendBootString(lvd, command)
	}

	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func sendBootString(d libvirt.VirDomain, original string) {
	//	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"
	var err error
	var ok bool
	var key uint

	for len(original) > 0 {
		//		var keyCode uint
		//		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		if strings.HasPrefix(original, "<esc>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<esc>"]}, 0)
			original = original[len("<esc>"):]
		}
		if strings.HasPrefix(original, "<enter>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<enter>"]}, 0)
			original = original[len("<enter>"):]
		}

		r, size := utf8.DecodeRuneInString(original)
		original = original[size:]
		key = ecodes[string(r)]
		log.Printf("find code for char %s %d", string(r), key)
		//VIR_KEYCODE_SET_LINUX, VIR_KEYCODE_SET_USB, VIR_KEYCODE_SET_RFB, VIR_KEYCODE_SET_WIN32, VIR_KEYCODE_SET_XT_KBD
		log.Printf("send code %d", key)
		if err = d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{key}, 0); err != nil {
			log.Printf("Sending code %d failed: %s", key, err.Error())
		}
	}

}
<commit_msg>fix<commit_after>package libvirt

import (
	"fmt"
	"log"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"gopkg.in/alexzorin/libvirt-go.v2"
)

const KeyLeftShift uint32 = 0xFFE1

type bootCommandTemplateData struct {
	HTTPIP   string
	HTTPPort uint
	Name     string
}

// This step "types" the boot command into the VM over VNC.
//
// Uses:
//   config *config
//   http_port int
//   ui     packer.Ui
//
// Produces:
//   <nothing>
type stepTypeBootCommand struct{}

func (s *stepTypeBootCommand) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	//	httpPort := state.Get("http_port").(uint)
	//	hostIp := state.Get("host_ip").(string)
	ui := state.Get("ui").(packer.Ui)

	var lvd libvirt.VirDomain
	lv, err := libvirt.NewVirConnection(config.LibvirtUrl)
	if err != nil {
		err := fmt.Errorf("Error connecting to libvirt: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lv.CloseConnection()
	if lvd, err = lv.LookupDomainByName(config.VMName); err != nil {
		err := fmt.Errorf("Error lookup domain: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer lvd.Free()

	//	tplData := &bootCommandTemplateData{
	//		hostIp,
	//		httpPort,
	//		config.VMName,
	//	}

	ui.Say("Typing the boot command...")
	for _, command := range config.BootCommand {
		//		command, err := config.tpl.Process(command, tplData)
		//		if err != nil {
		//			err := fmt.Errorf("Error preparing boot command: %s", err)
		//			state.Put("error", err)
		//			ui.Error(err.Error())
		//			return multistep.ActionHalt
		//		}

		// Check for interrupts between typing things so we can cancel
		// since this isn't the fastest thing.
		if _, ok := state.GetOk(multistep.StateCancelled); ok {
			return multistep.ActionHalt
		}

		sendBootString(lvd, command)
	}

	return multistep.ActionContinue
}

func (*stepTypeBootCommand) Cleanup(multistep.StateBag) {}

func sendBootString(d libvirt.VirDomain, original string) {
	//	shiftedChars := "~!@#$%^&*()_+{}|:\"<>?"
	var err error
	var key uint

	for len(original) > 0 {
		//		var keyCode uint
		//		keyShift := false

		if strings.HasPrefix(original, "<wait>") {
			log.Printf("Special code '<wait>' found, sleeping one second")
			time.Sleep(1 * time.Second)
			original = original[len("<wait>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait5>") {
			log.Printf("Special code '<wait5>' found, sleeping 5 seconds")
			time.Sleep(5 * time.Second)
			original = original[len("<wait5>"):]
			continue
		}

		if strings.HasPrefix(original, "<wait10>") {
			log.Printf("Special code '<wait10>' found, sleeping 10 seconds")
			time.Sleep(10 * time.Second)
			original = original[len("<wait10>"):]
			continue
		}

		if strings.HasPrefix(original, "<esc>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<esc>"]}, 0)
			original = original[len("<esc>"):]
		}
		if strings.HasPrefix(original, "<enter>") {
			d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{ecodes["<enter>"]}, 0)
			original = original[len("<enter>"):]
		}

		r, size := utf8.DecodeRuneInString(original)
		original = original[size:]
		key = ecodes[string(r)]
		log.Printf("find code for char %s %d", string(r), key)
		//VIR_KEYCODE_SET_LINUX, VIR_KEYCODE_SET_USB, VIR_KEYCODE_SET_RFB, VIR_KEYCODE_SET_WIN32, VIR_KEYCODE_SET_XT_KBD
		log.Printf("send code %d", key)
		if err = d.SendKey(libvirt.VIR_KEYCODE_SET_RFB, 400, []uint{key}, 0); err != nil {
			log.Printf("Sending code %d failed: %s", key, err.Error())
		}
	}

}
<|endoftext|>


<commit_before>package uhost

import (
	"context"
	"fmt"
	"github.com/hashicorp/packer/common/retry"
	"time"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/ucloud/ucloud-sdk-go/services/uhost"
	"github.com/ucloud/ucloud-sdk-go/ucloud"
)

type stepCreateImage struct {
	image *uhost.UHostImageSet
}

func (s *stepCreateImage) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*UCloudClient)
	conn := client.uhostconn
	instance := state.Get("instance").(*uhost.UHostInstanceSet)
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)

	ui.Say(fmt.Sprintf("Creating image %s", config.ImageName))

	req := conn.NewCreateCustomImageRequest()
	req.ImageName = ucloud.String(config.ImageName)
	req.ImageDescription = ucloud.String(config.ImageDescription)
	req.UHostId = ucloud.String(instance.UHostId)

	resp, err := conn.CreateCustomImage(req)
	if err != nil {
		return halt(state, err, "Error on creating image")
	}

	err = retry.Config{
		Tries: 200,
		ShouldRetry: func(err error) bool {
			return isExpectedStateError(err)
		},
		RetryDelay: (&retry.Backoff{InitialBackoff: 2 * time.Second, MaxBackoff: 12 * time.Second, Multiplier: 2}).Linear,
	}.Run(ctx, func(ctx context.Context) error {
		inst, err := client.DescribeImageById(resp.ImageId)
		if err != nil {
			return err
		}
		if inst == nil || inst.State != imageStateAvailable {
			return newExpectedStateError("image", resp.ImageId)
		}

		return nil
	})

	if err != nil {
		return halt(state, err, fmt.Sprintf("Error on waiting for image %q available", resp.ImageId))
	}

	imageSet, err := client.DescribeImageById(resp.ImageId)
	if err != nil {
		return halt(state, err, fmt.Sprintf("Error on reading image when creating %q", resp.ImageId))
	}

	s.image = imageSet
	state.Put("image_id", imageSet.ImageId)

	images := []imageInfo{
		{
			ImageId:   imageSet.ImageId,
			ProjectId: config.ProjectId,
			Region:    config.Region,
		},
	}

	state.Put("ucloud_images", newImageInfoSet(images))
	ui.Message(fmt.Sprintf("Creating image %q complete", imageSet.ImageId))
	return multistep.ActionContinue
}

func (s *stepCreateImage) Cleanup(state multistep.StateBag) {
	if s.image == nil {
		return
	}
	_, cancelled := state.GetOk(multistep.StateCancelled)
	_, halted := state.GetOk(multistep.StateHalted)
	if !cancelled && !halted {
		return
	}

	client := state.Get("client").(*UCloudClient)
	conn := client.uhostconn
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deleting image because of cancellation or error...")
	req := conn.NewTerminateCustomImageRequest()
	req.ImageId = ucloud.String(s.image.ImageId)
	_, err := conn.TerminateCustomImage(req)
	if err != nil {
		ui.Error(fmt.Sprintf("Error on deleting image %q", s.image.ImageId))
	}
	ui.Message(fmt.Sprintf("Deleting image %q complete", s.image.ImageId))
}
<commit_msg>Update builder/ucloud/uhost/step_create_image.go<commit_after>package uhost

import (
	"context"
	"fmt"
	"github.com/hashicorp/packer/common/retry"
	"time"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/ucloud/ucloud-sdk-go/services/uhost"
	"github.com/ucloud/ucloud-sdk-go/ucloud"
)

type stepCreateImage struct {
	image *uhost.UHostImageSet
}

func (s *stepCreateImage) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*UCloudClient)
	conn := client.uhostconn
	instance := state.Get("instance").(*uhost.UHostInstanceSet)
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)

	ui.Say(fmt.Sprintf("Creating image %s", config.ImageName))

	req := conn.NewCreateCustomImageRequest()
	req.ImageName = ucloud.String(config.ImageName)
	req.ImageDescription = ucloud.String(config.ImageDescription)
	req.UHostId = ucloud.String(instance.UHostId)

	resp, err := conn.CreateCustomImage(req)
	if err != nil {
		return halt(state, err, "Error on creating image")
	}

	err = retry.Config{
		Tries: 200,
		ShouldRetry: func(err error) bool {
			return isExpectedStateError(err)
		},
		RetryDelay: (&retry.Backoff{InitialBackoff: 2 * time.Second, MaxBackoff: 12 * time.Second, Multiplier: 2}).Linear,
	}.Run(ctx, func(ctx context.Context) error {
		inst, err := client.DescribeImageById(resp.ImageId)
		if err != nil {
			return err
		}
		if inst == nil || inst.State != imageStateAvailable {
			return newExpectedStateError("image", resp.ImageId)
		}

		return nil
	})

	if err != nil {
		return halt(state, err, fmt.Sprintf("Error on waiting for image %q to become available", resp.ImageId))
	}

	imageSet, err := client.DescribeImageById(resp.ImageId)
	if err != nil {
		return halt(state, err, fmt.Sprintf("Error on reading image when creating %q", resp.ImageId))
	}

	s.image = imageSet
	state.Put("image_id", imageSet.ImageId)

	images := []imageInfo{
		{
			ImageId:   imageSet.ImageId,
			ProjectId: config.ProjectId,
			Region:    config.Region,
		},
	}

	state.Put("ucloud_images", newImageInfoSet(images))
	ui.Message(fmt.Sprintf("Creating image %q complete", imageSet.ImageId))
	return multistep.ActionContinue
}

func (s *stepCreateImage) Cleanup(state multistep.StateBag) {
	if s.image == nil {
		return
	}
	_, cancelled := state.GetOk(multistep.StateCancelled)
	_, halted := state.GetOk(multistep.StateHalted)
	if !cancelled && !halted {
		return
	}

	client := state.Get("client").(*UCloudClient)
	conn := client.uhostconn
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deleting image because of cancellation or error...")
	req := conn.NewTerminateCustomImageRequest()
	req.ImageId = ucloud.String(s.image.ImageId)
	_, err := conn.TerminateCustomImage(req)
	if err != nil {
		ui.Error(fmt.Sprintf("Error on deleting image %q", s.image.ImageId))
	}
	ui.Message(fmt.Sprintf("Deleting image %q complete", s.image.ImageId))
}
<|endoftext|>