<commit_before>package main

import (
	"bufio"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"net"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
)

// https://github.com/memcached/memcached/blob/master/doc/protocol.txt
var graphdef = map[string](mp.Graphs){
	"memcached.connections": mp.Graphs{
		Label: "Memcached Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "curr_connections", Label: "Connections", Diff: false},
		},
	},
	"memcached.cmd": mp.Graphs{
		Label: "Memcached Command",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "cmd_get", Label: "Get", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "cmd_set", Label: "Set", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "cmd_flush", Label: "Flush", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "cmd_touch", Label: "Touch", Diff: true, Type: "uint64"},
		},
	},
	"memcached.hitmiss": mp.Graphs{
		Label: "Memcached Hits/Misses",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "get_hits", Label: "Get Hits", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "get_misses", Label: "Get Misses", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "delete_hits", Label: "Delete Hits", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "delete_misses", Label: "Delete Misses", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "incr_hits", Label: "Incr Hits", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "incr_misses", Label: "Incr Misses", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "cas_hits", Label: "Cas Hits", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "cas_misses", Label: "Cas Misses", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "touch_hits", Label: "Touch Hits", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "touch_misses", Label: "Touch Misses", Diff: true, Type: "uint64"},
		},
	},
	"memcached.evictions": mp.Graphs{
		Label: "Memcached Evictions",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "evictions", Label: "Evictions", Diff: true, Type: "uint64"},
		},
	},
	"memcached.unfetched": mp.Graphs{
		Label: "Memcached Unfetched",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "expired_unfetched", Label: "Expired unfetched", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "evicted_unfetched", Label: "Evicted unfetched", Diff: true, Type: "uint64"},
		},
	},
	"memcached.rusage": mp.Graphs{
		Label: "Memcached Resouce Usage",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "rusage_user", Label: "User", Diff: true},
			mp.Metrics{Name: "rusage_system", Label: "System", Diff: true},
		},
	},
	"memcached.bytes": mp.Graphs{
		Label: "Memcached Traffics",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "bytes_read", Label: "Read", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "bytes_written", Label: "Write", Diff: true, Type: "uint64"},
		},
	},
	"memcached.cachesize": mp.Graphs{
		Label: "Memcached Cache Size",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "limit_maxbytes", Label: "Total", Diff: false},
			mp.Metrics{Name: "bytes", Label: "Used", Diff: false, Type: "uint64"},
		},
	},
}

// MemcachedPlugin mackerel plugin for memchached
type MemcachedPlugin struct {
	Target   string
	Socket   string
	Tempfile string
}

// FetchMetrics interface for mackerelplugin
func (m MemcachedPlugin) FetchMetrics() (map[string]interface{}, error) {
	network := "tcp"
	target := m.Target
	if m.Socket != "" {
		network = "unix"
		target = m.Socket
	}
	conn, err := net.Dial(network, target)
	if err != nil {
		return nil, err
	}
	fmt.Fprintln(conn, "stats")
	return m.parseStats(conn)
}

func (m MemcachedPlugin) parseStats(conn io.Reader) (map[string]interface{}, error) {
	scanner := bufio.NewScanner(conn)
	stat := make(map[string]interface{})

	for scanner.Scan() {
		line := scanner.Text()
		s := string(line)
		if s == "END" {
			return stat, nil
		}

		res := strings.Split(s, " ")
		if res[0] == "STAT" {
			stat[res[1]] = res[2]
		}
	}
	if err := scanner.Err(); err != nil {
		return stat, err
	}
	return nil, nil
}

// GraphDefinition interface for mackerelplugin
func (m MemcachedPlugin) GraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "11211", "Port")
	optSocket := flag.String("socket", "", "Server socket (overrides hosts and port)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var memcached MemcachedPlugin
	if *optSocket != "" {
		memcached.Socket = *optSocket
	} else {
		memcached.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	}
	helper := mp.NewMackerelPlugin(memcached)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		if memcached.Socket != "" {
			helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s", fmt.Sprintf("%x", md5.Sum([]byte(memcached.Socket))))
		} else {
			helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s-%s", *optHost, *optPort)
		}
	}
	helper.Run()
}
<commit_msg>move graph definition into GraphDefinition<commit_after>package main

import (
	"bufio"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"net"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
)

// MemcachedPlugin mackerel plugin for memchached
type MemcachedPlugin struct {
	Target   string
	Socket   string
	Tempfile string
}

// FetchMetrics interface for mackerelplugin
func (m MemcachedPlugin) FetchMetrics() (map[string]interface{}, error) {
	network := "tcp"
	target := m.Target
	if m.Socket != "" {
		network = "unix"
		target = m.Socket
	}
	conn, err := net.Dial(network, target)
	if err != nil {
		return nil, err
	}
	fmt.Fprintln(conn, "stats")
	return m.parseStats(conn)
}

func (m MemcachedPlugin) parseStats(conn io.Reader) (map[string]interface{}, error) {
	scanner := bufio.NewScanner(conn)
	stat := make(map[string]interface{})

	for scanner.Scan() {
		line := scanner.Text()
		s := string(line)
		if s == "END" {
			return stat, nil
		}

		res := strings.Split(s, " ")
		if res[0] == "STAT" {
			stat[res[1]] = res[2]
		}
	}
	if err := scanner.Err(); err != nil {
		return stat, err
	}
	return nil, nil
}

// GraphDefinition interface for mackerelplugin
func (m MemcachedPlugin) GraphDefinition() map[string](mp.Graphs) {
	// https://github.com/memcached/memcached/blob/master/doc/protocol.txt
	var graphdef = map[string](mp.Graphs){
		"memcached.connections": mp.Graphs{
			Label: "Memcached Connections",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "curr_connections", Label: "Connections", Diff: false},
			},
		},
		"memcached.cmd": mp.Graphs{
			Label: "Memcached Command",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "cmd_get", Label: "Get", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "cmd_set", Label: "Set", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "cmd_flush", Label: "Flush", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "cmd_touch", Label: "Touch", Diff: true, Type: "uint64"},
			},
		},
		"memcached.hitmiss": mp.Graphs{
			Label: "Memcached Hits/Misses",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "get_hits", Label: "Get Hits", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "get_misses", Label: "Get Misses", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "delete_hits", Label: "Delete Hits", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "delete_misses", Label: "Delete Misses", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "incr_hits", Label: "Incr Hits", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "incr_misses", Label: "Incr Misses", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "cas_hits", Label: "Cas Hits", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "cas_misses", Label: "Cas Misses", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "touch_hits", Label: "Touch Hits", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "touch_misses", Label: "Touch Misses", Diff: true, Type: "uint64"},
			},
		},
		"memcached.evictions": mp.Graphs{
			Label: "Memcached Evictions",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "evictions", Label: "Evictions", Diff: true, Type: "uint64"},
			},
		},
		"memcached.unfetched": mp.Graphs{
			Label: "Memcached Unfetched",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "expired_unfetched", Label: "Expired unfetched", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "evicted_unfetched", Label: "Evicted unfetched", Diff: true, Type: "uint64"},
			},
		},
		"memcached.rusage": mp.Graphs{
			Label: "Memcached Resouce Usage",
			Unit:  "float",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "rusage_user", Label: "User", Diff: true},
				mp.Metrics{Name: "rusage_system", Label: "System", Diff: true},
			},
		},
		"memcached.bytes": mp.Graphs{
			Label: "Memcached Traffics",
			Unit:  "bytes",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "bytes_read", Label: "Read", Diff: true, Type: "uint64"},
				mp.Metrics{Name: "bytes_written", Label: "Write", Diff: true, Type: "uint64"},
			},
		},
		"memcached.cachesize": mp.Graphs{
			Label: "Memcached Cache Size",
			Unit:  "bytes",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "limit_maxbytes", Label: "Total", Diff: false},
				mp.Metrics{Name: "bytes", Label: "Used", Diff: false, Type: "uint64"},
			},
		},
	}
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "11211", "Port")
	optSocket := flag.String("socket", "", "Server socket (overrides hosts and port)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var memcached MemcachedPlugin
	if *optSocket != "" {
		memcached.Socket = *optSocket
	} else {
		memcached.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	}
	helper := mp.NewMackerelPlugin(memcached)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		if memcached.Socket != "" {
			helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s", fmt.Sprintf("%x", md5.Sum([]byte(memcached.Socket))))
		} else {
			helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s-%s", *optHost, *optPort)
		}
	}
	helper.Run()
}
<|endoftext|>


<commit_before>package benchmarks

import (
	sflag "veyron.io/veyron/veyron/security/flag"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/vlog"
)

type impl struct {
}

func (i *impl) Echo(ctx ipc.ServerCall, payload []byte) ([]byte, error) {
	return payload, nil
}

func (i *impl) EchoStream(ctx ipc.ServerCall, stream BenchmarkServiceEchoStreamStream) error {
	rStream := stream.RecvStream()
	sender := stream.SendStream()
	for rStream.Advance() {
		chunk := rStream.Value()
		if err := sender.Send(chunk); err != nil {
			return err
		}
	}

	return rStream.Err()
}

// StartServer starts a server that implements the Benchmark service. The
// server listens to the given protocol and address, and returns the veyron
// address of the server and a callback function to stop the server.
func StartServer(runtime veyron2.Runtime, listenSpec ipc.ListenSpec) (string, func()) {
	server, err := runtime.NewServer()
	if err != nil {
		vlog.Fatalf("NewServer failed: %v", err)
	}
	ep, err := server.Listen(listenSpec)
	if err != nil {
		vlog.Fatalf("Listen failed: %v", err)
	}
	if err := server.Serve("", &impl{}, sflag.NewAuthorizerOrDie()); err != nil {
		vlog.Fatalf("Serve failed: %v", err)
	}
	return naming.JoinAddressName(ep.String(), ""), func() {
		if err := server.Stop(); err != nil {
			vlog.Fatalf("Stop() failed: %v", err)
		}
	}
}
<commit_msg>ipc/benchmarks: fix for https://veyron-review.googlesource.com/#/c/6373/<commit_after>package benchmarks

import (
	"io"

	sflag "veyron.io/veyron/veyron/security/flag"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/vlog"
)

type impl struct {
}

func (i *impl) Echo(ctx ipc.ServerCall, payload []byte) ([]byte, error) {
	return payload, nil
}

func (i *impl) EchoStream(ctx ipc.ServerCall) error {
	for {
		var chunk []byte
		if err := ctx.Recv(&chunk); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if err := ctx.Send(chunk); err != nil {
			return err
		}
	}
	return nil
}

