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

package cache

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"sync"

	"github.com/google/trillian/merkle/compact"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage/storagepb"
	"github.com/google/trillian/storage/tree"
	"google.golang.org/protobuf/proto"
)

// TODO(al): move this up the stack
var populateConcurrency = flag.Int("populate_subtree_concurrency", 256, "Max number of concurrent workers concurrently populating subtrees")

// TODO(pavelkalinnikov): Rename subtrees to tiles.

// GetSubtreesFunc describes a function which can return a number of Subtrees from storage.
type GetSubtreesFunc func(ids [][]byte) ([]*storagepb.SubtreeProto, error)

// SetSubtreesFunc describes a function which can store a collection of Subtrees into storage.
type SetSubtreesFunc func(ctx context.Context, s []*storagepb.SubtreeProto) error

// SubtreeCache provides a caching access to Subtree storage. Currently there are assumptions
// in the code that all subtrees are multiple of 8 in depth and that log subtrees are always
// of depth 8. It is not possible to just change the constants above and have things still
// work. This is because of issues like byte packing of node IDs.
//
// The cache is optimized for the following use-cases (see sync.Map type for more details):
//  1. Parallel readers/writers working on non-intersecting subsets of subtrees/nodes.
//  2. Subtrees/nodes are rarely written, and mostly read.
type SubtreeCache struct {
	hasher hashers.LogHasher

	// subtrees contains the Subtree data read from storage, and is updated by
	// calls to SetNodes.
	subtrees sync.Map
	// dirtyPrefixes keeps track of all Subtrees which need to be written back
	// to storage.
	dirtyPrefixes sync.Map

	// populateConcurrency sets the amount of concurrency when repopulating subtrees.
	populateConcurrency int
}

// NewLogSubtreeCache creates and returns a SubtreeCache appropriate for use with a log
// tree. The caller must supply a suitable LogHasher.
func NewLogSubtreeCache(hasher hashers.LogHasher) *SubtreeCache {
	if *populateConcurrency <= 0 {
		panic(fmt.Errorf("populate_subtree_concurrency must be set to >= 1"))
	}
	return &SubtreeCache{
		hasher:              hasher,
		populateConcurrency: *populateConcurrency,
	}
}

// preload calculates the set of subtrees required to know the hashes of the
// passed in node IDs, uses getSubtrees to retrieve them, and finally populates
// the cache structures with the data. Returns the list of tile IDs not found.
func (s *SubtreeCache) preload(ids []compact.NodeID, getSubtrees GetSubtreesFunc) ([]string, error) {
	// Figure out the set of subtrees we need.
	want := make(map[string]bool)
	for _, id := range ids {
		subID := string(getTileID(id))
		if _, ok := want[subID]; ok {
			// No need to check s.subtrees map twice.
			continue
		}
		if _, ok := s.subtrees.Load(subID); !ok {
			want[subID] = true
		}
	}
	// Note: At this point multiple parallel preload invocations can happen to
	// getSubtrees with overlapping sets of IDs. It's okay because we collapse
	// results further below.

	// Don't make a read request for zero subtrees.
	if len(want) == 0 {
		return nil, nil
	}

	list := make([][]byte, 0, len(want))
	for id := range want {
		list = append(list, []byte(id))
	}
	subtrees, err := getSubtrees(list)
	if err != nil {
		return nil, err
	}
	if got, max := len(subtrees), len(want); got > max {
		return nil, fmt.Errorf("too many subtrees: %d, want <= %d", got, max)
	}

	ch := make(chan *storagepb.SubtreeProto, len(subtrees))
	workTokens := make(chan bool, s.populateConcurrency)
	for i := 0; i < s.populateConcurrency; i++ {
		workTokens <- true
	}
	wg := &sync.WaitGroup{}

	for _, t := range subtrees {
		t := t
		wg.Add(1)
		go func() {
			defer wg.Done()
			// wait for a token before starting work
			<-workTokens
			// return it when done
			defer func() { workTokens <- true }()

			PopulateLogTile(t, s.hasher)
			ch <- t // Note: This never blocks because len(ch) == len(subtrees).
		}()
	}

	go func() {
		wg.Wait()
		close(ch)
		close(workTokens)
	}()

	for t := range ch {
		if err := s.cacheSubtree(t); err != nil {
			return nil, err
		}
		delete(want, string(t.Prefix))
	}
	notFound := make([]string, len(want))
	for id := range want {
		notFound = append(notFound, id)
	}
	return notFound, nil
}

func (s *SubtreeCache) cacheSubtree(t *storagepb.SubtreeProto) error {
	raw, loaded := s.subtrees.LoadOrStore(string(t.Prefix), t)
	if loaded { // There is already something with this key.
		if subtree, ok := raw.(*storagepb.SubtreeProto); !ok {
			return fmt.Errorf("at %x: not a subtree: %T", t.Prefix, raw)
		} else if !proto.Equal(t, subtree) {
			return fmt.Errorf("at %x: subtree mismatch", t.Prefix)
		}
	}
	return nil
}

// GetNodes returns the requested nodes, calling the getSubtrees function if
// they are not already cached.
func (s *SubtreeCache) GetNodes(ids []compact.NodeID, getSubtrees GetSubtreesFunc) ([]tree.Node, error) {
	if notFound, err := s.preload(ids, getSubtrees); err != nil {
		return nil, err
	} else if r := len(notFound); r != 0 {
		return nil, fmt.Errorf("preload did not get all tiles: %d not found", r)
	}

	ret := make([]tree.Node, 0, len(ids))
	for _, id := range ids {
		if h, err := s.getNodeHash(id); err != nil {
			return nil, fmt.Errorf("getNodeHash(%+v): %v", id, err)
		} else if h != nil {
			ret = append(ret, tree.Node{ID: id, Hash: h})
		}
	}
	return ret, nil
}

func (s *SubtreeCache) getCachedSubtree(id []byte) *storagepb.SubtreeProto {
	raw, found := s.subtrees.Load(string(id))
	if !found || raw == nil {
		return nil
	}
	// Note: If type assertion fails, nil is returned. Should not happen though.
	res, _ := raw.(*storagepb.SubtreeProto)
	return res
}

func (s *SubtreeCache) prefixIsDirty(prefixKey string) bool {
	_, found := s.dirtyPrefixes.Load(prefixKey)
	return found
}

// getNodeHash returns a single node hash from the cache.
func (s *SubtreeCache) getNodeHash(id compact.NodeID) ([]byte, error) {
	subID, sx := splitID(id)
	c := s.getCachedSubtree(subID)
	if c == nil {
		return nil, fmt.Errorf("tile %x not found", subID)
	}

	// Look up the hash in the appropriate map.
	// The leaf hashes are stored in a separate map to the internal nodes so that
	// we can easily dump (and later reconstruct) the internal nodes. As log subtrees
	// have a fixed depth if the suffix has the same number of significant bits as the
	// subtree depth then this is a leaf. For example if the subtree is depth 8 its leaves
	// have 8 significant suffix bits.
	if int32(sx.Bits()) == c.Depth {
		return c.Leaves[sx.String()], nil
	}
	return c.InternalNodes[sx.String()], nil
}

// SetNodes sets hashes for the given nodes in the cache.
func (s *SubtreeCache) SetNodes(nodes []tree.Node, getSubtrees GetSubtreesFunc) error {
	ids := make([]compact.NodeID, len(nodes))
	for i, n := range nodes {
		ids[i] = n.ID
	}
	notFound, err := s.preload(ids, getSubtrees)
	if err != nil {
		return err
	}
	for _, id := range notFound {
		c := newEmptyTile([]byte(id))
		s.subtrees.Store(string(id), c)
	}

	for _, n := range nodes {
		subID, sx := splitID(n.ID)
		c := s.getCachedSubtree(subID)
		if c == nil {
			return fmt.Errorf("tile %x not found", subID)
		}

		// Store the hash to the containing tile, and mark it as dirty if the hash
		// differs from the previously stored one.
		sfxKey := sx.String()
		if int32(sx.Bits()) == c.Depth { // This is a leaf node.
			if !bytes.Equal(c.Leaves[sfxKey], n.Hash) {
				c.Leaves[sfxKey] = n.Hash
				s.dirtyPrefixes.Store(string(subID), nil)
			}
		} else { // This is an internal node.
			if !bytes.Equal(c.InternalNodes[sfxKey], n.Hash) {
				c.InternalNodes[sfxKey] = n.Hash
				s.dirtyPrefixes.Store(string(subID), nil)
			}
		}
	}

	return nil
}

// Flush causes the cache to write all dirty Subtrees back to storage.
func (s *SubtreeCache) Flush(ctx context.Context, setSubtrees SetSubtreesFunc) error {
	treesToWrite := make([]*storagepb.SubtreeProto, 0)
	var rangeErr error
	s.subtrees.Range(func(rawK, rawV interface{}) bool {
		k, ok := rawK.(string)
		if !ok {
			rangeErr = fmt.Errorf("unknown key type: %t", rawV)
			return false
		}
		v, ok := rawV.(*storagepb.SubtreeProto)
		if !ok {
			rangeErr = fmt.Errorf("unknown value type: %t", rawV)
			return false
		}
		if s.prefixIsDirty(k) {
			bk := []byte(k)
			if !bytes.Equal(bk, v.Prefix) {
				rangeErr = fmt.Errorf("inconsistent cache: prefix key is %v, but cached object claims %v", bk, v.Prefix)
				return false
			}

			if len(v.Leaves) > 0 {
				if err := prepareLogTile(v); err != nil {
					rangeErr = err
					return false
				}
				treesToWrite = append(treesToWrite, v)
			}
		}
		return true
	})
	if rangeErr != nil {
		return rangeErr
	}
	if len(treesToWrite) == 0 {
		return nil
	}
	err := setSubtrees(ctx, treesToWrite)
	return err
}
<commit_msg>Fix slice allocation size/capacity (#2598)<commit_after>// Copyright 2016 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cache

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"sync"

	"github.com/google/trillian/merkle/compact"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage/storagepb"
	"github.com/google/trillian/storage/tree"
	"google.golang.org/protobuf/proto"
)

// TODO(al): move this up the stack
var populateConcurrency = flag.Int("populate_subtree_concurrency", 256, "Max number of concurrent workers concurrently populating subtrees")

// TODO(pavelkalinnikov): Rename subtrees to tiles.

// GetSubtreesFunc describes a function which can return a number of Subtrees from storage.
type GetSubtreesFunc func(ids [][]byte) ([]*storagepb.SubtreeProto, error)

// SetSubtreesFunc describes a function which can store a collection of Subtrees into storage.
type SetSubtreesFunc func(ctx context.Context, s []*storagepb.SubtreeProto) error

// SubtreeCache provides a caching access to Subtree storage. Currently there are assumptions
// in the code that all subtrees are multiple of 8 in depth and that log subtrees are always
// of depth 8. It is not possible to just change the constants above and have things still
// work. This is because of issues like byte packing of node IDs.
//
// The cache is optimized for the following use-cases (see sync.Map type for more details):
//  1. Parallel readers/writers working on non-intersecting subsets of subtrees/nodes.
//  2. Subtrees/nodes are rarely written, and mostly read.
type SubtreeCache struct {
	hasher hashers.LogHasher

	// subtrees contains the Subtree data read from storage, and is updated by
	// calls to SetNodes.
	subtrees sync.Map
	// dirtyPrefixes keeps track of all Subtrees which need to be written back
	// to storage.
	dirtyPrefixes sync.Map

	// populateConcurrency sets the amount of concurrency when repopulating subtrees.
	populateConcurrency int
}

