<commit_before>package raft

import (
	"fmt"
	"io"
	"sync"
	"time"
)

// NewInmemAddr returns a new in-memory addr with
// a randomly generate UUID as the ID.
func NewInmemAddr() ServerAddress {
	return ServerAddress(generateUUID())
}

// inmemPipeline is used to pipeline requests for the in-mem transport.
type inmemPipeline struct {
	trans    *InmemTransport
	peer     *InmemTransport
	peerAddr ServerAddress

	doneCh       chan AppendFuture
	inprogressCh chan *inmemPipelineInflight

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

type inmemPipelineInflight struct {
	future *appendFuture
	respCh <-chan RPCResponse
}

// InmemTransport Implements the Transport interface, to allow Raft to be
// tested in-memory without going over a network.
type InmemTransport struct {
	sync.RWMutex
	consumerCh chan RPC
	localAddr  ServerAddress
	peers      map[ServerAddress]*InmemTransport
	pipelines  []*inmemPipeline
	timeout    time.Duration
}

// NewInmemTransport is used to initialize a new transport
// and generates a random local address if none is specified
func NewInmemTransport(addr ServerAddress) (ServerAddress, *InmemTransport) {
	if string(addr) == "" {
		addr = NewInmemAddr()
	}
	trans := &InmemTransport{
		consumerCh: make(chan RPC, 16),
		localAddr:  addr,
		peers:      make(map[ServerAddress]*InmemTransport),
		timeout:    50 * time.Millisecond,
	}
	return addr, trans
}

// SetHeartbeatHandler is used to set optional fast-path for
// heartbeats, not supported for this transport.
func (i *InmemTransport) SetHeartbeatHandler(cb func(RPC)) {
}

// Consumer implements the Transport interface.
func (i *InmemTransport) Consumer() <-chan RPC {
	return i.consumerCh
}

// LocalAddr implements the Transport interface.
func (i *InmemTransport) LocalAddr() ServerAddress {
	return i.localAddr
}

// AppendEntriesPipeline returns an interface that can be used to pipeline
// AppendEntries requests.
func (i *InmemTransport) AppendEntriesPipeline(id ServerID, target ServerAddress) (AppendPipeline, error) {
	i.RLock()
	peer, ok := i.peers[target]
	i.RUnlock()
	if !ok {
		return nil, fmt.Errorf("failed to connect to peer: %v", target)
	}
	pipeline := newInmemPipeline(i, peer, target)
	i.Lock()
	i.pipelines = append(i.pipelines, pipeline)
	i.Unlock()
	return pipeline, nil
}

// AppendEntries implements the Transport interface.
func (i *InmemTransport) AppendEntries(id ServerID, target ServerAddress, args *AppendEntriesRequest, resp *AppendEntriesResponse) error {
	rpcResp, err := i.makeRPC(target, args, nil, i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*AppendEntriesResponse)
	*resp = *out
	return nil
}

// RequestVote implements the Transport interface.
func (i *InmemTransport) RequestVote(id ServerID, target ServerAddress, args *RequestVoteRequest, resp *RequestVoteResponse) error {
	rpcResp, err := i.makeRPC(target, args, nil, i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*RequestVoteResponse)
	*resp = *out
	return nil
}

// InstallSnapshot implements the Transport interface.
func (i *InmemTransport) InstallSnapshot(id ServerID, target ServerAddress, args *InstallSnapshotRequest, resp *InstallSnapshotResponse, data io.Reader) error {
	rpcResp, err := i.makeRPC(target, args, data, 10*i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*InstallSnapshotResponse)
	*resp = *out
	return nil
}

func (i *InmemTransport) makeRPC(target ServerAddress, args interface{}, r io.Reader, timeout time.Duration) (rpcResp RPCResponse, err error) {
	i.RLock()
	peer, ok := i.peers[target]
	i.RUnlock()

	if !ok {
		err = fmt.Errorf("failed to connect to peer: %v", target)
		return
	}

	// Send the RPC over
	respCh := make(chan RPCResponse)
	peer.consumerCh <- RPC{
		Command:  args,
		Reader:   r,
		RespChan: respCh,
	}

	// Wait for a response
	select {
	case rpcResp = <-respCh:
		if rpcResp.Error != nil {
			err = rpcResp.Error
		}
	case <-time.After(timeout):
		err = fmt.Errorf("command timed out")
	}
	return
}

// EncodePeer implements the Transport interface.
func (i *InmemTransport) EncodePeer(id ServerID, p ServerAddress) []byte {
	return []byte(p)
}

// DecodePeer implements the Transport interface.
func (i *InmemTransport) DecodePeer(buf []byte) ServerAddress {
	return ServerAddress(buf)
}

// Connect is used to connect this transport to another transport for
// a given peer name. This allows for local routing.
func (i *InmemTransport) Connect(peer ServerAddress, t Transport) {
	trans := t.(*InmemTransport)
	i.Lock()
	defer i.Unlock()
	i.peers[peer] = trans
}

// Disconnect is used to remove the ability to route to a given peer.
func (i *InmemTransport) Disconnect(peer ServerAddress) {
	i.Lock()
	defer i.Unlock()
	delete(i.peers, peer)

	// Disconnect any pipelines
	n := len(i.pipelines)
	for idx := 0; idx < n; idx++ {
		if i.pipelines[idx].peerAddr == peer {
			i.pipelines[idx].Close()
			i.pipelines[idx], i.pipelines[n-1] = i.pipelines[n-1], nil
			idx--
			n--
		}
	}
	i.pipelines = i.pipelines[:n]
}

// DisconnectAll is used to remove all routes to peers.
func (i *InmemTransport) DisconnectAll() {
	i.Lock()
	defer i.Unlock()
	i.peers = make(map[ServerAddress]*InmemTransport)

	// Handle pipelines
	for _, pipeline := range i.pipelines {
		pipeline.Close()
	}
	i.pipelines = nil
}

// Close is used to permanently disable the transport
func (i *InmemTransport) Close() error {
	i.DisconnectAll()
	return nil
}

func newInmemPipeline(trans *InmemTransport, peer *InmemTransport, addr ServerAddress) *inmemPipeline {
	i := &inmemPipeline{
		trans:        trans,
		peer:         peer,
		peerAddr:     addr,
		doneCh:       make(chan AppendFuture, 16),
		inprogressCh: make(chan *inmemPipelineInflight, 16),
		shutdownCh:   make(chan struct{}),
	}
	go i.decodeResponses()
	return i
}

func (i *inmemPipeline) decodeResponses() {
	timeout := i.trans.timeout
	for {
		select {
		case inp := <-i.inprogressCh:
			var timeoutCh <-chan time.Time
			if timeout > 0 {
				timeoutCh = time.After(timeout)
			}

			select {
			case rpcResp := <-inp.respCh:
				// Copy the result back
				*inp.future.resp = *rpcResp.Response.(*AppendEntriesResponse)
				inp.future.respond(rpcResp.Error)

				select {
				case i.doneCh <- inp.future:
				case <-i.shutdownCh:
					return
				}

			case <-timeoutCh:
				inp.future.respond(fmt.Errorf("command timed out"))
				select {
				case i.doneCh <- inp.future:
				case <-i.shutdownCh:
					return
				}

			case <-i.shutdownCh:
				return
			}
		case <-i.shutdownCh:
			return
		}
	}
}

func (i *inmemPipeline) AppendEntries(args *AppendEntriesRequest, resp *AppendEntriesResponse) (AppendFuture, error) {
	// Create a new future
	future := &appendFuture{
		start: time.Now(),
		args:  args,
		resp:  resp,
	}
	future.init()

	// Handle a timeout
	var timeout <-chan time.Time
	if i.trans.timeout > 0 {
		timeout = time.After(i.trans.timeout)
	}

	// Send the RPC over
	respCh := make(chan RPCResponse, 1)
	rpc := RPC{
		Command:  args,
		RespChan: respCh,
	}
	select {
	case i.peer.consumerCh <- rpc:
	case <-timeout:
		return nil, fmt.Errorf("command enqueue timeout")
	case <-i.shutdownCh:
		return nil, ErrPipelineShutdown
	}

	// Send to be decoded
	select {
	case i.inprogressCh <- &inmemPipelineInflight{future, respCh}:
		return future, nil
	case <-i.shutdownCh:
		return nil, ErrPipelineShutdown
	}
}

func (i *inmemPipeline) Consumer() <-chan AppendFuture {
	return i.doneCh
}

func (i *inmemPipeline) Close() error {
	i.shutdownLock.Lock()
	defer i.shutdownLock.Unlock()
	if i.shutdown {
		return nil
	}

	i.shutdown = true
	close(i.shutdownCh)
	return nil
}
<commit_msg>Fix race between AppendEntriesPipeline and Disconnect in InmemTransport<commit_after>package raft

import (
	"fmt"
	"io"
	"sync"
	"time"
)

// NewInmemAddr returns a new in-memory addr with
// a randomly generate UUID as the ID.
func NewInmemAddr() ServerAddress {
	return ServerAddress(generateUUID())
}

// inmemPipeline is used to pipeline requests for the in-mem transport.
type inmemPipeline struct {
	trans    *InmemTransport
	peer     *InmemTransport
	peerAddr ServerAddress

	doneCh       chan AppendFuture
	inprogressCh chan *inmemPipelineInflight

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

type inmemPipelineInflight struct {
	future *appendFuture
	respCh <-chan RPCResponse
}

// InmemTransport Implements the Transport interface, to allow Raft to be
// tested in-memory without going over a network.
type InmemTransport struct {
	sync.RWMutex
	consumerCh chan RPC
	localAddr  ServerAddress
	peers      map[ServerAddress]*InmemTransport
	pipelines  []*inmemPipeline
	timeout    time.Duration
}

// NewInmemTransport is used to initialize a new transport
// and generates a random local address if none is specified
func NewInmemTransport(addr ServerAddress) (ServerAddress, *InmemTransport) {
	if string(addr) == "" {
		addr = NewInmemAddr()
	}
	trans := &InmemTransport{
		consumerCh: make(chan RPC, 16),
		localAddr:  addr,
		peers:      make(map[ServerAddress]*InmemTransport),
		timeout:    50 * time.Millisecond,
	}
	return addr, trans
}

// SetHeartbeatHandler is used to set optional fast-path for
// heartbeats, not supported for this transport.
func (i *InmemTransport) SetHeartbeatHandler(cb func(RPC)) {
}

// Consumer implements the Transport interface.
func (i *InmemTransport) Consumer() <-chan RPC {
	return i.consumerCh
}

// LocalAddr implements the Transport interface.
func (i *InmemTransport) LocalAddr() ServerAddress {
	return i.localAddr
}

// AppendEntriesPipeline returns an interface that can be used to pipeline
// AppendEntries requests.
func (i *InmemTransport) AppendEntriesPipeline(id ServerID, target ServerAddress) (AppendPipeline, error) {
	i.Lock()
	defer i.Unlock()

	peer, ok := i.peers[target]
	if !ok {
		return nil, fmt.Errorf("failed to connect to peer: %v", target)
	}
	pipeline := newInmemPipeline(i, peer, target)
	i.pipelines = append(i.pipelines, pipeline)
	return pipeline, nil
}

// AppendEntries implements the Transport interface.
func (i *InmemTransport) AppendEntries(id ServerID, target ServerAddress, args *AppendEntriesRequest, resp *AppendEntriesResponse) error {
	rpcResp, err := i.makeRPC(target, args, nil, i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*AppendEntriesResponse)
	*resp = *out
	return nil
}

// RequestVote implements the Transport interface.
func (i *InmemTransport) RequestVote(id ServerID, target ServerAddress, args *RequestVoteRequest, resp *RequestVoteResponse) error {
	rpcResp, err := i.makeRPC(target, args, nil, i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*RequestVoteResponse)
	*resp = *out
	return nil
}

// InstallSnapshot implements the Transport interface.
func (i *InmemTransport) InstallSnapshot(id ServerID, target ServerAddress, args *InstallSnapshotRequest, resp *InstallSnapshotResponse, data io.Reader) error {
	rpcResp, err := i.makeRPC(target, args, data, 10*i.timeout)
	if err != nil {
		return err
	}

	// Copy the result back
	out := rpcResp.Response.(*InstallSnapshotResponse)
	*resp = *out
	return nil
}

func (i *InmemTransport) makeRPC(target ServerAddress, args interface{}, r io.Reader, timeout time.Duration) (rpcResp RPCResponse, err error) {
	i.RLock()
	peer, ok := i.peers[target]
	i.RUnlock()

	if !ok {
		err = fmt.Errorf("failed to connect to peer: %v", target)
		return
	}

	// Send the RPC over
	respCh := make(chan RPCResponse)
	peer.consumerCh <- RPC{
		Command:  args,
		Reader:   r,
		RespChan: respCh,
	}

	// Wait for a response
	select {
	case rpcResp = <-respCh:
		if rpcResp.Error != nil {
			err = rpcResp.Error
		}
	case <-time.After(timeout):
		err = fmt.Errorf("command timed out")
	}
	return
}

// EncodePeer implements the Transport interface.
func (i *InmemTransport) EncodePeer(id ServerID, p ServerAddress) []byte {
	return []byte(p)
}

// DecodePeer implements the Transport interface.
func (i *InmemTransport) DecodePeer(buf []byte) ServerAddress {
	return ServerAddress(buf)
}

// Connect is used to connect this transport to another transport for
// a given peer name. This allows for local routing.
func (i *InmemTransport) Connect(peer ServerAddress, t Transport) {
	trans := t.(*InmemTransport)
	i.Lock()
	defer i.Unlock()
	i.peers[peer] = trans
}

// Disconnect is used to remove the ability to route to a given peer.
func (i *InmemTransport) Disconnect(peer ServerAddress) {
	i.Lock()
	defer i.Unlock()
	delete(i.peers, peer)

	// Disconnect any pipelines
	n := len(i.pipelines)
	for idx := 0; idx < n; idx++ {
		if i.pipelines[idx].peerAddr == peer {
			i.pipelines[idx].Close()
			i.pipelines[idx], i.pipelines[n-1] = i.pipelines[n-1], nil
			idx--
			n--
		}
	}
	i.pipelines = i.pipelines[:n]
}

// DisconnectAll is used to remove all routes to peers.
func (i *InmemTransport) DisconnectAll() {
	i.Lock()
	defer i.Unlock()
	i.peers = make(map[ServerAddress]*InmemTransport)

	// Handle pipelines
	for _, pipeline := range i.pipelines {
		pipeline.Close()
	}
	i.pipelines = nil
}

// Close is used to permanently disable the transport
func (i *InmemTransport) Close() error {
	i.DisconnectAll()
	return nil
}

func newInmemPipeline(trans *InmemTransport, peer *InmemTransport, addr ServerAddress) *inmemPipeline {
	i := &inmemPipeline{
		trans:        trans,
		peer:         peer,
		peerAddr:     addr,
		doneCh:       make(chan AppendFuture, 16),
		inprogressCh: make(chan *inmemPipelineInflight, 16),
		shutdownCh:   make(chan struct{}),
	}
	go i.decodeResponses()
	return i
}

func (i *inmemPipeline) decodeResponses() {
	timeout := i.trans.timeout
	for {
		select {
		case inp := <-i.inprogressCh:
			var timeoutCh <-chan time.Time
			if timeout > 0 {
				timeoutCh = time.After(timeout)
			}

			select {
			case rpcResp := <-inp.respCh:
				// Copy the result back
				*inp.future.resp = *rpcResp.Response.(*AppendEntriesResponse)
				inp.future.respond(rpcResp.Error)

				select {
				case i.doneCh <- inp.future:
				case <-i.shutdownCh:
					return
				}

			case <-timeoutCh:
				inp.future.respond(fmt.Errorf("command timed out"))
				select {
				case i.doneCh <- inp.future:
				case <-i.shutdownCh:
					return
				}

			case <-i.shutdownCh:
				return
			}
		case <-i.shutdownCh:
			return
		}
	}
}

func (i *inmemPipeline) AppendEntries(args *AppendEntriesRequest, resp *AppendEntriesResponse) (AppendFuture, error) {
	// Create a new future
	future := &appendFuture{
		start: time.Now(),
		args:  args,
		resp:  resp,
	}
	future.init()

	// Handle a timeout
	var timeout <-chan time.Time
	if i.trans.timeout > 0 {
		timeout = time.After(i.trans.timeout)
	}

	// Send the RPC over
	respCh := make(chan RPCResponse, 1)
	rpc := RPC{
		Command:  args,
		RespChan: respCh,
	}
	select {
	case i.peer.consumerCh <- rpc:
	case <-timeout:
		return nil, fmt.Errorf("command enqueue timeout")
	case <-i.shutdownCh:
		return nil, ErrPipelineShutdown
	}

	// Send to be decoded
	select {
	case i.inprogressCh <- &inmemPipelineInflight{future, respCh}:
		return future, nil
	case <-i.shutdownCh:
		return nil, ErrPipelineShutdown
	}
}

func (i *inmemPipeline) Consumer() <-chan AppendFuture {
	return i.doneCh
}

func (i *inmemPipeline) Close() error {
	i.shutdownLock.Lock()
	defer i.shutdownLock.Unlock()
	if i.shutdown {
		return nil
	}

	i.shutdown = true
	close(i.shutdownCh)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/raintank/raintank-metric/dur"
	"github.com/raintank/raintank-metric/metricdef"
	"github.com/raintank/raintank-metric/schema"
	"log"
	"os"
	"time"
)

func perror(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var esAddr = flag.String("es-addr", "localhost:9200", "elasticsearch address")
var esIndex = flag.String("es-index", "metric", "elasticsearch index to query")
var format = flag.String("format", "list", "format: list|vegeta-graphite|vegeta-mt")
var maxAge = flag.Int("max-age", 23400, "max age (last update diff with now) of metricdefs. defaults to 6.5hr. use 0 to disable")
var from = flag.String("from", "30min", "from. eg '30min', '5h', '14d', etc")
var silent = flag.Bool("silent", false, "silent mode (don't print number of metrics loaded to stderr)")
var fromS uint32
var total int

func showList(ds []*schema.MetricDefinition) {
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Println(d.OrgId, d.Name)
		}
	}
}
func showVegetaGraphite(ds []*schema.MetricDefinition) {
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Printf("GET http://localhost:8888/render?target=%s&from=-%s\nX-Org-Id: %d\n\n", d.Name, *from, d.OrgId)
		}
	}
}
func showVegetaMT(ds []*schema.MetricDefinition) {
	from := time.Now().Add(-time.Duration(fromS) * time.Second)
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Printf("GET http://localhost:18763/get?target=%s&from=%d\n", d.Id, from.Unix())
		}
	}
}