// StartServer starts a server that implements the Benchmark service. The
// server listens to the given protocol and address, and returns the veyron
// address of the server and a callback function to stop the server.
func StartServer(runtime veyron2.Runtime, listenSpec ipc.ListenSpec) (string, func()) {
	server, err := runtime.NewServer()
	if err != nil {
		vlog.Fatalf("NewServer failed: %v", err)
	}
	ep, err := server.Listen(listenSpec)
	if err != nil {
		vlog.Fatalf("Listen failed: %v", err)
	}
	if err := server.Serve("", &impl{}, sflag.NewAuthorizerOrDie()); err != nil {
		vlog.Fatalf("Serve failed: %v", err)
	}
	return naming.JoinAddressName(ep.String(), ""), func() {
		if err := server.Stop(); err != nil {
			vlog.Fatalf("Stop() failed: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package canoe

import (
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"os"

	"github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/snap"
	"github.com/coreos/etcd/wal"
	"github.com/coreos/etcd/wal/walpb"
)

type walMetadata struct {
	NodeID    uint64 `json:"node_id"`
	ClusterID uint64 `json:"cluster_id"`
}

func (rn *Node) initPersistentStorage() error {
	if err := rn.initSnap(); err != nil {
		return errors.Wrap(err, "Error initializing snapshot")
	}

	raftSnap, err := rn.ss.Load()
	if err != nil {
		if err != snap.ErrNoSnapshot && err != snap.ErrEmptySnapshot {
			return errors.Wrap(err, "Error loading latest snapshot")
		}
	}

	var walSnap walpb.Snapshot

	if raftSnap != nil {
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term
	}

	if err := rn.initWAL(walSnap); err != nil {
		return errors.Wrap(err, "Error initializing WAL")
	}

	return nil
}

// Correct order of ops
// 1: Restore Metadata from WAL
// 2: Apply any persisted snapshot to FSM
// 3: Apply any Snapshot to raft storage
// 4: Apply any hardstate to raft storage
// 5: Apply and WAL Entries to raft storage
func (rn *Node) restoreRaft() error {
	raftSnap, err := rn.ss.Load()
	if err != nil {
		if err != snap.ErrNoSnapshot && err != snap.ErrEmptySnapshot {
			return errors.Wrap(err, "Error loading latest snapshot")
		}
	}

	var walSnap walpb.Snapshot

	if raftSnap != nil {
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term
	} else {
		raftSnap = &raftpb.Snapshot{}
	}

	wMetadata, hState, ents, err := rn.wal.ReadAll()
	if err != nil {
		return errors.Wrap(err, "Error reading WAL")
	}

	// NOTE: Step 1
	if err := rn.restoreMetadata(wMetadata); err != nil {
		return errors.Wrap(err, "Error restoring from WAL metadata")
	}

	// We can do this now that we restored the metadata
	if err := rn.attachTransport(); err != nil {
		return errors.Wrap(err, "Error attaching raft Transport layer")
	}

	if err := rn.transport.Start(); err != nil {
		return errors.Wrap(err, "Error starting raft transport layer")
	}

	// NOTE: Step 2
	if err := rn.restoreFSMFromSnapshot(*raftSnap); err != nil {
		return errors.Wrap(err, "Error restoring FSM from snapshot")
	}

	// NOTE: Step 3, 4, 5
	if err := rn.restoreMemoryStorage(*raftSnap, hState, ents); err != nil {
		return errors.Wrap(err, "Error restoring raft memory storage")
	}

	// NOTE: Step 6
	/*if err := rn.restoreFSMFromWAL(ents); err != nil {
		return errors.Wrap(err, "Error restoring FSM from WAL")
	}*/

	return nil
}

func (rn *Node) initSnap() error {
	if rn.snapDir() == "" {
		return nil
	}

	if err := os.MkdirAll(rn.snapDir(), 0750); err != nil && !os.IsExist(err) {
		return errors.Wrap(err, "Error trying to create directory for snapshots")
	}

	rn.ss = snap.New(rn.snapDir())

	return nil
}

func (rn *Node) persistSnapshot(raftSnap raftpb.Snapshot) error {

	if rn.ss != nil {
		if err := rn.ss.SaveSnap(raftSnap); err != nil {
			return errors.Wrap(err, "Error saving snapshot to persistent storage")
		}
	}

	if rn.wal != nil {
		var walSnap walpb.Snapshot
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term

		if err := rn.wal.SaveSnapshot(walSnap); err != nil {
			return errors.Wrap(err, "Error updating WAL with snapshot")
		}
	}
	return nil
}

func (rn *Node) initWAL(walSnap walpb.Snapshot) error {
	if rn.walDir() == "" {
		return nil
	}

	if !wal.Exist(rn.walDir()) {

		if err := os.MkdirAll(rn.walDir(), 0750); err != nil && !os.IsExist(err) {
			return errors.Wrap(err, "Error creating directory for raft WAL")
		}

		metaStruct := &walMetadata{
			NodeID:    rn.id,
			ClusterID: rn.cid,
		}

		metaData, err := json.Marshal(metaStruct)
		if err != nil {
			return errors.Wrap(err, "Error marshaling WAL metadata")
		}

		w, err := wal.Create(rn.walDir(), metaData)
		if err != nil {
			return errors.Wrap(err, "Error creating new WAL")
		}
		rn.wal = w
	} else {
		// This assumes we WILL be reading this once elsewhere
		w, err := wal.Open(rn.walDir(), walSnap)
		if err != nil {
			return errors.Wrap(err, "Error opening existing WAL")
		}
		rn.wal = w
	}

	return nil
}

func (rn *Node) restoreMetadata(wMetadata []byte) error {
	var metaData walMetadata
	if err := json.Unmarshal(wMetadata, &metaData); err != nil {
		return errors.Wrap(err, "Error unmarshaling WAL metadata")
	}

	rn.id, rn.cid = metaData.NodeID, metaData.ClusterID
	rn.raftConfig.ID = metaData.NodeID
	return nil
}

// restores FSM AND it sets the NodeID and ClusterID if present in Metadata
func (rn *Node) restoreFSMFromWAL(ents []raftpb.Entry) error {
	if rn.wal == nil {
		return nil
	}

	if err := rn.publishEntries(ents); err != nil {
		return errors.Wrap(err, "Error publishing entries from WAL")
	}

	return nil
}

func (rn *Node) restoreMemoryStorage(raftSnap raftpb.Snapshot, hState raftpb.HardState, ents []raftpb.Entry) error {
	if !raft.IsEmptySnap(raftSnap) {
		if err := rn.raftStorage.ApplySnapshot(raftSnap); err != nil {
			return errors.Wrap(err, "Error applying snapshot to raft memory storage")
		}
	}

	if rn.wal != nil {
		if err := rn.raftStorage.SetHardState(hState); err != nil {
			return errors.Wrap(err, "Error setting memory hardstate")
		}

		if err := rn.raftStorage.Append(ents); err != nil {
			return errors.Wrap(err, "Error appending entries to memory storage")
		}
	}

	return nil
}

func (rn *Node) deletePersistentData() error {
	if rn.snapDir() != "" {
		if err := os.RemoveAll(rn.snapDir()); err != nil {
			return errors.Wrap(err, "Error deleting snapshot directory")
		}
	}
	if rn.walDir() != "" {
		//TODO: Should be delete walDir or snapDir()?
		if err := os.RemoveAll(rn.walDir()); err != nil {
			return errors.Wrap(err, "Error deleting WAL directory")
		}
	}
	return nil
}

func (rn *Node) walDir() string {
	if rn.dataDir == "" {
		return ""
	}
	return fmt.Sprintf("%s%s", rn.dataDir, walDirExtension)
}

func (rn *Node) snapDir() string {
	if rn.dataDir == "" {
		return ""
	}
	return fmt.Sprintf("%s%s", rn.dataDir, snapDirExtension)
}
<commit_msg>remove dead code<commit_after>package canoe

import (
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"os"

	"github.com/coreos/etcd/raft"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/snap"
	"github.com/coreos/etcd/wal"
	"github.com/coreos/etcd/wal/walpb"
)

type walMetadata struct {
	NodeID    uint64 `json:"node_id"`
	ClusterID uint64 `json:"cluster_id"`
}

func (rn *Node) initPersistentStorage() error {
	if err := rn.initSnap(); err != nil {
		return errors.Wrap(err, "Error initializing snapshot")
	}

	raftSnap, err := rn.ss.Load()
	if err != nil {
		if err != snap.ErrNoSnapshot && err != snap.ErrEmptySnapshot {
			return errors.Wrap(err, "Error loading latest snapshot")
		}
	}

	var walSnap walpb.Snapshot

	if raftSnap != nil {
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term
	}

	if err := rn.initWAL(walSnap); err != nil {
		return errors.Wrap(err, "Error initializing WAL")
	}

	return nil
}

// Correct order of ops
// 1: Restore Metadata from WAL
// 2: Apply any persisted snapshot to FSM
// 3: Apply any Snapshot to raft storage
// 4: Apply any hardstate to raft storage
// 5: Apply and WAL Entries to raft storage
func (rn *Node) restoreRaft() error {
	raftSnap, err := rn.ss.Load()
	if err != nil {
		if err != snap.ErrNoSnapshot && err != snap.ErrEmptySnapshot {
			return errors.Wrap(err, "Error loading latest snapshot")
		}
	}

	var walSnap walpb.Snapshot

	if raftSnap != nil {
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term
	} else {
		raftSnap = &raftpb.Snapshot{}
	}

	wMetadata, hState, ents, err := rn.wal.ReadAll()
	if err != nil {
		return errors.Wrap(err, "Error reading WAL")
	}

	// NOTE: Step 1
	if err := rn.restoreMetadata(wMetadata); err != nil {
		return errors.Wrap(err, "Error restoring from WAL metadata")
	}

	// We can do this now that we restored the metadata
	if err := rn.attachTransport(); err != nil {
		return errors.Wrap(err, "Error attaching raft Transport layer")
	}

	if err := rn.transport.Start(); err != nil {
		return errors.Wrap(err, "Error starting raft transport layer")
	}

	// NOTE: Step 2
	if err := rn.restoreFSMFromSnapshot(*raftSnap); err != nil {
		return errors.Wrap(err, "Error restoring FSM from snapshot")
	}

	// NOTE: Step 3, 4, 5
	if err := rn.restoreMemoryStorage(*raftSnap, hState, ents); err != nil {
		return errors.Wrap(err, "Error restoring raft memory storage")
	}

	return nil
}

func (rn *Node) initSnap() error {
	if rn.snapDir() == "" {
		return nil
	}

	if err := os.MkdirAll(rn.snapDir(), 0750); err != nil && !os.IsExist(err) {
		return errors.Wrap(err, "Error trying to create directory for snapshots")
	}

	rn.ss = snap.New(rn.snapDir())

	return nil
}

func (rn *Node) persistSnapshot(raftSnap raftpb.Snapshot) error {

	if rn.ss != nil {
		if err := rn.ss.SaveSnap(raftSnap); err != nil {
			return errors.Wrap(err, "Error saving snapshot to persistent storage")
		}
	}

	if rn.wal != nil {
		var walSnap walpb.Snapshot
		walSnap.Index, walSnap.Term = raftSnap.Metadata.Index, raftSnap.Metadata.Term

		if err := rn.wal.SaveSnapshot(walSnap); err != nil {
			return errors.Wrap(err, "Error updating WAL with snapshot")
		}
	}
	return nil
}

func (rn *Node) initWAL(walSnap walpb.Snapshot) error {
	if rn.walDir() == "" {
		return nil
	}

	if !wal.Exist(rn.walDir()) {

		if err := os.MkdirAll(rn.walDir(), 0750); err != nil && !os.IsExist(err) {
			return errors.Wrap(err, "Error creating directory for raft WAL")
		}

		metaStruct := &walMetadata{
			NodeID:    rn.id,
			ClusterID: rn.cid,
		}

		metaData, err := json.Marshal(metaStruct)
		if err != nil {
			return errors.Wrap(err, "Error marshaling WAL metadata")
		}

		w, err := wal.Create(rn.walDir(), metaData)
		if err != nil {
			return errors.Wrap(err, "Error creating new WAL")
		}
		rn.wal = w
	} else {
		// This assumes we WILL be reading this once elsewhere
		w, err := wal.Open(rn.walDir(), walSnap)
		if err != nil {
			return errors.Wrap(err, "Error opening existing WAL")
		}
		rn.wal = w
	}

	return nil
}

func (rn *Node) restoreMetadata(wMetadata []byte) error {
	var metaData walMetadata
	if err := json.Unmarshal(wMetadata, &metaData); err != nil {
		return errors.Wrap(err, "Error unmarshaling WAL metadata")
	}

	rn.id, rn.cid = metaData.NodeID, metaData.ClusterID
	rn.raftConfig.ID = metaData.NodeID
	return nil
}

func (rn *Node) restoreMemoryStorage(raftSnap raftpb.Snapshot, hState raftpb.HardState, ents []raftpb.Entry) error {
	if !raft.IsEmptySnap(raftSnap) {
		if err := rn.raftStorage.ApplySnapshot(raftSnap); err != nil {
			return errors.Wrap(err, "Error applying snapshot to raft memory storage")
		}
	}

	if rn.wal != nil {
		if err := rn.raftStorage.SetHardState(hState); err != nil {
			return errors.Wrap(err, "Error setting memory hardstate")
		}

		if err := rn.raftStorage.Append(ents); err != nil {
			return errors.Wrap(err, "Error appending entries to memory storage")
		}
	}

	return nil
}

func (rn *Node) deletePersistentData() error {
	if rn.snapDir() != "" {
		if err := os.RemoveAll(rn.snapDir()); err != nil {
			return errors.Wrap(err, "Error deleting snapshot directory")
		}
	}
	if rn.walDir() != "" {
		//TODO: Should be delete walDir or snapDir()?
		if err := os.RemoveAll(rn.walDir()); err != nil {
			return errors.Wrap(err, "Error deleting WAL directory")
		}
	}
	return nil
}

func (rn *Node) walDir() string {
	if rn.dataDir == "" {
		return ""
	}
	return fmt.Sprintf("%s%s", rn.dataDir, walDirExtension)
}

func (rn *Node) snapDir() string {
	if rn.dataDir == "" {
		return ""
	}
	return fmt.Sprintf("%s%s", rn.dataDir, snapDirExtension)
}
<|endoftext|>


<commit_before>// Copyright 2022 PingCAP, 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 checker

import (
	"container/list"
	"context"
	"testing"

	"github.com/stretchr/testify/require"
)

type parseTestData struct {
	sql                 string
	parseSucceeded      bool
	tableNeededExist    []string
	tableNeededNonExist []string
	executeSucceeded    bool
}

func setUpTestData() *list.List {
	testData := list.New()

	testData.PushBack(parseTestData{sql: "drop table if exists t1,t2,t3,t4,t5;", parseSucceeded: true, tableNeededExist: []string{"t1", "t2", "t3", "t4", "t5"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop database if exists mysqltest;", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t1 (b char(0));", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "insert into t1 values (''),(null);", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "select * from t1;", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop table if exists t1;", parseSucceeded: true, tableNeededExist: []string{"t1"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t1 (b char(0) not null);", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table if not exists t1 (b char(0) not null);", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "insert into t1 values (''),(null);", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: false})
	testData.PushBack(parseTestData{sql: "select * from t1;", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop table t1;", parseSucceeded: true, tableNeededExist: []string{"t1"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t(a int comment '[[range=1,10]]');", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t"}, executeSucceeded: true})
	return testData
}

func TestParse(t *testing.T) {
	ec, err := NewExecutableChecker()
	require.NoError(t, err)
	defer ec.Close()
	testData := setUpTestData()
	for e := testData.Front(); e != nil; e = e.Next() {
		data := e.Value.(parseTestData)
		stmt, err := ec.Parse(data.sql)
		if err != nil {
			require.False(t, data.parseSucceeded)
			continue
		}
		tableNeededExist, _ := GetTablesNeededExist(stmt)
		tableNeededNonExist, _ := GetTablesNeededNonExist(stmt)
		require.True(t, data.parseSucceeded)
		require.Equal(t, data.tableNeededExist, tableNeededExist)
		require.Equal(t, data.tableNeededNonExist, tableNeededNonExist)
	}
}

func TestExecute(t *testing.T) {
	ec, err := NewExecutableChecker()
	require.NoError(t, err)
	defer ec.Close()
	testData := setUpTestData()
	tidbContext := context.Background()
	err = ec.Execute(tidbContext, "use test;")
	require.NoError(t, err)
	for e := testData.Front(); e != nil; e = e.Next() {
		data := e.Value.(parseTestData)
		err := ec.Execute(tidbContext, data.sql)
		require.Equal(t, data.executeSucceeded, err == nil)
	}
}
<commit_msg>util: refactor TestExecute to fix data race (#33702)<commit_after>// Copyright 2022 PingCAP, 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 checker

import (
	"container/list"
	"testing"

	"github.com/pingcap/tidb/testkit"
	"github.com/stretchr/testify/require"
)

type parseTestData struct {
	sql                 string
	parseSucceeded      bool
	tableNeededExist    []string
	tableNeededNonExist []string
	executeSucceeded    bool
}

func setUpTestData() *list.List {
	testData := list.New()

	testData.PushBack(parseTestData{sql: "drop table if exists t1,t2,t3,t4,t5;", parseSucceeded: true, tableNeededExist: []string{"t1", "t2", "t3", "t4", "t5"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop database if exists mysqltest;", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t1 (b char(0));", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "insert into t1 values (''),(null);", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "select * from t1;", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop table if exists t1;", parseSucceeded: true, tableNeededExist: []string{"t1"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t1 (b char(0) not null);", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table if not exists t1 (b char(0) not null);", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t1"}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "insert into t1 values (''),(null);", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: false})
	testData.PushBack(parseTestData{sql: "select * from t1;", parseSucceeded: true, tableNeededExist: nil, tableNeededNonExist: nil, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "drop table t1;", parseSucceeded: true, tableNeededExist: []string{"t1"}, tableNeededNonExist: []string{}, executeSucceeded: true})
	testData.PushBack(parseTestData{sql: "create table t(a int comment '[[range=1,10]]');", parseSucceeded: true, tableNeededExist: []string{}, tableNeededNonExist: []string{"t"}, executeSucceeded: true})
	return testData
}

func TestParse(t *testing.T) {
	ec, err := NewExecutableChecker()
	require.NoError(t, err)
	defer ec.Close()
	testData := setUpTestData()
	for e := testData.Front(); e != nil; e = e.Next() {
		data := e.Value.(parseTestData)
		stmt, err := ec.Parse(data.sql)
		if err != nil {
			require.False(t, data.parseSucceeded)
			continue
		}
		tableNeededExist, _ := GetTablesNeededExist(stmt)
		tableNeededNonExist, _ := GetTablesNeededNonExist(stmt)
		require.True(t, data.parseSucceeded)
		require.Equal(t, data.tableNeededExist, tableNeededExist)
		require.Equal(t, data.tableNeededNonExist, tableNeededNonExist)
	}
}

func TestExecute(t *testing.T) {
	store, clean := testkit.CreateMockStore(t)
	defer clean()
	tk := testkit.NewTestKit(t, store)
	testData := setUpTestData()
	tk.MustExec("use test;")
	for e := testData.Front(); e != nil; e = e.Next() {
		data := e.Value.(parseTestData)
		_, err := tk.Exec(data.sql)
		require.Equal(t, data.executeSucceeded, err == nil)
	}
}
<|endoftext|>


<commit_before>package machine

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/docker"
)

type machineProvider struct {
	machine docker_helpers.Machine
	details machinesDetails
	lock    sync.RWMutex
	// provider stores a real executor that is used to start run the builds
	provider common.ExecutorProvider
}

func (m *machineProvider) machineDetails(name string, acquire bool) *machineDetails {
	m.lock.Lock()
	defer m.lock.Unlock()

	details, ok := m.details[name]
	if !ok {
		details = &machineDetails{
			Name:    name,
			Created: time.Now(),
			Used:    time.Now(),
			State:   machineStateIdle,
		}
		m.details[name] = details
	}

	if acquire {
		if details.isUsed() {
			return nil
		}
		details.State = machineStateAcquired
	}

	return details
}

func (m *machineProvider) create(config *common.RunnerConfig, state machineState) (details *machineDetails, errCh chan error) {
	name := newMachineName(machineFilter(config))
	details = m.machineDetails(name, true)
	details.State = machineStateCreating
	errCh = make(chan error, 1)

	// Create machine asynchronously
	go func() {
		started := time.Now()
		err := m.machine.Create(config.Machine.MachineDriver, details.Name, config.Machine.MachineOptions...)
		for i := 0; i < 3 && err != nil; i++ {
			logrus.WithField("name", details.Name).
				Warningln("Machine creation failed, trying to provision", err)
			time.Sleep(provisionRetryInterval)
			err = m.machine.Provision(details.Name)
		}

		if err != nil {
			m.remove(details.Name, "Failed to create")
		} else {
			details.State = state
			details.Used = time.Now()
			logrus.WithField("time", time.Since(started)).
				WithField("name", details.Name).
				Infoln("Machine created")
		}
		errCh <- err
	}()
	return
}

func (m *machineProvider) findFreeMachine(machines ...string) (details *machineDetails) {
	// Enumerate all machines
	for _, name := range machines {
		details := m.machineDetails(name, true)
		if details == nil {
			continue
		}

		// Check if node is running
		canConnect := m.machine.CanConnect(name)
		if !canConnect {
			m.remove(name, "machine is unavailable")
			continue
		}
		return details
	}

	return nil
}

func (m *machineProvider) useMachine(config *common.RunnerConfig) (details *machineDetails, err error) {
	machines, err := m.loadMachines(config)
	if err != nil {
		return
	}
	details = m.findFreeMachine(machines...)
	if details == nil {
		var errCh chan error
		details, errCh = m.create(config, machineStateAcquired)
		err = <-errCh
	}
	return
}

func (m *machineProvider) retryUseMachine(config *common.RunnerConfig) (details *machineDetails, err error) {
	// Try to find a machine
	for i := 0; i < 3; i++ {
		details, err = m.useMachine(config)
		if err == nil {
			break
		}
		time.Sleep(provisionRetryInterval)
	}
	return
}

func (m *machineProvider) finalizeRemoval(details *machineDetails) {
	for {
		if !m.machine.Exist(details.Name) {
			logrus.WithField("name", details.Name).
				WithField("created", time.Since(details.Created)).
				WithField("used", time.Since(details.Used)).
				WithField("reason", details.Reason).
				Warningln("Skipping machine removal, because it doesn't exist")
			break
		}

		err := m.machine.Remove(details.Name)
		if err == nil {
			break
		}
		time.Sleep(30 * time.Second)
		logrus.WithField("name", details.Name).
			WithField("created", time.Since(details.Created)).
			WithField("used", time.Since(details.Used)).
			WithField("reason", details.Reason).
			Warningln("Retrying removal")
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.details, details.Name)
}

func (m *machineProvider) remove(machineName string, reason ...interface{}) {
	m.lock.Lock()
	defer m.lock.Unlock()

	details, _ := m.details[machineName]
	if details == nil {
		return
	}

	details.Reason = fmt.Sprint(reason...)
	details.State = machineStateRemoving
	logrus.WithField("name", machineName).
		WithField("created", time.Since(details.Created)).
		WithField("used", time.Since(details.Used)).
		WithField("reason", details.Reason).
		Warningln("Removing machine")
	details.Used = time.Now()
	details.writeDebugInformation()

	go m.finalizeRemoval(details)
}

func (m *machineProvider) updateMachine(config *common.RunnerConfig, data *machinesData, details *machineDetails) error {
	if details.State != machineStateIdle {
		return nil
	}

	if config.Machine.MaxBuilds > 0 && details.UsedCount >= config.Machine.MaxBuilds {
		// Limit number of builds
		return errors.New("Too many builds")
	}

	if data.Total() >= config.Limit && config.Limit > 0 {
		// Limit maximum number of machines
		return errors.New("Too many machines")
	}

	if time.Since(details.Used) > time.Second*time.Duration(config.Machine.IdleTime) {
		if data.Idle >= config.Machine.IdleCount {
			// Remove machine that are way over the idle time
			return errors.New("Too many idle machines")
		}
	}
	return nil
}

func (m *machineProvider) updateMachines(machines []string, config *common.RunnerConfig) (data machinesData) {
	data.Runner = config.ShortDescription()

	for _, name := range machines {
		details := m.machineDetails(name, false)
		err := m.updateMachine(config, &data, details)
		if err != nil {
			m.remove(details.Name, err)
		}

		data.Add(details.State)
	}
	return
}

func (m *machineProvider) createMachines(config *common.RunnerConfig, data *machinesData) {
	// Create a new machines and mark them as Idle
	for {
		if data.Available() >= config.Machine.IdleCount {
			// Limit maximum number of idle machines
			break
		}
		if data.Total() >= config.Limit && config.Limit > 0 {
			// Limit maximum number of machines
			break
		}
		m.create(config, machineStateIdle)
		data.Creating++
	}
}

func (m *machineProvider) loadMachines(config *common.RunnerConfig) ([]string, error) {
	// Find a new machine
	return m.machine.List(machineFilter(config))
}

func (m *machineProvider) Acquire(config *common.RunnerConfig) (data common.ExecutorData, err error) {
	if config.Machine == nil || config.Machine.MachineName == "" {
		err = fmt.Errorf("Missing Machine options")
		return
	}

	machines, err := m.loadMachines(config)
	if err != nil {
		return
	}

	// Update a list of currently configured machines
	machinesData := m.updateMachines(machines, config)

	// Pre-create machines
	m.createMachines(config, &machinesData)

	logrus.WithFields(machinesData.Fields()).
		WithField("runner", config.ShortDescription()).
		WithField("minIdleCount", config.Machine.IdleCount).
		WithField("maxMachines", config.Limit).
		WithField("time", time.Now()).
		Debugln("Docker Machine Details")
	machinesData.writeDebugInformation()

	// Try to find a free machine
	details := m.findFreeMachine(machines...)
	if details != nil {
		data = details
		return
	}

	// If we have a free machines we can process a build
	if config.Machine.IdleCount != 0 && machinesData.Idle == 0 {
		err = errors.New("No free machines that can process builds")
	}
	return
}

func (m *machineProvider) Use(config *common.RunnerConfig, data common.ExecutorData) (newConfig common.RunnerConfig, newData common.ExecutorData, err error) {
	// Find a new machine
	details, _ := data.(*machineDetails)
	if details == nil {
		details, err = m.retryUseMachine(config)
		if err != nil {
			return
		}

		// Return details only if this is a new instance
		newData = details
	}

	// Get machine credentials
	dc, err := m.machine.Credentials(details.Name)
	if err != nil {
		if newData != nil {
			m.Release(config, newData)
		}
		return
	}

	// Create shallow copy of config and store in it docker credentials
	newConfig = *config
	newConfig.Docker = &common.DockerConfig{}
	if config.Docker != nil {
		*newConfig.Docker = *config.Docker
	}
	newConfig.Docker.DockerCredentials = dc

	// Mark machine as used
	details.State = machineStateUsed
	return
}

func (m *machineProvider) Release(config *common.RunnerConfig, data common.ExecutorData) error {
	// Release machine
	details, ok := data.(*machineDetails)
	if ok {
		// Mark last used time when is Used
		if details.State == machineStateUsed {
			details.Used = time.Now()
			details.UsedCount++
		}
		details.State = machineStateIdle
	}
	return nil
}

func (m *machineProvider) CanCreate() bool {
	return m.provider.CanCreate()
}

func (m *machineProvider) GetFeatures(features *common.FeaturesInfo) {
	m.provider.GetFeatures(features)
}

func (m *machineProvider) Create() common.Executor {
	return &machineExecutor{
		provider: m,
	}
}

func newMachineProvider(executor string) *machineProvider {
	provider := common.GetExecutor(executor)
	if provider == nil {
		logrus.Panicln("Missing", executor)
	}

	return &machineProvider{
		details:  make(machinesDetails),
		machine:  docker_helpers.NewMachineCommand(),
		provider: provider,
	}
}
<commit_msg>Lock updating machines when acquire happens<commit_after>package machine

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers/docker"
)

type machineProvider struct {
	machine     docker_helpers.Machine
	details     machinesDetails
	lock        sync.RWMutex
	acquireLock sync.Mutex
	// provider stores a real executor that is used to start run the builds
	provider common.ExecutorProvider
}

func (m *machineProvider) machineDetails(name string, acquire bool) *machineDetails {
	m.lock.Lock()
	defer m.lock.Unlock()

	details, ok := m.details[name]
	if !ok {
		details = &machineDetails{
			Name:    name,
			Created: time.Now(),
			Used:    time.Now(),
			State:   machineStateIdle,
		}
		m.details[name] = details
	}

	if acquire {
		if details.isUsed() {
			return nil
		}
		details.State = machineStateAcquired
	}

	return details
}

func (m *machineProvider) create(config *common.RunnerConfig, state machineState) (details *machineDetails, errCh chan error) {
	name := newMachineName(machineFilter(config))
	details = m.machineDetails(name, true)
	details.State = machineStateCreating
	errCh = make(chan error, 1)

	// Create machine asynchronously
	go func() {
		started := time.Now()
		err := m.machine.Create(config.Machine.MachineDriver, details.Name, config.Machine.MachineOptions...)
		for i := 0; i < 3 && err != nil; i++ {
			logrus.WithField("name", details.Name).
				Warningln("Machine creation failed, trying to provision", err)
			time.Sleep(provisionRetryInterval)
			err = m.machine.Provision(details.Name)
		}

		if err != nil {
			m.remove(details.Name, "Failed to create")
		} else {
			details.State = state
			details.Used = time.Now()
			logrus.WithField("time", time.Since(started)).
				WithField("name", details.Name).
				Infoln("Machine created")
		}
		errCh <- err
	}()
	return
}

func (m *machineProvider) findFreeMachine(machines ...string) (details *machineDetails) {
	// Enumerate all machines
	for _, name := range machines {
		details := m.machineDetails(name, true)
		if details == nil {
			continue
		}

		// Check if node is running
		canConnect := m.machine.CanConnect(name)
		if !canConnect {
			m.remove(name, "machine is unavailable")
			continue
		}
		return details
	}

	return nil
}

func (m *machineProvider) useMachine(config *common.RunnerConfig) (details *machineDetails, err error) {
	machines, err := m.loadMachines(config)
	if err != nil {
		return
	}
	details = m.findFreeMachine(machines...)
	if details == nil {
		var errCh chan error
		details, errCh = m.create(config, machineStateAcquired)
		err = <-errCh
	}
	return
}

func (m *machineProvider) retryUseMachine(config *common.RunnerConfig) (details *machineDetails, err error) {
	// Try to find a machine
	for i := 0; i < 3; i++ {
		details, err = m.useMachine(config)
		if err == nil {
			break
		}
		time.Sleep(provisionRetryInterval)
	}
	return
}

func (m *machineProvider) finalizeRemoval(details *machineDetails) {
	for {
		if !m.machine.Exist(details.Name) {
			logrus.WithField("name", details.Name).
				WithField("created", time.Since(details.Created)).
				WithField("used", time.Since(details.Used)).
				WithField("reason", details.Reason).
				Warningln("Skipping machine removal, because it doesn't exist")
			break
		}

		err := m.machine.Remove(details.Name)
		if err == nil {
			break
		}
		time.Sleep(30 * time.Second)
		logrus.WithField("name", details.Name).
			WithField("created", time.Since(details.Created)).
			WithField("used", time.Since(details.Used)).
			WithField("reason", details.Reason).
			Warningln("Retrying removal")
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.details, details.Name)
}

func (m *machineProvider) remove(machineName string, reason ...interface{}) {
	m.lock.Lock()
	defer m.lock.Unlock()

	details, _ := m.details[machineName]
	if details == nil {
		return
	}

	details.Reason = fmt.Sprint(reason...)
	details.State = machineStateRemoving
	logrus.WithField("name", machineName).
		WithField("created", time.Since(details.Created)).
		WithField("used", time.Since(details.Used)).
		WithField("reason", details.Reason).
		Warningln("Removing machine")
	details.Used = time.Now()
	details.writeDebugInformation()

	go m.finalizeRemoval(details)
}

func (m *machineProvider) updateMachine(config *common.RunnerConfig, data *machinesData, details *machineDetails) error {
	if details.State != machineStateIdle {
		return nil
	}

	if config.Machine.MaxBuilds > 0 && details.UsedCount >= config.Machine.MaxBuilds {
		// Limit number of builds
		return errors.New("Too many builds")
	}

	if data.Total() >= config.Limit && config.Limit > 0 {
		// Limit maximum number of machines
		return errors.New("Too many machines")
	}

	if time.Since(details.Used) > time.Second*time.Duration(config.Machine.IdleTime) {
		if data.Idle >= config.Machine.IdleCount {
			// Remove machine that are way over the idle time
			return errors.New("Too many idle machines")
		}
	}
	return nil
}

func (m *machineProvider) updateMachines(machines []string, config *common.RunnerConfig) (data machinesData) {
	data.Runner = config.ShortDescription()

	for _, name := range machines {
		details := m.machineDetails(name, false)
		err := m.updateMachine(config, &data, details)
		if err != nil {
			m.remove(details.Name, err)
		}

		data.Add(details.State)
	}
	return
}

func (m *machineProvider) createMachines(config *common.RunnerConfig, data *machinesData) {
	// Create a new machines and mark them as Idle
	for {
		if data.Available() >= config.Machine.IdleCount {
			// Limit maximum number of idle machines
			break
		}
		if data.Total() >= config.Limit && config.Limit > 0 {
			// Limit maximum number of machines
			break
		}
		m.create(config, machineStateIdle)
		data.Creating++
	}
}

func (m *machineProvider) loadMachines(config *common.RunnerConfig) ([]string, error) {
	// Find a new machine
	return m.machine.List(machineFilter(config))
}

func (m *machineProvider) Acquire(config *common.RunnerConfig) (data common.ExecutorData, err error) {
	if config.Machine == nil || config.Machine.MachineName == "" {
		err = fmt.Errorf("Missing Machine options")
		return
	}

	machines, err := m.loadMachines(config)
	if err != nil {
		return
	}

	// Lock updating machines, because two Acquires can be run at the same time
	m.acquireLock.Lock()

	// Update a list of currently configured machines
	machinesData := m.updateMachines(machines, config)

	// Pre-create machines
	m.createMachines(config, &machinesData)

	m.acquireLock.Unlock()

	logrus.WithFields(machinesData.Fields()).
		WithField("runner", config.ShortDescription()).
		WithField("minIdleCount", config.Machine.IdleCount).
		WithField("maxMachines", config.Limit).
		WithField("time", time.Now()).
		Debugln("Docker Machine Details")
	machinesData.writeDebugInformation()

	// Try to find a free machine
	details := m.findFreeMachine(machines...)
	if details != nil {
		data = details
		return
	}

	// If we have a free machines we can process a build
	if config.Machine.IdleCount != 0 && machinesData.Idle == 0 {
		err = errors.New("No free machines that can process builds")
	}
	return
}

func (m *machineProvider) Use(config *common.RunnerConfig, data common.ExecutorData) (newConfig common.RunnerConfig, newData common.ExecutorData, err error) {
	// Find a new machine
	details, _ := data.(*machineDetails)
	if details == nil {
		details, err = m.retryUseMachine(config)
		if err != nil {
			return
		}

		// Return details only if this is a new instance
		newData = details
	}

	// Get machine credentials
	dc, err := m.machine.Credentials(details.Name)
	if err != nil {
		if newData != nil {
			m.Release(config, newData)
		}
		return
	}

	// Create shallow copy of config and store in it docker credentials
	newConfig = *config
	newConfig.Docker = &common.DockerConfig{}
	if config.Docker != nil {
		*newConfig.Docker = *config.Docker
	}
	newConfig.Docker.DockerCredentials = dc

	// Mark machine as used
	details.State = machineStateUsed
	return
}

func (m *machineProvider) Release(config *common.RunnerConfig, data common.ExecutorData) error {
	// Release machine
	details, ok := data.(*machineDetails)
	if ok {
		// Mark last used time when is Used
		if details.State == machineStateUsed {
			details.Used = time.Now()
			details.UsedCount++
		}
		details.State = machineStateIdle
	}
	return nil
}

func (m *machineProvider) CanCreate() bool {
	return m.provider.CanCreate()
}

func (m *machineProvider) GetFeatures(features *common.FeaturesInfo) {
	m.provider.GetFeatures(features)
}

func (m *machineProvider) Create() common.Executor {
	return &machineExecutor{
		provider: m,
	}
}

func newMachineProvider(executor string) *machineProvider {
	provider := common.GetExecutor(executor)
	if provider == nil {
		logrus.Panicln("Missing", executor)
	}

	return &machineProvider{
		details:  make(machinesDetails),
		machine:  docker_helpers.NewMachineCommand(),
		provider: provider,
	}
}
<|endoftext|>


<commit_before>package router

import (
	wt "github.com/weaveworks/weave/testing"
	"testing"
	"time"
)

// TODO test gossip unicast and broadcast; atm we only test topology
// gossip, which does not employ unicast or broadcast.

type mockChannelConnection struct {
	RemoteConnection
	dest *Router
}

// Construct a "passive" Router, i.e. without any goroutines, except
// for Routes and GossipSenders.
func NewTestRouter(name PeerName) *Router {
	router := NewRouter(RouterConfig{}, name, "")
	// need to create a dummy channel otherwise tests hang on nil
	// channels when the Router invoked ConnectionMaker.Refresh
	router.ConnectionMaker.actionChan = make(chan ConnectionMakerAction, ChannelSize)
	router.Routes.Start()
	return router
}

func (conn *mockChannelConnection) SendProtocolMsg(protocolMsg ProtocolMsg) {
	if err := conn.dest.handleGossip(protocolMsg.tag, protocolMsg.msg); err != nil {
		panic(err)
	}
	conn.dest.sendPendingGossip()
}

func (router *Router) AddTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer := NewPeer(fromName, "", router.Ourself.Peer.UID, 0)
	toPeer := NewPeer(toName, "", r.Ourself.Peer.UID, 0)

	r.Peers.FetchWithDefault(fromPeer)    // Has side-effect of incrementing refcount
	router.Peers.FetchWithDefault(toPeer) //

	conn := &mockChannelConnection{RemoteConnection{router.Ourself.Peer, toPeer, "", false, true}, r}
	router.Ourself.handleAddConnection(conn)
	router.Ourself.handleConnectionEstablished(conn)
	router.sendPendingGossip()
}

func (router *Router) DeleteTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer, _ := r.Peers.Fetch(fromName)
	toPeer, _ := router.Peers.Fetch(toName)

	r.Peers.Dereference(fromPeer)
	router.Peers.Dereference(toPeer)

	conn, _ := router.Ourself.ConnectionTo(toName)
	router.Ourself.handleDeleteConnection(conn)
	router.sendPendingGossip()
}

func TestGossipTopology(t *testing.T) {
	wt.RunWithTimeout(t, 1*time.Second, func() {
		implTestGossipTopology(t)
	})
}

// Create a Peer representing the receiver router, with connections to
// the routers supplied as arguments, carrying across all UID and
// version information.
func (router *Router) tp(routers ...*Router) *Peer {
	peer := NewPeer(router.Ourself.Peer.Name, "", router.Ourself.Peer.UID, 0)
	connections := make(map[PeerName]Connection)
	for _, r := range routers {
		p := NewPeer(r.Ourself.Peer.Name, "", r.Ourself.Peer.UID, r.Ourself.Peer.version)
		connections[r.Ourself.Peer.Name] = newMockConnection(peer, p)
	}
	peer.version = router.Ourself.Peer.version
	peer.connections = connections
	return peer
}

// Check that the topology of router matches the peers and all of their connections
func checkTopology(t *testing.T, router *Router, wantedPeers ...*Peer) {
	checkTopologyPeers(t, true, router.Peers.allPeers(), wantedPeers...)
}

func implTestGossipTopology(t *testing.T) {
	// Create some peers that will talk to each other
	peer1Name, _ := PeerNameFromString("01:00:00:01:00:00")
	peer2Name, _ := PeerNameFromString("02:00:00:02:00:00")
	peer3Name, _ := PeerNameFromString("03:00:00:03:00:00")
	r1 := NewTestRouter(peer1Name)
	r2 := NewTestRouter(peer2Name)
	r3 := NewTestRouter(peer3Name)

	// Check state when they have no connections
	checkTopology(t, r1, r1.tp())
	checkTopology(t, r2, r2.tp())

	// Now try adding some connections
	r1.AddTestChannelConnection(r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp())
	r2.AddTestChannelConnection(r1)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))

	// Currently, the connection from 2 to 3 is one-way only
	r2.AddTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	// When r2 gossiped to r3, 1 was unreachable from r3 so it got removed from the
	// list of peers, but remains referenced in the connection from 1 to 3.
	checkTopology(t, r3, r2.tp(r1, r3), r3.tp())

	// Add a connection from 3 to 1 and now r1 is reachable.
	r3.AddTestChannelConnection(r1)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(), r2.tp(r1, r3), r3.tp(r1))

	r1.AddTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))

	// Drop the connection from 2 to 3
	r2.DeleteTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// Drop the connection from 1 to 3
	r1.DeleteTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))

	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))
	// r3 still thinks r1 has a connection to it
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// On a timer, r3 will gossip to r1
	r3.SendAllGossip()
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
}
<commit_msg>fixing build on master (broken by 68ef0e)<commit_after>package router

