<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 socket

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"

	"github.com/google/testimony/go/protocol"
	"github.com/google/testimony/go/testimonyd/internal/vlog"
)

// Testimony is the configuration parsed from the config file.
type Testimony []SocketConfig

const protocolVersion = 2

// SocketConfig defines how an individual socket should be set up.
type SocketConfig struct {
	SocketName         string // filename for the socket
	Interface          string // interface to sniff packets on
	BlockSize          int    // block size (in bytes) of a single packet block
	NumBlocks          int    // number of packet blocks in the memory region
	BlockTimeoutMillis int    // timeout for filling up a single block
	FanoutType         int    // which type of fanout to use (see linux/if_packet.h)
	FanoutSize         int    // number of threads to fan out to
	User               string // user to provide the socket to (will chown it)
	Filter             string // BPF filter to apply to this socket
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

// RunTestimony runs the testimonyd server given the passed in configuration.
func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		// Check for duplicate socket names
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true

		// Set up FanoutSize sockets and start goroutines to manage each.
		var socks []*socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}

		// Set up UNIX socket to serve these AF_PACKET sockets on, and start
		// goroutine to manage its connections.
		_ignore_error_ := os.Remove(sc.SocketName)
		_ = _ignore_error_
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		} else if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {} // Block (serving) forever.
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	vlog.V(1, "chowning %q to %d", sc.SocketName, uid)
	if err := syscall.Chown(sc.SocketName, uid, 0); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	connStr := c.RemoteAddr().String()
	log.Printf("Received new connection %q", connStr)
	var version [1]byte
	version[0] = protocolVersion
	if _, err := c.Write(version[:]); err != nil {
		log.Printf("new conn %q failed to write version: %v", connStr, err)
		return
	}
	conf := socks[0].conf
	if err := protocol.SendUint32(c, protocol.TypeFanoutSize, uint32(len(socks))); err != nil {
		log.Printf("new conn %q failed to send fanout size: %v", connStr, err)
		return
	}
	if err := protocol.SendUint32(c, protocol.TypeBlockSize, uint32(conf.BlockSize)); err != nil {
		log.Printf("new conn %q failed to send block size: %v", connStr, err)
		return
	}
	if err := protocol.SendUint32(c, protocol.TypeNumBlocks, uint32(conf.NumBlocks)); err != nil {
		log.Printf("new conn %q failed to send number of blocks: %v", connStr, err)
		return
	}
	if err := protocol.SendType(c, protocol.TypeWaitingForFanoutIndex); err != nil {
		log.Printf("new conn %q failed to send wait: %v", connStr, err)
		return
	}
	var fanoutMsg [8]byte
	if _, err := io.ReadFull(c, fanoutMsg[:]); err == io.EOF {
		log.Printf("new conn %q closed early, probably just gathering connection data", connStr)
		return
	} else if err != nil {
		log.Printf("new conn %q failed to read fanout index: %v", connStr, err)
		return
	}
	valA, valB := binary.BigEndian.Uint32(fanoutMsg[:4]), binary.BigEndian.Uint32(fanoutMsg[4:])
	if valA != protocol.ToTL(protocol.TypeFanoutIndex, 4) {
		log.Printf("new conn %q got unexpected type/value waiting for fanout message: %d/%d", connStr, valA>>16, valA&0xFFFF)
		return
	}
	idx := int(valB)
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn %q invalid index %v", connStr, idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn %q failed to send file descriptor: %v", connStr, err)
		return
	}
	vlog.V(2, "new conn %q spun up, passing off to socket", connStr)
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}
<commit_msg>Add Group to config, to chown sockets to a group.<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 socket

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"
	"unsafe"

	"github.com/google/testimony/go/protocol"
	"github.com/google/testimony/go/testimonyd/internal/vlog"
)

// #include <grp.h>
// #include <stdlib.h>
import "C"

// Testimony is the configuration parsed from the config file.
type Testimony []SocketConfig

const protocolVersion = 2

// SocketConfig defines how an individual socket should be set up.
type SocketConfig struct {
	SocketName         string // filename for the socket
	Interface          string // interface to sniff packets on
	BlockSize          int    // block size (in bytes) of a single packet block
	NumBlocks          int    // number of packet blocks in the memory region
	BlockTimeoutMillis int    // timeout for filling up a single block
	FanoutType         int    // which type of fanout to use (see linux/if_packet.h)
	FanoutSize         int    // number of threads to fan out to
	User, Group        string // user/group to provide the socket to (will chown it)
	Filter             string // BPF filter to apply to this socket
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

func (s SocketConfig) gid() (int, error) {
	// Sadly, at present Go doesn't have group functions to match its os/user
	// functions... so we jump down into C.
	if s.Group == "" {
		return 0, nil
	}
	groupName := C.CString(s.Group)
	defer C.free(unsafe.Pointer(groupName))
	var buf [2048]byte
	var grp C.struct_group
	var grpPtr *C.struct_group
	if _, err := C.getgrnam_r(groupName, &grp, (*C.char)(unsafe.Pointer(&buf[0])), C.size_t(len(buf)), &grpPtr); err != nil {
		return -1, err
	} else if grpPtr == nil {
		return -1, fmt.Errorf("group %q not found", s.Group)
	}
	return int(grpPtr.gr_gid), nil
}

// RunTestimony runs the testimonyd server given the passed in configuration.
func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		// Check for duplicate socket names
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true

		// Set up FanoutSize sockets and start goroutines to manage each.
		var socks []*socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}

		// Set up UNIX socket to serve these AF_PACKET sockets on, and start
		// goroutine to manage its connections.
		_ignore_error_ := os.Remove(sc.SocketName)
		_ = _ignore_error_
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		} else if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {} // Block (serving) forever.
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	gid, err := sc.gid()
	if err != nil {
		return fmt.Errorf("could not get gid to change to: %v", err)
	}
	vlog.V(1, "chowning %q to %d/%d", sc.SocketName, uid, gid)
	if err := syscall.Chown(sc.SocketName, uid, gid); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	connStr := c.RemoteAddr().String()
	log.Printf("Received new connection %q", connStr)
	var version [1]byte
	version[0] = protocolVersion
	if _, err := c.Write(version[:]); err != nil {
		log.Printf("new conn %q failed to write version: %v", connStr, err)
		return
	}
	conf := socks[0].conf
	if err := protocol.SendUint32(c, protocol.TypeFanoutSize, uint32(len(socks))); err != nil {
		log.Printf("new conn %q failed to send fanout size: %v", connStr, err)
		return
	}
	if err := protocol.SendUint32(c, protocol.TypeBlockSize, uint32(conf.BlockSize)); err != nil {
		log.Printf("new conn %q failed to send block size: %v", connStr, err)
		return
	}
	if err := protocol.SendUint32(c, protocol.TypeNumBlocks, uint32(conf.NumBlocks)); err != nil {
		log.Printf("new conn %q failed to send number of blocks: %v", connStr, err)
		return
	}
	if err := protocol.SendType(c, protocol.TypeWaitingForFanoutIndex); err != nil {
		log.Printf("new conn %q failed to send wait: %v", connStr, err)
		return
	}
	var fanoutMsg [8]byte
	if _, err := io.ReadFull(c, fanoutMsg[:]); err == io.EOF {
		log.Printf("new conn %q closed early, probably just gathering connection data", connStr)
		return
	} else if err != nil {
		log.Printf("new conn %q failed to read fanout index: %v", connStr, err)
		return
	}
	valA, valB := binary.BigEndian.Uint32(fanoutMsg[:4]), binary.BigEndian.Uint32(fanoutMsg[4:])
	if valA != protocol.ToTL(protocol.TypeFanoutIndex, 4) {
		log.Printf("new conn %q got unexpected type/value waiting for fanout message: %d/%d", connStr, valA>>16, valA&0xFFFF)
		return
	}
	idx := int(valB)
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn %q invalid index %v", connStr, idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn %q failed to send file descriptor: %v", connStr, err)
		return
	}
	vlog.V(2, "new conn %q spun up, passing off to socket", connStr)
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess 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.
*/

package reparentutil

import (
	"context"
	"sync"
	"time"

	"k8s.io/apimachinery/pkg/util/sets"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/logutil"
	replicationdatapb "vitess.io/vitess/go/vt/proto/replicationdata"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools/events"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tmclient"
)

// FindValidEmergencyReparentCandidates will find candidates for an emergency
// reparent, and, if successful, return a mapping of those tablet aliases (as
// raw strings) to their replication positions for later comparison.
func FindValidEmergencyReparentCandidates(
	statusMap map[string]*replicationdatapb.StopReplicationStatus,
	primaryStatusMap map[string]*replicationdatapb.PrimaryStatus,
) (map[string]mysql.Position, error) {
	replicationStatusMap := make(map[string]*mysql.ReplicationStatus, len(statusMap))
	positionMap := make(map[string]mysql.Position)

	// Build out replication status list from proto types.
	for alias, statuspb := range statusMap {
		status := mysql.ProtoToReplicationStatus(statuspb.After)
		replicationStatusMap[alias] = &status
	}

	// Determine if we're GTID-based. If we are, we'll need to look for errant
	// GTIDs below.
	var (
		isGTIDBased                bool
		isNonGTIDBased             bool
		emptyRelayPosErrorRecorder concurrency.FirstErrorRecorder
	)

	for alias, status := range replicationStatusMap {
		if _, ok := status.RelayLogPosition.GTIDSet.(mysql.Mysql56GTIDSet); ok {
			isGTIDBased = true
		} else {
			isNonGTIDBased = true
		}

		if status.RelayLogPosition.IsZero() {
			// Potentially bail. If any other tablet is detected to have
			// GTID-based relay log positions, we will return the error recorded
			// here.
			emptyRelayPosErrorRecorder.RecordError(vterrors.Errorf(vtrpc.Code_UNAVAILABLE, "encountered tablet %v with no relay log position, when at least one other tablet in the status map has GTID based relay log positions", alias))
		}
	}

	if isGTIDBased && emptyRelayPosErrorRecorder.HasErrors() {
		return nil, emptyRelayPosErrorRecorder.Error()
	}

	if isGTIDBased && isNonGTIDBased {
		return nil, vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "encountered mix of GTID-based and non GTID-based relay logs")
	}

	// Create relevant position list of errant GTID-based positions for later
	// comparison.
	for alias, status := range replicationStatusMap {
		// If we're not GTID-based, no need to search for errant GTIDs, so just
		// add the position to the map and continue.
		if !isGTIDBased {
			positionMap[alias] = status.Position

			continue
		}

		// This condition should really never happen, since we did the same cast
		// in the earlier loop, but let's be doubly sure.
		relayLogGTIDSet, ok := status.RelayLogPosition.GTIDSet.(mysql.Mysql56GTIDSet)
		if !ok {
			return nil, vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "we got a filled-in relay log position, but it's not of type Mysql56GTIDSet, even though we've determined we need to use GTID based assesment")
		}

		// We need to remove this alias's status from the list, otherwise the
		// GTID diff will always be empty.
		statusList := make([]*mysql.ReplicationStatus, 0, len(replicationStatusMap)-1)

		for a, s := range replicationStatusMap {
			if a != alias {
				statusList = append(statusList, s)
			}
		}

		errantGTIDs, err := status.FindErrantGTIDs(statusList)
		switch {
		case err != nil:
			// Could not look up GTIDs to determine if we have any. It's not
			// safe to continue.
			return nil, err
		case len(errantGTIDs) != 0:
			// This tablet has errant GTIDs. It's not a valid candidate for
			// reparent, so don't insert it into the final mapping.
			continue
		}

		pos := mysql.Position{GTIDSet: relayLogGTIDSet}
		positionMap[alias] = pos
	}

	for alias, primaryStatus := range primaryStatusMap {
		executedPosition, err := mysql.DecodePosition(primaryStatus.Position)
		if err != nil {
			return nil, vterrors.Wrapf(err, "could not decode a primary status executed position for tablet %v: %v", alias, err)
		}

		positionMap[alias] = executedPosition
	}

	return positionMap, nil
}