// NewLogSubtreeCache creates and returns a SubtreeCache appropriate for use with a log
// tree. The caller must supply a suitable LogHasher.
func NewLogSubtreeCache(hasher hashers.LogHasher) *SubtreeCache {
	if *populateConcurrency <= 0 {
		panic(fmt.Errorf("populate_subtree_concurrency must be set to >= 1"))
	}
	return &SubtreeCache{
		hasher:              hasher,
		populateConcurrency: *populateConcurrency,
	}
}

// preload calculates the set of subtrees required to know the hashes of the
// passed in node IDs, uses getSubtrees to retrieve them, and finally populates
// the cache structures with the data. Returns the list of tile IDs not found.
func (s *SubtreeCache) preload(ids []compact.NodeID, getSubtrees GetSubtreesFunc) ([]string, error) {
	// Figure out the set of subtrees we need.
	want := make(map[string]bool)
	for _, id := range ids {
		subID := string(getTileID(id))
		if _, ok := want[subID]; ok {
			// No need to check s.subtrees map twice.
			continue
		}
		if _, ok := s.subtrees.Load(subID); !ok {
			want[subID] = true
		}
	}
	// Note: At this point multiple parallel preload invocations can happen to
	// getSubtrees with overlapping sets of IDs. It's okay because we collapse
	// results further below.

	// Don't make a read request for zero subtrees.
	if len(want) == 0 {
		return nil, nil
	}

	list := make([][]byte, 0, len(want))
	for id := range want {
		list = append(list, []byte(id))
	}
	subtrees, err := getSubtrees(list)
	if err != nil {
		return nil, err
	}
	if got, max := len(subtrees), len(want); got > max {
		return nil, fmt.Errorf("too many subtrees: %d, want <= %d", got, max)
	}

	ch := make(chan *storagepb.SubtreeProto, len(subtrees))
	workTokens := make(chan bool, s.populateConcurrency)
	for i := 0; i < s.populateConcurrency; i++ {
		workTokens <- true
	}
	wg := &sync.WaitGroup{}

	for _, t := range subtrees {
		t := t
		wg.Add(1)
		go func() {
			defer wg.Done()
			// wait for a token before starting work
			<-workTokens
			// return it when done
			defer func() { workTokens <- true }()

			PopulateLogTile(t, s.hasher)
			ch <- t // Note: This never blocks because len(ch) == len(subtrees).
		}()
	}

	go func() {
		wg.Wait()
		close(ch)
		close(workTokens)
	}()

	for t := range ch {
		if err := s.cacheSubtree(t); err != nil {
			return nil, err
		}
		delete(want, string(t.Prefix))
	}
	notFound := make([]string, 0, len(want))
	for id := range want {
		notFound = append(notFound, id)
	}
	return notFound, nil
}

func (s *SubtreeCache) cacheSubtree(t *storagepb.SubtreeProto) error {
	raw, loaded := s.subtrees.LoadOrStore(string(t.Prefix), t)
	if loaded { // There is already something with this key.
		if subtree, ok := raw.(*storagepb.SubtreeProto); !ok {
			return fmt.Errorf("at %x: not a subtree: %T", t.Prefix, raw)
		} else if !proto.Equal(t, subtree) {
			return fmt.Errorf("at %x: subtree mismatch", t.Prefix)
		}
	}
	return nil
}

// GetNodes returns the requested nodes, calling the getSubtrees function if
// they are not already cached.
func (s *SubtreeCache) GetNodes(ids []compact.NodeID, getSubtrees GetSubtreesFunc) ([]tree.Node, error) {
	if notFound, err := s.preload(ids, getSubtrees); err != nil {
		return nil, err
	} else if r := len(notFound); r != 0 {
		return nil, fmt.Errorf("preload did not get all tiles: %d not found", r)
	}

	ret := make([]tree.Node, 0, len(ids))
	for _, id := range ids {
		if h, err := s.getNodeHash(id); err != nil {
			return nil, fmt.Errorf("getNodeHash(%+v): %v", id, err)
		} else if h != nil {
			ret = append(ret, tree.Node{ID: id, Hash: h})
		}
	}
	return ret, nil
}

func (s *SubtreeCache) getCachedSubtree(id []byte) *storagepb.SubtreeProto {
	raw, found := s.subtrees.Load(string(id))
	if !found || raw == nil {
		return nil
	}
	// Note: If type assertion fails, nil is returned. Should not happen though.
	res, _ := raw.(*storagepb.SubtreeProto)
	return res
}

func (s *SubtreeCache) prefixIsDirty(prefixKey string) bool {
	_, found := s.dirtyPrefixes.Load(prefixKey)
	return found
}

// getNodeHash returns a single node hash from the cache.
func (s *SubtreeCache) getNodeHash(id compact.NodeID) ([]byte, error) {
	subID, sx := splitID(id)
	c := s.getCachedSubtree(subID)
	if c == nil {
		return nil, fmt.Errorf("tile %x not found", subID)
	}

	// Look up the hash in the appropriate map.
	// The leaf hashes are stored in a separate map to the internal nodes so that
	// we can easily dump (and later reconstruct) the internal nodes. As log subtrees
	// have a fixed depth if the suffix has the same number of significant bits as the
	// subtree depth then this is a leaf. For example if the subtree is depth 8 its leaves
	// have 8 significant suffix bits.
	if int32(sx.Bits()) == c.Depth {
		return c.Leaves[sx.String()], nil
	}
	return c.InternalNodes[sx.String()], nil
}

// SetNodes sets hashes for the given nodes in the cache.
func (s *SubtreeCache) SetNodes(nodes []tree.Node, getSubtrees GetSubtreesFunc) error {
	ids := make([]compact.NodeID, len(nodes))
	for i, n := range nodes {
		ids[i] = n.ID
	}
	notFound, err := s.preload(ids, getSubtrees)
	if err != nil {
		return err
	}
	for _, id := range notFound {
		c := newEmptyTile([]byte(id))
		s.subtrees.Store(string(id), c)
	}

	for _, n := range nodes {
		subID, sx := splitID(n.ID)
		c := s.getCachedSubtree(subID)
		if c == nil {
			return fmt.Errorf("tile %x not found", subID)
		}

		// Store the hash to the containing tile, and mark it as dirty if the hash
		// differs from the previously stored one.
		sfxKey := sx.String()
		if int32(sx.Bits()) == c.Depth { // This is a leaf node.
			if !bytes.Equal(c.Leaves[sfxKey], n.Hash) {
				c.Leaves[sfxKey] = n.Hash
				s.dirtyPrefixes.Store(string(subID), nil)
			}
		} else { // This is an internal node.
			if !bytes.Equal(c.InternalNodes[sfxKey], n.Hash) {
				c.InternalNodes[sfxKey] = n.Hash
				s.dirtyPrefixes.Store(string(subID), nil)
			}
		}
	}

	return nil
}

// Flush causes the cache to write all dirty Subtrees back to storage.
func (s *SubtreeCache) Flush(ctx context.Context, setSubtrees SetSubtreesFunc) error {
	treesToWrite := make([]*storagepb.SubtreeProto, 0)
	var rangeErr error
	s.subtrees.Range(func(rawK, rawV interface{}) bool {
		k, ok := rawK.(string)
		if !ok {
			rangeErr = fmt.Errorf("unknown key type: %t", rawV)
			return false
		}
		v, ok := rawV.(*storagepb.SubtreeProto)
		if !ok {
			rangeErr = fmt.Errorf("unknown value type: %t", rawV)
			return false
		}
		if s.prefixIsDirty(k) {
			bk := []byte(k)
			if !bytes.Equal(bk, v.Prefix) {
				rangeErr = fmt.Errorf("inconsistent cache: prefix key is %v, but cached object claims %v", bk, v.Prefix)
				return false
			}

			if len(v.Leaves) > 0 {
				if err := prepareLogTile(v); err != nil {
					rangeErr = err
					return false
				}
				treesToWrite = append(treesToWrite, v)
			}
		}
		return true
	})
	if rangeErr != nil {
		return rangeErr
	}
	if len(treesToWrite) == 0 {
		return nil
	}
	err := setSubtrees(ctx, treesToWrite)
	return err
}
<|endoftext|>


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

// This file contains replica methods related to range leases.

package storage

import (
	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/protoutil"
)

// pendingLeaseRequest coalesces RequestLease requests and lets callers
// join an in-progress one and wait for the result.
// The actual execution of the RequestLease Raft request is delegated to a
// replica.
// Methods are not thread-safe; a pendingLeaseRequest is logically part of
// a replica, so replica.mu should be used to synchronize all calls.
type pendingLeaseRequest struct {
	// Slice of channels to send on after lease acquisition.
	// If empty, then no request is in progress.
	llChans []chan *roachpb.Error
	// nextLease is the pending RequestLease request, if any. It can be used to
	// figure out if we're in the process of extending our own lease, or
	// transferring it to another replica.
	nextLease roachpb.Lease
}

// RequestPending returns the pending Lease, if one is in progress.
// Otherwise, nil.
func (p *pendingLeaseRequest) RequestPending() *roachpb.Lease {
	pending := len(p.llChans) > 0
	if pending {
		return &p.nextLease
	}
	return nil
}

// InitOrJoinRequest executes a RequestLease command asynchronously and returns a
// channel on which the result will be posted. If there's already a request in
// progress, we join in waiting for the results of that request.
// It is an error to call InitOrJoinRequest() while a request is in progress
// naming another replica as lease holder.
//
// replica is used to schedule and execute async work (proposing a RequestLease
// command). replica.mu is locked when delivering results, so calls from the
// replica happen either before or after a result for a pending request has
// happened.
//
// transfer needs to be set if the request represents a lease transfer (as
// opposed to an extension, or acquiring the lease when none is held).
//
// Note: Once this function gets a context to be used for cancellation, instead
// of replica.store.Stopper().ShouldQuiesce(), care will be needed for cancelling
// the Raft command, similar to replica.addWriteCmd.
func (p *pendingLeaseRequest) InitOrJoinRequest(
	replica *Replica,
	nextLeaseHolder roachpb.ReplicaDescriptor,
	timestamp hlc.Timestamp,
	startKey roachpb.Key,
	transfer bool,
) <-chan *roachpb.Error {
	if nextLease := p.RequestPending(); nextLease != nil {
		if nextLease.Replica.ReplicaID == nextLeaseHolder.ReplicaID {
			// Join a pending request asking for the same replica to become lease
			// holder.
			return p.JoinRequest()
		}
		llChan := make(chan *roachpb.Error, 1)
		// We can't join the request in progress.
		llChan <- roachpb.NewErrorf("request for different replica in progress "+
			"(requesting: %+v, in progress: %+v)",
			nextLeaseHolder.ReplicaID, nextLease.Replica.ReplicaID)
		return llChan
	}
	llChan := make(chan *roachpb.Error, 1)
	// No request in progress. Let's propose a Lease command asynchronously.
	// TODO(tschottdorf): get duration from configuration, either as a
	// config flag or, later, dynamically adjusted.
	startStasis := timestamp.Add(int64(replica.store.ctx.rangeLeaseActiveDuration), 0)
	expiration := startStasis.Add(int64(replica.store.Clock().MaxOffset()), 0)
	reqSpan := roachpb.Span{
		Key: startKey,
	}
	var leaseReq roachpb.Request
	reqLease := roachpb.Lease{
		Start:       timestamp,
		StartStasis: startStasis,
		Expiration:  expiration,
		Replica:     nextLeaseHolder,
	}
	if transfer {
		leaseReq = &roachpb.TransferLeaseRequest{
			Span:  reqSpan,
			Lease: reqLease,
		}
	} else {
		leaseReq = &roachpb.RequestLeaseRequest{
			Span:  reqSpan,
			Lease: reqLease,
		}
	}
	if replica.store.Stopper().RunAsyncTask(func() {
		// Propose a RequestLease command and wait for it to apply.
		var execPErr *roachpb.Error
		ba := roachpb.BatchRequest{}
		ba.Timestamp = replica.store.Clock().Now()
		ba.RangeID = replica.RangeID
		ba.Add(leaseReq)
		// Send lease request directly to raft in order to skip unnecessary
		// checks from normal request machinery, (e.g. the command queue).
		// Note that the command itself isn't traced, but usually the caller
		// waiting for the result has an active Trace.
		ch, _, err := replica.proposeRaftCommand(context.Background(), ba)
		if err != nil {
			execPErr = roachpb.NewError(err)
		} else {
			// If the command was committed, wait for the range to apply it.
			select {
			case c := <-ch:
				if c.Err != nil {
					if log.V(1) {
						log.Infof("failed to acquire lease for replica %s: %s", replica.store, c.Err)
					}
					execPErr = c.Err
				}
			case <-replica.store.Stopper().ShouldQuiesce():
				execPErr = roachpb.NewError(
					replica.newNotLeaseHolderError(nil, replica.store.StoreID(), replica.Desc()))
			}
		}

		// Send result of lease to all waiter channels.
		replica.mu.Lock()
		defer replica.mu.Unlock()
		for i, llChan := range p.llChans {
			// Don't send the same pErr object twice; this can lead to races. We could
			// clone every time but it's more efficient to send pErr itself to one of
			// the channels (the last one; if we send it earlier the race can still
			// happen).
			if i == len(p.llChans)-1 {
				llChan <- execPErr
			} else {
				llChan <- protoutil.Clone(execPErr).(*roachpb.Error) // works with `nil`
			}
		}
		p.llChans = p.llChans[:0]
		p.nextLease = roachpb.Lease{}
	}) != nil {
		// We failed to start the asynchronous task. Send a blank NotLeaseHolderError
		// back to indicate that we have no idea who the range lease holder might
		// be; we've withdrawn from active duty.
		llChan <- roachpb.NewError(
			replica.newNotLeaseHolderError(nil, replica.store.StoreID(), replica.mu.state.Desc))
		return llChan
	}
	p.llChans = append(p.llChans, llChan)
	p.nextLease = reqLease
	return llChan
}