import (
	"testing"
	"time"

	wt "github.com/weaveworks/weave/testing"
)

// TODO test gossip unicast and broadcast; atm we only test topology
// gossip, which does not employ unicast or broadcast.

type mockChannelConnection struct {
	RemoteConnection
	dest *Router
}

// Construct a "passive" Router, i.e. without any goroutines, except
// for Routes and GossipSenders.
func NewTestRouter(name PeerName) *Router {
	router := NewRouter(Config{}, name, "")
	// need to create a dummy channel otherwise tests hang on nil
	// channels when the Router invoked ConnectionMaker.Refresh
	router.ConnectionMaker.actionChan = make(chan ConnectionMakerAction, ChannelSize)
	router.Routes.Start()
	return router
}

func (conn *mockChannelConnection) SendProtocolMsg(protocolMsg ProtocolMsg) {
	if err := conn.dest.handleGossip(protocolMsg.tag, protocolMsg.msg); err != nil {
		panic(err)
	}
	conn.dest.sendPendingGossip()
}

func (router *Router) AddTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer := NewPeer(fromName, "", router.Ourself.Peer.UID, 0)
	toPeer := NewPeer(toName, "", r.Ourself.Peer.UID, 0)

	r.Peers.FetchWithDefault(fromPeer)    // Has side-effect of incrementing refcount
	router.Peers.FetchWithDefault(toPeer) //

	conn := &mockChannelConnection{RemoteConnection{router.Ourself.Peer, toPeer, "", false, true}, r}
	router.Ourself.handleAddConnection(conn)
	router.Ourself.handleConnectionEstablished(conn)
	router.sendPendingGossip()
}