// ReplicaWasRunning returns true if a StopReplicationStatus indicates that the
// replica had running replication threads before being stopped. It returns an
// error if the Before state of replication is nil.
func ReplicaWasRunning(stopStatus *replicationdatapb.StopReplicationStatus) (bool, error) {
	if stopStatus == nil || stopStatus.Before == nil {
		return false, vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "could not determine Before state of StopReplicationStatus %v", stopStatus)
	}

	return stopStatus.Before.IoThreadRunning || stopStatus.Before.SqlThreadRunning, nil
}

// StopReplicationAndBuildStatusMaps stops replication on all replicas, then
// collects and returns a mapping of TabletAlias (as string) to their current
// replication positions.
func StopReplicationAndBuildStatusMaps(
	ctx context.Context,
	tmc tmclient.TabletManagerClient,
	ev *events.Reparent,
	tabletMap map[string]*topo.TabletInfo,
	waitReplicasTimeout time.Duration,
	ignoredTablets sets.String,
	tabletToWaitFor *topodatapb.TabletAlias,
	logger logutil.Logger,
) (map[string]*replicationdatapb.StopReplicationStatus, map[string]*replicationdatapb.PrimaryStatus, error) {
	event.DispatchUpdate(ev, "stop replication on all replicas")

	var (
		statusMap        = map[string]*replicationdatapb.StopReplicationStatus{}
		primaryStatusMap = map[string]*replicationdatapb.PrimaryStatus{}
		m                sync.Mutex
		errChan          = make(chan concurrency.Error)
	)

	groupCtx, groupCancel := context.WithTimeout(ctx, waitReplicasTimeout)
	defer groupCancel()

	fillStatus := func(alias string, tabletInfo *topo.TabletInfo, mustWaitForTablet bool) {
		var concurrencyErr concurrency.Error
		var err error
		defer func() {
			concurrencyErr.Err = err
			concurrencyErr.MustWaitFor = mustWaitForTablet
			errChan <- concurrencyErr
		}()

		logger.Infof("getting replication position from %v", alias)

		_, stopReplicationStatus, err := tmc.StopReplicationAndGetStatus(groupCtx, tabletInfo.Tablet, replicationdatapb.StopReplicationMode_IOTHREADONLY)
		sqlErr, isSQLErr := mysql.NewSQLErrorFromError(err).(*mysql.SQLError)
		switch {
		case isSQLErr && sqlErr != nil && sqlErr.Number() == mysql.ERNotReplica:
			var primaryStatus *replicationdatapb.PrimaryStatus

			primaryStatus, err = tmc.DemotePrimary(groupCtx, tabletInfo.Tablet)
			if err != nil {
				msg := "replica %v thinks it's primary but we failed to demote it"
				err = vterrors.Wrapf(err, msg+": %v", alias, err)

				logger.Warningf(msg, alias)
				return
			}

			m.Lock()
			primaryStatusMap[alias] = primaryStatus
			m.Unlock()
		case err == nil:
			m.Lock()
			statusMap[alias] = stopReplicationStatus
			m.Unlock()
		default:
			logger.Warningf("failed to get replication status from %v: %v", alias, err)

			err = vterrors.Wrapf(err, "error when getting replication status for alias %v: %v", alias, err)
		}
	}

	tabletAliasToWaitFor := ""
	numErrosToWaitFor := 0
	if tabletToWaitFor != nil {
		tabletAliasToWaitFor = topoproto.TabletAliasString(tabletToWaitFor)
	}
	for alias, tabletInfo := range tabletMap {
		if !ignoredTablets.Has(alias) {
			mustWaitFor := tabletAliasToWaitFor == alias
			if mustWaitFor {
				numErrosToWaitFor++
			}
			go fillStatus(alias, tabletInfo, mustWaitFor)
		}
	}

	errgroup := concurrency.ErrorGroup{
		NumGoroutines:        len(tabletMap) - ignoredTablets.Len(),
		NumRequiredSuccesses: len(tabletMap) - ignoredTablets.Len() - 1,
		NumAllowedErrors:     1,
		NumErrorsToWaitFor:   numErrosToWaitFor,
	}

	errRecorder := errgroup.Wait(groupCancel, errChan)
	if len(errRecorder.Errors) > 1 {
		return nil, nil, vterrors.Wrapf(errRecorder.Error(), "encountered more than one error when trying to stop replication and get positions: %v", errRecorder.Error())
	}

	return statusMap, primaryStatusMap, nil
}

// WaitForRelayLogsToApply blocks execution waiting for the given tablet's relay
// logs to apply, unless the specified context is canceled or exceeded.
// Typically a caller will set a timeout of WaitReplicasTimeout on a context and
// use that context with this function.
func WaitForRelayLogsToApply(ctx context.Context, tmc tmclient.TabletManagerClient, tabletInfo *topo.TabletInfo, status *replicationdatapb.StopReplicationStatus) error {
	switch status.After.RelayLogPosition {
	case "":
		return tmc.WaitForPosition(ctx, tabletInfo.Tablet, status.After.FileRelayLogPosition)
	default:
		return tmc.WaitForPosition(ctx, tabletInfo.Tablet, status.After.RelayLogPosition)
	}
}
<commit_msg>feat: replace switch case with if else conditions<commit_after>/*
Copyright 2021 The Vitess 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.
*/

package reparentutil

import (
	"context"
	"sync"
	"time"

	"k8s.io/apimachinery/pkg/util/sets"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/logutil"
	replicationdatapb "vitess.io/vitess/go/vt/proto/replicationdata"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools/events"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tmclient"
)

// FindValidEmergencyReparentCandidates will find candidates for an emergency
// reparent, and, if successful, return a mapping of those tablet aliases (as
// raw strings) to their replication positions for later comparison.
func FindValidEmergencyReparentCandidates(
	statusMap map[string]*replicationdatapb.StopReplicationStatus,
	primaryStatusMap map[string]*replicationdatapb.PrimaryStatus,
) (map[string]mysql.Position, error) {
	replicationStatusMap := make(map[string]*mysql.ReplicationStatus, len(statusMap))
	positionMap := make(map[string]mysql.Position)

	// Build out replication status list from proto types.
	for alias, statuspb := range statusMap {
		status := mysql.ProtoToReplicationStatus(statuspb.After)
		replicationStatusMap[alias] = &status
	}

	// Determine if we're GTID-based. If we are, we'll need to look for errant
	// GTIDs below.
	var (
		isGTIDBased                bool
		isNonGTIDBased             bool
		emptyRelayPosErrorRecorder concurrency.FirstErrorRecorder
	)

	for alias, status := range replicationStatusMap {
		if _, ok := status.RelayLogPosition.GTIDSet.(mysql.Mysql56GTIDSet); ok {
			isGTIDBased = true
		} else {
			isNonGTIDBased = true
		}

		if status.RelayLogPosition.IsZero() {
			// Potentially bail. If any other tablet is detected to have
			// GTID-based relay log positions, we will return the error recorded
			// here.
			emptyRelayPosErrorRecorder.RecordError(vterrors.Errorf(vtrpc.Code_UNAVAILABLE, "encountered tablet %v with no relay log position, when at least one other tablet in the status map has GTID based relay log positions", alias))
		}
	}

	if isGTIDBased && emptyRelayPosErrorRecorder.HasErrors() {
		return nil, emptyRelayPosErrorRecorder.Error()
	}

	if isGTIDBased && isNonGTIDBased {
		return nil, vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "encountered mix of GTID-based and non GTID-based relay logs")
	}

	// Create relevant position list of errant GTID-based positions for later
	// comparison.
	for alias, status := range replicationStatusMap {
		// If we're not GTID-based, no need to search for errant GTIDs, so just
		// add the position to the map and continue.
		if !isGTIDBased {
			positionMap[alias] = status.Position

			continue
		}

		// This condition should really never happen, since we did the same cast
		// in the earlier loop, but let's be doubly sure.
		relayLogGTIDSet, ok := status.RelayLogPosition.GTIDSet.(mysql.Mysql56GTIDSet)
		if !ok {
			return nil, vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "we got a filled-in relay log position, but it's not of type Mysql56GTIDSet, even though we've determined we need to use GTID based assesment")
		}

		// We need to remove this alias's status from the list, otherwise the
		// GTID diff will always be empty.
		statusList := make([]*mysql.ReplicationStatus, 0, len(replicationStatusMap)-1)

		for a, s := range replicationStatusMap {
			if a != alias {
				statusList = append(statusList, s)
			}
		}

		errantGTIDs, err := status.FindErrantGTIDs(statusList)
		switch {
		case err != nil:
			// Could not look up GTIDs to determine if we have any. It's not
			// safe to continue.
			return nil, err
		case len(errantGTIDs) != 0:
			// This tablet has errant GTIDs. It's not a valid candidate for
			// reparent, so don't insert it into the final mapping.
			continue
		}

		pos := mysql.Position{GTIDSet: relayLogGTIDSet}
		positionMap[alias] = pos
	}

	for alias, primaryStatus := range primaryStatusMap {
		executedPosition, err := mysql.DecodePosition(primaryStatus.Position)
		if err != nil {
			return nil, vterrors.Wrapf(err, "could not decode a primary status executed position for tablet %v: %v", alias, err)
		}

		positionMap[alias] = executedPosition
	}

	return positionMap, nil
}

// ReplicaWasRunning returns true if a StopReplicationStatus indicates that the
// replica had running replication threads before being stopped. It returns an
// error if the Before state of replication is nil.
func ReplicaWasRunning(stopStatus *replicationdatapb.StopReplicationStatus) (bool, error) {
	if stopStatus == nil || stopStatus.Before == nil {
		return false, vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "could not determine Before state of StopReplicationStatus %v", stopStatus)
	}

	return stopStatus.Before.IoThreadRunning || stopStatus.Before.SqlThreadRunning, nil
}

// StopReplicationAndBuildStatusMaps stops replication on all replicas, then
// collects and returns a mapping of TabletAlias (as string) to their current
// replication positions.
func StopReplicationAndBuildStatusMaps(
	ctx context.Context,
	tmc tmclient.TabletManagerClient,
	ev *events.Reparent,
	tabletMap map[string]*topo.TabletInfo,
	waitReplicasTimeout time.Duration,
	ignoredTablets sets.String,
	tabletToWaitFor *topodatapb.TabletAlias,
	logger logutil.Logger,
) (map[string]*replicationdatapb.StopReplicationStatus, map[string]*replicationdatapb.PrimaryStatus, error) {
	event.DispatchUpdate(ev, "stop replication on all replicas")

	var (
		statusMap        = map[string]*replicationdatapb.StopReplicationStatus{}
		primaryStatusMap = map[string]*replicationdatapb.PrimaryStatus{}
		m                sync.Mutex
		errChan          = make(chan concurrency.Error)
	)

	groupCtx, groupCancel := context.WithTimeout(ctx, waitReplicasTimeout)
	defer groupCancel()

	fillStatus := func(alias string, tabletInfo *topo.TabletInfo, mustWaitForTablet bool) {
		var concurrencyErr concurrency.Error
		var err error
		defer func() {
			concurrencyErr.Err = err
			concurrencyErr.MustWaitFor = mustWaitForTablet
			errChan <- concurrencyErr
		}()

		logger.Infof("getting replication position from %v", alias)

		_, stopReplicationStatus, err := tmc.StopReplicationAndGetStatus(groupCtx, tabletInfo.Tablet, replicationdatapb.StopReplicationMode_IOTHREADONLY)
		if err != nil {
			sqlErr, isSQLErr := mysql.NewSQLErrorFromError(err).(*mysql.SQLError)
			if isSQLErr && sqlErr != nil && sqlErr.Number() == mysql.ERNotReplica {
				var primaryStatus *replicationdatapb.PrimaryStatus

				primaryStatus, err = tmc.DemotePrimary(groupCtx, tabletInfo.Tablet)
				if err != nil {
					msg := "replica %v thinks it's primary but we failed to demote it"
					err = vterrors.Wrapf(err, msg+": %v", alias, err)

					logger.Warningf(msg, alias)
					return
				}

				m.Lock()
				primaryStatusMap[alias] = primaryStatus
				m.Unlock()
			} else {
				logger.Warningf("failed to get replication status from %v: %v", alias, err)
				err = vterrors.Wrapf(err, "error when getting replication status for alias %v: %v", alias, err)
			}
		} else {
			m.Lock()
			statusMap[alias] = stopReplicationStatus
			m.Unlock()
		}
	}

	tabletAliasToWaitFor := ""
	numErrosToWaitFor := 0
	if tabletToWaitFor != nil {
		tabletAliasToWaitFor = topoproto.TabletAliasString(tabletToWaitFor)
	}
	for alias, tabletInfo := range tabletMap {
		if !ignoredTablets.Has(alias) {
			mustWaitFor := tabletAliasToWaitFor == alias
			if mustWaitFor {
				numErrosToWaitFor++
			}
			go fillStatus(alias, tabletInfo, mustWaitFor)
		}
	}

	errgroup := concurrency.ErrorGroup{
		NumGoroutines:        len(tabletMap) - ignoredTablets.Len(),
		NumRequiredSuccesses: len(tabletMap) - ignoredTablets.Len() - 1,
		NumAllowedErrors:     1,
		NumErrorsToWaitFor:   numErrosToWaitFor,
	}

	errRecorder := errgroup.Wait(groupCancel, errChan)
	if len(errRecorder.Errors) > 1 {
		return nil, nil, vterrors.Wrapf(errRecorder.Error(), "encountered more than one error when trying to stop replication and get positions: %v", errRecorder.Error())
	}

	return statusMap, primaryStatusMap, nil
}