// JoinRequest adds one more waiter to the currently pending request.
// It is the caller's responsibility to ensure that there is a pending request,
// and that the request is compatible with whatever the caller is currently
// wanting to do (i.e. the request is naming the intended node as the next
// lease holder).
func (p *pendingLeaseRequest) JoinRequest() <-chan *roachpb.Error {
	llChan := make(chan *roachpb.Error, 1)
	if len(p.llChans) == 0 {
		llChan <- roachpb.NewErrorf("no request in progress")
		return llChan
	}
	p.llChans = append(p.llChans, llChan)
	return llChan
}

// TransferInProgress returns the next lease, if the replica is in the process
// of transferring aways its range lease. This next lease indicates the next
// lease holder. Returns nil if no lease transfer is in progress.
//
// It is assumed that the replica owning this pendingLeaseRequest owns the
// LeaderLease.
//
// replicaID is the ID of the parent replica.
func (p *pendingLeaseRequest) TransferInProgress(
	replicaID roachpb.ReplicaID,
) *roachpb.Lease {
	if nextLease := p.RequestPending(); nextLease != nil {
		// Is the lease being transferred? (as opposed to just extended)
		if replicaID != nextLease.Replica.ReplicaID {
			return nextLease
		}
	}
	return nil
}

// requestLeaseLocked executes a request to obtain or extend a lease
// asynchronously and returns a channel on which the result will be posted. If
// there's already a request in progress, we join in waiting for the results of
// that request. Unless an error is returned, the obtained lease will be valid
// for a time interval containing the requested timestamp.
// If a transfer is in progress, a NotLeaderError directing to the recipient is
// sent on the returned chan.
func (r *Replica) requestLeaseLocked(timestamp hlc.Timestamp) <-chan *roachpb.Error {
	// Propose a Raft command to get a lease for this replica.
	repDesc, err := r.getReplicaDescriptorLocked()
	if err != nil {
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(err)
		return llChan
	}
	transferLease := r.mu.pendingLeaseRequest.TransferInProgress(repDesc.ReplicaID)
	if transferLease != nil {
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(
			r.newNotLeaseHolderError(transferLease, r.store.StoreID(), r.mu.state.Desc))
		return llChan
	}
	if r.store.IsDrainingLeases() {
		// We've retired from active duty.
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(r.newNotLeaseHolderError(nil, r.store.StoreID(), r.mu.state.Desc))
		return llChan
	}
	return r.mu.pendingLeaseRequest.InitOrJoinRequest(
		r, repDesc, timestamp, r.mu.state.Desc.StartKey.AsRawKey(), false /* transfer */)
}

// AdminTransferLease transfers the LeaderLease to another replica. Only the
// current holder of the LeaderLease can do a transfer, because it needs to
// stop serving reads and proposing Raft commands (CPut is a read) after
// sending the transfer command. If it did not stop serving reads immediately,
// it would potentially serve reads with timestamps greater than the start
// timestamp of the new (transferred) lease. More subtly, the replica can't
// even serve reads or propose commands with timestamps lower than the start of
// the new lease because it could lead to read your own write violations (see
// comments on the stasis period in the Lease proto). We could, in principle,
// serve reads more than the maximum clock offset in the past.
//
// The method waits for any in-progress lease extension to be done, and it also
// blocks until the transfer is done. If a transfer is already in progress,
// this method joins in waiting for it to complete if it's transferring to the
// same replica. Otherwise, a NotLeaderError is returned.
//
// TODO(andrei): figure out how to persist the "not serving" state across node
// restarts.
func (r *Replica) AdminTransferLease(nextLeader roachpb.ReplicaDescriptor) error {
	// initTransferHelper inits a transfer if no extension is in progress.
	// It returns a channel for waiting for the result of a pending
	// extension (if any is in progress) and a channel for waiting for the
	// transfer (if it was successfully initiated).
	initTransferHelper := func() (<-chan *roachpb.Error, <-chan *roachpb.Error, error) {
		r.mu.Lock()
		defer r.mu.Unlock()
		lease := r.mu.state.Lease
		if !lease.OwnedBy(r.store.StoreID()) {
			return nil, nil, r.newNotLeaseHolderError(lease, r.store.StoreID(), r.Desc())
		}
		nextLease := r.mu.pendingLeaseRequest.RequestPending()
		if nextLease != nil && nextLease.Replica != nextLeader {
			repDesc, err := r.getReplicaDescriptorLocked()
			if err != nil {
				return nil, nil, err
			}
			if nextLease.Replica == repDesc {
				// There's an extension in progress. Let's wait for it to succeed and
				// try again.
				return r.mu.pendingLeaseRequest.JoinRequest(), nil, nil
			}
			// Another transfer is in progress, and it's not transferring to the
			// same replica we'd like.
			return nil, nil, r.newNotLeaseHolderError(nextLease, r.store.StoreID(), r.mu.state.Desc)
		}
		// No extension in progress; start a transfer.
		nextLeaseBegin := r.store.Clock().Now()
		// Don't transfer immediately after a node restart since we might have
		// served higher timestamps before the restart.
		nextLeaseBegin.Forward(
			hlc.ZeroTimestamp.Add(r.store.startedAt+int64(r.store.Clock().MaxOffset()), 0))
		transfer := r.mu.pendingLeaseRequest.InitOrJoinRequest(
			r, nextLeader, nextLeaseBegin,
			r.mu.state.Desc.StartKey.AsRawKey(), true /* transfer */)
		return nil, transfer, nil
	}

	// Loop while there's an extension in progress.
	for {
		// See if there's an extension in progress that we have to wait for.
		// If there isn't, request a transfer.
		extension, transfer, err := initTransferHelper()
		if err != nil {
			return err
		}
		if extension == nil {
			return (<-transfer).GoError()
		}
		// Wait for the in-progress extension without holding the mutex.
		if r.store.TestingKnobs().LeaseTransferBlockedOnExtensionEvent != nil {
			r.store.TestingKnobs().LeaseTransferBlockedOnExtensionEvent(nextLeader)
		}
		<-extension
	}
}
<commit_msg>storage: fixes a deadlock in which the process tries to lock itself<commit_after>// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Andrei Matei (andreimatei1@gmail.com)
// Author: Tobias Schottdorf (tobias.schottdorf@gmail.com)

// This file contains replica methods related to range leases.

package storage

import (
	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/protoutil"
)

// pendingLeaseRequest coalesces RequestLease requests and lets callers
// join an in-progress one and wait for the result.
// The actual execution of the RequestLease Raft request is delegated to a
// replica.
// Methods are not thread-safe; a pendingLeaseRequest is logically part of
// a replica, so replica.mu should be used to synchronize all calls.
type pendingLeaseRequest struct {
	// Slice of channels to send on after lease acquisition.
	// If empty, then no request is in progress.
	llChans []chan *roachpb.Error
	// nextLease is the pending RequestLease request, if any. It can be used to
	// figure out if we're in the process of extending our own lease, or
	// transferring it to another replica.
	nextLease roachpb.Lease
}

// RequestPending returns the pending Lease, if one is in progress.
// Otherwise, nil.
func (p *pendingLeaseRequest) RequestPending() *roachpb.Lease {
	pending := len(p.llChans) > 0
	if pending {
		return &p.nextLease
	}
	return nil
}