func main() {
	flag.Parse()
	var show func(ds []*schema.MetricDefinition)
	switch *format {
	case "list":
		show = showList
	case "vegeta-graphite":
		show = showVegetaGraphite
	case "vegeta-mt":
		show = showVegetaMT
	default:
		log.Fatal("invalid format")
	}
	var err error
	fromS, err = dur.ParseUNsec(*from)
	perror(err)
	defs, err := metricdef.NewDefsEs(*esAddr, "", "", *esIndex)
	perror(err)
	met, scroll_id, err := defs.GetMetrics("")
	perror(err)
	show(met)
	for scroll_id != "" {
		met, scroll_id, err = defs.GetMetrics(scroll_id)
		perror(err)
		show(met)
	}
	if !*silent {
		fmt.Fprintf(os.Stderr, "listed %d metrics\n", total)
	}
}
<commit_msg>update inspect-es for new graphite 'native' format for MT<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/raintank/raintank-metric/dur"
	"github.com/raintank/raintank-metric/metricdef"
	"github.com/raintank/raintank-metric/schema"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"
)

func perror(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var esAddr = flag.String("es-addr", "localhost:9200", "elasticsearch address")
var esIndex = flag.String("es-index", "metric", "elasticsearch index to query")
var format = flag.String("format", "list", "format: list|vegeta-graphite|vegeta-mt|vegeta-mt-graphite")
var maxAge = flag.Int("max-age", 23400, "max age (last update diff with now) of metricdefs. defaults to 6.5hr. use 0 to disable")
var from = flag.String("from", "30min", "from. eg '30min', '5h', '14d', etc")
var silent = flag.Bool("silent", false, "silent mode (don't print number of metrics loaded to stderr)")
var fromS uint32
var total int

func showList(ds []*schema.MetricDefinition) {
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Println(d.OrgId, d.Name)
		}
	}
}
func showVegetaGraphite(ds []*schema.MetricDefinition) {
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Printf("GET http://localhost:8888/render?target=%s&from=-%s\nX-Org-Id: %d\n\n", d.Name, *from, d.OrgId)
		}
	}
}

func showVegetaMT(ds []*schema.MetricDefinition) {
	from := time.Now().Add(-time.Duration(fromS) * time.Second)
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			fmt.Printf("GET http://localhost:6063/get?target=%s&from=%d\n", d.Id, from.Unix())
		}
	}
}

func showVegetaMTGraphite(ds []*schema.MetricDefinition) {
	for _, d := range ds {
		if *maxAge != 0 && d.LastUpdate > time.Now().Unix()-int64(*maxAge) {
			total += 1
			mode := rand.Intn(3)
			name := d.Name
			if mode == 0 {
				// in this mode, replaces a node with a dot
				which := rand.Intn(d.NodeCount)
				parts := strings.Split(d.Name, ".")
				parts[which] = "*"
				name = strings.Join(parts, ".")
			} else if mode == 1 {
				// randomly replace chars with a *
				// note that in 1/5 cases, nothing happens
				// and otherwise, sometimes valid patterns are produced,
				// but it's also possible to produce patterns that won't match anything (if '.' was taken out)
				chars := rand.Intn(5)
				pos := rand.Intn(len(d.Name) - chars)
				name = name[0:pos] + "*" + name[pos+chars:]
			}
			// mode 3: do nothing :)

			fmt.Printf("GET http://localhost:6063/render?target=%s&from=-%s\nX-Org-Id: %d\n\n", name, *from, d.OrgId)
		}
	}
}