// WaitForRelayLogsToApply blocks execution waiting for the given tablet's relay
// logs to apply, unless the specified context is canceled or exceeded.
// Typically a caller will set a timeout of WaitReplicasTimeout on a context and
// use that context with this function.
func WaitForRelayLogsToApply(ctx context.Context, tmc tmclient.TabletManagerClient, tabletInfo *topo.TabletInfo, status *replicationdatapb.StopReplicationStatus) error {
	switch status.After.RelayLogPosition {
	case "":
		return tmc.WaitForPosition(ctx, tabletInfo.Tablet, status.After.FileRelayLogPosition)
	default:
		return tmc.WaitForPosition(ctx, tabletInfo.Tablet, status.After.RelayLogPosition)
	}
}
<|endoftext|>


<commit_before>/*
	Original copyright by GitHub as follows. Additions by the Vitess authors as follows.
*/
/*
   Copyright 2016 GitHub Inc.
	 See https://github.com/github/gh-ost/blob/master/LICENSE
*/
/*
Copyright 2021 The Vitess 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.
*/

package vrepl

import (
	"fmt"
	"reflect"
	"strings"
)

// ColumnType indicated some MySQL data types
type ColumnType int

const (
	UnknownColumnType ColumnType = iota
	TimestampColumnType
	DateTimeColumnType
	EnumColumnType
	MediumIntColumnType
	JSONColumnType
	FloatColumnType
	BinaryColumnType
	StringColumnType
	IntegerColumnType
)

// Column represents a table column
type Column struct {
	Name                 string
	IsUnsigned           bool
	Charset              string
	Collation            string
	Type                 ColumnType
	EnumValues           string
	EnumToTextConversion bool

	IsNullable    bool
	IsDefaultNull bool

	// add Octet length for binary type, fix bytes with suffix "00" get clipped in mysql binlog.
	// https://github.com/github/gh-ost/issues/909
	BinaryOctetLength uint64
}

// SetTypeIfUnknown will set a new column type only if the current type is unknown, otherwise silently skip
func (c *Column) SetTypeIfUnknown(t ColumnType) {
	if c.Type == UnknownColumnType {
		c.Type = t
	}
}

// NewColumns creates a new column array from non empty names
func NewColumns(names []string) []Column {
	result := []Column{}
	for _, name := range names {
		if name == "" {
			continue
		}
		result = append(result, Column{Name: name})
	}
	return result
}

// ParseColumns creates a new column array fby parsing comma delimited names list
func ParseColumns(names string) []Column {
	namesArray := strings.Split(names, ",")
	return NewColumns(namesArray)
}

// ColumnsMap maps a column name onto its ordinal position
type ColumnsMap map[string]int

// NewEmptyColumnsMap creates an empty map
func NewEmptyColumnsMap() ColumnsMap {
	columnsMap := make(map[string]int)
	return ColumnsMap(columnsMap)
}

// NewColumnsMap creates a column map based on ordered list of columns
func NewColumnsMap(orderedColumns []Column) ColumnsMap {
	columnsMap := NewEmptyColumnsMap()
	for i, column := range orderedColumns {
		columnsMap[column.Name] = i
	}
	return columnsMap
}

// ColumnList makes for a named list of columns
type ColumnList struct {
	columns  []Column
	Ordinals ColumnsMap
}

// NewColumnList creates an object given ordered list of column names
func NewColumnList(names []string) *ColumnList {
	result := &ColumnList{
		columns: NewColumns(names),
	}
	result.Ordinals = NewColumnsMap(result.columns)
	return result
}

// ParseColumnList parses a comma delimited list of column names
func ParseColumnList(names string) *ColumnList {
	result := &ColumnList{
		columns: ParseColumns(names),
	}
	result.Ordinals = NewColumnsMap(result.columns)
	return result
}

// Columns returns the list of columns
func (l *ColumnList) Columns() []Column {
	return l.columns
}

// Names returns list of column names
func (l *ColumnList) Names() []string {
	names := make([]string, len(l.columns))
	for i := range l.columns {
		names[i] = l.columns[i].Name
	}
	return names
}

// GetColumn gets a column by name
func (l *ColumnList) GetColumn(columnName string) *Column {
	if ordinal, ok := l.Ordinals[columnName]; ok {
		return &l.columns[ordinal]
	}
	return nil
}

// ColumnExists returns true if this column list has a column by a given name
func (l *ColumnList) ColumnExists(columnName string) bool {
	_, ok := l.Ordinals[columnName]
	return ok
}

// String returns a comma separated list of column names
func (l *ColumnList) String() string {
	return strings.Join(l.Names(), ",")
}

// Equals checks for complete (deep) identities of columns, in order.
func (l *ColumnList) Equals(other *ColumnList) bool {
	return reflect.DeepEqual(l.Columns, other.Columns)
}

// EqualsByNames chcks if the names in this list equals the names of another list, in order. Type is ignored.
func (l *ColumnList) EqualsByNames(other *ColumnList) bool {
	return reflect.DeepEqual(l.Names(), other.Names())
}

// IsSubsetOf returns 'true' when column names of this list are a subset of
// another list, in arbitrary order (order agnostic)
func (l *ColumnList) IsSubsetOf(other *ColumnList) bool {
	for _, column := range l.columns {
		if _, exists := other.Ordinals[column.Name]; !exists {
			return false
		}
	}
	return true
}

// Difference returns a (new copy) subset of this column list, consisting of all
// column NOT in given list.
// The result is never nil, even if the difference is empty
func (l *ColumnList) Difference(other *ColumnList) (diff *ColumnList) {
	names := []string{}
	for _, column := range l.columns {
		if !other.ColumnExists(column.Name) {
			names = append(names, column.Name)
		}
	}
	return NewColumnList(names)
}

// Len returns the length of this list
func (l *ColumnList) Len() int {
	return len(l.columns)
}

// MappedNamesColumnList returns a column list based on this list, with names possibly mapped by given map
func (l *ColumnList) MappedNamesColumnList(columnNamesMap map[string]string) *ColumnList {
	names := l.Names()
	for i := range names {
		if mappedName, ok := columnNamesMap[names[i]]; ok {
			names[i] = mappedName
		}
	}
	return NewColumnList(names)
}

// SetEnumToTextConversion tells this column list that an enum is conveted to text
func (l *ColumnList) SetEnumToTextConversion(columnName string, enumValues string) {
	l.GetColumn(columnName).EnumToTextConversion = true
	l.GetColumn(columnName).EnumValues = enumValues
}

// IsEnumToTextConversion tells whether an enum was converted to text
func (l *ColumnList) IsEnumToTextConversion(columnName string) bool {
	return l.GetColumn(columnName).EnumToTextConversion
}

// UniqueKey is the combination of a key's name and columns
type UniqueKey struct {
	Name            string
	Columns         ColumnList
	HasNullable     bool
	HasFloat        bool
	IsAutoIncrement bool
}

// IsPrimary checks if this unique key is primary
func (k *UniqueKey) IsPrimary() bool {
	return k.Name == "PRIMARY"
}

// Len returns the length of this list
func (k *UniqueKey) Len() int {
	return k.Columns.Len()
}

// String returns a visual representation of this key
func (k *UniqueKey) String() string {
	description := k.Name
	if k.IsAutoIncrement {
		description = fmt.Sprintf("%s (auto_increment)", description)
	}
	return fmt.Sprintf("%s: %s; has nullable: %+v", description, k.Columns.Names(), k.HasNullable)
}
<commit_msg>helper function<commit_after>/*
	Original copyright by GitHub as follows. Additions by the Vitess authors as follows.
*/
/*
   Copyright 2016 GitHub Inc.
	 See https://github.com/github/gh-ost/blob/master/LICENSE
*/
/*
Copyright 2021 The Vitess 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.
*/

package vrepl

import (
	"fmt"
	"reflect"
	"strings"
)

// ColumnType indicated some MySQL data types
type ColumnType int

const (
	UnknownColumnType ColumnType = iota
	TimestampColumnType
	DateTimeColumnType
	EnumColumnType
	MediumIntColumnType
	JSONColumnType
	FloatColumnType
	BinaryColumnType
	StringColumnType
	IntegerColumnType
)

// Column represents a table column
type Column struct {
	Name                 string
	IsUnsigned           bool
	Charset              string
	Collation            string
	Type                 ColumnType
	EnumValues           string
	EnumToTextConversion bool

	IsNullable    bool
	IsDefaultNull bool

	// add Octet length for binary type, fix bytes with suffix "00" get clipped in mysql binlog.
	// https://github.com/github/gh-ost/issues/909
	BinaryOctetLength uint64
}

// SetTypeIfUnknown will set a new column type only if the current type is unknown, otherwise silently skip
func (c *Column) SetTypeIfUnknown(t ColumnType) {
	if c.Type == UnknownColumnType {
		c.Type = t
	}
}

// HasDefault returns true if the column at all has a default value (possibly NULL)
func (c *Column) HasDefault() bool {
	if c.IsDefaultNull && !c.IsNullable {
		// based on INFORMATION_SCHEMA.COLUMNS, this is the indicator for a 'NOT NULL' column with no default value.
		return false
	}
	return true
}

// NewColumns creates a new column array from non empty names
func NewColumns(names []string) []Column {
	result := []Column{}
	for _, name := range names {
		if name == "" {
			continue
		}
		result = append(result, Column{Name: name})
	}
	return result
}

// ParseColumns creates a new column array fby parsing comma delimited names list
func ParseColumns(names string) []Column {
	namesArray := strings.Split(names, ",")
	return NewColumns(namesArray)
}

// ColumnsMap maps a column name onto its ordinal position
type ColumnsMap map[string]int

// NewEmptyColumnsMap creates an empty map
func NewEmptyColumnsMap() ColumnsMap {
	columnsMap := make(map[string]int)
	return ColumnsMap(columnsMap)
}

// NewColumnsMap creates a column map based on ordered list of columns
func NewColumnsMap(orderedColumns []Column) ColumnsMap {
	columnsMap := NewEmptyColumnsMap()
	for i, column := range orderedColumns {
		columnsMap[column.Name] = i
	}
	return columnsMap
}

// ColumnList makes for a named list of columns
type ColumnList struct {
	columns  []Column
	Ordinals ColumnsMap
}

// NewColumnList creates an object given ordered list of column names
func NewColumnList(names []string) *ColumnList {
	result := &ColumnList{
		columns: NewColumns(names),
	}
	result.Ordinals = NewColumnsMap(result.columns)
	return result
}