// InitOrJoinRequest executes a RequestLease command asynchronously and returns a
// channel on which the result will be posted. If there's already a request in
// progress, we join in waiting for the results of that request.
// It is an error to call InitOrJoinRequest() while a request is in progress
// naming another replica as lease holder.
//
// replica is used to schedule and execute async work (proposing a RequestLease
// command). replica.mu is locked when delivering results, so calls from the
// replica happen either before or after a result for a pending request has
// happened.
//
// transfer needs to be set if the request represents a lease transfer (as
// opposed to an extension, or acquiring the lease when none is held).
//
// Note: Once this function gets a context to be used for cancellation, instead
// of replica.store.Stopper().ShouldQuiesce(), care will be needed for cancelling
// the Raft command, similar to replica.addWriteCmd.
func (p *pendingLeaseRequest) InitOrJoinRequest(
	replica *Replica,
	nextLeaseHolder roachpb.ReplicaDescriptor,
	timestamp hlc.Timestamp,
	startKey roachpb.Key,
	transfer bool,
) <-chan *roachpb.Error {
	if nextLease := p.RequestPending(); nextLease != nil {
		if nextLease.Replica.ReplicaID == nextLeaseHolder.ReplicaID {
			// Join a pending request asking for the same replica to become lease
			// holder.
			return p.JoinRequest()
		}
		llChan := make(chan *roachpb.Error, 1)
		// We can't join the request in progress.
		llChan <- roachpb.NewErrorf("request for different replica in progress "+
			"(requesting: %+v, in progress: %+v)",
			nextLeaseHolder.ReplicaID, nextLease.Replica.ReplicaID)
		return llChan
	}
	llChan := make(chan *roachpb.Error, 1)
	// No request in progress. Let's propose a Lease command asynchronously.
	// TODO(tschottdorf): get duration from configuration, either as a
	// config flag or, later, dynamically adjusted.
	startStasis := timestamp.Add(int64(replica.store.ctx.rangeLeaseActiveDuration), 0)
	expiration := startStasis.Add(int64(replica.store.Clock().MaxOffset()), 0)
	reqSpan := roachpb.Span{
		Key: startKey,
	}
	var leaseReq roachpb.Request
	reqLease := roachpb.Lease{
		Start:       timestamp,
		StartStasis: startStasis,
		Expiration:  expiration,
		Replica:     nextLeaseHolder,
	}
	if transfer {
		leaseReq = &roachpb.TransferLeaseRequest{
			Span:  reqSpan,
			Lease: reqLease,
		}
	} else {
		leaseReq = &roachpb.RequestLeaseRequest{
			Span:  reqSpan,
			Lease: reqLease,
		}
	}
	if replica.store.Stopper().RunAsyncTask(func() {
		// Propose a RequestLease command and wait for it to apply.
		var execPErr *roachpb.Error
		ba := roachpb.BatchRequest{}
		ba.Timestamp = replica.store.Clock().Now()
		ba.RangeID = replica.RangeID
		ba.Add(leaseReq)
		// Send lease request directly to raft in order to skip unnecessary
		// checks from normal request machinery, (e.g. the command queue).
		// Note that the command itself isn't traced, but usually the caller
		// waiting for the result has an active Trace.
		ch, _, err := replica.proposeRaftCommand(context.Background(), ba)
		if err != nil {
			execPErr = roachpb.NewError(err)
		} else {
			// If the command was committed, wait for the range to apply it.
			select {
			case c := <-ch:
				if c.Err != nil {
					if log.V(1) {
						log.Infof("failed to acquire lease for replica %s: %s", replica.store, c.Err)
					}
					execPErr = c.Err
				}
			case <-replica.store.Stopper().ShouldQuiesce():
				execPErr = roachpb.NewError(
					replica.newNotLeaseHolderError(nil, replica.store.StoreID(), replica.Desc()))
			}
		}

		// Send result of lease to all waiter channels.
		replica.mu.Lock()
		defer replica.mu.Unlock()
		for i, llChan := range p.llChans {
			// Don't send the same pErr object twice; this can lead to races. We could
			// clone every time but it's more efficient to send pErr itself to one of
			// the channels (the last one; if we send it earlier the race can still
			// happen).
			if i == len(p.llChans)-1 {
				llChan <- execPErr
			} else {
				llChan <- protoutil.Clone(execPErr).(*roachpb.Error) // works with `nil`
			}
		}
		p.llChans = p.llChans[:0]
		p.nextLease = roachpb.Lease{}
	}) != nil {
		// We failed to start the asynchronous task. Send a blank NotLeaseHolderError
		// back to indicate that we have no idea who the range lease holder might
		// be; we've withdrawn from active duty.
		llChan <- roachpb.NewError(
			replica.newNotLeaseHolderError(nil, replica.store.StoreID(), replica.mu.state.Desc))
		return llChan
	}
	p.llChans = append(p.llChans, llChan)
	p.nextLease = reqLease
	return llChan
}

// JoinRequest adds one more waiter to the currently pending request.
// It is the caller's responsibility to ensure that there is a pending request,
// and that the request is compatible with whatever the caller is currently
// wanting to do (i.e. the request is naming the intended node as the next
// lease holder).
func (p *pendingLeaseRequest) JoinRequest() <-chan *roachpb.Error {
	llChan := make(chan *roachpb.Error, 1)
	if len(p.llChans) == 0 {
		llChan <- roachpb.NewErrorf("no request in progress")
		return llChan
	}
	p.llChans = append(p.llChans, llChan)
	return llChan
}

// TransferInProgress returns the next lease, if the replica is in the process
// of transferring aways its range lease. This next lease indicates the next
// lease holder. Returns nil if no lease transfer is in progress.
//
// It is assumed that the replica owning this pendingLeaseRequest owns the
// LeaderLease.
//
// replicaID is the ID of the parent replica.
func (p *pendingLeaseRequest) TransferInProgress(
	replicaID roachpb.ReplicaID,
) *roachpb.Lease {
	if nextLease := p.RequestPending(); nextLease != nil {
		// Is the lease being transferred? (as opposed to just extended)
		if replicaID != nextLease.Replica.ReplicaID {
			return nextLease
		}
	}
	return nil
}

// requestLeaseLocked executes a request to obtain or extend a lease
// asynchronously and returns a channel on which the result will be posted. If
// there's already a request in progress, we join in waiting for the results of
// that request. Unless an error is returned, the obtained lease will be valid
// for a time interval containing the requested timestamp.
// If a transfer is in progress, a NotLeaderError directing to the recipient is
// sent on the returned chan.
func (r *Replica) requestLeaseLocked(timestamp hlc.Timestamp) <-chan *roachpb.Error {
	// Propose a Raft command to get a lease for this replica.
	repDesc, err := r.getReplicaDescriptorLocked()
	if err != nil {
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(err)
		return llChan
	}
	transferLease := r.mu.pendingLeaseRequest.TransferInProgress(repDesc.ReplicaID)
	if transferLease != nil {
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(
			r.newNotLeaseHolderError(transferLease, r.store.StoreID(), r.mu.state.Desc))
		return llChan
	}
	if r.store.IsDrainingLeases() {
		// We've retired from active duty.
		llChan := make(chan *roachpb.Error, 1)
		llChan <- roachpb.NewError(r.newNotLeaseHolderError(nil, r.store.StoreID(), r.mu.state.Desc))
		return llChan
	}
	return r.mu.pendingLeaseRequest.InitOrJoinRequest(
		r, repDesc, timestamp, r.mu.state.Desc.StartKey.AsRawKey(), false /* transfer */)
}

// AdminTransferLease transfers the LeaderLease to another replica. Only the
// current holder of the LeaderLease can do a transfer, because it needs to
// stop serving reads and proposing Raft commands (CPut is a read) after
// sending the transfer command. If it did not stop serving reads immediately,
// it would potentially serve reads with timestamps greater than the start
// timestamp of the new (transferred) lease. More subtly, the replica can't
// even serve reads or propose commands with timestamps lower than the start of
// the new lease because it could lead to read your own write violations (see
// comments on the stasis period in the Lease proto). We could, in principle,
// serve reads more than the maximum clock offset in the past.
//
// The method waits for any in-progress lease extension to be done, and it also
// blocks until the transfer is done. If a transfer is already in progress,
// this method joins in waiting for it to complete if it's transferring to the
// same replica. Otherwise, a NotLeaderError is returned.
//
// TODO(andrei): figure out how to persist the "not serving" state across node
// restarts.
func (r *Replica) AdminTransferLease(nextLeader roachpb.ReplicaDescriptor) error {
	// initTransferHelper inits a transfer if no extension is in progress.
	// It returns a channel for waiting for the result of a pending
	// extension (if any is in progress) and a channel for waiting for the
	// transfer (if it was successfully initiated).
	initTransferHelper := func() (<-chan *roachpb.Error, <-chan *roachpb.Error, error) {
		r.mu.Lock()
		defer r.mu.Unlock()
		lease := r.mu.state.Lease
		if !lease.OwnedBy(r.store.StoreID()) {
			return nil, nil, r.newNotLeaseHolderError(lease, r.store.StoreID(), r.mu.state.Desc)
		}
		nextLease := r.mu.pendingLeaseRequest.RequestPending()
		if nextLease != nil && nextLease.Replica != nextLeader {
			repDesc, err := r.getReplicaDescriptorLocked()
			if err != nil {
				return nil, nil, err
			}
			if nextLease.Replica == repDesc {
				// There's an extension in progress. Let's wait for it to succeed and
				// try again.
				return r.mu.pendingLeaseRequest.JoinRequest(), nil, nil
			}
			// Another transfer is in progress, and it's not transferring to the
			// same replica we'd like.
			return nil, nil, r.newNotLeaseHolderError(nextLease, r.store.StoreID(), r.mu.state.Desc)
		}
		// No extension in progress; start a transfer.
		nextLeaseBegin := r.store.Clock().Now()
		// Don't transfer immediately after a node restart since we might have
		// served higher timestamps before the restart.
		nextLeaseBegin.Forward(
			hlc.ZeroTimestamp.Add(r.store.startedAt+int64(r.store.Clock().MaxOffset()), 0))
		transfer := r.mu.pendingLeaseRequest.InitOrJoinRequest(
			r, nextLeader, nextLeaseBegin,
			r.mu.state.Desc.StartKey.AsRawKey(), true /* transfer */)
		return nil, transfer, nil
	}

	// Loop while there's an extension in progress.
	for {
		// See if there's an extension in progress that we have to wait for.
		// If there isn't, request a transfer.
		extension, transfer, err := initTransferHelper()
		if err != nil {
			return err
		}
		if extension == nil {
			return (<-transfer).GoError()
		}
		// Wait for the in-progress extension without holding the mutex.
		if r.store.TestingKnobs().LeaseTransferBlockedOnExtensionEvent != nil {
			r.store.TestingKnobs().LeaseTransferBlockedOnExtensionEvent(nextLeader)
		}
		<-extension
	}
}
<|endoftext|>


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

package db

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"regexp"

	"github.com/couchbaselabs/go-couchbase"

	"github.com/couchbaselabs/sync_gateway/auth"
	"github.com/couchbaselabs/sync_gateway/base"
	"github.com/couchbaselabs/sync_gateway/channels"
)

var kDBNameMatch = regexp.MustCompile("[-%+()$_a-z0-9]+")

type DatabaseContext struct {
	Name          string
	Bucket        *couchbase.Bucket
	sequences     *sequenceAllocator
	ChannelMapper *channels.ChannelMapper
	Validator     *Validator
}

// Represents a simulated CouchDB database.
type Database struct {
	*DatabaseContext
	user *auth.User
}

// Helper function to open a Couchbase connection and return a specific bucket.
func ConnectToBucket(couchbaseURL, poolName, bucketName string) (bucket *couchbase.Bucket, err error) {
	bucket, err = couchbase.GetBucket(couchbaseURL, poolName, bucketName)
	if err != nil {
		return
	}
	log.Printf("Connected to <%s>, pool %s, bucket %s", couchbaseURL, poolName, bucketName)
	err = installViews(bucket)
	return
}

func NewDatabaseContext(dbName string, bucket *couchbase.Bucket) (*DatabaseContext, error) {
	sequences, err := newSequenceAllocator(bucket)
	if err != nil {
		return nil, err
	}
	return &DatabaseContext{Name: dbName, Bucket: bucket, sequences: sequences}, nil
}

// Makes a Database object given its name and bucket.
func GetDatabase(context *DatabaseContext, user *auth.User) (*Database, error) {
	return &Database{context, user}, nil
}

func CreateDatabase(context *DatabaseContext) (*Database, error) {
	return &Database{context, nil}, nil
}

func (db *Database) SameAs(otherdb *Database) bool {
	return db != nil && otherdb != nil &&
		db.Bucket == otherdb.Bucket
}

// Sets the database object's channelMapper and validator based on the JS code in _design/channels
func (db *Database) ReadDesignDocument() error {
	body, err := db.Get("_design/channels")
	if err != nil {
		if status, _ := base.ErrorAsHTTPStatus(err); status == http.StatusNotFound {
			err = nil // missing design document is not an error
		}
		return err
	}
	src, ok := body["channelmap"].(string)
	if ok {
		log.Printf("Channel mapper = %s", src)
		db.ChannelMapper, err = channels.NewChannelMapper(src)
		if err != nil {
			log.Printf("WARNING: Error loading channel mapper: %s", err)
			return err
		}
	}
	src, ok = body["validate_doc_update"].(string)
	if ok {
		log.Printf("Validator = %s", src)
		db.Validator, err = NewValidator(src)
		if err != nil {
			log.Printf("WARNING: Error loading validator: %s", err)
			return err
		}
	}
	return nil
}