func main() {
	flag.Parse()
	var show func(ds []*schema.MetricDefinition)
	switch *format {
	case "list":
		show = showList
	case "vegeta-graphite":
		show = showVegetaGraphite
	case "vegeta-mt":
		show = showVegetaMT
	case "vegeta-mt-graphite":
		show = showVegetaMTGraphite
	default:
		log.Fatal("invalid format")
	}
	var err error
	fromS, err = dur.ParseUNsec(*from)
	perror(err)
	defs, err := metricdef.NewDefsEs(*esAddr, "", "", *esIndex)
	perror(err)
	met, scroll_id, err := defs.GetMetrics("")
	perror(err)
	show(met)
	for scroll_id != "" {
		met, scroll_id, err = defs.GetMetrics(scroll_id)
		perror(err)
		show(met)
	}
	if !*silent {
		fmt.Fprintf(os.Stderr, "listed %d metrics\n", total)
	}
}
<|endoftext|>


<commit_before>package internal

import (
	"context"

	"strings"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach-go/v2/testserver"

	"github.com/ory/x/configx"

	"github.com/stretchr/testify/require"

	"github.com/ory/hydra/x"
	"github.com/ory/x/contextx"
	"github.com/ory/x/sqlcon/dockertest"

	"github.com/ory/hydra/driver"
	"github.com/ory/hydra/driver/config"
	"github.com/ory/hydra/jwk"
	"github.com/ory/x/logrusx"
)

func resetConfig(p *config.DefaultProvider) {
	p.MustSet(context.Background(), config.KeyBCryptCost, "4")
	p.MustSet(context.Background(), config.KeySubjectIdentifierAlgorithmSalt, "00000000")
	p.MustSet(context.Background(), config.KeyGetSystemSecret, []string{"000000000000000000000000000000000000000000000000"})
	p.MustSet(context.Background(), config.KeyGetCookieSecrets, []string{"000000000000000000000000000000000000000000000000"})
	p.MustSet(context.Background(), config.KeyLogLevel, "trace")
}

func NewConfigurationWithDefaults() *config.DefaultProvider {
	p := config.MustNew(context.Background(), logrusx.New("", ""), configx.SkipValidation())
	resetConfig(p)
	p.MustSet(context.Background(), "dangerous-force-http", true)
	return p
}

func NewConfigurationWithDefaultsAndHTTPS() *config.DefaultProvider {
	p := config.MustNew(context.Background(), logrusx.New("", ""), configx.SkipValidation())
	resetConfig(p)
	p.MustSet(context.Background(), "dangerous-force-http", false)
	return p
}

func NewRegistryMemory(t *testing.T, c *config.DefaultProvider, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, "memory", c, true, ctxer)
}

func NewMockedRegistry(t *testing.T, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, "memory", NewConfigurationWithDefaults(), true, ctxer)
}

func NewRegistrySQLFromURL(t *testing.T, url string, migrate bool, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, url, NewConfigurationWithDefaults(), migrate, ctxer)
}

func newRegistryDefault(t *testing.T, url string, c *config.DefaultProvider, migrate bool, ctxer contextx.Contextualizer) driver.Registry {
	ctx := context.Background()
	c.MustSet(ctx, config.KeyLogLevel, "trace")
	c.MustSet(ctx, config.KeyDSN, url)

	r, err := driver.NewRegistryFromDSN(ctx, c, logrusx.New("test_hydra", "master"), false, migrate, ctxer)
	require.NoError(t, err)

	return r
}

func CleanAndMigrate(reg driver.Registry) func(*testing.T) {
	return func(t *testing.T) {
		x.CleanSQLPop(t, reg.Persister().Connection(context.Background()))
		require.NoError(t, reg.Persister().MigrateUp(context.Background()))
		t.Log("clean and migrate done")
	}
}

func ConnectToMySQL(t *testing.T) string {
	c := dockertest.ConnectToTestMySQLPop(t)
	url := c.URL()
	if !strings.HasPrefix(url, "mysql://") {
		url = "mysql://" + url
	}
	require.NoError(t, c.Close())
	return url
}

func ConnectToPG(t *testing.T) string {
	c := dockertest.ConnectToTestPostgreSQLPop(t)
	require.NoError(t, c.Close())
	return c.URL()
}

func ConnectToCRDB(t *testing.T) string {
	c := dockertest.ConnectToTestCockroachDBPop(t)
	url := c.URL()
	if !strings.HasPrefix(url, "cockroach") {
		url = "cockroach://" + strings.Split(url, "://")[1]
	}
	require.NoError(t, c.Close())
	return url
}

func NewEmptyTestCRDBServer(t *testing.T) string {
	ts, err := testserver.NewTestServer()
	require.NoError(t, err)
	t.Cleanup(ts.Stop)

	ts.PGURL().Scheme = "cockroach"
	return ts.PGURL().String()
}

func ConnectDatabases(t *testing.T, migrate bool, ctxer contextx.Contextualizer) (pg, mysql, crdb driver.Registry, clean func(*testing.T)) {
	var pgURL, mysqlURL, crdbURL string
	wg := sync.WaitGroup{}

	wg.Add(3)
	go func() {
		pgURL = ConnectToPG(t)
		t.Log("Pg done")
		wg.Done()
	}()
	go func() {
		mysqlURL = ConnectToMySQL(t)
		t.Log("myssql done")
		wg.Done()
	}()
	go func() {
		crdbURL = NewEmptyTestCRDBServer(t)
		t.Log("crdb done")
		wg.Done()
	}()
	t.Log("beginning to wait")
	wg.Wait()
	t.Log("done waiting")

	pg = NewRegistrySQLFromURL(t, pgURL, migrate, ctxer)
	mysql = NewRegistrySQLFromURL(t, mysqlURL, migrate, ctxer)
	crdb = NewRegistrySQLFromURL(t, crdbURL, migrate, ctxer)
	dbs := []driver.Registry{pg, mysql, crdb}

	clean = func(t *testing.T) {
		wg := sync.WaitGroup{}

		wg.Add(len(dbs))
		for _, db := range dbs {
			go func(db driver.Registry) {
				defer wg.Done()
				CleanAndMigrate(db)(t)
			}(db)
		}
		wg.Wait()
	}
	clean(t)
	return
}

func MustEnsureRegistryKeys(r driver.Registry, key string) {
	if err := jwk.EnsureAsymmetricKeypairExists(context.Background(), r, "RS256", key); err != nil {
		panic(err)
	}
}
<commit_msg>fix: dont close crdb for reuse purposes<commit_after>package internal

import (
	"context"

	"strings"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach-go/v2/testserver"

	"github.com/ory/x/configx"

	"github.com/stretchr/testify/require"

	"github.com/ory/hydra/x"
	"github.com/ory/x/contextx"
	"github.com/ory/x/sqlcon/dockertest"

	"github.com/ory/hydra/driver"
	"github.com/ory/hydra/driver/config"
	"github.com/ory/hydra/jwk"
	"github.com/ory/x/logrusx"
)

func resetConfig(p *config.DefaultProvider) {
	p.MustSet(context.Background(), config.KeyBCryptCost, "4")
	p.MustSet(context.Background(), config.KeySubjectIdentifierAlgorithmSalt, "00000000")
	p.MustSet(context.Background(), config.KeyGetSystemSecret, []string{"000000000000000000000000000000000000000000000000"})
	p.MustSet(context.Background(), config.KeyGetCookieSecrets, []string{"000000000000000000000000000000000000000000000000"})
	p.MustSet(context.Background(), config.KeyLogLevel, "trace")
}

func NewConfigurationWithDefaults() *config.DefaultProvider {
	p := config.MustNew(context.Background(), logrusx.New("", ""), configx.SkipValidation())
	resetConfig(p)
	p.MustSet(context.Background(), "dangerous-force-http", true)
	return p
}

func NewConfigurationWithDefaultsAndHTTPS() *config.DefaultProvider {
	p := config.MustNew(context.Background(), logrusx.New("", ""), configx.SkipValidation())
	resetConfig(p)
	p.MustSet(context.Background(), "dangerous-force-http", false)
	return p
}

func NewRegistryMemory(t *testing.T, c *config.DefaultProvider, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, "memory", c, true, ctxer)
}

func NewMockedRegistry(t *testing.T, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, "memory", NewConfigurationWithDefaults(), true, ctxer)
}

func NewRegistrySQLFromURL(t *testing.T, url string, migrate bool, ctxer contextx.Contextualizer) driver.Registry {
	return newRegistryDefault(t, url, NewConfigurationWithDefaults(), migrate, ctxer)
}

func newRegistryDefault(t *testing.T, url string, c *config.DefaultProvider, migrate bool, ctxer contextx.Contextualizer) driver.Registry {
	ctx := context.Background()
	c.MustSet(ctx, config.KeyLogLevel, "trace")
	c.MustSet(ctx, config.KeyDSN, url)

	r, err := driver.NewRegistryFromDSN(ctx, c, logrusx.New("test_hydra", "master"), false, migrate, ctxer)
	require.NoError(t, err)

	return r
}

func CleanAndMigrate(reg driver.Registry) func(*testing.T) {
	return func(t *testing.T) {
		x.CleanSQLPop(t, reg.Persister().Connection(context.Background()))
		require.NoError(t, reg.Persister().MigrateUp(context.Background()))
		t.Log("clean and migrate done")
	}
}

func ConnectToMySQL(t *testing.T) string {
	c := dockertest.ConnectToTestMySQLPop(t)
	url := c.URL()
	if !strings.HasPrefix(url, "mysql://") {
		url = "mysql://" + url
	}
	require.NoError(t, c.Close())
	return url
}

func ConnectToPG(t *testing.T) string {
	c := dockertest.ConnectToTestPostgreSQLPop(t)
	require.NoError(t, c.Close())
	return c.URL()
}

func ConnectToCRDB(t *testing.T) string {
	c := dockertest.ConnectToTestCockroachDBPop(t)
	url := c.URL()
	if !strings.HasPrefix(url, "cockroach") {
		url = "cockroach://" + strings.Split(url, "://")[1]
	}
	require.NoError(t, c.Close())
	return url
}

func NewEmptyTestCRDBServer(t *testing.T) string {
	ts, err := testserver.NewTestServer()
	require.NoError(t, err)

	ts.PGURL().Scheme = "cockroach"
	return ts.PGURL().String()
}