// ParseColumnList parses a comma delimited list of column names
func ParseColumnList(names string) *ColumnList {
	result := &ColumnList{
		columns: ParseColumns(names),
	}
	result.Ordinals = NewColumnsMap(result.columns)
	return result
}

// Columns returns the list of columns
func (l *ColumnList) Columns() []Column {
	return l.columns
}

// Names returns list of column names
func (l *ColumnList) Names() []string {
	names := make([]string, len(l.columns))
	for i := range l.columns {
		names[i] = l.columns[i].Name
	}
	return names
}

// GetColumn gets a column by name
func (l *ColumnList) GetColumn(columnName string) *Column {
	if ordinal, ok := l.Ordinals[columnName]; ok {
		return &l.columns[ordinal]
	}
	return nil
}

// ColumnExists returns true if this column list has a column by a given name
func (l *ColumnList) ColumnExists(columnName string) bool {
	_, ok := l.Ordinals[columnName]
	return ok
}

// String returns a comma separated list of column names
func (l *ColumnList) String() string {
	return strings.Join(l.Names(), ",")
}

// Equals checks for complete (deep) identities of columns, in order.
func (l *ColumnList) Equals(other *ColumnList) bool {
	return reflect.DeepEqual(l.Columns, other.Columns)
}

// EqualsByNames chcks if the names in this list equals the names of another list, in order. Type is ignored.
func (l *ColumnList) EqualsByNames(other *ColumnList) bool {
	return reflect.DeepEqual(l.Names(), other.Names())
}

// IsSubsetOf returns 'true' when column names of this list are a subset of
// another list, in arbitrary order (order agnostic)
func (l *ColumnList) IsSubsetOf(other *ColumnList) bool {
	for _, column := range l.columns {
		if _, exists := other.Ordinals[column.Name]; !exists {
			return false
		}
	}
	return true
}

// Difference returns a (new copy) subset of this column list, consisting of all
// column NOT in given list.
// The result is never nil, even if the difference is empty
func (l *ColumnList) Difference(other *ColumnList) (diff *ColumnList) {
	names := []string{}
	for _, column := range l.columns {
		if !other.ColumnExists(column.Name) {
			names = append(names, column.Name)
		}
	}
	return NewColumnList(names)
}

// Len returns the length of this list
func (l *ColumnList) Len() int {
	return len(l.columns)
}

// MappedNamesColumnList returns a column list based on this list, with names possibly mapped by given map
func (l *ColumnList) MappedNamesColumnList(columnNamesMap map[string]string) *ColumnList {
	names := l.Names()
	for i := range names {
		if mappedName, ok := columnNamesMap[names[i]]; ok {
			names[i] = mappedName
		}
	}
	return NewColumnList(names)
}

// SetEnumToTextConversion tells this column list that an enum is conveted to text
func (l *ColumnList) SetEnumToTextConversion(columnName string, enumValues string) {
	l.GetColumn(columnName).EnumToTextConversion = true
	l.GetColumn(columnName).EnumValues = enumValues
}

// IsEnumToTextConversion tells whether an enum was converted to text
func (l *ColumnList) IsEnumToTextConversion(columnName string) bool {
	return l.GetColumn(columnName).EnumToTextConversion
}

// UniqueKey is the combination of a key's name and columns
type UniqueKey struct {
	Name            string
	Columns         ColumnList
	HasNullable     bool
	HasFloat        bool
	IsAutoIncrement bool
}

// IsPrimary checks if this unique key is primary
func (k *UniqueKey) IsPrimary() bool {
	return k.Name == "PRIMARY"
}

// Len returns the length of this list
func (k *UniqueKey) Len() int {
	return k.Columns.Len()
}

// String returns a visual representation of this key
func (k *UniqueKey) String() string {
	description := k.Name
	if k.IsAutoIncrement {
		description = fmt.Sprintf("%s (auto_increment)", description)
	}
	return fmt.Sprintf("%s: %s; has nullable: %+v", description, k.Columns.Names(), k.HasNullable)
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"

	"github.com/hashicorp/terraform/helper/schema"
	computeBeta "google.golang.org/api/compute/v0.beta"
)

func expandAliasIpRanges(ranges []interface{}) []*computeBeta.AliasIpRange {
	ipRanges := make([]*computeBeta.AliasIpRange, 0, len(ranges))
	for _, raw := range ranges {
		data := raw.(map[string]interface{})
		ipRanges = append(ipRanges, &computeBeta.AliasIpRange{
			IpCidrRange:         data["ip_cidr_range"].(string),
			SubnetworkRangeName: data["subnetwork_range_name"].(string),
		})
	}
	return ipRanges
}

func flattenAliasIpRange(ranges []*computeBeta.AliasIpRange) []map[string]interface{} {
	rangesSchema := make([]map[string]interface{}, 0, len(ranges))
	for _, ipRange := range ranges {
		rangesSchema = append(rangesSchema, map[string]interface{}{
			"ip_cidr_range":         ipRange.IpCidrRange,
			"subnetwork_range_name": ipRange.SubnetworkRangeName,
		})
	}
	return rangesSchema
}

func flattenScheduling(scheduling *computeBeta.Scheduling) []map[string]interface{} {
	result := make([]map[string]interface{}, 0, 1)
	schedulingMap := map[string]interface{}{
		"on_host_maintenance": scheduling.OnHostMaintenance,
		"preemptible":         scheduling.Preemptible,
	}
	if scheduling.AutomaticRestart != nil {
		schedulingMap["automatic_restart"] = *scheduling.AutomaticRestart
	}
	result = append(result, schedulingMap)
	return result
}

func flattenAccessConfigs(accessConfigs []*computeBeta.AccessConfig) ([]map[string]interface{}, string) {
	flattened := make([]map[string]interface{}, len(accessConfigs))
	natIP := ""
	for i, ac := range accessConfigs {
		flattened[i] = map[string]interface{}{
			"nat_ip":       ac.NatIP,
			"network_tier": ac.NetworkTier,
		}
		if ac.SetPublicPtr {
			flattened[i]["public_ptr_domain_name"] = ac.PublicPtrDomainName
		}
		if natIP == "" {
			natIP = ac.NatIP
		}
	}
	return flattened, natIP
}

func flattenNetworkInterfaces(d *schema.ResourceData, config *Config, networkInterfaces []*computeBeta.NetworkInterface) ([]map[string]interface{}, string, string, string, error) {
	flattened := make([]map[string]interface{}, len(networkInterfaces))
	var region, internalIP, externalIP string

	for i, iface := range networkInterfaces {
		var ac []map[string]interface{}
		ac, externalIP = flattenAccessConfigs(iface.AccessConfigs)

		subnet, err := ParseSubnetworkFieldValue(iface.Subnetwork, d, config)
		if err != nil {
			return nil, "", "", "", err
		}
		region = subnet.Region

		flattened[i] = map[string]interface{}{
			"network_ip":         iface.NetworkIP,
			"network":            ConvertSelfLinkToV1(iface.Network),
			"subnetwork":         ConvertSelfLinkToV1(iface.Subnetwork),
			"subnetwork_project": subnet.Project,
			"access_config":      ac,
			"alias_ip_range":     flattenAliasIpRange(iface.AliasIpRanges),
		}
		// Instance template interfaces never have names, so they're absent
		// in the instance template network_interface schema. We want to use the
		// same flattening code for both resource types, so we avoid trying to
		// set the name field when it's not set at the GCE end.
		if iface.Name != "" {
			flattened[i]["name"] = iface.Name
		}
		if internalIP == "" {
			internalIP = iface.NetworkIP
		}
	}
	return flattened, region, internalIP, externalIP, nil
}

func expandAccessConfigs(configs []interface{}) []*computeBeta.AccessConfig {
	acs := make([]*computeBeta.AccessConfig, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})
		acs[i] = &computeBeta.AccessConfig{
			Type:        "ONE_TO_ONE_NAT",
			NatIP:       data["nat_ip"].(string),
			NetworkTier: data["network_tier"].(string),
		}
		if ptr, ok := data["public_ptr_domain_name"]; ok && ptr != "" {
			acs[i].SetPublicPtr = true
			acs[i].PublicPtrDomainName = ptr.(string)
		}
	}
	return acs
}

func expandNetworkInterfaces(d TerraformResourceData, config *Config) ([]*computeBeta.NetworkInterface, error) {
	configs := d.Get("network_interface").([]interface{})
	ifaces := make([]*computeBeta.NetworkInterface, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})

		network := data["network"].(string)
		subnetwork := data["subnetwork"].(string)
		if (network == "" && subnetwork == "") || (network != "" && subnetwork != "") {
			return nil, fmt.Errorf("exactly one of network or subnetwork must be provided")
		}

		nf, err := ParseNetworkFieldValue(network, d, config)
		if err != nil {
			return nil, fmt.Errorf("cannot determine self_link for network %q: %s", network, err)
		}

		subnetProjectField := fmt.Sprintf("network_interface.%d.subnetwork_project", i)
		sf, err := ParseSubnetworkFieldValueWithProjectField(subnetwork, subnetProjectField, d, config)
		if err != nil {
			return nil, fmt.Errorf("cannot determine self_link for subnetwork %q: %s", subnetwork, err)
		}

		ifaces[i] = &computeBeta.NetworkInterface{
			NetworkIP:     data["network_ip"].(string),
			Network:       nf.RelativeLink(),
			Subnetwork:    sf.RelativeLink(),
			AccessConfigs: expandAccessConfigs(data["access_config"].([]interface{})),
			AliasIpRanges: expandAliasIpRanges(data["alias_ip_range"].([]interface{})),
		}

	}
	return ifaces, nil
}

func flattenServiceAccounts(serviceAccounts []*computeBeta.ServiceAccount) []map[string]interface{} {
	result := make([]map[string]interface{}, len(serviceAccounts))
	for i, serviceAccount := range serviceAccounts {
		result[i] = map[string]interface{}{
			"email":  serviceAccount.Email,
			"scopes": schema.NewSet(stringScopeHashcode, convertStringArrToInterface(serviceAccount.Scopes)),
		}
	}
	return result
}

func expandServiceAccounts(configs []interface{}) []*computeBeta.ServiceAccount {
	accounts := make([]*computeBeta.ServiceAccount, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})

		accounts[i] = &computeBeta.ServiceAccount{
			Email:  data["email"].(string),
			Scopes: canonicalizeServiceScopes(convertStringSet(data["scopes"].(*schema.Set))),
		}

		if accounts[i].Email == "" {
			accounts[i].Email = "default"
		}
	}
	return accounts
}

func flattenGuestAccelerators(accelerators []*computeBeta.AcceleratorConfig) []map[string]interface{} {
	acceleratorsSchema := make([]map[string]interface{}, len(accelerators))
	for i, accelerator := range accelerators {
		acceleratorsSchema[i] = map[string]interface{}{
			"count": accelerator.AcceleratorCount,
			"type":  accelerator.AcceleratorType,
		}
	}
	return acceleratorsSchema
}