func (router *Router) DeleteTestChannelConnection(r *Router) {
	fromName := router.Ourself.Peer.Name
	toName := r.Ourself.Peer.Name

	fromPeer, _ := r.Peers.Fetch(fromName)
	toPeer, _ := router.Peers.Fetch(toName)

	r.Peers.Dereference(fromPeer)
	router.Peers.Dereference(toPeer)

	conn, _ := router.Ourself.ConnectionTo(toName)
	router.Ourself.handleDeleteConnection(conn)
	router.sendPendingGossip()
}

func TestGossipTopology(t *testing.T) {
	wt.RunWithTimeout(t, 1*time.Second, func() {
		implTestGossipTopology(t)
	})
}

// Create a Peer representing the receiver router, with connections to
// the routers supplied as arguments, carrying across all UID and
// version information.
func (router *Router) tp(routers ...*Router) *Peer {
	peer := NewPeer(router.Ourself.Peer.Name, "", router.Ourself.Peer.UID, 0)
	connections := make(map[PeerName]Connection)
	for _, r := range routers {
		p := NewPeer(r.Ourself.Peer.Name, "", r.Ourself.Peer.UID, r.Ourself.Peer.version)
		connections[r.Ourself.Peer.Name] = newMockConnection(peer, p)
	}
	peer.version = router.Ourself.Peer.version
	peer.connections = connections
	return peer
}