func ConnectDatabases(t *testing.T, migrate bool, ctxer contextx.Contextualizer) (pg, mysql, crdb driver.Registry, clean func(*testing.T)) {
	var pgURL, mysqlURL, crdbURL string
	wg := sync.WaitGroup{}

	wg.Add(3)
	go func() {
		pgURL = ConnectToPG(t)
		t.Log("Pg done")
		wg.Done()
	}()
	go func() {
		mysqlURL = ConnectToMySQL(t)
		t.Log("myssql done")
		wg.Done()
	}()
	go func() {
		crdbURL = NewEmptyTestCRDBServer(t)
		t.Log("crdb done")
		wg.Done()
	}()
	t.Log("beginning to wait")
	wg.Wait()
	t.Log("done waiting")

	pg = NewRegistrySQLFromURL(t, pgURL, migrate, ctxer)
	mysql = NewRegistrySQLFromURL(t, mysqlURL, migrate, ctxer)
	crdb = NewRegistrySQLFromURL(t, crdbURL, migrate, ctxer)
	dbs := []driver.Registry{pg, mysql, crdb}

	clean = func(t *testing.T) {
		wg := sync.WaitGroup{}

		wg.Add(len(dbs))
		for _, db := range dbs {
			go func(db driver.Registry) {
				defer wg.Done()
				CleanAndMigrate(db)(t)
			}(db)
		}
		wg.Wait()
	}
	clean(t)
	return
}

func MustEnsureRegistryKeys(r driver.Registry, key string) {
	if err := jwk.EnsureAsymmetricKeypairExists(context.Background(), r, "RS256", key); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package chlib

import (
	"fmt"
	"os/user"
	"path"

	"github.com/boltdb/bolt"
	"github.com/kfeofantov/chkit-v2/helpers"
)

const configDir = ".containerum"
const configFile = "config.db"

var configDb *bolt.DB

var initializers map[string]helpers.MappedStruct = make(map[string]helpers.MappedStruct)

func init() {
	currentUser, err := user.Current()
	if err != nil {
		panic(fmt.Errorf("get current user: %s", err))
	}

	configPath := path.Join(currentUser.HomeDir, configDir, configFile)
	configDb, err = bolt.Open(configPath, 0600, nil)
	if err != nil {
		panic(fmt.Errorf("config db open: %s", err))
	}
	if err := runInitializers(); err != nil {
		panic(err)
	}
}

func recursiveInitialize(startBucket *bolt.Bucket, initializer helpers.MappedStruct) error {
	for key, value := range initializer {
		switch value.(type) {
		case []byte:
			err := startBucket.Put([]byte(key), value.([]byte))
			if err != nil {
				return err
			}
		case helpers.MappedStruct:
			newBucket, err := startBucket.CreateBucketIfNotExists([]byte(key))
			if err != nil {
				return fmt.Errorf("bucket create: %s", err)
			}
			err = recursiveInitialize(newBucket, value.(helpers.MappedStruct))
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func runInitializers() error {
	return configDb.Update(func(tx *bolt.Tx) error {
		for bucketName, initializer := range initializers {
			bucket := tx.Bucket([]byte(bucketName))
			if bucket == nil {
				bucket, err := tx.CreateBucket([]byte(bucketName))
				if err != nil {
					return fmt.Errorf("bucket create: %s", err)
				}
				recursiveInitialize(bucket, initializer)
			}
		}
		return nil
	})
}

func recursivePush(startBucket *bolt.Bucket, data helpers.MappedStruct) error {
	for key, value := range data {
		var err error
		switch value.(type) {
		case []byte:
			err = startBucket.Put([]byte(key), value.([]byte))
		case helpers.MappedStruct:
			newBucket, err := startBucket.CreateBucket([]byte(key))
			if err != nil {
				return err
			}
			err = recursivePush(newBucket, value.(helpers.MappedStruct))
		}
		if err != nil {
			return fmt.Errorf("value for key %s push: %s", key, err)
		}
	}
	return nil
}

func pushToBucket(bucketName string, data helpers.MappedStruct) error {
	return configDb.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(bucketName))
		return recursivePush(bucket, data)
	})
}

func recursiveRead(startBucket *bolt.Bucket) (data helpers.MappedStruct, err error) {
	data = make(helpers.MappedStruct)
	err = startBucket.ForEach(func(k []byte, v []byte) error {
		if v != nil {
			data[string(k)] = v
		} else {
			// it`s a bucket
			newBucket := startBucket.Bucket(k)
			var err error
			data[string(k)], err = recursiveRead(newBucket)
			if err != nil {
				return err
			}
		}
		return nil
	})
	return
}

func readFromBucket(bucketName string) (data helpers.MappedStruct, err error) {
	err = configDb.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(bucketName))
		if bucket == nil {
			return fmt.Errorf("bucket %s not exist", bucketName)
		}
		data, err = recursiveRead(bucket)
		return err
	})
	return
}

func Close() {
	configDb.Close()
}
<commit_msg>expose config directory path to package<commit_after>package chlib

import (
	"fmt"
	"os/user"
	"path"

	"github.com/boltdb/bolt"
	"github.com/kfeofantov/chkit-v2/helpers"
)

const configDir = ".containerum"
const configFile = "config.db"

var configPath string

var configDb *bolt.DB

var initializers map[string]helpers.MappedStruct = make(map[string]helpers.MappedStruct)

func init() {
	currentUser, err := user.Current()
	if err != nil {
		panic(fmt.Errorf("get current user: %s", err))
	}

	configPath = path.Join(currentUser.HomeDir, configDir)
	configDb, err = bolt.Open(path.Join(configPath, configFile), 0600, nil)
	if err != nil {
		panic(fmt.Errorf("config db open: %s", err))
	}
	if err := runInitializers(); err != nil {
		panic(err)
	}
}

func recursiveInitialize(startBucket *bolt.Bucket, initializer helpers.MappedStruct) error {
	for key, value := range initializer {
		switch value.(type) {
		case []byte:
			err := startBucket.Put([]byte(key), value.([]byte))
			if err != nil {
				return err
			}
		case helpers.MappedStruct:
			newBucket, err := startBucket.CreateBucketIfNotExists([]byte(key))
			if err != nil {
				return fmt.Errorf("bucket create: %s", err)
			}
			err = recursiveInitialize(newBucket, value.(helpers.MappedStruct))
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func runInitializers() error {
	return configDb.Update(func(tx *bolt.Tx) error {
		for bucketName, initializer := range initializers {
			bucket := tx.Bucket([]byte(bucketName))
			if bucket == nil {
				bucket, err := tx.CreateBucket([]byte(bucketName))
				if err != nil {
					return fmt.Errorf("bucket create: %s", err)
				}
				recursiveInitialize(bucket, initializer)
			}
		}
		return nil
	})
}

func recursivePush(startBucket *bolt.Bucket, data helpers.MappedStruct) error {
	for key, value := range data {
		var err error
		switch value.(type) {
		case []byte:
			err = startBucket.Put([]byte(key), value.([]byte))
		case helpers.MappedStruct:
			newBucket, err := startBucket.CreateBucket([]byte(key))
			if err != nil {
				return err
			}
			err = recursivePush(newBucket, value.(helpers.MappedStruct))
		}
		if err != nil {
			return fmt.Errorf("value for key %s push: %s", key, err)
		}
	}
	return nil
}

func pushToBucket(bucketName string, data helpers.MappedStruct) error {
	return configDb.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(bucketName))
		return recursivePush(bucket, data)
	})
}

func recursiveRead(startBucket *bolt.Bucket) (data helpers.MappedStruct, err error) {
	data = make(helpers.MappedStruct)
	err = startBucket.ForEach(func(k []byte, v []byte) error {
		if v != nil {
			data[string(k)] = v
		} else {
			// it`s a bucket
			newBucket := startBucket.Bucket(k)
			var err error
			data[string(k)], err = recursiveRead(newBucket)
			if err != nil {
				return err
			}
		}
		return nil
	})
	return
}

func readFromBucket(bucketName string) (data helpers.MappedStruct, err error) {
	err = configDb.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(bucketName))
		if bucket == nil {
			return fmt.Errorf("bucket %s not exist", bucketName)
		}
		data, err = recursiveRead(bucket)
		return err
	})
	return
}

func Close() {
	configDb.Close()
}
<|endoftext|>


<commit_before>// Copyright 2018 NetApp, Inc. All Rights Reserved.

package cmd

import (
	"archive/zip"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/netapp/trident/config"
	"github.com/spf13/cobra"
)

const (
	LogLimitBytes = 10485760 // 10 MiB

	logNameTrident         = "trident"
	logNameTridentPrevious = "trident-previous"
	logNameEtcd            = "etcd"
	logNameEtcdPrevious    = "etcd-previous"

	logTypeAuto    = "auto"
	logTypeTrident = "trident"
	logTypeEtcd    = "etcd"
	logTypeAll     = "all"

	archiveFilenameFormat = "support-2006-01-02T15-04-05-MST.zip"
)

var (
	logType  string
	archive  bool
	previous bool
)

func init() {
	RootCmd.AddCommand(logsCmd)
	logsCmd.Flags().StringVarP(&logType, "log", "l", logTypeAuto, "Trident log to display. One of trident|etcd|auto|all")
	logsCmd.Flags().BoolVarP(&archive, "archive", "a", false, "Create a support archive with all logs unless otherwise specified.")
	logsCmd.Flags().BoolVarP(&previous, "previous", "p", false, "Get the logs for the previous container instance if it exists.")
}

var logsCmd = &cobra.Command{
	Use:   "logs",
	Short: "Print the logs from Trident",
	Long:  "Print the logs from the Trident storage orchestrator for Kubernetes",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		err := discoverOperatingMode(cmd)
		return err
	},
	RunE: func(cmd *cobra.Command, args []string) error {

		err := checkValidLog()
		if err != nil {
			return err
		}

		if archive {
			return archiveLogs()
		} else {
			return consoleLogs()
		}
	},
}