func resourceInstanceTags(d TerraformResourceData) *computeBeta.Tags {
	// Calculate the tags
	var tags *computeBeta.Tags
	if v := d.Get("tags"); v != nil {
		vs := v.(*schema.Set)
		tags = new(computeBeta.Tags)
		tags.Items = make([]string, vs.Len())
		for i, v := range vs.List() {
			tags.Items[i] = v.(string)
		}

		tags.Fingerprint = d.Get("tags_fingerprint").(string)
	}

	return tags
}
<commit_msg>Remove validation in compute_instance expanders (#525)<commit_after>package google

import (
	"fmt"

	"github.com/hashicorp/terraform/helper/schema"
	computeBeta "google.golang.org/api/compute/v0.beta"
)

func expandAliasIpRanges(ranges []interface{}) []*computeBeta.AliasIpRange {
	ipRanges := make([]*computeBeta.AliasIpRange, 0, len(ranges))
	for _, raw := range ranges {
		data := raw.(map[string]interface{})
		ipRanges = append(ipRanges, &computeBeta.AliasIpRange{
			IpCidrRange:         data["ip_cidr_range"].(string),
			SubnetworkRangeName: data["subnetwork_range_name"].(string),
		})
	}
	return ipRanges
}

func flattenAliasIpRange(ranges []*computeBeta.AliasIpRange) []map[string]interface{} {
	rangesSchema := make([]map[string]interface{}, 0, len(ranges))
	for _, ipRange := range ranges {
		rangesSchema = append(rangesSchema, map[string]interface{}{
			"ip_cidr_range":         ipRange.IpCidrRange,
			"subnetwork_range_name": ipRange.SubnetworkRangeName,
		})
	}
	return rangesSchema
}

func flattenScheduling(scheduling *computeBeta.Scheduling) []map[string]interface{} {
	result := make([]map[string]interface{}, 0, 1)
	schedulingMap := map[string]interface{}{
		"on_host_maintenance": scheduling.OnHostMaintenance,
		"preemptible":         scheduling.Preemptible,
	}
	if scheduling.AutomaticRestart != nil {
		schedulingMap["automatic_restart"] = *scheduling.AutomaticRestart
	}
	result = append(result, schedulingMap)
	return result
}

func flattenAccessConfigs(accessConfigs []*computeBeta.AccessConfig) ([]map[string]interface{}, string) {
	flattened := make([]map[string]interface{}, len(accessConfigs))
	natIP := ""
	for i, ac := range accessConfigs {
		flattened[i] = map[string]interface{}{
			"nat_ip":       ac.NatIP,
			"network_tier": ac.NetworkTier,
		}
		if ac.SetPublicPtr {
			flattened[i]["public_ptr_domain_name"] = ac.PublicPtrDomainName
		}
		if natIP == "" {
			natIP = ac.NatIP
		}
	}
	return flattened, natIP
}

func flattenNetworkInterfaces(d *schema.ResourceData, config *Config, networkInterfaces []*computeBeta.NetworkInterface) ([]map[string]interface{}, string, string, string, error) {
	flattened := make([]map[string]interface{}, len(networkInterfaces))
	var region, internalIP, externalIP string

	for i, iface := range networkInterfaces {
		var ac []map[string]interface{}
		ac, externalIP = flattenAccessConfigs(iface.AccessConfigs)

		subnet, err := ParseSubnetworkFieldValue(iface.Subnetwork, d, config)
		if err != nil {
			return nil, "", "", "", err
		}
		region = subnet.Region

		flattened[i] = map[string]interface{}{
			"network_ip":         iface.NetworkIP,
			"network":            ConvertSelfLinkToV1(iface.Network),
			"subnetwork":         ConvertSelfLinkToV1(iface.Subnetwork),
			"subnetwork_project": subnet.Project,
			"access_config":      ac,
			"alias_ip_range":     flattenAliasIpRange(iface.AliasIpRanges),
		}
		// Instance template interfaces never have names, so they're absent
		// in the instance template network_interface schema. We want to use the
		// same flattening code for both resource types, so we avoid trying to
		// set the name field when it's not set at the GCE end.
		if iface.Name != "" {
			flattened[i]["name"] = iface.Name
		}
		if internalIP == "" {
			internalIP = iface.NetworkIP
		}
	}
	return flattened, region, internalIP, externalIP, nil
}

func expandAccessConfigs(configs []interface{}) []*computeBeta.AccessConfig {
	acs := make([]*computeBeta.AccessConfig, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})
		acs[i] = &computeBeta.AccessConfig{
			Type:        "ONE_TO_ONE_NAT",
			NatIP:       data["nat_ip"].(string),
			NetworkTier: data["network_tier"].(string),
		}
		if ptr, ok := data["public_ptr_domain_name"]; ok && ptr != "" {
			acs[i].SetPublicPtr = true
			acs[i].PublicPtrDomainName = ptr.(string)
		}
	}
	return acs
}

func expandNetworkInterfaces(d TerraformResourceData, config *Config) ([]*computeBeta.NetworkInterface, error) {
	configs := d.Get("network_interface").([]interface{})
	ifaces := make([]*computeBeta.NetworkInterface, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})

		network := data["network"].(string)
		subnetwork := data["subnetwork"].(string)
		if network == "" && subnetwork == "" {
			return nil, fmt.Errorf("exactly one of network or subnetwork must be provided")
		}

		nf, err := ParseNetworkFieldValue(network, d, config)
		if err != nil {
			return nil, fmt.Errorf("cannot determine self_link for network %q: %s", network, err)
		}

		subnetProjectField := fmt.Sprintf("network_interface.%d.subnetwork_project", i)
		sf, err := ParseSubnetworkFieldValueWithProjectField(subnetwork, subnetProjectField, d, config)
		if err != nil {
			return nil, fmt.Errorf("cannot determine self_link for subnetwork %q: %s", subnetwork, err)
		}

		ifaces[i] = &computeBeta.NetworkInterface{
			NetworkIP:     data["network_ip"].(string),
			Network:       nf.RelativeLink(),
			Subnetwork:    sf.RelativeLink(),
			AccessConfigs: expandAccessConfigs(data["access_config"].([]interface{})),
			AliasIpRanges: expandAliasIpRanges(data["alias_ip_range"].([]interface{})),
		}

	}
	return ifaces, nil
}

func flattenServiceAccounts(serviceAccounts []*computeBeta.ServiceAccount) []map[string]interface{} {
	result := make([]map[string]interface{}, len(serviceAccounts))
	for i, serviceAccount := range serviceAccounts {
		result[i] = map[string]interface{}{
			"email":  serviceAccount.Email,
			"scopes": schema.NewSet(stringScopeHashcode, convertStringArrToInterface(serviceAccount.Scopes)),
		}
	}
	return result
}

func expandServiceAccounts(configs []interface{}) []*computeBeta.ServiceAccount {
	accounts := make([]*computeBeta.ServiceAccount, len(configs))
	for i, raw := range configs {
		data := raw.(map[string]interface{})

		accounts[i] = &computeBeta.ServiceAccount{
			Email:  data["email"].(string),
			Scopes: canonicalizeServiceScopes(convertStringSet(data["scopes"].(*schema.Set))),
		}

		if accounts[i].Email == "" {
			accounts[i].Email = "default"
		}
	}
	return accounts
}

func flattenGuestAccelerators(accelerators []*computeBeta.AcceleratorConfig) []map[string]interface{} {
	acceleratorsSchema := make([]map[string]interface{}, len(accelerators))
	for i, accelerator := range accelerators {
		acceleratorsSchema[i] = map[string]interface{}{
			"count": accelerator.AcceleratorCount,
			"type":  accelerator.AcceleratorType,
		}
	}
	return acceleratorsSchema
}

func resourceInstanceTags(d TerraformResourceData) *computeBeta.Tags {
	// Calculate the tags
	var tags *computeBeta.Tags
	if v := d.Get("tags"); v != nil {
		vs := v.(*schema.Set)
		tags = new(computeBeta.Tags)
		tags.Items = make([]string, vs.Len())
		for i, v := range vs.List() {
			tags.Items[i] = v.(string)
		}

		tags.Fingerprint = d.Get("tags_fingerprint").(string)
	}

	return tags
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/integration/checker"
	"github.com/go-check/check"
)

// This is a heisen-test.  Because the created timestamp of images and the behavior of
// sort is not predictable it doesn't always fail.
func (s *DockerSuite) TestBuildHistory(c *check.C) {
	testRequires(c, DaemonIsLinux)
	name := "testbuildhistory"
	_, err := buildImage(name, `FROM busybox
RUN echo "A"
RUN echo "B"
RUN echo "C"
RUN echo "D"
RUN echo "E"
RUN echo "F"
RUN echo "G"
RUN echo "H"
RUN echo "I"
RUN echo "J"
RUN echo "K"
RUN echo "L"
RUN echo "M"
RUN echo "N"
RUN echo "O"
RUN echo "P"
RUN echo "Q"
RUN echo "R"
RUN echo "S"
RUN echo "T"
RUN echo "U"
RUN echo "V"
RUN echo "W"
RUN echo "X"
RUN echo "Y"
RUN echo "Z"`,
		true)

	c.Assert(err, checker.IsNil)

	out, _ := dockerCmd(c, "history", "testbuildhistory")
	actualValues := strings.Split(out, "\n")[1:27]
	expectedValues := [26]string{"Z", "Y", "X", "W", "V", "U", "T", "S", "R", "Q", "P", "O", "N", "M", "L", "K", "J", "I", "H", "G", "F", "E", "D", "C", "B", "A"}

	for i := 0; i < 26; i++ {
		echoValue := fmt.Sprintf("echo \"%s\"", expectedValues[i])
		actualValue := actualValues[i]
		c.Assert(actualValue, checker.Contains, echoValue)
	}

}

func (s *DockerSuite) TestHistoryExistentImage(c *check.C) {
	testRequires(c, DaemonIsLinux)
	dockerCmd(c, "history", "busybox")
}

func (s *DockerSuite) TestHistoryNonExistentImage(c *check.C) {
	_, _, err := dockerCmdWithError("history", "testHistoryNonExistentImage")
	c.Assert(err, checker.NotNil, check.Commentf("history on a non-existent image should fail."))
}

func (s *DockerSuite) TestHistoryImageWithComment(c *check.C) {
	testRequires(c, DaemonIsLinux)
	name := "testhistoryimagewithcomment"

	// make a image through docker commit <container id> [ -m messages ]

	dockerCmd(c, "run", "--name", name, "busybox", "true")
	dockerCmd(c, "wait", name)

	comment := "This_is_a_comment"
	dockerCmd(c, "commit", "-m="+comment, name, name)

	// test docker history <image id> to check comment messages

	out, _ := dockerCmd(c, "history", name)
	outputTabs := strings.Fields(strings.Split(out, "\n")[1])
	actualValue := outputTabs[len(outputTabs)-1]
	c.Assert(actualValue, checker.Contains, comment)
}

func (s *DockerSuite) TestHistoryHumanOptionFalse(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "history", "--human=false", "busybox")
	lines := strings.Split(out, "\n")
	sizeColumnRegex, _ := regexp.Compile("SIZE +")
	indices := sizeColumnRegex.FindStringIndex(lines[0])
	startIndex := indices[0]
	endIndex := indices[1]
	for i := 1; i < len(lines)-1; i++ {
		if endIndex > len(lines[i]) {
			endIndex = len(lines[i])
		}
		sizeString := lines[i][startIndex:endIndex]

		_, err := strconv.Atoi(strings.TrimSpace(sizeString))
		c.Assert(err, checker.IsNil, check.Commentf("The size '%s' was not an Integer", sizeString))
	}
}

func (s *DockerSuite) TestHistoryHumanOptionTrue(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "history", "--human=true", "busybox")
	lines := strings.Split(out, "\n")
	sizeColumnRegex, _ := regexp.Compile("SIZE +")
	humanSizeRegexRaw := "\\d+.*B" // Matches human sizes like 10 MB, 3.2 KB, etc
	indices := sizeColumnRegex.FindStringIndex(lines[0])
	startIndex := indices[0]
	endIndex := indices[1]
	for i := 1; i < len(lines)-1; i++ {
		if endIndex > len(lines[i]) {
			endIndex = len(lines[i])
		}
		sizeString := lines[i][startIndex:endIndex]
		c.Assert(strings.TrimSpace(sizeString), checker.Matches, humanSizeRegexRaw, check.Commentf("The size '%s' was not in human format", sizeString))
	}
}
<commit_msg>Windows CI: Enable TestHistory* cli tests<commit_after>package main

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/integration/checker"
	"github.com/go-check/check"
)