// Check that the topology of router matches the peers and all of their connections
func checkTopology(t *testing.T, router *Router, wantedPeers ...*Peer) {
	checkTopologyPeers(t, true, router.Peers.allPeers(), wantedPeers...)
}

func implTestGossipTopology(t *testing.T) {
	// Create some peers that will talk to each other
	peer1Name, _ := PeerNameFromString("01:00:00:01:00:00")
	peer2Name, _ := PeerNameFromString("02:00:00:02:00:00")
	peer3Name, _ := PeerNameFromString("03:00:00:03:00:00")
	r1 := NewTestRouter(peer1Name)
	r2 := NewTestRouter(peer2Name)
	r3 := NewTestRouter(peer3Name)

	// Check state when they have no connections
	checkTopology(t, r1, r1.tp())
	checkTopology(t, r2, r2.tp())

	// Now try adding some connections
	r1.AddTestChannelConnection(r2)
	checkTopology(t, r1, r1.tp(r2), r2.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp())
	r2.AddTestChannelConnection(r1)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))

	// Currently, the connection from 2 to 3 is one-way only
	r2.AddTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp())
	// When r2 gossiped to r3, 1 was unreachable from r3 so it got removed from the
	// list of peers, but remains referenced in the connection from 1 to 3.
	checkTopology(t, r3, r2.tp(r1, r3), r3.tp())

	// Add a connection from 3 to 1 and now r1 is reachable.
	r3.AddTestChannelConnection(r1)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(), r2.tp(r1, r3), r3.tp(r1))

	r1.AddTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1))

	// Drop the connection from 2 to 3
	r2.DeleteTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1))
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// Drop the connection from 1 to 3
	r1.DeleteTestChannelConnection(r3)
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))

	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))
	// r3 still thinks r1 has a connection to it
	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1))

	// On a timer, r3 will gossip to r1
	r3.SendAllGossip()
	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1))
}
<|endoftext|>