func archiveLogs() error {

	// In archive mode, "auto" means to attempt to get all logs (current & previous).
	if logType == logTypeAuto {
		logType = logTypeAll
		previous = true
	}

	logMap := make(map[string][]byte)
	getLogs(logMap)

	// If there aren't any logs, bail out.
	anyLogs := false
	for log := range logMap {
		if log != "error" {
			anyLogs = true
			break
		}
	}
	if !anyLogs {
		return errors.New("no Trident-related logs found")
	}

	// Create archive file.
	zipFileName := time.Now().Format(archiveFilenameFormat)
	zipFile, err := os.Create(zipFileName)
	if err != nil {
		return err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	// Write to the archive file.
	for log, logBytes := range logMap {
		entry, err := zipWriter.Create(log)
		if err != nil {
			return err
		}
		_, err = entry.Write(logBytes)
		if err != nil {
			return err
		}
		fmt.Printf("Wrote %s log to %s archive file.\n", log, zipFileName)
	}

	return nil
}

func consoleLogs() error {

	logMap := make(map[string][]byte)
	err := getLogs(logMap)

	SetExitCodeFromError(err)
	if err != nil {
		// Preserve anything written to stdout/stderr
		logMessage := strings.TrimSuffix(strings.TrimSpace(string(logMap["error"])), ".")
		if len(logMessage) > 0 {
			errMessage := strings.TrimSuffix(strings.TrimSpace(err.Error()), ".")
			return fmt.Errorf("%s. %s", errMessage, logMessage)
		} else {
			return err
		}
	}

	// Print to the console
	anyLogs := false
	for log, logBytes := range logMap {
		if log != "error" {
			fmt.Printf("%s log:\n", log)
			fmt.Printf("%s\n", string(logBytes))
			anyLogs = true
		}
	}
	if !anyLogs {
		return errors.New("no Trident-related logs found")
	}

	return nil
}

func getLogs(logMap map[string][]byte) error {

	var err error

	if OperatingMode != ModeTunnel {
		return errors.New("'tridentctl logs' only supports Trident running in a Kubernetes pod")
	}

	switch logType {
	case logTypeTrident, logTypeAuto:
		err = getTridentLogs(logNameTrident, logMap)
	case logTypeEtcd:
		err = getTridentLogs(logNameEtcd, logMap)
	case logTypeAll:
		getTridentLogs(logNameTrident, logMap)
		getTridentLogs(logNameEtcd, logMap)
	}

	if previous {
		switch logType {
		case logTypeTrident, logTypeAuto:
			getTridentLogs(logNameTridentPrevious, logMap)
		case logTypeEtcd:
			getTridentLogs(logNameEtcdPrevious, logMap)
		case logTypeAll:
			getTridentLogs(logNameTridentPrevious, logMap)
			getTridentLogs(logNameEtcdPrevious, logMap)
		}
	}

	return err
}

func checkValidLog() error {
	switch logType {
	case logTypeTrident, logTypeEtcd, logTypeAuto, logTypeAll:
		return nil
	default:
		return fmt.Errorf("%s is not a valid Trident log", logType)
	}
}

func getTridentLogs(logName string, logMap map[string][]byte) error {

	var container string
	var prev bool

	switch logName {
	case logNameTrident:
		container, prev = config.ContainerTrident, false
	case logNameTridentPrevious:
		container, prev = config.ContainerTrident, true
	case logNameEtcd:
		container, prev = config.ContainerEtcd, false
	case logNameEtcdPrevious:
		container, prev = config.ContainerEtcd, true
	default:
		return fmt.Errorf("%s is not a valid Trident log", logName)
	}

	// Build command to get K8S logs
	limitArg := fmt.Sprintf("--limit-bytes=%d", LogLimitBytes)
	prevArg := fmt.Sprintf("--previous=%v", prev)
	logsCommand := []string{"logs", TridentPodName, "-n", TridentPodNamespace, "-c", container, limitArg, prevArg}

	if Debug {
		fmt.Printf("Invoking command: %s %v\n", KubernetesCLI, strings.Join(logsCommand, " "))
	}

	// Get logs
	logBytes, err := exec.Command(KubernetesCLI, logsCommand...).CombinedOutput()
	if err != nil {
		logMap["error"] = appendError(logMap["error"], logBytes)
	} else {
		logMap[logName] = logBytes
	}
	return err
}

func appendError(oldErrors, newError []byte) []byte {

	if len(oldErrors) == 0 {
		return newError
	} else {
		oldErrorsStr := string(oldErrors)
		oldErrorsStr = strings.TrimSpace(oldErrorsStr)
		oldErrorsStr = strings.TrimSuffix(oldErrorsStr, ".")
		oldErrorsStr += ". "
		oldErrors = append([]byte(oldErrorsStr), newError...)
		return oldErrors
	}
}
<commit_msg>Remove size limit for `tridentctl logs`<commit_after>// Copyright 2019 NetApp, Inc. All Rights Reserved.

package cmd

import (
	"archive/zip"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/netapp/trident/config"
	"github.com/spf13/cobra"
)

const (
	logNameTrident         = "trident"
	logNameTridentPrevious = "trident-previous"
	logNameEtcd            = "etcd"
	logNameEtcdPrevious    = "etcd-previous"

	logTypeAuto    = "auto"
	logTypeTrident = "trident"
	logTypeEtcd    = "etcd"
	logTypeAll     = "all"

	archiveFilenameFormat = "support-2006-01-02T15-04-05-MST.zip"
)

var (
	logType  string
	archive  bool
	previous bool
)

func init() {
	RootCmd.AddCommand(logsCmd)
	logsCmd.Flags().StringVarP(&logType, "log", "l", logTypeAuto, "Trident log to display. One of trident|etcd|auto|all")
	logsCmd.Flags().BoolVarP(&archive, "archive", "a", false, "Create a support archive with all logs unless otherwise specified.")
	logsCmd.Flags().BoolVarP(&previous, "previous", "p", false, "Get the logs for the previous container instance if it exists.")
}

var logsCmd = &cobra.Command{
	Use:   "logs",
	Short: "Print the logs from Trident",
	Long:  "Print the logs from the Trident storage orchestrator for Kubernetes",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		err := discoverOperatingMode(cmd)
		return err
	},
	RunE: func(cmd *cobra.Command, args []string) error {

		err := checkValidLog()
		if err != nil {
			return err
		}

		if archive {
			return archiveLogs()
		} else {
			return consoleLogs()
		}
	},
}

func archiveLogs() error {

	// In archive mode, "auto" means to attempt to get all logs (current & previous).
	if logType == logTypeAuto {
		logType = logTypeAll
		previous = true
	}

	logMap := make(map[string][]byte)
	getLogs(logMap)

	// If there aren't any logs, bail out.
	anyLogs := false
	for log := range logMap {
		if log != "error" {
			anyLogs = true
			break
		}
	}
	if !anyLogs {
		return errors.New("no Trident-related logs found")
	}

	// Create archive file.
	zipFileName := time.Now().Format(archiveFilenameFormat)
	zipFile, err := os.Create(zipFileName)
	if err != nil {
		return err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	// Write to the archive file.
	for log, logBytes := range logMap {
		entry, err := zipWriter.Create(log)
		if err != nil {
			return err
		}
		_, err = entry.Write(logBytes)
		if err != nil {
			return err
		}
		fmt.Printf("Wrote %s log to %s archive file.\n", log, zipFileName)
	}

	return nil
}

func consoleLogs() error {

	logMap := make(map[string][]byte)
	err := getLogs(logMap)

	SetExitCodeFromError(err)
	if err != nil {
		// Preserve anything written to stdout/stderr
		logMessage := strings.TrimSuffix(strings.TrimSpace(string(logMap["error"])), ".")
		if len(logMessage) > 0 {
			errMessage := strings.TrimSuffix(strings.TrimSpace(err.Error()), ".")
			return fmt.Errorf("%s. %s", errMessage, logMessage)
		} else {
			return err
		}
	}

	// Print to the console
	anyLogs := false
	for log, logBytes := range logMap {
		if log != "error" {
			fmt.Printf("%s log:\n", log)
			fmt.Printf("%s\n", string(logBytes))
			anyLogs = true
		}
	}
	if !anyLogs {
		return errors.New("no Trident-related logs found")
	}

	return nil
}

func getLogs(logMap map[string][]byte) error {

	var err error

	if OperatingMode != ModeTunnel {
		return errors.New("'tridentctl logs' only supports Trident running in a Kubernetes pod")
	}

	switch logType {
	case logTypeTrident, logTypeAuto:
		err = getTridentLogs(logNameTrident, logMap)
	case logTypeEtcd:
		err = getTridentLogs(logNameEtcd, logMap)
	case logTypeAll:
		getTridentLogs(logNameTrident, logMap)
		getTridentLogs(logNameEtcd, logMap)
	}

	if previous {
		switch logType {
		case logTypeTrident, logTypeAuto:
			getTridentLogs(logNameTridentPrevious, logMap)
		case logTypeEtcd:
			getTridentLogs(logNameEtcdPrevious, logMap)
		case logTypeAll:
			getTridentLogs(logNameTridentPrevious, logMap)
			getTridentLogs(logNameEtcdPrevious, logMap)
		}
	}

	return err
}

func checkValidLog() error {
	switch logType {
	case logTypeTrident, logTypeEtcd, logTypeAuto, logTypeAll:
		return nil
	default:
		return fmt.Errorf("%s is not a valid Trident log", logType)
	}
}

func getTridentLogs(logName string, logMap map[string][]byte) error {

	var container string
	var prev bool

	switch logName {
	case logNameTrident:
		container, prev = config.ContainerTrident, false
	case logNameTridentPrevious:
		container, prev = config.ContainerTrident, true
	case logNameEtcd:
		container, prev = config.ContainerEtcd, false
	case logNameEtcdPrevious:
		container, prev = config.ContainerEtcd, true
	default:
		return fmt.Errorf("%s is not a valid Trident log", logName)
	}

	// Build command to get K8S logs
	prevArg := fmt.Sprintf("--previous=%v", prev)
	logsCommand := []string{"logs", TridentPodName, "-n", TridentPodNamespace, "-c", container, prevArg}

	if Debug {
		fmt.Printf("Invoking command: %s %v\n", KubernetesCLI, strings.Join(logsCommand, " "))
	}

	// Get logs
	logBytes, err := exec.Command(KubernetesCLI, logsCommand...).CombinedOutput()
	if err != nil {
		logMap["error"] = appendError(logMap["error"], logBytes)
	} else {
		logMap[logName] = logBytes
	}
	return err
}

func appendError(oldErrors, newError []byte) []byte {

	if len(oldErrors) == 0 {
		return newError
	} else {
		oldErrorsStr := string(oldErrors)
		oldErrorsStr = strings.TrimSpace(oldErrorsStr)
		oldErrorsStr = strings.TrimSuffix(oldErrorsStr, ".")
		oldErrorsStr += ". "
		oldErrors = append([]byte(oldErrorsStr), newError...)
		return oldErrors
	}
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"github.com/gopherjs/jquery"
	//"honnef.co/go/js/console"
)

//EventFunc is something that can handle an event.  It should not
//return anything, it should use the event object to stop default
//event handling and similar.
type EventFunc func(jquery.Event)

type EventName int

const (
	//keys
	BLUR EventName = iota
	CHANGE
	CLICK
	DBLCLICK
	FOCUS
	FOCUSIN
	FOCUSOUT
	HOVER
	KEYDOWN
	KEYPRESS
	KEYUP
	INPUT_EVENT
	MOUSE_ENTER
	MOUSE_LEAVE
)

func (self EventName) String() string {
	switch self {
	case BLUR:
		return jquery.BLUR
	case CHANGE:
		return jquery.CHANGE
	case CLICK:
		return jquery.CLICK
	case DBLCLICK:
		return jquery.DBLCLICK
	case FOCUS:
		return jquery.FOCUS
	case FOCUSIN:
		return jquery.FOCUSIN
	case FOCUSOUT:
		return jquery.FOCUSOUT
	case HOVER:
		return jquery.HOVER
	case KEYDOWN:
		return jquery.KEYDOWN
	case KEYPRESS:
		return jquery.KEYPRESS
	case KEYUP:
		return jquery.KEYUP
	case INPUT_EVENT:
		return "input"
	case MOUSE_ENTER:
		return jquery.MOUSEENTER
	case MOUSE_LEAVE:
		return jquery.MOUSELEAVE
	}
	panic(fmt.Sprintf("unknown event name %v", self))
}

type eventHandler struct {
	name EventName
	t    NarrowDom
	fn   EventFunc
}

func (self eventHandler) register() {
	self.t.On(self.name, self.handle)
}

func (self eventHandler) handle(event jquery.Event) {
	self.fn(event)
	DrainEagerQueue()
}
<commit_msg>safety<commit_after>package client

import (
	"fmt"
	"github.com/gopherjs/jquery"
	//"honnef.co/go/js/console"
)

//EventFunc is something that can handle an event.  It should not
//return anything, it should use the event object to stop default
//event handling and similar.
type EventFunc func(jquery.Event)

type EventName int

const (
	//keys
	BLUR EventName = iota
	CHANGE
	CLICK
	DBLCLICK
	FOCUS
	FOCUSIN
	FOCUSOUT
	HOVER
	KEYDOWN
	KEYPRESS
	KEYUP
	INPUT_EVENT
	MOUSE_ENTER
	MOUSE_LEAVE
	SELECT
)

func (self EventName) String() string {
	switch self {
	case BLUR:
		return jquery.BLUR
	case CHANGE:
		return jquery.CHANGE
	case CLICK:
		return jquery.CLICK
	case DBLCLICK:
		return jquery.DBLCLICK
	case FOCUS:
		return jquery.FOCUS
	case SELECT:
		return jquery.SELECT
	case FOCUSIN:
		return jquery.FOCUSIN
	case FOCUSOUT:
		return jquery.FOCUSOUT
	case HOVER:
		return jquery.HOVER
	case KEYDOWN:
		return jquery.KEYDOWN
	case KEYPRESS:
		return jquery.KEYPRESS
	case KEYUP:
		return jquery.KEYUP
	case INPUT_EVENT:
		return "input"
	case MOUSE_ENTER:
		return jquery.MOUSEENTER
	case MOUSE_LEAVE:
		return jquery.MOUSELEAVE
	}
	panic(fmt.Sprintf("unknown event name %v", self))
}

type eventHandler struct {
	name EventName
	t    NarrowDom
	fn   EventFunc
}

func (self eventHandler) register() {
	self.t.On(self.name, self.handle)
}

func (self eventHandler) handle(event jquery.Event) {
	self.fn(event)
	DrainEagerQueue()
}
<|endoftext|>


<commit_before>package deploystack

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"text/template"

	domains "cloud.google.com/go/domains/apiv1beta1"
	"google.golang.org/api/iterator"
	domainspb "google.golang.org/genproto/googleapis/cloud/domains/v1beta1"
	"google.golang.org/genproto/googleapis/type/postaladdress"
	"gopkg.in/yaml.v2"
)