//////// ALL DOCUMENTS:

// The number of documents in the database.
func (db *Database) DocCount() int {
	vres, err := db.queryAllDocs(true)
	if err != nil {
		return -1
	}
	if len(vres.Rows) == 0 {
		return 0
	}
	return int(vres.Rows[0].Value.(float64))
}

func installViews(bucket *couchbase.Bucket) error {
	node := bucket.Nodes[rand.Intn(len(bucket.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		fmt.Printf("Failed to parse %s", node.CouchAPIBase)
		return err
	}
	u.Path = fmt.Sprintf("/%s/_design/%s", bucket.Name, "sync_gateway")

	// View for finding every Couchbase doc used by a database (for cleanup upon deletion)
	allbits_map := `function (doc, meta) {
                     var pieces = meta.id.split(":", 2);
                     if (pieces.length < 2)
                       return;
					 var type = pieces[0];
					 if (type == "doc" || type == "ldoc" || type == "seq")
                       emit(meta.id, null); }`
	// View for _all_docs
	alldocs_map := `function (doc, meta) {
                     var pieces = meta.id.split(":", 2);
                     if (pieces.length < 2 || pieces[0] != "doc")
                       return;
                     if (doc.deleted || doc.id===undefined)
                       return;
                     emit(doc.id, doc.rev); }`
	// View for _changes feed, i.e. a by-sequence index
	changes_map := `function (doc, meta) {
                    if (doc.sequence === undefined)
                        return;
                    var pieces = meta.id.split(":", 2);
                    if (pieces.length < 2 || pieces[0] != "doc" || doc.id===undefined)
                        return;
                    var value = [doc.id, doc.rev];
                    if (doc.deleted)
                        value.push(true);
                    emit(doc.sequence, value); }`
	// By-channels view
	channels_map := `function (doc, meta) {
						var sequence = doc.sequence;
	                    if (sequence === undefined)
	                        return;
	                    var pieces = meta.id.split(":", 2);
	                    if (pieces.length < 2 || pieces[0] != "doc" || doc.id===undefined)
	                        return;
	                    var value = [doc.id, doc.rev];
	                    if (doc.deleted)
	                        value.push(true);
						emit(["*", sequence], value);
						var channels = doc["channels"];
						if (channels) {
							for (var name in channels) {
								removed = channels[name];
								if (!removed)
									emit([name, sequence], value);
								else
									emit([name, removed.seq], [doc.id, removed.rev, false]);
							}
						}
					}`

	access_map := `function (doc, meta) {
	                    var sequence = doc.sequence;
	                    if (doc.deleted || sequence === undefined)
	                        return;
	                    var access = doc["access"];
	                    if (access) {
	                        for (var name in access) {
	                            emit(name, access[name]);
	                        }
	                    }
	               }`
	// View for mapping revision IDs to documents (for vacuuming)
	revs_map := `function (doc, meta) {
					  var pieces = meta.id.split(":", 2);
					  if (pieces.length < 2)
					    return;
					  if (pieces[0] == "rev") {
					    emit(pieces[1], true);
					  } else if (pieces[0] == "doc") {
					    var keys = doc.history.keys;
					    for (var i = 0; i < keys.length; ++i) {
					      if (keys[i] != "")
					        emit(keys[i], false);
					    }
					  }
					}`
	// View for mapping attachment IDs to revisions (for vacuuming)
	atts_map := `function (doc, meta) {
				  var pieces = meta.id.split(":", 2);
				  if (pieces.length < 2)
				    return;
				  if (pieces[0] == "att") {
				    emit(pieces[1], true);
				  } else if (pieces[0] == "rev") {
				    var atts = doc._attachments;
				    if (atts) {
				      for (var i = 0; i < atts.length; ++i) {
				        var digest = atts[i].digest
				        if (digest)
				          emit(digest, false);
				      }
				    }
				  }
				}`
	// Reduce function for revs_map and atts_map
	revs_or_atts_reduce := `function(keys, values, rereduce) {
							  for (var i = 0; i < values.length; ++i) {
							    if (!values[i])
							      return false;
							  }
							  return true;
							}`

	ddoc := Body{
		"language": "javascript",
		"views": Body{
			"all_bits": Body{"map": allbits_map},
			"all_docs": Body{"map": alldocs_map, "reduce": "_count"},
			"channels": Body{"map": channels_map},
			"access":   Body{"map": access_map},
			"revs":     Body{"map": revs_map, "reduce": revs_or_atts_reduce},
			"atts":     Body{"map": atts_map, "reduce": revs_or_atts_reduce},
			"changes":  Body{"map": changes_map}}}
	payload, err := json.Marshal(ddoc)
	rq, err := http.NewRequest("PUT", u.String(), bytes.NewBuffer(payload))
	rq.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(rq)

	if err == nil && response.StatusCode > 299 {
		err = &base.HTTPError{Status: response.StatusCode, Message: response.Status}
	}
	if err != nil {
		log.Printf("WARNING: Error installing Couchbase design doc: %v", err)
	}
	return err
}

type IDAndRev struct {
	DocID string
	RevID string
}

// Returns all document IDs as an array.
func (db *Database) AllDocIDs() ([]IDAndRev, error) {
	vres, err := db.queryAllDocs(false)
	if err != nil {
		return nil, err
	}

	rows := vres.Rows
	result := make([]IDAndRev, 0, len(rows))
	for _, row := range rows {
		result = append(result, IDAndRev{DocID: row.Key.(string), RevID: row.Value.(string)})
	}
	return result, nil
}

func (db *Database) queryAllDocs(reduce bool) (couchbase.ViewResult, error) {
	opts := Body{"stale": false, "reduce": reduce}
	vres, err := db.Bucket.View("sync_gateway", "all_docs", opts)
	if err != nil {
		log.Printf("WARNING: all_docs got error: %v", err)
	}
	return vres, err
}

// Deletes a database (and all documents)
func (db *Database) Delete() error {
	opts := Body{"stale": false}
	vres, err := db.Bucket.View("sync_gateway", "all_bits", opts)
	if err != nil {
		log.Printf("WARNING: all_bits view returned %v", err)
		return err
	}

	//FIX: Is there a way to do this in one operation?
	log.Printf("Deleting %d documents of %q ...", len(vres.Rows), db.Name)
	for _, row := range vres.Rows {
		if base.Logging {
			log.Printf("\tDeleting %q", row.ID)
		}
		if err := db.Bucket.Delete(row.ID); err != nil {
			log.Printf("WARNING: Error deleting %q: %v", row.ID, err)
		}
	}
	return nil
}

func vacuumContent(bucket *couchbase.Bucket, viewName string, docPrefix string) (int, error) {
	vres, err := bucket.View("sync_gateway", viewName, Body{"stale": false, "group_level": 1})
	if err != nil {
		log.Printf("WARNING: %s view returned %v", viewName, err)
		return 0, err
	}
	deletions := 0
	for _, row := range vres.Rows {
		if row.Value.(bool) {
			deletions += vacIt(bucket, docPrefix+row.Key.(string))
		}
	}
	return deletions, nil
}

func vacIt(bucket *couchbase.Bucket, docid string) int {
	if err := bucket.Delete(docid); err != nil {
		log.Printf("\tError vacuuming %q: %v", docid, err)
		return 0
	}
	if base.Logging {
		log.Printf("\tVacuumed %q", docid)
	}
	return 1
}

// Deletes all orphaned CouchDB revisions not used by any documents.
func VacuumRevisions(bucket *couchbase.Bucket) (int, error) {
	return vacuumContent(bucket, "revs", "rev:")
}

// Deletes all orphaned CouchDB attachments not used by any revisions.
func VacuumAttachments(bucket *couchbase.Bucket) (int, error) {
	return vacuumContent(bucket, "atts", "att:")
}

// Re-runs the channelMapper on every document in the database.
// To be used when the JavaScript channelmap function changes.
func (db *Database) UpdateAllDocChannels() error {
	log.Printf("Recomputing document channels...")
	vres, err := db.Bucket.View("sync_gateway", "all_docs", Body{"stale": false, "reduce": false})
	if err != nil {
		return err
	}
	for _, row := range vres.Rows {
		docid := row.Key.(string)
		key := db.realDocID(docid)
		err := db.Bucket.Update(key, 0, func(currentValue []byte) ([]byte, error) {
			// Be careful: this block can be invoked multiple times if there are races!
			if currentValue == nil {
				return nil, couchbase.UpdateCancel // someone deleted it?!
			}
			doc := newDocument()
			if err := json.Unmarshal(currentValue, doc); err != nil {
				return nil, err
			}
			body, err := db.getRevFromDoc(doc, "", false)
			if err != nil {
				return nil, err
			}
			channels, access := db.getChannelsAndAccess(body)
			if !db.updateDocChannels(doc, channels) {
				return nil, couchbase.UpdateCancel // unchanged
			}
			if !db.updateDocAccess(doc, access) {
				return nil, couchbase.UpdateCancel // unchanged
			}
			log.Printf("\tSaving updated channels of %q", docid)
			return json.Marshal(doc)
		})
		if err != nil && err != couchbase.UpdateCancel {
			log.Printf("WARNING: Error updating doc %q: %v", docid, err)
		}
	}
	return nil
}
<commit_msg>avoid cancel<commit_after>//  Copyright (c) 2012 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package db

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"regexp"

	"github.com/couchbaselabs/go-couchbase"

	"github.com/couchbaselabs/sync_gateway/auth"
	"github.com/couchbaselabs/sync_gateway/base"
	"github.com/couchbaselabs/sync_gateway/channels"
)

var kDBNameMatch = regexp.MustCompile("[-%+()$_a-z0-9]+")

type DatabaseContext struct {
	Name          string
	Bucket        *couchbase.Bucket
	sequences     *sequenceAllocator
	ChannelMapper *channels.ChannelMapper
	Validator     *Validator
}

// Represents a simulated CouchDB database.
type Database struct {
	*DatabaseContext
	user *auth.User
}

// Helper function to open a Couchbase connection and return a specific bucket.
func ConnectToBucket(couchbaseURL, poolName, bucketName string) (bucket *couchbase.Bucket, err error) {
	bucket, err = couchbase.GetBucket(couchbaseURL, poolName, bucketName)
	if err != nil {
		return
	}
	log.Printf("Connected to <%s>, pool %s, bucket %s", couchbaseURL, poolName, bucketName)
	err = installViews(bucket)
	return
}

func NewDatabaseContext(dbName string, bucket *couchbase.Bucket) (*DatabaseContext, error) {
	sequences, err := newSequenceAllocator(bucket)
	if err != nil {
		return nil, err
	}
	return &DatabaseContext{Name: dbName, Bucket: bucket, sequences: sequences}, nil
}

// Makes a Database object given its name and bucket.
func GetDatabase(context *DatabaseContext, user *auth.User) (*Database, error) {
	return &Database{context, user}, nil
}

func CreateDatabase(context *DatabaseContext) (*Database, error) {
	return &Database{context, nil}, nil
}

func (db *Database) SameAs(otherdb *Database) bool {
	return db != nil && otherdb != nil &&
		db.Bucket == otherdb.Bucket
}

// Sets the database object's channelMapper and validator based on the JS code in _design/channels
func (db *Database) ReadDesignDocument() error {
	body, err := db.Get("_design/channels")
	if err != nil {
		if status, _ := base.ErrorAsHTTPStatus(err); status == http.StatusNotFound {
			err = nil // missing design document is not an error
		}
		return err
	}
	src, ok := body["channelmap"].(string)
	if ok {
		log.Printf("Channel mapper = %s", src)
		db.ChannelMapper, err = channels.NewChannelMapper(src)
		if err != nil {
			log.Printf("WARNING: Error loading channel mapper: %s", err)
			return err
		}
	}
	src, ok = body["validate_doc_update"].(string)
	if ok {
		log.Printf("Validator = %s", src)
		db.Validator, err = NewValidator(src)
		if err != nil {
			log.Printf("WARNING: Error loading validator: %s", err)
			return err
		}
	}
	return nil
}

//////// ALL DOCUMENTS:

// The number of documents in the database.
func (db *Database) DocCount() int {
	vres, err := db.queryAllDocs(true)
	if err != nil {
		return -1
	}
	if len(vres.Rows) == 0 {
		return 0
	}
	return int(vres.Rows[0].Value.(float64))
}

func installViews(bucket *couchbase.Bucket) error {
	node := bucket.Nodes[rand.Intn(len(bucket.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		fmt.Printf("Failed to parse %s", node.CouchAPIBase)
		return err
	}
	u.Path = fmt.Sprintf("/%s/_design/%s", bucket.Name, "sync_gateway")

	// View for finding every Couchbase doc used by a database (for cleanup upon deletion)
	allbits_map := `function (doc, meta) {
                     var pieces = meta.id.split(":", 2);
                     if (pieces.length < 2)
                       return;
					 var type = pieces[0];
					 if (type == "doc" || type == "ldoc" || type == "seq")
                       emit(meta.id, null); }`
	// View for _all_docs
	alldocs_map := `function (doc, meta) {
                     var pieces = meta.id.split(":", 2);
                     if (pieces.length < 2 || pieces[0] != "doc")
                       return;
                     if (doc.deleted || doc.id===undefined)
                       return;
                     emit(doc.id, doc.rev); }`
	// View for _changes feed, i.e. a by-sequence index
	changes_map := `function (doc, meta) {
                    if (doc.sequence === undefined)
                        return;
                    var pieces = meta.id.split(":", 2);
                    if (pieces.length < 2 || pieces[0] != "doc" || doc.id===undefined)
                        return;
                    var value = [doc.id, doc.rev];
                    if (doc.deleted)
                        value.push(true);
                    emit(doc.sequence, value); }`
	// By-channels view
	channels_map := `function (doc, meta) {
						var sequence = doc.sequence;
	                    if (sequence === undefined)
	                        return;
	                    var pieces = meta.id.split(":", 2);
	                    if (pieces.length < 2 || pieces[0] != "doc" || doc.id===undefined)
	                        return;
	                    var value = [doc.id, doc.rev];
	                    if (doc.deleted)
	                        value.push(true);
						emit(["*", sequence], value);
						var channels = doc["channels"];
						if (channels) {
							for (var name in channels) {
								removed = channels[name];
								if (!removed)
									emit([name, sequence], value);
								else
									emit([name, removed.seq], [doc.id, removed.rev, false]);
							}
						}
					}`

	access_map := `function (doc, meta) {
	                    var sequence = doc.sequence;
	                    if (doc.deleted || sequence === undefined)
	                        return;
	                    var access = doc["access"];
	                    if (access) {
	                        for (var name in access) {
	                            emit(name, access[name]);
	                        }
	                    }
	               }`
	// View for mapping revision IDs to documents (for vacuuming)
	revs_map := `function (doc, meta) {
					  var pieces = meta.id.split(":", 2);
					  if (pieces.length < 2)
					    return;
					  if (pieces[0] == "rev") {
					    emit(pieces[1], true);
					  } else if (pieces[0] == "doc") {
					    var keys = doc.history.keys;
					    for (var i = 0; i < keys.length; ++i) {
					      if (keys[i] != "")
					        emit(keys[i], false);
					    }
					  }
					}`
	// View for mapping attachment IDs to revisions (for vacuuming)
	atts_map := `function (doc, meta) {
				  var pieces = meta.id.split(":", 2);
				  if (pieces.length < 2)
				    return;
				  if (pieces[0] == "att") {
				    emit(pieces[1], true);
				  } else if (pieces[0] == "rev") {
				    var atts = doc._attachments;
				    if (atts) {
				      for (var i = 0; i < atts.length; ++i) {
				        var digest = atts[i].digest
				        if (digest)
				          emit(digest, false);
				      }
				    }
				  }
				}`
	// Reduce function for revs_map and atts_map
	revs_or_atts_reduce := `function(keys, values, rereduce) {
							  for (var i = 0; i < values.length; ++i) {
							    if (!values[i])
							      return false;
							  }
							  return true;
							}`

	ddoc := Body{
		"language": "javascript",
		"views": Body{
			"all_bits": Body{"map": allbits_map},
			"all_docs": Body{"map": alldocs_map, "reduce": "_count"},
			"channels": Body{"map": channels_map},
			"access":   Body{"map": access_map},
			"revs":     Body{"map": revs_map, "reduce": revs_or_atts_reduce},
			"atts":     Body{"map": atts_map, "reduce": revs_or_atts_reduce},
			"changes":  Body{"map": changes_map}}}
	payload, err := json.Marshal(ddoc)
	rq, err := http.NewRequest("PUT", u.String(), bytes.NewBuffer(payload))
	rq.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(rq)

	if err == nil && response.StatusCode > 299 {
		err = &base.HTTPError{Status: response.StatusCode, Message: response.Status}
	}
	if err != nil {
		log.Printf("WARNING: Error installing Couchbase design doc: %v", err)
	}
	return err
}

type IDAndRev struct {
	DocID string
	RevID string
}

// Returns all document IDs as an array.
func (db *Database) AllDocIDs() ([]IDAndRev, error) {
	vres, err := db.queryAllDocs(false)
	if err != nil {
		return nil, err
	}

	rows := vres.Rows
	result := make([]IDAndRev, 0, len(rows))
	for _, row := range rows {
		result = append(result, IDAndRev{DocID: row.Key.(string), RevID: row.Value.(string)})
	}
	return result, nil
}

func (db *Database) queryAllDocs(reduce bool) (couchbase.ViewResult, error) {
	opts := Body{"stale": false, "reduce": reduce}
	vres, err := db.Bucket.View("sync_gateway", "all_docs", opts)
	if err != nil {
		log.Printf("WARNING: all_docs got error: %v", err)
	}
	return vres, err
}

// Deletes a database (and all documents)
func (db *Database) Delete() error {
	opts := Body{"stale": false}
	vres, err := db.Bucket.View("sync_gateway", "all_bits", opts)
	if err != nil {
		log.Printf("WARNING: all_bits view returned %v", err)
		return err
	}

	//FIX: Is there a way to do this in one operation?
	log.Printf("Deleting %d documents of %q ...", len(vres.Rows), db.Name)
	for _, row := range vres.Rows {
		if base.Logging {
			log.Printf("\tDeleting %q", row.ID)
		}
		if err := db.Bucket.Delete(row.ID); err != nil {
			log.Printf("WARNING: Error deleting %q: %v", row.ID, err)
		}
	}
	return nil
}

func vacuumContent(bucket *couchbase.Bucket, viewName string, docPrefix string) (int, error) {
	vres, err := bucket.View("sync_gateway", viewName, Body{"stale": false, "group_level": 1})
	if err != nil {
		log.Printf("WARNING: %s view returned %v", viewName, err)
		return 0, err
	}
	deletions := 0
	for _, row := range vres.Rows {
		if row.Value.(bool) {
			deletions += vacIt(bucket, docPrefix+row.Key.(string))
		}
	}
	return deletions, nil
}

func vacIt(bucket *couchbase.Bucket, docid string) int {
	if err := bucket.Delete(docid); err != nil {
		log.Printf("\tError vacuuming %q: %v", docid, err)
		return 0
	}
	if base.Logging {
		log.Printf("\tVacuumed %q", docid)
	}
	return 1
}

// Deletes all orphaned CouchDB revisions not used by any documents.
func VacuumRevisions(bucket *couchbase.Bucket) (int, error) {
	return vacuumContent(bucket, "revs", "rev:")
}

// Deletes all orphaned CouchDB attachments not used by any revisions.
func VacuumAttachments(bucket *couchbase.Bucket) (int, error) {
	return vacuumContent(bucket, "atts", "att:")
}

// Re-runs the channelMapper on every document in the database.
// To be used when the JavaScript channelmap function changes.
func (db *Database) UpdateAllDocChannels() error {
	log.Printf("Recomputing document channels...")
	vres, err := db.Bucket.View("sync_gateway", "all_docs", Body{"stale": false, "reduce": false})
	if err != nil {
		return err
	}
	for _, row := range vres.Rows {
		docid := row.Key.(string)
		key := db.realDocID(docid)
		err := db.Bucket.Update(key, 0, func(currentValue []byte) ([]byte, error) {
			// Be careful: this block can be invoked multiple times if there are races!
			if currentValue == nil {
				return nil, couchbase.UpdateCancel // someone deleted it?!
			}
			doc := newDocument()
			if err := json.Unmarshal(currentValue, doc); err != nil {
				return nil, err
			}
			body, err := db.getRevFromDoc(doc, "", false)
			if err != nil {
				return nil, err
			}
			channels, access := db.getChannelsAndAccess(body)
			db.updateDocAccess(doc, access)
			db.updateDocChannels(doc, channels)
			log.Printf("\tSaving updated channels and access grants of %q", docid)
			return json.Marshal(doc)
		})
		if err != nil && err != couchbase.UpdateCancel {
			log.Printf("WARNING: Error updating doc %q: %v", docid, err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package moneybird

import "testing"

func TestContactGatewayCRUD(t *testing.T) {
	var err error
	var testContact *Contact
	gateway := testClient.Contact()

	// 1. Create
	testContact, err = gateway.Create(&Contact{
		Email:     "johndoe@email.com",
		FirstName: "John",
		LastName:  "Doe",
	})
	if err != nil {
		t.Fatal(err)
	}

	// 2. Scheduled Delete
	defer func() {
		err = gateway.Delete(testContact)
		if err != nil {
			t.Error(err)
		}
	}()

	// 3. Update
	testContact.FirstName = "Peter"
	testContact, err = gateway.Update(testContact)
	if err != nil {
		t.Error(err)
	}

	if testContact.FirstName != "Peter" {
		t.Errorf("Contact was not properly updated.")
	}

	// 4. Get
	testContact, err = gateway.Get(testContact.ID)
	if err != nil {
		t.Error(err)
	}

	if testContact.LastName != "Doe" {
		t.Errorf("Invalid Contact.LastName: %#v", testContact.LastName)
	}

}
<commit_msg>add test for ContactGateway.List() func<commit_after>package moneybird

import "testing"

func TestContactGatewayCRUD(t *testing.T) {
	var err error
	var testContact *Contact
	gateway := testClient.Contact()

	// 1. Create
	testContact, err = gateway.Create(&Contact{
		Email:     "johndoe@email.com",
		FirstName: "John",
		LastName:  "Doe",
	})
	if err != nil {
		t.Fatal(err)
	}

	// 2. Scheduled Delete
	defer func() {
		err = gateway.Delete(testContact)
		if err != nil {
			t.Error(err)
		}
	}()

	// 3. Update
	testContact.FirstName = "Peter"
	testContact, err = gateway.Update(testContact)
	if err != nil {
		t.Error(err)
	}

	if testContact.FirstName != "Peter" {
		t.Errorf("Contact was not properly updated.")
	}

	// 4. Get
	testContact, err = gateway.Get(testContact.ID)
	if err != nil {
		t.Error(err)
	}

	if testContact.LastName != "Doe" {
		t.Errorf("Invalid Contact.LastName: %#v", testContact.LastName)
	}

}

func TestContactGatewayList(t *testing.T) {
	_, err := testClient.Contact().List()
	if err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package reporter

import (
	"fmt"
	"io"
	"os"

	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/formatter"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/gauge"
)

// SimpleConsoleOutput represents if coloring should be removed from the Console output
var SimpleConsoleOutput bool

// Verbose represents level of console Reporting. If true its at step level, else at scenario level.
var Verbose bool

const newline = "\n"

// Reporter reports the progress of spec execution. It reports
// 1. Which spec / scenarion / step (if verbose) is currently executing.
// 2. Status (pass/fail) of the spec / scenario / step (if verbose) once its executed.
type Reporter interface {
	SpecStart(string)
	SpecEnd()
	ScenarioStart(string)
	ScenarioEnd(bool)
	StepStart(string)
	StepEnd(bool)
	ConceptStart(string)
	ConceptEnd(bool)
	DataTable(string)

	Errorf(string, ...interface{})

	io.Writer
}

var currentReporter Reporter

// Current returns the current instance of Reporter, if present. Else, it returns a new Reporter.
func Current() Reporter {
	if currentReporter == nil {
		if SimpleConsoleOutput {
			currentReporter = newSimpleConsole(os.Stdout)
		} else {
			currentReporter = newColoredConsole(os.Stdout)
		}
	}
	return currentReporter
}

type parallelReportWriter struct {
	nRunner int
}

func (p *parallelReportWriter) Write(b []byte) (int, error) {
	return fmt.Printf("[runner: %d] %s", p.nRunner, string(b))
}

// NewParallelConsole returns the instance of parallel console reporter
func NewParallelConsole(n int) Reporter {
	writer := &parallelReportWriter{nRunner: n}
	return newSimpleConsole(writer)
}

// ListenExecutionEvents listens to all execution events for reporting on console
func ListenExecutionEvents() {
	ch := make(chan event.ExecutionEvent, 0)
	event.Register(ch, event.SpecStart, event.SpecEnd, event.ScenarioStart, event.ScenarioEnd, event.StepStart, event.StepEnd)

	go func() {
		for {
			e := <-ch
			switch e.Topic {
			case event.SpecStart:
				Current().SpecStart(e.Item.(*gauge.Specification).Heading.Value)
			case event.SpecEnd:
				Current().SpecEnd()
			case event.ScenarioStart:
				Current().ScenarioStart(e.Item.(*gauge.Scenario).Heading.Value)
			case event.ScenarioEnd:
				Current().ScenarioEnd(e.Result.(*result.ScenarioResult).GetFailed())
			case event.StepStart:
				Current().StepStart(formatter.FormatStep(e.Item.(*gauge.Step)))
			case event.StepEnd:
				Current().StepEnd(e.Result.(*result.StepResult).GetFailed())
			}
		}
	}()
}
<commit_msg>Reordering events in ConsoleReporter<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package reporter

import (
	"fmt"
	"io"
	"os"

	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/formatter"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/gauge"
)

// SimpleConsoleOutput represents if coloring should be removed from the Console output
var SimpleConsoleOutput bool

// Verbose represents level of console Reporting. If true its at step level, else at scenario level.
var Verbose bool

const newline = "\n"

// Reporter reports the progress of spec execution. It reports
// 1. Which spec / scenarion / step (if verbose) is currently executing.
// 2. Status (pass/fail) of the spec / scenario / step (if verbose) once its executed.
type Reporter interface {
	SpecStart(string)
	SpecEnd()
	ScenarioStart(string)
	ScenarioEnd(bool)
	StepStart(string)
	StepEnd(bool)
	ConceptStart(string)
	ConceptEnd(bool)
	DataTable(string)

	Errorf(string, ...interface{})

	io.Writer
}

var currentReporter Reporter

// Current returns the current instance of Reporter, if present. Else, it returns a new Reporter.
func Current() Reporter {
	if currentReporter == nil {
		if SimpleConsoleOutput {
			currentReporter = newSimpleConsole(os.Stdout)
		} else {
			currentReporter = newColoredConsole(os.Stdout)
		}
	}
	return currentReporter
}

type parallelReportWriter struct {
	nRunner int
}

func (p *parallelReportWriter) Write(b []byte) (int, error) {
	return fmt.Printf("[runner: %d] %s", p.nRunner, string(b))
}

// NewParallelConsole returns the instance of parallel console reporter
func NewParallelConsole(n int) Reporter {
	writer := &parallelReportWriter{nRunner: n}
	return newSimpleConsole(writer)
}

// ListenExecutionEvents listens to all execution events for reporting on console
func ListenExecutionEvents() {
	ch := make(chan event.ExecutionEvent, 0)
	event.Register(ch, event.SpecStart, event.SpecEnd, event.ScenarioStart, event.ScenarioEnd, event.StepStart, event.StepEnd)

	go func() {
		for {
			e := <-ch
			switch e.Topic {
			case event.SpecStart:
				Current().SpecStart(e.Item.(*gauge.Specification).Heading.Value)
			case event.ScenarioStart:
				Current().ScenarioStart(e.Item.(*gauge.Scenario).Heading.Value)
			case event.StepStart:
				Current().StepStart(formatter.FormatStep(e.Item.(*gauge.Step)))
			case event.StepEnd:
				Current().StepEnd(e.Result.(*result.StepResult).GetFailed())
			case event.ScenarioEnd:
				Current().ScenarioEnd(e.Result.(*result.ScenarioResult).GetFailed())
			case event.SpecEnd:
				Current().SpecEnd()
			}
		}
	}()
}
<|endoftext|>


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

// Package reqtrace contains a very simple request tracing framework.
package reqtrace

import "golang.org/x/net/context"

// A function that must be called exactly once to report the outcome of an
// operation represented by a span.
type ReportFunc func(error)

// Return a context descending from the supplied parent that contains the
// smarts necessary to be used with the other functions in this package. If ctx
// is already the result of calling Trace, do nothing.
//
// This function starts a root span. The returned report function must be
// called when the overall operation completes.
func Trace(
	parent context.Context,
	desc string) (ctx context.Context, report ReportFunc)

// If ctx is the result of calling Trace, begin a span in the trace with the
// supplied description and return a report function that must be called to
// report the outcome of the span. Otherwise return a function that does
// nothing.
func Start(
	ctx context.Context,
	desc string) (report ReportFunc)

// Call Start, then return a function that reports the value of *err at the
// time it is invoked. Intended to be used with defer at the start of a
// function with a named error return value:
//
//     func DoSomething(ctx context.Context) (err error) {
//       defer reqtrace.StartWithError(ctx, &err, "DoSomething")
//       [...]
//     }
//
func StartWithError(
	ctx context.Context,
	err *error,
	desc string) (f func())
<commit_msg>Added stub implementations.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package reqtrace contains a very simple request tracing framework.
package reqtrace

import "golang.org/x/net/context"

// A function that must be called exactly once to report the outcome of an
// operation represented by a span.
type ReportFunc func(error)

// Return a context descending from the supplied parent that contains the
// smarts necessary to be used with the other functions in this package. If ctx
// is already the result of calling Trace, do nothing.
//
// This function starts a root span. The returned report function must be
// called when the overall operation completes.
func Trace(
	parent context.Context,
	desc string) (ctx context.Context, report ReportFunc) {
	// TODO(jacobsa): Do something interesting.
	ctx = parent
	report = func(err error) {}
	return
}

// If ctx is the result of calling Trace, begin a span in the trace with the
// supplied description and return a report function that must be called to
// report the outcome of the span. Otherwise return a function that does
// nothing.
func Start(
	ctx context.Context,
	desc string) (report ReportFunc) {
	// TODO(jacobsa): Do something interesting.
	report = func(err error) {}
	return
}

// Call Start, then return a function that reports the value of *err at the
// time it is invoked. Intended to be used with defer at the start of a
// function with a named error return value:
//
//     func DoSomething(ctx context.Context) (err error) {
//       defer reqtrace.StartWithError(ctx, &err, "DoSomething")
//       [...]
//     }
//
func StartWithError(
	ctx context.Context,
	err *error,
	desc string) (f func()) {
	// TODO(jacobsa): Do something interesting.
	f = func() {}
	return
}
<|endoftext|>


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

// Command brio-gen generates (un)marshaler code for types.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/importer"
	"go/types"
	"io"
	"log"
	"os"
	"strings"
)

var (
	typeNames = flag.String("t", "", "comma-separated list of type names")
	pkgPath   = flag.String("p", "", "package import path")
	output    = flag.String("o", "brio_gen.go", "output file name")
)

func main() {
	flag.Parse()

	log.SetPrefix("brio: ")
	log.SetFlags(0)

	if *typeNames == "" {
		flag.Usage()
		os.Exit(2)
	}

	types := strings.Split(*typeNames, ",")
	g, err := newGenerator(*pkgPath)
	if err != nil {
		log.Fatal(err)
	}

	for _, t := range types {
		g.generate(t)
	}

	// FIXME(sbinet): debug
	io.Copy(os.Stdout, g.buf)
}

// Generator holds the state of the generation.
type Generator struct {
	buf *bytes.Buffer
	pkg *types.Package
}

func newGenerator(path string) (*Generator, error) {
	pkg, err := importer.Default().Import(path)
	if err != nil {
		return nil, err
	}

	return &Generator{
		buf: new(bytes.Buffer),
		pkg: pkg,
	}, nil
}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(g.buf, format, args...)
}

func (g *Generator) generate(typeName string) {
	scope := g.pkg.Scope()
	obj := scope.Lookup(typeName)
	if obj == nil {
		log.Fatalf("no such type %q in package %q\n", typeName, g.pkg.Path()+"/"+g.pkg.Name())
	}

	tn, ok := obj.(*types.TypeName)
	if !ok {
		log.Fatalf("%q is not a type (%v)\n", typeName, obj)
	}

	typ, ok := tn.Type().Underlying().(*types.Struct)
	if !ok {
		log.Fatalf("%q is not a named struct (%v)\n", typeName, tn)
	}
	log.Printf("typ: %+v\n", typ)

	for i := 0; i < typ.NumFields(); i++ {
		ft := typ.Field(i)
		log.Printf("field[%d]: %v\n", i, ft)
	}

}
<commit_msg>cmd/brio-gen: generate marshaling for (u)int,float,cplx<commit_after>// Copyright 2016 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Command brio-gen generates (un)marshaler code for types.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/importer"
	"go/types"
	"io"
	"log"
	"os"
	"strings"
)

var (
	typeNames = flag.String("t", "", "comma-separated list of type names")
	pkgPath   = flag.String("p", "", "package import path")
	output    = flag.String("o", "brio_gen.go", "output file name")
)

func main() {
	flag.Parse()

	log.SetPrefix("brio: ")
	log.SetFlags(0)

	if *typeNames == "" {
		flag.Usage()
		os.Exit(2)
	}

	types := strings.Split(*typeNames, ",")
	g, err := newGenerator(*pkgPath)
	if err != nil {
		log.Fatal(err)
	}

	for _, t := range types {
		g.generate(t)
	}

	// FIXME(sbinet): debug
	io.Copy(os.Stdout, g.buf)
}

// Generator holds the state of the generation.
type Generator struct {
	buf *bytes.Buffer
	pkg *types.Package

	// set of imported packages.
	// usually: "encoding/binary", "math"
	imps map[string]int
}

func newGenerator(path string) (*Generator, error) {
	pkg, err := importer.Default().Import(path)
	if err != nil {
		return nil, err
	}

	return &Generator{
		buf:  new(bytes.Buffer),
		pkg:  pkg,
		imps: map[string]int{"encoding/binary": 1},
	}, nil
}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(g.buf, format, args...)
}

func (g *Generator) generate(typeName string) {
	scope := g.pkg.Scope()
	obj := scope.Lookup(typeName)
	if obj == nil {
		log.Fatalf("no such type %q in package %q\n", typeName, g.pkg.Path()+"/"+g.pkg.Name())
	}

	tn, ok := obj.(*types.TypeName)
	if !ok {
		log.Fatalf("%q is not a type (%v)\n", typeName, obj)
	}

	typ, ok := tn.Type().Underlying().(*types.Struct)
	if !ok {
		log.Fatalf("%q is not a named struct (%v)\n", typeName, tn)
	}
	log.Printf("typ: %+v\n", typ)

	g.Printf(`// MarshalBinary implements encoding.BinaryMarshaler
func (o *%[1]s) MarshalBinary() (data []byte, err error) {
	var buf [8]byte
`,
		typeName,
	)

	for i := 0; i < typ.NumFields(); i++ {
		ft := typ.Field(i)
		log.Printf("field[%d]: %v\n", i, ft)
		ut := ft.Type().Underlying()
		switch ut := ut.(type) {
		case *types.Basic:
			switch kind := ut.Kind(); kind {
			case types.Uint:
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], uint64(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			case types.Uint8:
				g.Printf("\tdata = append(data, byte(o.%s))\n", ft.Name())
			case types.Uint16:
				g.Printf(
					"\tbinary.LittleEndian.PutUint16(buf[:2], o.%s)\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:2])\n")
			case types.Uint32:
				g.Printf(
					"\tbinary.LittleEndian.PutUint32(buf[:4], o.%s)\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:4])\n")
			case types.Uint64:
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], o.%s)\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			case types.Int:
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], int64(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			case types.Int8:
				g.Printf("\tdata = append(data, byte(o.%s))\n", ft.Name())
			case types.Int16:
				g.Printf(
					"\tbinary.LittleEndian.PutUint16(buf[:2], uint16(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:2])\n")
			case types.Int32:
				g.Printf(
					"\tbinary.LittleEndian.PutUint32(buf[:4], uint32(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:4])\n")
			case types.Int64:
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], uint64(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			case types.Float32:
				g.imps["math"]++
				g.Printf(
					"\tbinary.LittleEndian.PutUint32(buf[:4], math.Float32bits(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:4])\n")
			case types.Float64:
				g.imps["math"]++
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], math.Float64bits(o.%s))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			case types.Complex64:
				g.imps["math"]++
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:4], math.Float32bits(real(o.%s)))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:4])\n")
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:4], math.Float32bits(imag(o.%s)))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:4])\n")
			case types.Complex128:
				g.imps["math"]++
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], math.Float64bits(real(o.%s)))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
				g.Printf(
					"\tbinary.LittleEndian.PutUint64(buf[:8], math.Float64bits(imag(o.%s)))\n",
					ft.Name(),
				)
				g.Printf("\tdata = append(data, buf[:8])\n")
			}
		}
	}

	g.Printf("\treturn data, err\n}\n\n")
}
<|endoftext|>


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

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path/filepath"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v3"

	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
)