<commit_before>package cluster

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/rancher/kontainer-engine/service"
	"github.com/rancher/norman/api/access"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	gaccess "github.com/rancher/rancher/pkg/api/customization/globalnamespaceaccess"
	"github.com/rancher/rancher/pkg/controllers/management/k3supgrade"
	"github.com/rancher/rancher/pkg/controllers/user/cis"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	mgmtSchema "github.com/rancher/types/apis/management.cattle.io/v3/schema"
	mgmtclient "github.com/rancher/types/client/management/v3"
	"github.com/robfig/cron"
	"github.com/sirupsen/logrus"
)

type Validator struct {
	ClusterLister                 v3.ClusterLister
	ClusterTemplateLister         v3.ClusterTemplateLister
	ClusterTemplateRevisionLister v3.ClusterTemplateRevisionLister
	Users                         v3.UserInterface
	GrbLister                     v3.GlobalRoleBindingLister
	GrLister                      v3.GlobalRoleLister
	CisConfigClient               v3.CisConfigInterface
	CisConfigLister               v3.CisConfigLister
	CisBenchmarkVersionClient     v3.CisBenchmarkVersionInterface
	CisBenchmarkVersionLister     v3.CisBenchmarkVersionLister
}

func (v *Validator) Validator(request *types.APIContext, schema *types.Schema, data map[string]interface{}) error {
	var clusterSpec v3.ClusterSpec
	var clientClusterSpec mgmtclient.Cluster
	logrus.Tracef("Validator: data: %+v\n", data)
	if err := convert.ToObj(data, &clusterSpec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Cluster spec conversion error")
	}

	if err := convert.ToObj(data, &clientClusterSpec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Cluster spec conversion error")
	}

	logrus.Tracef("Validator: clusterSpec: %+v\n", &clusterSpec)
	logrus.Tracef("Validator: clientClusterSpec: %+v\n", &clientClusterSpec)

	if err := v.validateEnforcement(request, data); err != nil {
		return err
	}

	if err := v.validateLocalClusterAuthEndpoint(request, &clusterSpec); err != nil {
		return err
	}

	if err := v.validateK3sVersionUpgrade(request, &clusterSpec); err != nil {
		return err
	}

	if err := v.validateScheduledClusterScan(&clientClusterSpec); err != nil {
		return err
	}

	if err := v.validateGenericEngineConfig(request, &clusterSpec); err != nil {
		return err
	}
	return nil
}

func (v *Validator) validateScheduledClusterScan(spec *mgmtclient.Cluster) error {
	// If this cluster is created using a template, we dont have the version in the provided data, skip
	if spec.ClusterTemplateRevisionID != "" {
		return nil
	}

	// If CIS scan is not present/enabled, skip
	if spec.ScheduledClusterScan == nil ||
		(spec.ScheduledClusterScan != nil && !spec.ScheduledClusterScan.Enabled) {
		return nil
	}
	currentK8sVersion := spec.RancherKubernetesEngineConfig.Version
	overrideBenchmarkVersion := ""
	if spec.ScheduledClusterScan.ScanConfig.CisScanConfig != nil {
		overrideBenchmarkVersion = spec.ScheduledClusterScan.ScanConfig.CisScanConfig.OverrideBenchmarkVersion
	}
	_, _, err := cis.GetBenchmarkVersionToUse(overrideBenchmarkVersion, currentK8sVersion,
		v.CisConfigLister, v.CisConfigClient,
		v.CisBenchmarkVersionLister, v.CisBenchmarkVersionClient,
	)
	if err != nil {
		return httperror.NewAPIError(httperror.InvalidBodyContent, err.Error())
	}
	return validateScheduledClusterScan(spec)
}

func validateScheduledClusterScan(spec *mgmtclient.Cluster) error {
	// If this cluster is created using a template, we dont have the version in the provided data, skip
	if spec.ClusterTemplateRevisionID != "" {
		return nil
	}

	if spec.ScheduledClusterScan.ScanConfig != nil &&
		spec.ScheduledClusterScan.ScanConfig.CisScanConfig != nil {
		profile := spec.ScheduledClusterScan.ScanConfig.CisScanConfig.Profile
		if profile != string(v3.CisScanProfileTypePermissive) &&
			profile != string(v3.CisScanProfileTypeHardened) {
			return httperror.NewFieldAPIError(httperror.InvalidOption, "ScheduledClusterScan.ScanConfig.CisScanConfig.Profile", "profile can be either permissive or hardened")
		}
	}

	if spec.ScheduledClusterScan.ScheduleConfig != nil {
		if spec.ScheduledClusterScan.ScheduleConfig.Retention < 0 {
			return httperror.NewFieldAPIError(httperror.MinLimitExceeded, "ScheduledClusterScan.ScheduleConfig.Retention", "Retention count cannot be negative")
		}
		schedule, err := cron.ParseStandard(spec.ScheduledClusterScan.ScheduleConfig.CronSchedule)
		if err != nil {
			return httperror.NewFieldAPIError(httperror.InvalidFormat, "ScheduledClusterScan.ScheduleConfig.CronSchedule", fmt.Sprintf("error parsing cron schedule: %v", err))
		}
		now := time.Now().Round(time.Second)
		next1 := schedule.Next(now).Round(time.Second)
		next2 := schedule.Next(next1).Round(time.Second)
		timeAfter := next2.Sub(next1).Round(time.Second)

		if timeAfter < (1 * time.Hour) {
			if spec.ScheduledClusterScan.ScanConfig.CisScanConfig.DebugMaster ||
				spec.ScheduledClusterScan.ScanConfig.CisScanConfig.DebugWorker {
				return nil
			}
			return httperror.NewFieldAPIError(httperror.MinLimitExceeded, "ScheduledClusterScan.ScheduleConfig.CronSchedule", "minimum interval is one hour")
		}
	}
	return nil
}