// This is a heisen-test.  Because the created timestamp of images and the behavior of
// sort is not predictable it doesn't always fail.
func (s *DockerSuite) TestBuildHistory(c *check.C) {
	name := "testbuildhistory"
	_, err := buildImage(name, `FROM busybox
RUN echo "A"
RUN echo "B"
RUN echo "C"
RUN echo "D"
RUN echo "E"
RUN echo "F"
RUN echo "G"
RUN echo "H"
RUN echo "I"
RUN echo "J"
RUN echo "K"
RUN echo "L"
RUN echo "M"
RUN echo "N"
RUN echo "O"
RUN echo "P"
RUN echo "Q"
RUN echo "R"
RUN echo "S"
RUN echo "T"
RUN echo "U"
RUN echo "V"
RUN echo "W"
RUN echo "X"
RUN echo "Y"
RUN echo "Z"`,
		true)

	c.Assert(err, checker.IsNil)

	out, _ := dockerCmd(c, "history", "testbuildhistory")
	actualValues := strings.Split(out, "\n")[1:27]
	expectedValues := [26]string{"Z", "Y", "X", "W", "V", "U", "T", "S", "R", "Q", "P", "O", "N", "M", "L", "K", "J", "I", "H", "G", "F", "E", "D", "C", "B", "A"}

	for i := 0; i < 26; i++ {
		echoValue := fmt.Sprintf("echo \"%s\"", expectedValues[i])
		actualValue := actualValues[i]
		c.Assert(actualValue, checker.Contains, echoValue)
	}

}

func (s *DockerSuite) TestHistoryExistentImage(c *check.C) {
	dockerCmd(c, "history", "busybox")
}

func (s *DockerSuite) TestHistoryNonExistentImage(c *check.C) {
	_, _, err := dockerCmdWithError("history", "testHistoryNonExistentImage")
	c.Assert(err, checker.NotNil, check.Commentf("history on a non-existent image should fail."))
}

func (s *DockerSuite) TestHistoryImageWithComment(c *check.C) {
	name := "testhistoryimagewithcomment"

	// make a image through docker commit <container id> [ -m messages ]

	dockerCmd(c, "run", "--name", name, "busybox", "true")
	dockerCmd(c, "wait", name)

	comment := "This_is_a_comment"
	dockerCmd(c, "commit", "-m="+comment, name, name)

	// test docker history <image id> to check comment messages

	out, _ := dockerCmd(c, "history", name)
	outputTabs := strings.Fields(strings.Split(out, "\n")[1])
	actualValue := outputTabs[len(outputTabs)-1]
	c.Assert(actualValue, checker.Contains, comment)
}

func (s *DockerSuite) TestHistoryHumanOptionFalse(c *check.C) {
	out, _ := dockerCmd(c, "history", "--human=false", "busybox")
	lines := strings.Split(out, "\n")
	sizeColumnRegex, _ := regexp.Compile("SIZE +")
	indices := sizeColumnRegex.FindStringIndex(lines[0])
	startIndex := indices[0]
	endIndex := indices[1]
	for i := 1; i < len(lines)-1; i++ {
		if endIndex > len(lines[i]) {
			endIndex = len(lines[i])
		}
		sizeString := lines[i][startIndex:endIndex]

		_, err := strconv.Atoi(strings.TrimSpace(sizeString))
		c.Assert(err, checker.IsNil, check.Commentf("The size '%s' was not an Integer", sizeString))
	}
}

func (s *DockerSuite) TestHistoryHumanOptionTrue(c *check.C) {
	out, _ := dockerCmd(c, "history", "--human=true", "busybox")
	lines := strings.Split(out, "\n")
	sizeColumnRegex, _ := regexp.Compile("SIZE +")
	humanSizeRegexRaw := "\\d+.*B" // Matches human sizes like 10 MB, 3.2 KB, etc
	indices := sizeColumnRegex.FindStringIndex(lines[0])
	startIndex := indices[0]
	endIndex := indices[1]
	for i := 1; i < len(lines)-1; i++ {
		if endIndex > len(lines[i]) {
			endIndex = len(lines[i])
		}
		sizeString := lines[i][startIndex:endIndex]
		c.Assert(strings.TrimSpace(sizeString), checker.Matches, humanSizeRegexRaw, check.Commentf("The size '%s' was not in human format", sizeString))
	}
}
<|endoftext|>


<commit_before>// Copyright © 2018 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package generic

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/builtins/collector"
	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-agent/internal/release"
	"github.com/circonus-labs/circonus-agent/internal/tags"
	cgm "github.com/circonus-labs/circonus-gometrics/v3"
	"github.com/pkg/errors"
	"github.com/rs/zerolog"
	"github.com/shirou/gopsutil/cpu"
)

// CPU metrics from psutils
type CPU struct {
	gencommon
	reportAllCPUs bool // OPT report all cpus (vs just total) may be overridden in config file
}

// cpuOptions defines what elements can be overridden in a config file
type cpuOptions struct {
	// common
	ID     string `json:"id" toml:"id" yaml:"id"`
	RunTTL string `json:"run_ttl" toml:"run_ttl" yaml:"run_ttl"`

	// collector specific
	AllCPU string `json:"report_all_cpus" toml:"report_all_cpus" yaml:"report_all_cpus"`
}

// NewCPUCollector creates new psutils cpu collector
func NewCPUCollector(cfgBaseName string, parentLogger zerolog.Logger) (collector.Collector, error) {
	c := CPU{}
	c.id = NameCPU
	c.pkgID = PackageName + "." + c.id
	c.logger = parentLogger.With().Str("id", c.id).Logger()
	c.reportAllCPUs = false
	c.baseTags = tags.FromList(tags.GetBaseTags())

	var opts cpuOptions
	err := config.LoadConfigFile(cfgBaseName, &opts)
	if err != nil {
		if strings.Contains(err.Error(), "no config found matching") {
			return &c, nil
		}
		c.logger.Warn().Err(err).Str("file", cfgBaseName).Msg("loading config file")
		return nil, errors.Wrapf(err, "%s config", c.pkgID)
	}

	c.logger.Debug().Interface("config", opts).Msg("loaded config")

	if opts.AllCPU != "" {
		rpt, err := strconv.ParseBool(opts.AllCPU)
		if err != nil {
			return nil, errors.Wrapf(err, "%s parsing report_all_cpus", c.pkgID)
		}
		c.reportAllCPUs = rpt
	}

	if opts.ID != "" {
		c.id = opts.ID
	}

	if opts.RunTTL != "" {
		dur, err := time.ParseDuration(opts.RunTTL)
		if err != nil {
			return nil, errors.Wrapf(err, "%s parsing run_ttl", c.pkgID)
		}
		c.runTTL = dur
	}

	return &c, nil
}

// Collect cpu metrics
func (c *CPU) Collect() error {
	c.Lock()
	if c.runTTL > time.Duration(0) {
		if time.Since(c.lastEnd) < c.runTTL {
			c.logger.Warn().Msg(collector.ErrTTLNotExpired.Error())
			c.Unlock()
			return collector.ErrTTLNotExpired
		}
	}
	if c.running {
		c.logger.Warn().Msg(collector.ErrAlreadyRunning.Error())
		c.Unlock()
		return collector.ErrAlreadyRunning
	}

	c.running = true
	c.lastStart = time.Now()
	c.Unlock()

	moduleTags := tags.Tags{
		tags.Tag{Category: release.NAME + "-module", Value: c.id},
	}

	metrics := cgm.Metrics{}
	pcts, err := cpu.Percent(time.Duration(0), c.reportAllCPUs)
	if err != nil {
		c.logger.Warn().Err(err).Msg("collecting metrics, cpu%")
	} else {
		if !c.reportAllCPUs && len(pcts) == 1 {
			var tagList tags.Tags
			tagList = append(tagList, moduleTags...)
			tagList = append(tagList, tags.Tag{Category: "units", Value: "percent"})
			_ = c.addMetric(&metrics, "used", "n", pcts[0], tagList)
		} else {
			for idx, pct := range pcts {
				var tagList tags.Tags
				tagList = append(tagList, moduleTags...)
				tagList = append(tagList, tags.Tags{
					tags.Tag{Category: "cpu", Value: fmt.Sprintf("%d", idx)},
					tags.Tag{Category: "units", Value: "percent"},
				}...)
				_ = c.addMetric(&metrics, "used", "n", pct, tagList)
			}
		}
	}

	var cpuTags tags.Tags
	cpuTags = append(cpuTags, moduleTags...)
	cpuTags = append(cpuTags, tags.Tags{
		tags.Tag{Category: "units", Value: "jiffies"},
	}...)

	ts, err := cpu.Times(c.reportAllCPUs)
	if err != nil {
		c.logger.Warn().Err(err).Msg("collecting metrics, cpu times")
	} else {
		if !c.reportAllCPUs && len(ts) == 1 {
			_ = c.addMetric(&metrics, "user", "n", ts[0].User, cpuTags)
			_ = c.addMetric(&metrics, "system", "n", ts[0].System, cpuTags)
			_ = c.addMetric(&metrics, "idle", "n", ts[0].Idle, cpuTags)
			_ = c.addMetric(&metrics, "nice", "n", ts[0].Nice, cpuTags)
			_ = c.addMetric(&metrics, "iowait", "n", ts[0].Iowait, cpuTags)
			_ = c.addMetric(&metrics, "irq", "n", ts[0].Irq, cpuTags)
			_ = c.addMetric(&metrics, "soft_irq", "n", ts[0].Softirq, cpuTags)
			_ = c.addMetric(&metrics, "steal", "n", ts[0].Steal, cpuTags)
			_ = c.addMetric(&metrics, "guest", "n", ts[0].Guest, cpuTags)
			_ = c.addMetric(&metrics, "guest_nice", "n", ts[0].GuestNice, cpuTags)
			_ = c.addMetric(&metrics, "stolen", "n", ts[0].Stolen, cpuTags)
		} else {
			for idx, v := range ts {
				var tagList tags.Tags
				tagList = append(tagList, cpuTags...)
				tagList = append(tagList, tags.Tag{Category: "cpu", Value: fmt.Sprintf("%d", idx)})
				_ = c.addMetric(&metrics, "user", "n", v.User, tagList)
				_ = c.addMetric(&metrics, "system", "n", v.System, tagList)
				_ = c.addMetric(&metrics, "idle", "n", v.Idle, tagList)
				_ = c.addMetric(&metrics, "nice", "n", v.Nice, tagList)
				_ = c.addMetric(&metrics, "iowait", "n", v.Iowait, tagList)
				_ = c.addMetric(&metrics, "irq", "n", v.Irq, tagList)
				_ = c.addMetric(&metrics, "soft_irq", "n", v.Softirq, tagList)
				_ = c.addMetric(&metrics, "steal", "n", v.Steal, tagList)
				_ = c.addMetric(&metrics, "guest", "n", v.Guest, tagList)
				_ = c.addMetric(&metrics, "guest_nice", "n", v.GuestNice, tagList)
				_ = c.addMetric(&metrics, "stolen", "n", v.Stolen, tagList)
			}
		}
	}

	c.setStatus(metrics, nil)
	return nil
}
<commit_msg>upd: remove module specific tags, let common metric handler add them<commit_after>// Copyright © 2018 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package generic

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/builtins/collector"
	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-agent/internal/tags"
	cgm "github.com/circonus-labs/circonus-gometrics/v3"
	"github.com/pkg/errors"
	"github.com/rs/zerolog"
	"github.com/shirou/gopsutil/cpu"
)

// CPU metrics from psutils
type CPU struct {
	gencommon
	reportAllCPUs bool // OPT report all cpus (vs just total) may be overridden in config file
}

// cpuOptions defines what elements can be overridden in a config file
type cpuOptions struct {
	// common
	ID     string `json:"id" toml:"id" yaml:"id"`
	RunTTL string `json:"run_ttl" toml:"run_ttl" yaml:"run_ttl"`

	// collector specific
	AllCPU string `json:"report_all_cpus" toml:"report_all_cpus" yaml:"report_all_cpus"`
}