func New() *cobra.Command {
	deployCmd := &cobra.Command{
		Use:   "deploy <deployment yaml>",
		Short: "Deploy cluster.",
		RunE:  deployFn,
	}
	return deployCmd
}

type Cluster interface {
	Deploy(context.Context) error
}

type Ingress interface {
	Deploy(context.Context) error
	SetKClient(kubernetes.Interface)
	Healthy(context.Context) error
}

type CNI interface {
	Deploy(context.Context) error
	SetKClient(kubernetes.Interface)
	Healthy(context.Context) error
}

type Deployment struct {
	Cluster Cluster
	Ingress Ingress
	CNI     CNI
}

type DeploymentConfig struct {
	Cluster ClusterSpec `yaml:"cluster"`
	Ingress IngressSpec `yaml:"ingress"`
	CNI     CNISpec     `yaml:"cni"`
}

func NewDeployment(cfg *DeploymentConfig) (*Deployment, error) {
	d := &Deployment{}
	switch cfg.Cluster.Kind {
	case "Kind":
		log.Infof("Using kind scenario") // Entry point
		v := &KindSpec{}
		if err := cfg.Cluster.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.Cluster = v
	default:
		return nil, fmt.Errorf("cluster type not supported: %s", cfg.Cluster.Kind)
	}
	switch cfg.CNI.Kind {
	case "Meshnet":
		v := &MeshnetSpec{}
		if err := cfg.CNI.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.CNI = v
	default:
		return nil, fmt.Errorf("CNI type not supported: %s", cfg.CNI.Kind)
	}
	switch cfg.Ingress.Kind {
	case "MetalLB":
		v := &MetalLBSpec{}
		if err := cfg.Ingress.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.Ingress = v
	default:
		return nil, fmt.Errorf("ingress type not supported: %s", cfg.Ingress.Kind)
	}
	return d, nil
}