func (v *Validator) validateLocalClusterAuthEndpoint(request *types.APIContext, spec *v3.ClusterSpec) error {
	if !spec.LocalClusterAuthEndpoint.Enabled {
		return nil
	}

	var isValidCluster bool
	if request.ID == "" {
		isValidCluster = spec.RancherKubernetesEngineConfig != nil
	} else {
		cluster, err := v.ClusterLister.Get("", request.ID)
		if err != nil {
			return err
		}
		isValidCluster = cluster.Status.Driver == "" ||
			cluster.Status.Driver == v3.ClusterDriverRKE ||
			cluster.Status.Driver == v3.ClusterDriverImported
	}
	if !isValidCluster {
		return httperror.NewFieldAPIError(httperror.InvalidState, "LocalClusterAuthEndpoint.Enabled", "Can only enable LocalClusterAuthEndpoint with RKE")
	}

	if spec.LocalClusterAuthEndpoint.CACerts != "" && spec.LocalClusterAuthEndpoint.FQDN == "" {
		return httperror.NewFieldAPIError(httperror.MissingRequired, "LocalClusterAuthEndpoint.FQDN", "CACerts defined but FQDN is not defined")
	}

	return nil
}

func (v *Validator) validateEnforcement(request *types.APIContext, data map[string]interface{}) error {

	if !strings.EqualFold(settings.ClusterTemplateEnforcement.Get(), "true") {
		return nil
	}

	var spec mgmtclient.Cluster
	if err := convert.ToObj(data, &spec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Cluster spec conversion error")
	}

	if !v.checkClusterForEnforcement(&spec) {
		return nil
	}

	ma := gaccess.MemberAccess{
		Users:     v.Users,
		GrLister:  v.GrLister,
		GrbLister: v.GrbLister,
	}

	//if user is admin, no checks needed
	callerID := request.Request.Header.Get(gaccess.ImpersonateUserHeader)

	isAdmin, err := ma.IsAdmin(callerID)
	if err != nil {
		return err
	}
	if isAdmin {
		return nil
	}

	//enforcement is true, template is a must
	if spec.ClusterTemplateRevisionID == "" {
		return httperror.NewFieldAPIError(httperror.MissingRequired, "", "A clusterTemplateRevision to create a cluster")
	}

	err = v.accessTemplate(request, &spec)
	if err != nil {
		if httperror.IsForbidden(err) || httperror.IsNotFound(err) {
			return httperror.NewAPIError(httperror.NotFound, "The clusterTemplateRevision is not found")
		}
		return err
	}

	return nil
}

// TODO: test validator
// prevents downgrades, no-ops, and upgrading before versions have been set
func (v *Validator) validateK3sVersionUpgrade(request *types.APIContext, spec *v3.ClusterSpec) error {
	upgradeNotReadyErr := httperror.NewAPIError(httperror.Conflict, "k3s version upgrade is not ready, try again later")

	if request.Method == http.MethodPost {
		return nil
	}

	if spec.K3sConfig == nil {
		// only applies to k3s clusters
		return nil
	}

	// must wait for original spec version to be set
	if spec.K3sConfig.Version == "" {
		return upgradeNotReadyErr
	}

	cluster, err := v.ClusterLister.Get("", request.ID)
	if err != nil {
		return err
	}

	// must wait for original status version to be set
	if cluster.Status.Version == nil {
		return upgradeNotReadyErr
	}

	prevVersion := cluster.Status.Version.GitVersion
	updateVersion := spec.K3sConfig.Version

	if prevVersion == updateVersion {
		// no op
		return nil
	}

	isNewer, err := k3supgrade.IsNewerVersion(prevVersion, updateVersion)
	if err != nil {
		errMsg := fmt.Sprintf("unable to compare k3s version [%s]", spec.K3sConfig.Version)
		return httperror.NewAPIError(httperror.InvalidBodyContent, errMsg)
	}

	if !isNewer {
		// update version must be higher than previous version, downgrades are not supported
		errMsg := fmt.Sprintf("cannot upgrade k3s cluster version from [%s] to [%s]. New version must be higher.", prevVersion, updateVersion)
		return httperror.NewAPIError(httperror.InvalidBodyContent, errMsg)
	}

	return nil
}

func (v *Validator) checkClusterForEnforcement(spec *mgmtclient.Cluster) bool {
	if spec.RancherKubernetesEngineConfig != nil {
		return true
	}

	if spec.ClusterTemplateRevisionID != "" {
		return true
	}
	return false
}

func (v *Validator) accessTemplate(request *types.APIContext, spec *mgmtclient.Cluster) error {
	split := strings.SplitN(spec.ClusterTemplateRevisionID, ":", 2)
	if len(split) != 2 {
		return fmt.Errorf("error in splitting clusterTemplateRevision name %v", spec.ClusterTemplateRevisionID)
	}
	revName := split[1]
	clusterTempRev, err := v.ClusterTemplateRevisionLister.Get(namespace.GlobalNamespace, revName)
	if err != nil {
		return err
	}

	var ctMap map[string]interface{}
	if err := access.ByID(request, &mgmtSchema.Version, mgmtclient.ClusterTemplateType, clusterTempRev.Spec.ClusterTemplateName, &ctMap); err != nil {
		return err
	}

	return nil
}

// validateGenericEngineConfig allows for additional validation of clusters that depend on Kontainer Engine or Rancher Machine driver
func (v *Validator) validateGenericEngineConfig(request *types.APIContext, spec *v3.ClusterSpec) error {

	if request.Method == http.MethodPost {
		return nil
	}

	if spec.AmazonElasticContainerServiceConfig != nil {
		// compare with current cluster
		clusterName := request.ID
		prevCluster, err := v.ClusterLister.Get("", clusterName)
		if err != nil {
			return httperror.WrapAPIError(err, httperror.InvalidBodyContent, err.Error())
		}

		err = validateEKS(*prevCluster.Spec.GenericEngineConfig, *spec.AmazonElasticContainerServiceConfig)
		if err != nil {
			return err
		}
	}

	return nil

}

func validateEKS(prevCluster, newCluster map[string]interface{}) error {
	// check config is for EKS clusters
	if driver, ok := prevCluster["driverName"]; ok {
		if driver != service.AmazonElasticContainerServiceDriverName {
			return nil
		}
	}

	// don't allow for updating subnets
	if prev, ok := prevCluster["subnets"]; ok {
		if new, ok := newCluster["subnets"]; ok {
			if !reflect.DeepEqual(prev, new) {
				return httperror.NewAPIError(httperror.InvalidBodyContent, "cannot modify EKS subnets after creation")
			}
		}
	}
	return nil
}
<commit_msg>Address review comments on CIS cluster validation<commit_after>package cluster

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/rancher/kontainer-engine/service"
	"github.com/rancher/norman/api/access"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	gaccess "github.com/rancher/rancher/pkg/api/customization/globalnamespaceaccess"
	"github.com/rancher/rancher/pkg/controllers/management/k3supgrade"
	"github.com/rancher/rancher/pkg/controllers/user/cis"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	mgmtSchema "github.com/rancher/types/apis/management.cattle.io/v3/schema"
	mgmtclient "github.com/rancher/types/client/management/v3"
	"github.com/robfig/cron"
)

type Validator struct {
	ClusterLister                 v3.ClusterLister
	ClusterTemplateLister         v3.ClusterTemplateLister
	ClusterTemplateRevisionLister v3.ClusterTemplateRevisionLister
	Users                         v3.UserInterface
	GrbLister                     v3.GlobalRoleBindingLister
	GrLister                      v3.GlobalRoleLister
	CisConfigClient               v3.CisConfigInterface
	CisConfigLister               v3.CisConfigLister
	CisBenchmarkVersionClient     v3.CisBenchmarkVersionInterface
	CisBenchmarkVersionLister     v3.CisBenchmarkVersionLister
}

func (v *Validator) Validator(request *types.APIContext, schema *types.Schema, data map[string]interface{}) error {
	var clusterSpec v3.ClusterSpec
	var clientClusterSpec mgmtclient.Cluster
	if err := convert.ToObj(data, &clusterSpec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Cluster spec conversion error")
	}

	if err := convert.ToObj(data, &clientClusterSpec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Client cluster spec conversion error")
	}

	if err := v.validateEnforcement(request, data); err != nil {
		return err
	}

	if err := v.validateLocalClusterAuthEndpoint(request, &clusterSpec); err != nil {
		return err
	}

	if err := v.validateK3sVersionUpgrade(request, &clusterSpec); err != nil {
		return err
	}

	if err := v.validateScheduledClusterScan(&clientClusterSpec); err != nil {
		return err
	}

	if err := v.validateGenericEngineConfig(request, &clusterSpec); err != nil {
		return err
	}
	return nil
}

func (v *Validator) validateScheduledClusterScan(spec *mgmtclient.Cluster) error {
	// If this cluster is created using a template, we dont have the version in the provided data, skip
	if spec.ClusterTemplateRevisionID != "" {
		return nil
	}

	// If CIS scan is not present/enabled, skip
	if spec.ScheduledClusterScan == nil ||
		(spec.ScheduledClusterScan != nil && !spec.ScheduledClusterScan.Enabled) {
		return nil
	}
	currentK8sVersion := spec.RancherKubernetesEngineConfig.Version
	overrideBenchmarkVersion := ""
	if spec.ScheduledClusterScan.ScanConfig.CisScanConfig != nil {
		overrideBenchmarkVersion = spec.ScheduledClusterScan.ScanConfig.CisScanConfig.OverrideBenchmarkVersion
	}
	_, _, err := cis.GetBenchmarkVersionToUse(overrideBenchmarkVersion, currentK8sVersion,
		v.CisConfigLister, v.CisConfigClient,
		v.CisBenchmarkVersionLister, v.CisBenchmarkVersionClient,
	)
	if err != nil {
		return httperror.NewAPIError(httperror.InvalidBodyContent, err.Error())
	}
	return validateScheduledClusterScan(spec)
}