// NewCPUCollector creates new psutils cpu collector
func NewCPUCollector(cfgBaseName string, parentLogger zerolog.Logger) (collector.Collector, error) {
	c := CPU{}
	c.id = NameCPU
	c.pkgID = PackageName + "." + c.id
	c.logger = parentLogger.With().Str("id", c.id).Logger()
	c.reportAllCPUs = false
	c.baseTags = tags.FromList(tags.GetBaseTags())

	var opts cpuOptions
	err := config.LoadConfigFile(cfgBaseName, &opts)
	if err != nil {
		if strings.Contains(err.Error(), "no config found matching") {
			return &c, nil
		}
		c.logger.Warn().Err(err).Str("file", cfgBaseName).Msg("loading config file")
		return nil, errors.Wrapf(err, "%s config", c.pkgID)
	}

	c.logger.Debug().Interface("config", opts).Msg("loaded config")

	if opts.AllCPU != "" {
		rpt, err := strconv.ParseBool(opts.AllCPU)
		if err != nil {
			return nil, errors.Wrapf(err, "%s parsing report_all_cpus", c.pkgID)
		}
		c.reportAllCPUs = rpt
	}

	if opts.ID != "" {
		c.id = opts.ID
	}

	if opts.RunTTL != "" {
		dur, err := time.ParseDuration(opts.RunTTL)
		if err != nil {
			return nil, errors.Wrapf(err, "%s parsing run_ttl", c.pkgID)
		}
		c.runTTL = dur
	}

	return &c, nil
}

// Collect cpu metrics
func (c *CPU) Collect() error {
	c.Lock()
	if c.runTTL > time.Duration(0) {
		if time.Since(c.lastEnd) < c.runTTL {
			c.logger.Warn().Msg(collector.ErrTTLNotExpired.Error())
			c.Unlock()
			return collector.ErrTTLNotExpired
		}
	}
	if c.running {
		c.logger.Warn().Msg(collector.ErrAlreadyRunning.Error())
		c.Unlock()
		return collector.ErrAlreadyRunning
	}

	c.running = true
	c.lastStart = time.Now()
	c.Unlock()

	metrics := cgm.Metrics{}
	pcts, err := cpu.Percent(time.Duration(0), c.reportAllCPUs)
	if err != nil {
		c.logger.Warn().Err(err).Msg("collecting metrics, cpu%")
	} else {
		metricName := "cpu_used"
		metricType := "n"
		tagUnitsPercent := tags.Tag{Category: "units", Value: "percent"}
		if !c.reportAllCPUs && len(pcts) == 1 {
			tagList := tags.Tags{tagUnitsPercent}
			_ = c.addMetric(&metrics, metricName, metricType, pcts[0], tagList)
		} else {
			for idx, pct := range pcts {
				tagList := tags.Tags{
					tags.Tag{Category: "cpu", Value: fmt.Sprintf("%d", idx)},
				}
				tagList = append(tagList, tagUnitsPercent)
				_ = c.addMetric(&metrics, metricName, metricType, pct, tagList)
			}
		}
	}

	ts, err := cpu.Times(c.reportAllCPUs)
	if err != nil {
		c.logger.Warn().Err(err).Msg("collecting metrics, cpu times")
		c.setStatus(metrics, nil)
		return nil
	}

	tagUnitsCentiseconds := tags.Tag{Category: "units", Value: "centiseconds"} // aka jiffies
	if !c.reportAllCPUs && len(ts) == 1 {
		tagList := tags.Tags{tagUnitsCentiseconds}
		_ = c.addMetric(&metrics, "cpu_user", "n", ts[0].User, tagList)
		_ = c.addMetric(&metrics, "cpu_system", "n", ts[0].System, tagList)
		_ = c.addMetric(&metrics, "cpu_idle", "n", ts[0].Idle, tagList)
		_ = c.addMetric(&metrics, "cpu_nice", "n", ts[0].Nice, tagList)
		_ = c.addMetric(&metrics, "cpu_iowait", "n", ts[0].Iowait, tagList)
		_ = c.addMetric(&metrics, "cpu_irq", "n", ts[0].Irq, tagList)
		_ = c.addMetric(&metrics, "cpu_soft_irq", "n", ts[0].Softirq, tagList)
		_ = c.addMetric(&metrics, "cpu_steal", "n", ts[0].Steal, tagList)
		_ = c.addMetric(&metrics, "cpu_guest", "n", ts[0].Guest, tagList)
		_ = c.addMetric(&metrics, "cpu_guest_nice", "n", ts[0].GuestNice, tagList)
		_ = c.addMetric(&metrics, "cpu_stolen", "n", ts[0].Stolen, tagList)
	} else {
		for idx, v := range ts {
			tagList := tags.Tags{
				tags.Tag{Category: "cpu", Value: fmt.Sprintf("%d", idx)},
			}
			tagList = append(tagList, tagUnitsCentiseconds)
			_ = c.addMetric(&metrics, "cpu_user", "n", v.User, tagList)
			_ = c.addMetric(&metrics, "cpu_system", "n", v.System, tagList)
			_ = c.addMetric(&metrics, "cpu_idle", "n", v.Idle, tagList)
			_ = c.addMetric(&metrics, "cpu_nice", "n", v.Nice, tagList)
			_ = c.addMetric(&metrics, "cpu_iowait", "n", v.Iowait, tagList)
			_ = c.addMetric(&metrics, "cpu_irq", "n", v.Irq, tagList)
			_ = c.addMetric(&metrics, "cpu_soft_irq", "n", v.Softirq, tagList)
			_ = c.addMetric(&metrics, "cpu_steal", "n", v.Steal, tagList)
			_ = c.addMetric(&metrics, "cpu_guest", "n", v.Guest, tagList)
			_ = c.addMetric(&metrics, "cpu_guest_nice", "n", v.GuestNice, tagList)
			_ = c.addMetric(&metrics, "cpu_stolen", "n", v.Stolen, tagList)
		}
	}

	c.setStatus(metrics, nil)
	return nil
}
<|endoftext|>


<commit_before>package xmlutil