var (
	deploymentBasePath string
)

func deploymentFromArg(p string) (*DeploymentConfig, string, error) {
	bp, err := filepath.Abs(p)
	if err != nil {
		return nil, "", err
	}
	b, err := ioutil.ReadFile(bp)
	if err != nil {
		return nil, "", err
	}
	bp = filepath.Dir(bp)
	dCfg := &DeploymentConfig{}
	decoder := yaml.NewDecoder(bytes.NewBuffer(b))
	decoder.KnownFields(true)
	if err := decoder.Decode(dCfg); err != nil {
		return nil, "", err
	}
	return dCfg, bp, nil
}

func deployFn(cmd *cobra.Command, args []string) error {
	if len(args) != 1 {
		return fmt.Errorf("%s: missing args", cmd.Use)
	}
	if _, err := exec.LookPath("kubectl"); err != nil {
		return fmt.Errorf("install kubectl before running deploy: %v", err)
	}
	dCfg, bp, err := deploymentFromArg(args[0])
	if err != nil {
		return err
	}
	deploymentBasePath = bp
	d, err := NewDeployment(dCfg)
	if err != nil {
		return err
	}
	if err := d.Cluster.Deploy(cmd.Context()); err != nil {
		return err
	}
	// Once cluster is up set kClient
	kubecfg, err := cmd.Flags().GetString("kubecfg")
	if err != nil {
		return err
	}
	rCfg, err := clientcmd.BuildConfigFromFlags("", kubecfg)
	if err != nil {
		return err
	}
	kClient, err := kubernetes.NewForConfig(rCfg)
	if err != nil {
		return err
	}
	d.Ingress.SetKClient(kClient)
	log.Infof("Validating cluster health")
	if err := d.Ingress.Deploy(cmd.Context()); err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(cmd.Context(), 1*time.Minute)
	defer cancel()
	if err := d.Ingress.Healthy(ctx); err != nil {
		return err
	}
	if err := d.CNI.Deploy(cmd.Context()); err != nil {
		return err
	}
	d.CNI.SetKClient(kClient)
	ctx, cancel = context.WithTimeout(cmd.Context(), 1*time.Minute)
	defer cancel()
	if err := d.CNI.Healthy(ctx); err != nil {
		return err
	}
	log.Infof("Ready for topology")
	return nil
}
<commit_msg>bare bone controller server with unimplemented rpc's<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package deploy

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path/filepath"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v3"

	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
)