var (
	// ErrorDomainUntenable is returned when a domain isn't available for registration, but
	// is also not owned by the user. It can't be used in this app
	ErrorDomainUntenable = fmt.Errorf("domain is not available, and not owned by attempting user")
	// ErrorDomainUserDeny is returned when an user declines the choice to purchase.
	ErrorDomainUserDeny = fmt.Errorf("user said no to buying the domain")
)

var msgDomainRegisterContactExplanation = fmt.Sprintf(`Domain registration requires some contact data. This process only asks for the
absolutely mandatory ones. The domain will be registered with user privacy
enabled, so that your contact info will not be public. This will create a file,
so that you never have to do it again.
%sThis file will only exist locally, or in your Cloud Shell environment.%s
`, TERMCYAN, TERMCLEAR)

var (
	msgDomainRegisterHeader      = fmt.Sprintf("%sManage Domain Registration %s", TERMCYANREV, TERMCLEAR)
	msgDomainContactHeader       = fmt.Sprintf("Managing Domain contact information")
	msgDomainContactFileWrite    = fmt.Sprintf("Your information was recorded and placed in a file for your future use.")
	msgDomainContactFileRead     = fmt.Sprintf("Your information was read from the local contact file.")
	msgDomainAvailablityHeader   = fmt.Sprintf("Managing Domain Availability")
	msgDomainAvailablityVerified = fmt.Sprintf("Domain is unavailable for purchase, but records show you are verified as the owner, so it can be used for this application.")
	msgDomainPurchase            = fmt.Sprintf("%sBuying a domain is not reversable, saying 'y' will incur a charge.%s", TERMREDB, TERMCLEAR)
	msgDomainRegisterSuccess     = fmt.Sprintf("%sDomain Registered%s. \n", TERMCYANB, TERMCLEAR)
)

// ContactData represents the structure that we need for Registrar Contact
// Data
type ContactData struct {
	AllContacts DomainRegistrarContact `yaml:"allContacts"`
}

// DomainRegistrarContact represents the data required to register a domain
// with a public registrar.
type DomainRegistrarContact struct {
	Email         string        `yaml:"email"`
	Phone         string        `yaml:"phoneNumber"`
	PostalAddress PostalAddress `yaml:"postalAddress"`
}

// PostalAddress represents the mail address in a DomainRegistrarContact
type PostalAddress struct {
	RegionCode         string   `yaml:"regionCode"`
	PostalCode         string   `yaml:"postalCode"`
	AdministrativeArea string   `yaml:"administrativeArea"`
	Locality           string   `yaml:"locality"`
	AddressLines       []string `yaml:"addressLines"`
	Recipients         []string `yaml:"recipients"`
}

// YAML outputs the content of this structure into the contact format needed for
// domain registration
func (c ContactData) YAML() (string, error) {
	yaml := `allContacts:
  email: '{{ .AllContacts.Email}}'
  phoneNumber: '{{.AllContacts.Phone}}'
  postalAddress:
    regionCode: '{{ .AllContacts.PostalAddress.RegionCode}}'
    postalCode: '{{ .AllContacts.PostalAddress.PostalCode}}'
    administrativeArea: '{{ .AllContacts.PostalAddress.AdministrativeArea}}'
    locality: '{{ .AllContacts.PostalAddress.Locality}}'
    addressLines: [{{range $element := .AllContacts.PostalAddress.AddressLines}}'{{$element}}'{{end}}]
    recipients: [{{range $element := .AllContacts.PostalAddress.Recipients}}'{{$element}}'{{end}}]`

	t, err := template.New("yaml").Parse(yaml)
	if err != nil {
		return "", fmt.Errorf("error parsing the yaml template %s", err)
	}
	var tpl bytes.Buffer
	err = t.Execute(&tpl, c)
	if err != nil {
		return "", fmt.Errorf("error executing the yaml template %s", err)
	}

	return tpl.String(), nil
}

// DomainContact outputs a varible in the format that Domain Registration
// API needs.
func (c ContactData) DomainContact() (domainspb.ContactSettings, error) {
	dc := domainspb.ContactSettings{}

	pa := postaladdress.PostalAddress{
		RegionCode:         c.AllContacts.PostalAddress.RegionCode,
		PostalCode:         c.AllContacts.PostalAddress.PostalCode,
		AdministrativeArea: c.AllContacts.PostalAddress.AdministrativeArea,
		Locality:           c.AllContacts.PostalAddress.Locality,
		AddressLines:       c.AllContacts.PostalAddress.AddressLines,
		Recipients:         c.AllContacts.PostalAddress.Recipients,
	}

	all := domainspb.ContactSettings_Contact{
		Email:         c.AllContacts.Email,
		PhoneNumber:   c.AllContacts.Phone,
		PostalAddress: &pa,
	}

	dc.AdminContact = &all
	dc.RegistrantContact = &all
	dc.TechnicalContact = &all
	dc.Privacy = domainspb.ContactPrivacy_PRIVATE_CONTACT_DATA

	return dc, nil
}

func newContactData() ContactData {
	c := ContactData{}
	d := DomainRegistrarContact{}
	d.PostalAddress.AddressLines = []string{}
	d.PostalAddress.Recipients = []string{}
	c.AllContacts = d
	return c
}

func newContactDataFromFile(file string) (ContactData, error) {
	c := newContactData()

	dat, err := os.ReadFile(file)
	if err != nil {
		return c, err
	}

	err = yaml.Unmarshal(dat, &c)
	if err != nil {
		return c, err
	}

	return c, nil
}

// DomainManage walks a user through the porocess of collecting contact info and
// registering a domain.
func DomainManage(s *Stack) (string, error) {
	fmt.Println(Divider)
	fmt.Println(msgDomainRegisterHeader)
	fmt.Println(Divider)
	fmt.Println(msgDomainContactHeader)

	contactfile := "contact.yaml"
	contact := ContactData{}
	domain := ""
	project := s.GetSetting("project_id")

	item := Custom{Name: "domain", Description: "Enter a domain you wish to purchase and use for this application"}

	if err := item.Collect(); err != nil {
		return "", fmt.Errorf("trouble getting domain from keyboard: %s", err)
	}

	domain = item.Value

	fmt.Println(Divider)
	fmt.Println(msgDomainAvailablityHeader)

	domainInfo, err := domainIsAvailable(project, domain)
	if err != nil {
		return "", fmt.Errorf("error checking domain %s", err)
	}

	if domainInfo.Availability == domainspb.RegisterParameters_UNAVAILABLE {

		isVerified, err := domainsIsVerified(project, domain)
		if err != nil {
			return "", fmt.Errorf("error verifying domain %s", err)
		}
		// If not, fail and ask the user to repick
		if !isVerified {
			return "", ErrorDomainUntenable
		}

		fmt.Println(msgDomainAvailablityVerified)
		return domain, nil
	}

	if _, err := os.Stat(contactfile); errors.Is(err, os.ErrNotExist) {
		contact, err = RegistrarContactManage(contactfile)
		if err != nil {
			return "", fmt.Errorf("error getting contact data %s", err)
		}
	} else {
		contact, err = newContactDataFromFile(contactfile)
		if err != nil {
			contact, err = RegistrarContactManage(contactfile)
			if err != nil {
				return "", fmt.Errorf("error getting contact data %s", err)
			}
		} else {
			fmt.Println(msgDomainContactFileRead)
		}
	}

	fmt.Println(msgDomainPurchase)
	fmt.Printf("%sCost for %s%s%s%s will be %s%d%s%s%s.  Continue?%s\n",
		TERMREDB, TERMCYAN, domain, TERMCLEAR, TERMREDB,
		TERMCYAN,
		domainInfo.YearlyPrice.Units, domainInfo.YearlyPrice.CurrencyCode,
		TERMCLEAR, TERMREDB, TERMCLEAR)
	choice := Custom{Name: "choice", Description: "y or n?"}

	if err := choice.Collect(); err != nil {
		return "", fmt.Errorf("trouble getting domain from keyboard: %s", err)
	}

	if strings.ToLower(choice.Value) != "y" && strings.ToLower(choice.Value) != "yes" {
		return "", ErrorDomainUserDeny
	}

	if err := domainRegister(project, domainInfo, contact); err != nil {
		return "", fmt.Errorf("error registering domain %s", err)
	}
	fmt.Println(Divider)
	fmt.Println(msgDomainRegisterSuccess)
	return domain, nil
}