func validateScheduledClusterScan(spec *mgmtclient.Cluster) error {
	// If this cluster is created using a template, we dont have the version in the provided data, skip
	if spec.ClusterTemplateRevisionID != "" {
		return nil
	}

	if spec.ScheduledClusterScan.ScanConfig != nil &&
		spec.ScheduledClusterScan.ScanConfig.CisScanConfig != nil {
		profile := spec.ScheduledClusterScan.ScanConfig.CisScanConfig.Profile
		if profile != string(v3.CisScanProfileTypePermissive) &&
			profile != string(v3.CisScanProfileTypeHardened) {
			return httperror.NewFieldAPIError(httperror.InvalidOption, "ScheduledClusterScan.ScanConfig.CisScanConfig.Profile", "profile can be either permissive or hardened")
		}
	}

	if spec.ScheduledClusterScan.ScheduleConfig != nil {
		if spec.ScheduledClusterScan.ScheduleConfig.Retention < 0 {
			return httperror.NewFieldAPIError(httperror.MinLimitExceeded, "ScheduledClusterScan.ScheduleConfig.Retention", "Retention count cannot be negative")
		}
		schedule, err := cron.ParseStandard(spec.ScheduledClusterScan.ScheduleConfig.CronSchedule)
		if err != nil {
			return httperror.NewFieldAPIError(httperror.InvalidFormat, "ScheduledClusterScan.ScheduleConfig.CronSchedule", fmt.Sprintf("error parsing cron schedule: %v", err))
		}
		now := time.Now().Round(time.Second)
		next1 := schedule.Next(now).Round(time.Second)
		next2 := schedule.Next(next1).Round(time.Second)
		timeAfter := next2.Sub(next1).Round(time.Second)

		if timeAfter < (1 * time.Hour) {
			if spec.ScheduledClusterScan.ScanConfig.CisScanConfig.DebugMaster ||
				spec.ScheduledClusterScan.ScanConfig.CisScanConfig.DebugWorker {
				return nil
			}
			return httperror.NewFieldAPIError(httperror.MinLimitExceeded, "ScheduledClusterScan.ScheduleConfig.CronSchedule", "minimum interval is one hour")
		}
	}
	return nil
}

func (v *Validator) validateLocalClusterAuthEndpoint(request *types.APIContext, spec *v3.ClusterSpec) error {
	if !spec.LocalClusterAuthEndpoint.Enabled {
		return nil
	}

	var isValidCluster bool
	if request.ID == "" {
		isValidCluster = spec.RancherKubernetesEngineConfig != nil
	} else {
		cluster, err := v.ClusterLister.Get("", request.ID)
		if err != nil {
			return err
		}
		isValidCluster = cluster.Status.Driver == "" ||
			cluster.Status.Driver == v3.ClusterDriverRKE ||
			cluster.Status.Driver == v3.ClusterDriverImported
	}
	if !isValidCluster {
		return httperror.NewFieldAPIError(httperror.InvalidState, "LocalClusterAuthEndpoint.Enabled", "Can only enable LocalClusterAuthEndpoint with RKE")
	}

	if spec.LocalClusterAuthEndpoint.CACerts != "" && spec.LocalClusterAuthEndpoint.FQDN == "" {
		return httperror.NewFieldAPIError(httperror.MissingRequired, "LocalClusterAuthEndpoint.FQDN", "CACerts defined but FQDN is not defined")
	}

	return nil
}

func (v *Validator) validateEnforcement(request *types.APIContext, data map[string]interface{}) error {

	if !strings.EqualFold(settings.ClusterTemplateEnforcement.Get(), "true") {
		return nil
	}

	var spec mgmtclient.Cluster
	if err := convert.ToObj(data, &spec); err != nil {
		return httperror.WrapAPIError(err, httperror.InvalidBodyContent, "Cluster spec conversion error")
	}

	if !v.checkClusterForEnforcement(&spec) {
		return nil
	}

	ma := gaccess.MemberAccess{
		Users:     v.Users,
		GrLister:  v.GrLister,
		GrbLister: v.GrbLister,
	}

	//if user is admin, no checks needed
	callerID := request.Request.Header.Get(gaccess.ImpersonateUserHeader)

	isAdmin, err := ma.IsAdmin(callerID)
	if err != nil {
		return err
	}
	if isAdmin {
		return nil
	}

	//enforcement is true, template is a must
	if spec.ClusterTemplateRevisionID == "" {
		return httperror.NewFieldAPIError(httperror.MissingRequired, "", "A clusterTemplateRevision to create a cluster")
	}

	err = v.accessTemplate(request, &spec)
	if err != nil {
		if httperror.IsForbidden(err) || httperror.IsNotFound(err) {
			return httperror.NewAPIError(httperror.NotFound, "The clusterTemplateRevision is not found")
		}
		return err
	}

	return nil
}

// TODO: test validator
// prevents downgrades, no-ops, and upgrading before versions have been set
func (v *Validator) validateK3sVersionUpgrade(request *types.APIContext, spec *v3.ClusterSpec) error {
	upgradeNotReadyErr := httperror.NewAPIError(httperror.Conflict, "k3s version upgrade is not ready, try again later")

	if request.Method == http.MethodPost {
		return nil
	}

	if spec.K3sConfig == nil {
		// only applies to k3s clusters
		return nil
	}

	// must wait for original spec version to be set
	if spec.K3sConfig.Version == "" {
		return upgradeNotReadyErr
	}

	cluster, err := v.ClusterLister.Get("", request.ID)
	if err != nil {
		return err
	}

	// must wait for original status version to be set
	if cluster.Status.Version == nil {
		return upgradeNotReadyErr
	}

	prevVersion := cluster.Status.Version.GitVersion
	updateVersion := spec.K3sConfig.Version

	if prevVersion == updateVersion {
		// no op
		return nil
	}

	isNewer, err := k3supgrade.IsNewerVersion(prevVersion, updateVersion)
	if err != nil {
		errMsg := fmt.Sprintf("unable to compare k3s version [%s]", spec.K3sConfig.Version)
		return httperror.NewAPIError(httperror.InvalidBodyContent, errMsg)
	}

	if !isNewer {
		// update version must be higher than previous version, downgrades are not supported
		errMsg := fmt.Sprintf("cannot upgrade k3s cluster version from [%s] to [%s]. New version must be higher.", prevVersion, updateVersion)
		return httperror.NewAPIError(httperror.InvalidBodyContent, errMsg)
	}

	return nil
}

func (v *Validator) checkClusterForEnforcement(spec *mgmtclient.Cluster) bool {
	if spec.RancherKubernetesEngineConfig != nil {
		return true
	}

	if spec.ClusterTemplateRevisionID != "" {
		return true
	}
	return false
}

func (v *Validator) accessTemplate(request *types.APIContext, spec *mgmtclient.Cluster) error {
	split := strings.SplitN(spec.ClusterTemplateRevisionID, ":", 2)
	if len(split) != 2 {
		return fmt.Errorf("error in splitting clusterTemplateRevision name %v", spec.ClusterTemplateRevisionID)
	}
	revName := split[1]
	clusterTempRev, err := v.ClusterTemplateRevisionLister.Get(namespace.GlobalNamespace, revName)
	if err != nil {
		return err
	}

	var ctMap map[string]interface{}
	if err := access.ByID(request, &mgmtSchema.Version, mgmtclient.ClusterTemplateType, clusterTempRev.Spec.ClusterTemplateName, &ctMap); err != nil {
		return err
	}

	return nil
}

// validateGenericEngineConfig allows for additional validation of clusters that depend on Kontainer Engine or Rancher Machine driver
func (v *Validator) validateGenericEngineConfig(request *types.APIContext, spec *v3.ClusterSpec) error {

	if request.Method == http.MethodPost {
		return nil
	}

	if spec.AmazonElasticContainerServiceConfig != nil {
		// compare with current cluster
		clusterName := request.ID
		prevCluster, err := v.ClusterLister.Get("", clusterName)
		if err != nil {
			return httperror.WrapAPIError(err, httperror.InvalidBodyContent, err.Error())
		}

		err = validateEKS(*prevCluster.Spec.GenericEngineConfig, *spec.AmazonElasticContainerServiceConfig)
		if err != nil {
			return err
		}
	}

	return nil

}

func validateEKS(prevCluster, newCluster map[string]interface{}) error {
	// check config is for EKS clusters
	if driver, ok := prevCluster["driverName"]; ok {
		if driver != service.AmazonElasticContainerServiceDriverName {
			return nil
		}
	}

	// don't allow for updating subnets
	if prev, ok := prevCluster["subnets"]; ok {
		if new, ok := newCluster["subnets"]; ok {
			if !reflect.DeepEqual(prev, new) {
				return httperror.NewAPIError(httperror.InvalidBodyContent, "cannot modify EKS subnets after creation")
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package matchers_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/matchers"
)

var _ = Describe("AssignableTo", func() {
	Context("When asserting equality between types", func() {
		It("should do the right thing", func() {
			Ω(0).Should(FitTypeOf(0))
			Ω(5).Should(FitTypeOf(-1))
			Ω("foo").Should(FitTypeOf("bar"))
		})
	})

	Context("When asserting nil values", func() {
		It("should error", func() {
			success, _, err := (&AssignableToMatcher{Expected: nil}).Match(nil)

			Ω(success).Should(BeFalse())
			Ω(err).Should(HaveOccured())
		})
	})
})
<commit_msg>Update specs for AssignableTo matcher<commit_after>package matchers_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/matchers"
)

var _ = Describe("AssignableTo", func() {
	Context("When asserting equality between types", func() {
		It("should do the right thing", func() {
			Ω(0).Should(BeAssignableTo(0))
			Ω(5).Should(BeAssignableTo(-1))
			Ω("foo").Should(BeAssignableTo("bar"))

			Ω(0).ShouldNot(BeAssignableTo("bar"))
			Ω(5).ShouldNot(BeAssignableTo(struct{ Foo string }{}))
			Ω("foo").ShouldNot(BeAssignableTo(42))
		})
	})

	Context("When asserting nil values", func() {
		It("should error", func() {
			success, _, err := (&AssignableToMatcher{Expected: nil}).Match(nil)

			Ω(success).Should(BeFalse())
			Ω(err).Should(HaveOccured())
		})
	})
})
<|endoftext|>