func New() *cobra.Command {
	deployCmd := &cobra.Command{
		Use:   "deploy <deployment yaml>",
		Short: "Deploy cluster.",
		RunE:  deployFn,
	}
	return deployCmd
}

type Cluster interface {
	Deploy(context.Context) error
}

type Ingress interface {
	Deploy(context.Context) error
	SetKClient(kubernetes.Interface)
	Healthy(context.Context) error
}

type CNI interface {
	Deploy(context.Context) error
	SetKClient(kubernetes.Interface)
	Healthy(context.Context) error
}

type Deployment struct {
	Cluster Cluster
	Ingress Ingress
	CNI     CNI
}

type DeploymentConfig struct {
	Cluster ClusterSpec `yaml:"cluster"`
	Ingress IngressSpec `yaml:"ingress"`
	CNI     CNISpec     `yaml:"cni"`
}

func NewDeployment(cfg *DeploymentConfig) (*Deployment, error) {
	d := &Deployment{}
	switch cfg.Cluster.Kind {
	case "Kind":
		log.Infof("Using kind scenario")
		v := &KindSpec{}
		if err := cfg.Cluster.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.Cluster = v
	default:
		return nil, fmt.Errorf("cluster type not supported: %s", cfg.Cluster.Kind)
	}
	switch cfg.CNI.Kind {
	case "Meshnet":
		v := &MeshnetSpec{}
		if err := cfg.CNI.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.CNI = v
	default:
		return nil, fmt.Errorf("CNI type not supported: %s", cfg.CNI.Kind)
	}
	switch cfg.Ingress.Kind {
	case "MetalLB":
		v := &MetalLBSpec{}
		if err := cfg.Ingress.Spec.Decode(v); err != nil {
			return nil, err
		}
		d.Ingress = v
	default:
		return nil, fmt.Errorf("ingress type not supported: %s", cfg.Ingress.Kind)
	}
	return d, nil
}

var (
	deploymentBasePath string
)

func deploymentFromArg(p string) (*DeploymentConfig, string, error) {
	bp, err := filepath.Abs(p)
	if err != nil {
		return nil, "", err
	}
	b, err := ioutil.ReadFile(bp)
	if err != nil {
		return nil, "", err
	}
	bp = filepath.Dir(bp)
	dCfg := &DeploymentConfig{}
	decoder := yaml.NewDecoder(bytes.NewBuffer(b))
	decoder.KnownFields(true)
	if err := decoder.Decode(dCfg); err != nil {
		return nil, "", err
	}
	return dCfg, bp, nil
}

func deployFn(cmd *cobra.Command, args []string) error {
	if len(args) != 1 {
		return fmt.Errorf("%s: missing args", cmd.Use)
	}
	if _, err := exec.LookPath("kubectl"); err != nil {
		return fmt.Errorf("install kubectl before running deploy: %v", err)
	}
	dCfg, bp, err := deploymentFromArg(args[0])
	if err != nil {
		return err
	}
	deploymentBasePath = bp
	d, err := NewDeployment(dCfg)
	if err != nil {
		return err
	}
	if err := d.Cluster.Deploy(cmd.Context()); err != nil {
		return err
	}
	// Once cluster is up set kClient
	kubecfg, err := cmd.Flags().GetString("kubecfg")
	if err != nil {
		return err
	}
	rCfg, err := clientcmd.BuildConfigFromFlags("", kubecfg)
	if err != nil {
		return err
	}
	kClient, err := kubernetes.NewForConfig(rCfg)
	if err != nil {
		return err
	}
	d.Ingress.SetKClient(kClient)
	log.Infof("Validating cluster health")
	if err := d.Ingress.Deploy(cmd.Context()); err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(cmd.Context(), 1*time.Minute)
	defer cancel()
	if err := d.Ingress.Healthy(ctx); err != nil {
		return err
	}
	if err := d.CNI.Deploy(cmd.Context()); err != nil {
		return err
	}
	d.CNI.SetKClient(kClient)
	ctx, cancel = context.WithTimeout(cmd.Context(), 1*time.Minute)
	defer cancel()
	if err := d.CNI.Healthy(ctx); err != nil {
		return err
	}
	log.Infof("Ready for topology")
	return nil
}
<|endoftext|>