// RegistrarContactManage manages collecting domain registraton information
// from the user
func RegistrarContactManage(file string) (ContactData, error) {
	d := newContactData()

	fmt.Println(Divider)
	fmt.Printf(msgDomainRegisterContactExplanation)
	fmt.Println(Divider)

	items := Customs{
		{Name: "email", Description: "Enter an email address", Default: "person@example.com"},
		{Name: "phone", Description: "Enter a phone number. (Please enter with country code - +1 555 555 5555 for US for example)", Default: "+14155551234", Validation: "phonenumber"},
		{Name: "country", Description: "Enter a country code", Default: "US"},
		{Name: "postalcode", Description: "Enter a postal code", Default: "94502"},
		{Name: "state", Description: "Enter a state or administrative area", Default: "CA"},
		{Name: "city", Description: "Enter a city", Default: "San Francisco"},
		{Name: "address", Description: "Enter an address", Default: "345 Spear Street"},
		{Name: "name", Description: "Enter name", Default: "Googler"},
	}

	if err := items.Collect(); err != nil {
		return d, err
	}

	d.AllContacts.Email = items.Get("email").Value
	d.AllContacts.Phone = items.Get("phone").Value
	d.AllContacts.PostalAddress.RegionCode = items.Get("country").Value
	d.AllContacts.PostalAddress.PostalCode = items.Get("postalcode").Value
	d.AllContacts.PostalAddress.AdministrativeArea = items.Get("state").Value
	d.AllContacts.PostalAddress.Locality = items.Get("city").Value
	d.AllContacts.PostalAddress.AddressLines = append(d.AllContacts.PostalAddress.AddressLines, items.Get("address").Value)
	d.AllContacts.PostalAddress.Recipients = append(d.AllContacts.PostalAddress.Recipients, items.Get("name").Value)

	yaml, err := d.YAML()
	if err != nil {
		return d, err
	}

	if err := os.WriteFile(file, []byte(yaml), 0o644); err != nil {
		return d, err
	}
	fmt.Println(msgDomainContactFileWrite)

	return d, nil
}

func domainsSearch(project, domain string) ([]*domainspb.RegisterParameters, error) {
	ctx := context.Background()

	c, err := domains.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	defer c.Close()

	req := &domainspb.SearchDomainsRequest{
		Query:    domain,
		Location: fmt.Sprintf("projects/%s/locations/global", project),
	}
	resp, err := c.SearchDomains(ctx, req)
	if err != nil {
		return nil, err
	}

	return resp.RegisterParameters, nil
}

func domainIsAvailable(project, domain string) (*domainspb.RegisterParameters, error) {
	list, err := domainsSearch(project, domain)
	if err != nil {
		return nil, err
	}
	for _, v := range list {
		if v.DomainName == domain {
			return v, err
		}
	}

	return nil, err
}

func domainsIsVerified(project, domain string) (bool, error) {
	ctx := context.Background()

	c, err := domains.NewClient(ctx)
	if err != nil {
		return false, err
	}
	defer c.Close()

	req := &domainspb.ListRegistrationsRequest{
		Filter: fmt.Sprintf("domainName=\"%s\"", domain),
		Parent: fmt.Sprintf("projects/%s/locations/global", project),
	}
	it := c.ListRegistrations(ctx, req)
	for {
		resp, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return false, err
		}

		if resp.DomainName == domain {
			return true, nil
		}
	}

	return false, nil
}

func domainRegister(project string, domaininfo *domainspb.RegisterParameters, contact ContactData) error {
	ctx := context.Background()

	parent := fmt.Sprintf("projects/%s/locations/global", project)

	c, err := domains.NewClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	dnscontact, err := contact.DomainContact()
	if err != nil {
		return err
	}

	dnsprovider := domainspb.DnsSettings_CustomDns_{
		CustomDns: &domainspb.DnsSettings_CustomDns{
			NameServers: []string{
				"ns-cloud-e1.googledomains.com",
				"ns-cloud-e2.googledomains.com",
				"ns-cloud-e3.googledomains.com",
				"ns-cloud-e4.googledomains.com",
			},
		},
	}

	dnssettings := domainspb.DnsSettings{
		DnsProvider: &dnsprovider,
	}

	reg := domainspb.Registration{
		Name:            fmt.Sprintf("%s/registrations/%s", parent, domaininfo.DomainName),
		DomainName:      domaininfo.DomainName,
		DnsSettings:     &dnssettings,
		ContactSettings: &dnscontact,
	}

	req := &domainspb.RegisterDomainRequest{
		Registration: &reg,
		Parent:       parent,
		YearlyPrice:  domaininfo.YearlyPrice,
	}

	if _, err := c.RegisterDomain(ctx, req); err != nil {
		return err
	}

	return nil
}
<commit_msg>Cleaning up a bunch of code.<commit_after>package deploystack

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"text/template"

	domains "cloud.google.com/go/domains/apiv1beta1"
	"google.golang.org/api/iterator"
	domainspb "google.golang.org/genproto/googleapis/cloud/domains/v1beta1"
	"google.golang.org/genproto/googleapis/type/postaladdress"
	"gopkg.in/yaml.v2"
)

var (
	// ErrorDomainUntenable is returned when a domain isn't available for registration, but
	// is also not owned by the user. It can't be used in this app
	ErrorDomainUntenable = fmt.Errorf("domain is not available, and not owned by attempting user")
	// ErrorDomainUserDeny is returned when an user declines the choice to purchase.
	ErrorDomainUserDeny = fmt.Errorf("user said no to buying the domain")
)

var msgDomainRegisterContactExplanation = fmt.Sprintf(`Domain registration requires some contact data. This process only asks for the
absolutely mandatory ones. The domain will be registered with user privacy
enabled, so that your contact info will not be public. This will create a file,
so that you never have to do it again.
%sThis file will only exist locally, or in your Cloud Shell environment.%s
`, TERMCYAN, TERMCLEAR)

var (
	msgDomainRegisterHeader      = fmt.Sprintf("%sManage Domain Registration %s", TERMCYANREV, TERMCLEAR)
	msgDomainContactHeader       = fmt.Sprintf("Managing Domain contact information")
	msgDomainContactFileWrite    = fmt.Sprintf("Your information was recorded and placed in a file for your future use.")
	msgDomainContactFileRead     = fmt.Sprintf("Your information was read from the local contact file.")
	msgDomainAvailablityHeader   = fmt.Sprintf("Managing Domain Availability")
	msgDomainAvailablityVerified = fmt.Sprintf("Domain is unavailable for purchase, but records show you are verified as the owner, so it can be used for this application.")
	msgDomainPurchase            = fmt.Sprintf("%sBuying a domain is not reversable, saying 'y' will incur a charge.%s", TERMREDB, TERMCLEAR)
	msgDomainRegisterSuccess     = fmt.Sprintf("%sDomain Registered%s. \n", TERMCYANB, TERMCLEAR)
)

// ContactData represents the structure that we need for Registrar Contact
// Data
type ContactData struct {
	AllContacts DomainRegistrarContact `yaml:"allContacts"`
}

// DomainRegistrarContact represents the data required to register a domain
// with a public registrar.
type DomainRegistrarContact struct {
	Email         string        `yaml:"email"`
	Phone         string        `yaml:"phoneNumber"`
	PostalAddress PostalAddress `yaml:"postalAddress"`
}

// PostalAddress represents the mail address in a DomainRegistrarContact
type PostalAddress struct {
	RegionCode         string   `yaml:"regionCode"`
	PostalCode         string   `yaml:"postalCode"`
	AdministrativeArea string   `yaml:"administrativeArea"`
	Locality           string   `yaml:"locality"`
	AddressLines       []string `yaml:"addressLines"`
	Recipients         []string `yaml:"recipients"`
}

// YAML outputs the content of this structure into the contact format needed for
// domain registration
func (c ContactData) YAML() (string, error) {
	yaml := `allContacts:
  email: '{{ .AllContacts.Email}}'
  phoneNumber: '{{.AllContacts.Phone}}'
  postalAddress:
    regionCode: '{{ .AllContacts.PostalAddress.RegionCode}}'
    postalCode: '{{ .AllContacts.PostalAddress.PostalCode}}'
    administrativeArea: '{{ .AllContacts.PostalAddress.AdministrativeArea}}'
    locality: '{{ .AllContacts.PostalAddress.Locality}}'
    addressLines: [{{range $element := .AllContacts.PostalAddress.AddressLines}}'{{$element}}'{{end}}]
    recipients: [{{range $element := .AllContacts.PostalAddress.Recipients}}'{{$element}}'{{end}}]`

	t, err := template.New("yaml").Parse(yaml)
	if err != nil {
		return "", fmt.Errorf("error parsing the yaml template %s", err)
	}
	var tpl bytes.Buffer
	err = t.Execute(&tpl, c)
	if err != nil {
		return "", fmt.Errorf("error executing the yaml template %s", err)
	}

	return tpl.String(), nil
}

// DomainContact outputs a varible in the format that Domain Registration
// API needs.
func (c ContactData) DomainContact() (domainspb.ContactSettings, error) {
	dc := domainspb.ContactSettings{}

	pa := postaladdress.PostalAddress{
		RegionCode:         c.AllContacts.PostalAddress.RegionCode,
		PostalCode:         c.AllContacts.PostalAddress.PostalCode,
		AdministrativeArea: c.AllContacts.PostalAddress.AdministrativeArea,
		Locality:           c.AllContacts.PostalAddress.Locality,
		AddressLines:       c.AllContacts.PostalAddress.AddressLines,
		Recipients:         c.AllContacts.PostalAddress.Recipients,
	}

	all := domainspb.ContactSettings_Contact{
		Email:         c.AllContacts.Email,
		PhoneNumber:   c.AllContacts.Phone,
		PostalAddress: &pa,
	}

	dc.AdminContact = &all
	dc.RegistrantContact = &all
	dc.TechnicalContact = &all
	dc.Privacy = domainspb.ContactPrivacy_PRIVATE_CONTACT_DATA

	return dc, nil
}