import (
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func UnmarshalXML(v interface{}, d *xml.Decoder, wrapper string) error {
	n, _ := XMLToStruct(d, nil)
	if n.Children != nil {
		for _, root := range n.Children {
			for _, c := range root {
				if wrappedChild, ok := c.Children[wrapper]; ok {
					c = wrappedChild[0] // pull out wrapped element
				}

				err := parse(reflect.ValueOf(v), c, "")
				if err != nil {
					if err == io.EOF {
						return nil
					}
					return err
				}
			}
		}
		return nil
	}
	return nil
}

func parse(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	rtype := r.Type()
	if rtype.Kind() == reflect.Ptr {
		rtype = rtype.Elem() // check kind of actual element type
	}

	t := tag.Get("type")
	if t == "" {
		switch rtype.Kind() {
		case reflect.Struct:
			t = "structure"
		case reflect.Slice:
			t = "list"
		case reflect.Map:
			t = "map"
		}
	}

	switch t {
	case "structure":
		if field, ok := rtype.FieldByName("SDKShapeTraits"); ok {
			tag = field.Tag
		}
		return parseStruct(r, node, tag)
	case "list":
		return parseList(r, node, tag)
	case "map":
		return parseMap(r, node, tag)
	default:
		return parseScalar(r, node, tag)
	}
}

func parseStruct(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()
	if r.Kind() == reflect.Ptr {
		if r.IsNil() { // create the structure if it's nil
			s := reflect.New(r.Type().Elem())
			r.Set(s)
			r = s
		}

		r = r.Elem()
		t = t.Elem()
	}

	// unwrap any payloads
	if payload := tag.Get("payload"); payload != "" {
		field, _ := t.FieldByName(payload)
		return parseStruct(r.FieldByName(payload), node, field.Tag)
	}

	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)
		if c := field.Name[0:1]; strings.ToLower(c) == c {
			continue // ignore unexported fields
		}

		// figure out what this field is called
		name := field.Name
		if locName := field.Tag.Get("locationName"); locName != "" {
			name = locName
		} else if field.Tag.Get("flattened") != "" {
			if locName := field.Tag.Get("locationNameList"); locName != "" {
				name = locName
			}
		}

		// try to find the field by name in elements
		elems := node.Children[name]

		if elems == nil { // try to find the field in attributes
			for _, a := range node.Attr {
				if name == a.Name.Local {
					// turn this into a text node for de-serializing
					elems = []*XMLNode{&XMLNode{Text: a.Value}}
				}
			}
		}

		member := r.FieldByName(field.Name)
		for _, elem := range elems {
			err := parse(member, elem, field.Tag)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func parseList(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()

	if tag.Get("flattened") == "" { // look at all item entries
		mname := "member"
		if name := tag.Get("locationNameList"); name != "" {
			mname = name
		}

		if Children, ok := node.Children[mname]; ok {
			if r.IsNil() {
				r.Set(reflect.MakeSlice(t, len(Children), len(Children)))
			}

			for i, c := range Children {
				err := parse(r.Index(i), c, "")
				if err != nil {
					return err
				}
			}
		}
	} else { // flattened list means this is a single element
		if r.IsNil() {
			r.Set(reflect.MakeSlice(t, 0, 0))
		}

		childR := reflect.Zero(t.Elem())
		r.Set(reflect.Append(r, childR))
		err := parse(r.Index(r.Len()-1), node, "")
		if err != nil {
			return err
		}
	}

	return nil
}

func parseMap(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()
	if r.Kind() == reflect.Ptr {
		t = t.Elem()
		if r.IsNil() {
			r.Set(reflect.New(t))
			r.Elem().Set(reflect.MakeMap(t))
		}

		r = r.Elem()
	}

	if tag.Get("flattened") == "" { // look at all child entries
		for _, entry := range node.Children["entry"] {
			parseMapEntry(r, entry, tag)
		}
	} else { // this element is itself an entry
		parseMapEntry(r, node, tag)
	}

	return nil
}

func parseMapEntry(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	kname, vname := "key", "value"
	if n := tag.Get("locationNameKey"); n != "" {
		kname = n
	}
	if n := tag.Get("locationNameValue"); n != "" {
		vname = n
	}

	keys, ok := node.Children[kname]
	values := node.Children[vname]
	if ok {
		for i, key := range keys {
			keyR := reflect.ValueOf(key.Text)
			value := values[i]
			valueR := reflect.New(r.Type().Elem()).Elem()

			parse(valueR, value, "")
			r.SetMapIndex(keyR, valueR)
		}
	}
	return nil
}

func parseScalar(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	switch r.Interface().(type) {
	case *string:
		r.Set(reflect.ValueOf(&node.Text))
		return nil
	case []byte:
		b, err := base64.StdEncoding.DecodeString(node.Text)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(b))
	case *bool:
		v, err := strconv.ParseBool(node.Text)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *int64:
		v, err := strconv.ParseInt(node.Text, 10, 64)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *float64:
		v, err := strconv.ParseFloat(node.Text, 64)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *time.Time:
		const ISO8601UTC = "2006-01-02T15:04:05Z"
		t, err := time.Parse(ISO8601UTC, node.Text)
		if err != nil {
			return err
		} else {
			r.Set(reflect.ValueOf(&t))
		}
	default:
		return fmt.Errorf("unsupported value: %v (%s)", r.Interface(), r.Type())
	}
	return nil
}
<commit_msg>Prioritize locationNameList over locationName<commit_after>package xmlutil

import (
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func UnmarshalXML(v interface{}, d *xml.Decoder, wrapper string) error {
	n, _ := XMLToStruct(d, nil)
	if n.Children != nil {
		for _, root := range n.Children {
			for _, c := range root {
				if wrappedChild, ok := c.Children[wrapper]; ok {
					c = wrappedChild[0] // pull out wrapped element
				}

				err := parse(reflect.ValueOf(v), c, "")
				if err != nil {
					if err == io.EOF {
						return nil
					}
					return err
				}
			}
		}
		return nil
	}
	return nil
}

func parse(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	rtype := r.Type()
	if rtype.Kind() == reflect.Ptr {
		rtype = rtype.Elem() // check kind of actual element type
	}

	t := tag.Get("type")
	if t == "" {
		switch rtype.Kind() {
		case reflect.Struct:
			t = "structure"
		case reflect.Slice:
			t = "list"
		case reflect.Map:
			t = "map"
		}
	}

	switch t {
	case "structure":
		if field, ok := rtype.FieldByName("SDKShapeTraits"); ok {
			tag = field.Tag
		}
		return parseStruct(r, node, tag)
	case "list":
		return parseList(r, node, tag)
	case "map":
		return parseMap(r, node, tag)
	default:
		return parseScalar(r, node, tag)
	}
}

func parseStruct(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()
	if r.Kind() == reflect.Ptr {
		if r.IsNil() { // create the structure if it's nil
			s := reflect.New(r.Type().Elem())
			r.Set(s)
			r = s
		}

		r = r.Elem()
		t = t.Elem()
	}

	// unwrap any payloads
	if payload := tag.Get("payload"); payload != "" {
		field, _ := t.FieldByName(payload)
		return parseStruct(r.FieldByName(payload), node, field.Tag)
	}

	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)
		if c := field.Name[0:1]; strings.ToLower(c) == c {
			continue // ignore unexported fields
		}

		// figure out what this field is called
		name := field.Name
		if field.Tag.Get("flattened") != "" && field.Tag.Get("locationNameList") != "" {
			name = field.Tag.Get("locationNameList")
		} else if locName := field.Tag.Get("locationName"); locName != "" {
			name = locName
		}

		// try to find the field by name in elements
		elems := node.Children[name]

		if elems == nil { // try to find the field in attributes
			for _, a := range node.Attr {
				if name == a.Name.Local {
					// turn this into a text node for de-serializing
					elems = []*XMLNode{&XMLNode{Text: a.Value}}
				}
			}
		}

		member := r.FieldByName(field.Name)
		for _, elem := range elems {
			err := parse(member, elem, field.Tag)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func parseList(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()

	if tag.Get("flattened") == "" { // look at all item entries
		mname := "member"
		if name := tag.Get("locationNameList"); name != "" {
			mname = name
		}

		if Children, ok := node.Children[mname]; ok {
			if r.IsNil() {
				r.Set(reflect.MakeSlice(t, len(Children), len(Children)))
			}

			for i, c := range Children {
				err := parse(r.Index(i), c, "")
				if err != nil {
					return err
				}
			}
		}
	} else { // flattened list means this is a single element
		if r.IsNil() {
			r.Set(reflect.MakeSlice(t, 0, 0))
		}

		childR := reflect.Zero(t.Elem())
		r.Set(reflect.Append(r, childR))
		err := parse(r.Index(r.Len()-1), node, "")
		if err != nil {
			return err
		}
	}

	return nil
}

func parseMap(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	t := r.Type()
	if r.Kind() == reflect.Ptr {
		t = t.Elem()
		if r.IsNil() {
			r.Set(reflect.New(t))
			r.Elem().Set(reflect.MakeMap(t))
		}

		r = r.Elem()
	}

	if tag.Get("flattened") == "" { // look at all child entries
		for _, entry := range node.Children["entry"] {
			parseMapEntry(r, entry, tag)
		}
	} else { // this element is itself an entry
		parseMapEntry(r, node, tag)
	}

	return nil
}

func parseMapEntry(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	kname, vname := "key", "value"
	if n := tag.Get("locationNameKey"); n != "" {
		kname = n
	}
	if n := tag.Get("locationNameValue"); n != "" {
		vname = n
	}

	keys, ok := node.Children[kname]
	values := node.Children[vname]
	if ok {
		for i, key := range keys {
			keyR := reflect.ValueOf(key.Text)
			value := values[i]
			valueR := reflect.New(r.Type().Elem()).Elem()

			parse(valueR, value, "")
			r.SetMapIndex(keyR, valueR)
		}
	}
	return nil
}

func parseScalar(r reflect.Value, node *XMLNode, tag reflect.StructTag) error {
	switch r.Interface().(type) {
	case *string:
		r.Set(reflect.ValueOf(&node.Text))
		return nil
	case []byte:
		b, err := base64.StdEncoding.DecodeString(node.Text)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(b))
	case *bool:
		v, err := strconv.ParseBool(node.Text)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *int64:
		v, err := strconv.ParseInt(node.Text, 10, 64)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *float64:
		v, err := strconv.ParseFloat(node.Text, 64)
		if err != nil {
			return err
		}
		r.Set(reflect.ValueOf(&v))
	case *time.Time:
		const ISO8601UTC = "2006-01-02T15:04:05Z"
		t, err := time.Parse(ISO8601UTC, node.Text)
		if err != nil {
			return err
		} else {
			r.Set(reflect.ValueOf(&t))
		}
	default:
		return fmt.Errorf("unsupported value: %v (%s)", r.Interface(), r.Type())
	}
	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"github.com/MG-RAST/AWE/lib/conf"
	"os"
	"time"
)

const (
	WORK_STAT_QUEUED      = "queued"
	WORK_STAT_CHECKOUT    = "checkout"
	WORK_STAT_SUSPEND     = "suspend"
	WORK_STAT_DONE        = "done"
	WORK_STAT_FAIL        = "fail"
	WORK_STAT_PREPARED    = "prepared"
	WORK_STAT_COMPUTED    = "computed"
	WORK_STAT_DISCARDED   = "discarded"
	WORK_STAT_PROXYQUEUED = "proxyqueued"
)

type Workunit struct {
	Id           string            `bson:"wuid" json:"wuid"`
	Info         *Info             `bson:"info" json:"info"`
	Inputs       IOmap             `bson:"inputs" json:"inputs"`
	Outputs      IOmap             `bson:"outputs" json:"outputs"`
	Predata      IOmap             `bson:"predata" json:"predata"`
	Cmd          *Command          `bson:"cmd" json:"cmd"`
	Rank         int               `bson:"rank" json:"rank"`
	TotalWork    int               `bson:"totalwork" json:"totalwork"`
	Partition    *PartInfo         `bson:"part" json:"part"`
	State        string            `bson:"state" json:"state"`
	Failed       int               `bson:"failed" json:"failed"`
	CheckoutTime time.Time         `bson:"checkout_time" json:"checkout_time"`
	Client       string            `bson:"client" json:"client"`
	ComputeTime  int               `bson:"computetime" json:"computetime"`
	Notes        string            `bson:"-" json:"-"`
	UserAttr     map[string]string `bson:"userattr" json:"userattr"`
}

func NewWorkunit(task *Task, rank int) *Workunit {
	return &Workunit{
		Id:        fmt.Sprintf("%s_%d", task.Id, rank),
		Info:      task.Info,
		Inputs:    task.Inputs,
		Outputs:   task.Outputs,
		Predata:   task.Predata,
		Cmd:       task.Cmd,
		Rank:      rank,
		TotalWork: task.TotalWork, //keep this info in workunit for load balancing
		Partition: task.Partition,
		State:     WORK_STAT_QUEUED,
		Failed:    0,
		UserAttr:  task.UserAttr,
	}
}

func (work *Workunit) Mkdir() (err error) {
	RemoveAll(work.Path()) // delete workdir just in case it exists

	err = os.MkdirAll(work.Path(), 0777)
	if err != nil {
		return
	}
	return
}

func (work *Workunit) RemoveDir() (err error) {
	err = os.RemoveAll(work.Path())
	if err != nil {
		return
	}
	return
}

func (work *Workunit) Path() string {
	id := work.Id
	return fmt.Sprintf("%s/%s/%s/%s/%s", conf.WORK_PATH, id[0:2], id[2:4], id[4:6], id)
}

func (work *Workunit) CDworkpath() (err error) {
	return os.Chdir(work.Path())
}

func (work *Workunit) IndexType() (indextype string) {
	return work.Partition.Index
}

//calculate the range of data part
//algorithm: try to evenly distribute indexed parts to workunits
//e.g. totalWork=4, totalParts=10, then each workunits have parts 3,3,2,2
func (work *Workunit) Part() (part string) {
	if work.Rank == 0 {
		return ""
	}
	partsize := work.Partition.TotalIndex / work.TotalWork //floor
	remainder := work.Partition.TotalIndex % work.TotalWork
	var start, end int
	if work.Rank <= remainder {
		start = (partsize+1)*(work.Rank-1) + 1
		end = start + partsize
	} else {
		start = (partsize+1)*remainder + partsize*(work.Rank-remainder-1) + 1
		end = start + partsize - 1
	}
	if start == end {
		part = fmt.Sprintf("%d", start)
	} else {
		part = fmt.Sprintf("%d-%d", start, end)
	}
	return
}
<commit_msg>rmdir workdir<commit_after>package core

import (
	"fmt"
	"github.com/MG-RAST/AWE/lib/conf"
	"os"
	"time"
)

const (
	WORK_STAT_QUEUED      = "queued"
	WORK_STAT_CHECKOUT    = "checkout"
	WORK_STAT_SUSPEND     = "suspend"
	WORK_STAT_DONE        = "done"
	WORK_STAT_FAIL        = "fail"
	WORK_STAT_PREPARED    = "prepared"
	WORK_STAT_COMPUTED    = "computed"
	WORK_STAT_DISCARDED   = "discarded"
	WORK_STAT_PROXYQUEUED = "proxyqueued"
)

type Workunit struct {
	Id           string            `bson:"wuid" json:"wuid"`
	Info         *Info             `bson:"info" json:"info"`
	Inputs       IOmap             `bson:"inputs" json:"inputs"`
	Outputs      IOmap             `bson:"outputs" json:"outputs"`
	Predata      IOmap             `bson:"predata" json:"predata"`
	Cmd          *Command          `bson:"cmd" json:"cmd"`
	Rank         int               `bson:"rank" json:"rank"`
	TotalWork    int               `bson:"totalwork" json:"totalwork"`
	Partition    *PartInfo         `bson:"part" json:"part"`
	State        string            `bson:"state" json:"state"`
	Failed       int               `bson:"failed" json:"failed"`
	CheckoutTime time.Time         `bson:"checkout_time" json:"checkout_time"`
	Client       string            `bson:"client" json:"client"`
	ComputeTime  int               `bson:"computetime" json:"computetime"`
	Notes        string            `bson:"-" json:"-"`
	UserAttr     map[string]string `bson:"userattr" json:"userattr"`
}

func NewWorkunit(task *Task, rank int) *Workunit {
	return &Workunit{
		Id:        fmt.Sprintf("%s_%d", task.Id, rank),
		Info:      task.Info,
		Inputs:    task.Inputs,
		Outputs:   task.Outputs,
		Predata:   task.Predata,
		Cmd:       task.Cmd,
		Rank:      rank,
		TotalWork: task.TotalWork, //keep this info in workunit for load balancing
		Partition: task.Partition,
		State:     WORK_STAT_QUEUED,
		Failed:    0,
		UserAttr:  task.UserAttr,
	}
}

func (work *Workunit) Mkdir() (err error) {
	os.RemoveAll(work.Path()) // delete workdir just in case it exists

	err = os.MkdirAll(work.Path(), 0777)
	if err != nil {
		return
	}
	return
}

func (work *Workunit) RemoveDir() (err error) {
	err = os.RemoveAll(work.Path())
	if err != nil {
		return
	}
	return
}

func (work *Workunit) Path() string {
	id := work.Id
	return fmt.Sprintf("%s/%s/%s/%s/%s", conf.WORK_PATH, id[0:2], id[2:4], id[4:6], id)
}

func (work *Workunit) CDworkpath() (err error) {
	return os.Chdir(work.Path())
}

func (work *Workunit) IndexType() (indextype string) {
	return work.Partition.Index
}

//calculate the range of data part
//algorithm: try to evenly distribute indexed parts to workunits
//e.g. totalWork=4, totalParts=10, then each workunits have parts 3,3,2,2
func (work *Workunit) Part() (part string) {
	if work.Rank == 0 {
		return ""
	}
	partsize := work.Partition.TotalIndex / work.TotalWork //floor
	remainder := work.Partition.TotalIndex % work.TotalWork
	var start, end int
	if work.Rank <= remainder {
		start = (partsize+1)*(work.Rank-1) + 1
		end = start + partsize
	} else {
		start = (partsize+1)*remainder + partsize*(work.Rank-remainder-1) + 1
		end = start + partsize - 1
	}
	if start == end {
		part = fmt.Sprintf("%d", start)
	} else {
		part = fmt.Sprintf("%d-%d", start, end)
	}
	return
}
<|endoftext|>