func newContactData() ContactData {
	c := ContactData{}
	d := DomainRegistrarContact{}
	d.PostalAddress.AddressLines = []string{}
	d.PostalAddress.Recipients = []string{}
	c.AllContacts = d
	return c
}

func newContactDataFromFile(file string) (ContactData, error) {
	c := newContactData()

	dat, err := os.ReadFile(file)
	if err != nil {
		return c, err
	}

	err = yaml.Unmarshal(dat, &c)
	if err != nil {
		return c, err
	}

	return c, nil
}

// DomainManage walks a user through the porocess of collecting contact info and
// registering a domain.
func DomainManage(s *Stack) (string, error) {
	fmt.Println(Divider)
	fmt.Println(msgDomainRegisterHeader)
	fmt.Println(Divider)
	fmt.Println(msgDomainContactHeader)

	contactfile := "contact.yaml"
	contact := ContactData{}
	domain := ""
	project := s.GetSetting("project_id")

	item := Custom{Name: "domain", Description: "Enter a domain you wish to purchase and use for this application"}

	if err := item.Collect(); err != nil {
		return "", fmt.Errorf("trouble getting domain from keyboard: %s", err)
	}

	domain = item.Value

	fmt.Println(Divider)
	fmt.Println(msgDomainAvailablityHeader)

	domainInfo, err := domainIsAvailable(project, domain)
	if err != nil {
		return "", fmt.Errorf("error checking domain %s", err)
	}

	if domainInfo.Availability == domainspb.RegisterParameters_UNAVAILABLE {

		isVerified, err := domainsIsVerified(project, domain)
		if err != nil {
			return "", fmt.Errorf("error verifying domain %s", err)
		}
		// If not, fail and ask the user to repick
		if !isVerified {
			return "", ErrorDomainUntenable
		}

		fmt.Println(msgDomainAvailablityVerified)
		return domain, nil
	}

	if _, err := os.Stat(contactfile); errors.Is(err, os.ErrNotExist) {
		contact, err = RegistrarContactManage(contactfile)
		if err != nil {
			return "", fmt.Errorf("error getting contact data %s", err)
		}
	} else {
		contact, err = newContactDataFromFile(contactfile)
		if err != nil {
			contact, err = RegistrarContactManage(contactfile)
			if err != nil {
				return "", fmt.Errorf("error getting contact data %s", err)
			}
		} else {
			fmt.Println(msgDomainContactFileRead)
		}
	}

	fmt.Println(msgDomainPurchase)
	fmt.Printf("%sCost for %s%s%s%s will be %s%d%s%s%s.  Continue?%s\n",
		TERMREDB, TERMCYAN, domain, TERMCLEAR, TERMREDB,
		TERMCYAN,
		domainInfo.YearlyPrice.Units, domainInfo.YearlyPrice.CurrencyCode,
		TERMCLEAR, TERMREDB, TERMCLEAR)
	choice := Custom{Name: "choice", Description: "y or n?"}

	if err := choice.Collect(); err != nil {
		return "", fmt.Errorf("trouble getting domain from keyboard: %s", err)
	}

	if strings.ToLower(choice.Value) != "y" && strings.ToLower(choice.Value) != "yes" {
		return "", ErrorDomainUserDeny
	}

	if err := domainRegister(project, domainInfo, contact); err != nil {
		return "", fmt.Errorf("error registering domain %s", err)
	}
	fmt.Println(Divider)
	fmt.Println(msgDomainRegisterSuccess)
	return domain, nil
}

// RegistrarContactManage manages collecting domain registraton information
// from the user
func RegistrarContactManage(file string) (ContactData, error) {
	d := newContactData()

	fmt.Println(Divider)
	fmt.Printf(msgDomainRegisterContactExplanation)
	fmt.Println(Divider)

	items := Customs{
		{Name: "email", Description: "Enter an email address", Default: "person@example.com"},
		{Name: "phone", Description: "Enter a phone number. (Please enter with country code - +1 555 555 5555 for US for example)", Default: "+14155551234", Validation: "phonenumber"},
		{Name: "country", Description: "Enter a country code", Default: "US"},
		{Name: "postalcode", Description: "Enter a postal code", Default: "94502"},
		{Name: "state", Description: "Enter a state or administrative area", Default: "CA"},
		{Name: "city", Description: "Enter a city", Default: "San Francisco"},
		{Name: "address", Description: "Enter an address", Default: "345 Spear Street"},
		{Name: "name", Description: "Enter name", Default: "Googler"},
	}

	if err := items.Collect(); err != nil {
		return d, err
	}

	d.AllContacts.Email = items.Get("email").Value
	d.AllContacts.Phone = items.Get("phone").Value
	d.AllContacts.PostalAddress.RegionCode = items.Get("country").Value
	d.AllContacts.PostalAddress.PostalCode = items.Get("postalcode").Value
	d.AllContacts.PostalAddress.AdministrativeArea = items.Get("state").Value
	d.AllContacts.PostalAddress.Locality = items.Get("city").Value
	d.AllContacts.PostalAddress.AddressLines = append(d.AllContacts.PostalAddress.AddressLines, items.Get("address").Value)
	d.AllContacts.PostalAddress.Recipients = append(d.AllContacts.PostalAddress.Recipients, items.Get("name").Value)

	yaml, err := d.YAML()
	if err != nil {
		return d, err
	}

	if err := os.WriteFile(file, []byte(yaml), 0o644); err != nil {
		return d, err
	}
	fmt.Println(msgDomainContactFileWrite)

	return d, nil
}

func domainsSearch(project, domain string) ([]*domainspb.RegisterParameters, error) {
	ctx := context.Background()

	c, err := domains.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	defer c.Close()

	req := &domainspb.SearchDomainsRequest{
		Query:    domain,
		Location: fmt.Sprintf("projects/%s/locations/global", project),
	}
	resp, err := c.SearchDomains(ctx, req)
	if err != nil {
		return nil, err
	}

	return resp.RegisterParameters, nil
}

func domainIsAvailable(project, domain string) (*domainspb.RegisterParameters, error) {
	list, err := domainsSearch(project, domain)
	if err != nil {
		return nil, err
	}
	for _, v := range list {
		if v.DomainName == domain {
			return v, err
		}
	}

	return nil, err
}

func domainsIsVerified(project, domain string) (bool, error) {
	ctx := context.Background()

	c, err := domains.NewClient(ctx)
	if err != nil {
		return false, err
	}
	defer c.Close()

	req := &domainspb.ListRegistrationsRequest{
		Filter: fmt.Sprintf("domainName=\"%s\"", domain),
		Parent: fmt.Sprintf("projects/%s/locations/global", project),
	}
	it := c.ListRegistrations(ctx, req)
	for {
		resp, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return false, err
		}

		if resp.DomainName == domain {
			return true, nil
		}
	}

	return false, nil
}

func domainRegister(project string, domaininfo *domainspb.RegisterParameters, contact ContactData) error {
	ctx := context.Background()

	parent := fmt.Sprintf("projects/%s/locations/global", project)

	c, err := domains.NewClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	dnscontact, err := contact.DomainContact()
	if err != nil {
		return err
	}

	req := &domainspb.RegisterDomainRequest{
		Registration: &domainspb.Registration{
			Name:       fmt.Sprintf("%s/registrations/%s", parent, domaininfo.DomainName),
			DomainName: domaininfo.DomainName,
			DnsSettings: &domainspb.DnsSettings{
				DnsProvider: &domainspb.DnsSettings_CustomDns_{
					CustomDns: &domainspb.DnsSettings_CustomDns{
						NameServers: []string{
							"ns-cloud-e1.googledomains.com",
							"ns-cloud-e2.googledomains.com",
							"ns-cloud-e3.googledomains.com",
							"ns-cloud-e4.googledomains.com",
						},
					},
				},
			},
			ContactSettings: &dnscontact,
		},
		Parent:      parent,
		YearlyPrice: domaininfo.YearlyPrice,
	}

	if _, err := c.RegisterDomain(ctx, req); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string) *Node {
	e := &Node{
		Addr:       addr,
		Labels:     make(map[string]string),
		ch:         make(chan bool),
		containers: make(map[string]*Container),
		healthy:    true,
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch           chan bool
	containers   map[string]*Container
	client       dockerclient.Client
	eventHandler EventHandler
	healthy      bool
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	c, err := dockerclient.NewDockerClientTimeout(n.Addr, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", strings.Split(c.URL.Host, ":")[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) ForceRefreshContainer(c dockerclient.Container) error {
	return n.inspectContainer(c, n.containers, true)
}

func (n *Node) inspectContainer(c dockerclient.Container, containers map[string]*Container, lock bool) error {

	container := &Container{}
	container.Container = c
	container.Node = n

	info, err := n.client.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	container.Info = *info

	// real CpuShares -> nb of CPUs
	container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

	if lock {
		n.Lock()
		defer n.Unlock()
	}
	containers[container.Id] = container

	return nil
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		if err := n.inspectContainer(c, containers, false); err != nil {
			return nil, err
		}
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}

		if err != nil {
			n.healthy = false
			log.Errorf("[%s/%s] Flagging node as dead. Updated state failed: %v", n.ID, n.Name, err)
		} else {
			if !n.healthy {
				log.Infof("[%s/%s] Node came back to life. Hooray!", n.ID, n.Name)
			}
			n.healthy = true
		}
	}
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node: n,
	}
	event.Event = *ev

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<commit_msg>More canonical Event construction on node event handling<commit_after>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string) *Node {
	e := &Node{
		Addr:       addr,
		Labels:     make(map[string]string),
		ch:         make(chan bool),
		containers: make(map[string]*Container),
		healthy:    true,
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch           chan bool
	containers   map[string]*Container
	client       dockerclient.Client
	eventHandler EventHandler
	healthy      bool
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	c, err := dockerclient.NewDockerClientTimeout(n.Addr, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", strings.Split(c.URL.Host, ":")[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) ForceRefreshContainer(c dockerclient.Container) error {
	return n.inspectContainer(c, n.containers, true)
}

func (n *Node) inspectContainer(c dockerclient.Container, containers map[string]*Container, lock bool) error {

	container := &Container{}
	container.Container = c
	container.Node = n

	info, err := n.client.InspectContainer(c.Id)
	if err != nil {
		return err
	}
	container.Info = *info

	// real CpuShares -> nb of CPUs
	container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

	if lock {
		n.Lock()
		defer n.Unlock()
	}
	containers[container.Id] = container

	return nil
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		if err := n.inspectContainer(c, containers, false); err != nil {
			return nil, err
		}
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}

		if err != nil {
			n.healthy = false
			log.Errorf("[%s/%s] Flagging node as dead. Updated state failed: %v", n.ID, n.Name, err)
		} else {
			if !n.healthy {
				log.Infof("[%s/%s] Node came back to life. Hooray!", n.ID, n.Name)
			}
			n.healthy = true
		}
	}
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node:  n,
		Event: *ev,
	}

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<|endoftext|>