<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strconv"

	"time"

	"strings"

	. "github.com/claudetech/loggo/default"
	"github.com/mxk/go-flowrate/flowrate"
	"github.com/orcaman/concurrent-map"
)

var instances cmap.ConcurrentMap
var chunkPath string
var chunkSize int64
var chunkDirMaxSize int64
var speedLimit int64

func init() {
	instances = cmap.New()
}

// Buffer is a buffered stream
type Buffer struct {
	numberOfInstances int
	client            *http.Client
	object            *APIObject
	tempDir           string
	preload           bool
	chunks            cmap.ConcurrentMap
}

// GetBufferInstance gets a singleton instance of buffer
func GetBufferInstance(client *http.Client, object *APIObject) (*Buffer, error) {
	if !instances.Has(object.ObjectID) {
		i, err := newBuffer(client, object)
		if nil != err {
			return nil, err
		}

		instances.Set(object.ObjectID, i)
	}

	instance, ok := instances.Get(object.ObjectID)
	// if buffer allocation failed due to race conditions it will try to fetch a new one
	if !ok {
		i, err := GetBufferInstance(client, object)
		if nil != err {
			return nil, err
		}
		instance = i
	}
	instance.(*Buffer).numberOfInstances++
	return instance.(*Buffer), nil
}

// SetChunkPath sets the global chunk path
func SetChunkPath(path string) {
	chunkPath = path
}

// SetChunkSize sets the global chunk size
func SetChunkSize(size int64) {
	chunkSize = size
}

// SetChunkDirMaxSize sets the maximum size of the chunk directory
func SetChunkDirMaxSize(size int64) {
	chunkDirMaxSize = size
}

// SetDownloadSpeedLimit sets the download speed limit per chunk
func SetDownloadSpeedLimit(downloadSpeedLimit int64) {
	speedLimit = downloadSpeedLimit
}

// NewBuffer creates a new buffer instance
func newBuffer(client *http.Client, object *APIObject) (*Buffer, error) {
	Log.Infof("Starting playback of %v (%v)", object.ObjectID, object.Name)
	Log.Debugf("Creating buffer for object %v (%v)", object.ObjectID, object.Name)

	tempDir := filepath.Join(chunkPath, object.ObjectID)
	if err := os.MkdirAll(tempDir, 0777); nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not create temp path for object %v (%v)", object.ObjectID, object.Name)
	}

	if 0 == chunkSize {
		Log.Debugf("ChunkSize was 0, setting to default (5 MB)")
		chunkSize = 5 * 1024 * 1024
	}

	buffer := Buffer{
		numberOfInstances: 0,
		client:            client,
		object:            object,
		tempDir:           tempDir,
		preload:           true,
		chunks:            cmap.New(),
	}

	return &buffer, nil
}

// Close all handles
func (b *Buffer) Close() error {
	b.numberOfInstances--
	if 0 == b.numberOfInstances {
		Log.Infof("Stopping playback of %v (%v)", b.object.ObjectID, b.object.Name)
		Log.Debugf("Stop buffering for object %v (%v)", b.object.ObjectID, b.object.Name)

		b.preload = false
		instances.Remove(b.object.ObjectID)
	}
	return nil
}

// ReadBytes on a specific location
func (b *Buffer) ReadBytes(start, size int64, preload bool, delay int32) ([]byte, error) {
	fOffset := start % chunkSize
	offset := start - fOffset
	offsetEnd := offset + chunkSize

	Log.Tracef("Getting object %v (%v) - chunk %v - offset %v for %v bytes",
		b.object.ObjectID, b.object.Name, strconv.Itoa(int(offset)), fOffset, size)

	if !preload && b.preload && uint64(offsetEnd) < b.object.Size {
		defer func() {
			go func() {
				preloadStart := strconv.Itoa(int(offsetEnd))
				if !b.chunks.Has(preloadStart) {
					b.chunks.Set(preloadStart, true)
					b.ReadBytes(offsetEnd, size, true, 0)
				}
			}()
		}()
	}

	filename := filepath.Join(b.tempDir, strconv.Itoa(int(offset)))
	if f, err := os.Open(filename); nil == err {
		defer f.Close()

		buf := make([]byte, size)
		if n, err := f.ReadAt(buf, fOffset); n > 0 && (nil == err || io.EOF == err || io.ErrUnexpectedEOF == err) {
			Log.Tracef("Found file %s bytes %v - %v in cache", filename, offset, offsetEnd)

			// update the last modified time for files that are often in use
			if err := os.Chtimes(filename, time.Now(), time.Now()); nil != err {
				Log.Warningf("Could not update last modified time for %v", filename)
			}

			eOffset := int64(math.Min(float64(size), float64(len(buf))))
			return buf[:eOffset], nil
		}

		Log.Debugf("%v", err)
		Log.Debugf("Could not read file %s at %v", filename, fOffset)
	}

	if chunkDirMaxSize > 0 {
		go func() {
			if err := cleanChunkDir(chunkPath); nil != err {
				Log.Debugf("%v", err)
				Log.Warningf("Could not delete oldest chunk")
			}
		}()
	}

	// sleep if request is throttled
	if delay > 0 {
		time.Sleep(time.Duration(delay) * time.Second)
	}

	Log.Debugf("Requesting object %v (%v) bytes %v - %v from API", b.object.ObjectID, b.object.Name, offset, offsetEnd)
	req, err := http.NewRequest("GET", b.object.DownloadURL, nil)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not create request object %v (%v) from API", b.object.ObjectID, b.object.Name)
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%v-%v", offset, offsetEnd))

	Log.Tracef("Sending HTTP Request %v", req)

	res, err := b.client.Do(req)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not request object %v (%v) from API", b.object.ObjectID, b.object.Name)
	}
	defer res.Body.Close()

	reader := res.Body
	if speedLimit > 0 {
		reader = flowrate.NewReader(res.Body, speedLimit)
	}

	if res.StatusCode != 206 {
		if res.StatusCode != 403 {
			Log.Debugf("Request\n----------\n%v\n----------\n", req)
			Log.Debugf("Response\n----------\n%v\n----------\n", res)
			return nil, fmt.Errorf("Wrong status code %v", res.StatusCode)
		}

		// throttle requests
		if delay > 8 {
			return nil, fmt.Errorf("Maximum throttle interval has been reached")
		}
		bytes, err := ioutil.ReadAll(reader)
		if nil != err {
			Log.Debugf("%v", err)
			return nil, fmt.Errorf("Could not read body of 403 error")
		}
		body := string(bytes)
		if strings.Contains(body, "dailyLimitExceeded") ||
			strings.Contains(body, "userRateLimitExceeded") ||
			strings.Contains(body, "rateLimitExceeded") ||
			strings.Contains(body, "backendError") {
			if 0 == delay {
				delay = 1
			} else {
				delay = delay * 2
			}
			return b.ReadBytes(start, size, true, delay)
		}

		// return an error if other 403 error occurred
		Log.Debugf("%v", body)
		return nil, fmt.Errorf("Could not read file / StatusCode: %v", res.StatusCode)
	}

	bytes, err := ioutil.ReadAll(reader)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not read objects %v (%v) API response", b.object.ObjectID, b.object.Name)
	}

	if _, err := os.Stat(b.tempDir); os.IsNotExist(err) {
		if err := os.MkdirAll(b.tempDir, 0777); nil != err {
			Log.Debugf("%v", err)
			return nil, fmt.Errorf("Could not create chunk temp path for chunk %v", filename)
		}
	}

	if err := ioutil.WriteFile(filename, bytes, 0777); nil != err {
		Log.Debugf("%v", err)
		Log.Warningf("Could not write chunk temp file %v", filename)
	}

	sOffset := int64(math.Min(float64(fOffset), float64(len(bytes))))
	eOffset := int64(math.Min(float64(fOffset+size), float64(len(bytes))))
	return bytes[sOffset:eOffset], nil
}

// cleanChunkDir checks if the chunk folder is grown to big and clears the oldest file if necessary
func cleanChunkDir(chunkPath string) error {
	chunkDirSize, err := dirSize(chunkPath)
	if nil != err {
		return err
	}

	if chunkDirSize+chunkSize*2 > chunkDirMaxSize {
		if err := deleteOldestFile(chunkPath); nil != err {
			return err
		}
	}

	return nil
}

// deleteOldestFile deletes the oldest file in the directory
func deleteOldestFile(path string) error {
	var fpath string
	lastMod := time.Now()

	err := filepath.Walk(path, func(file string, info os.FileInfo, err error) error {
		if nil != err {
			Log.Tracef("%v", err)
			return filepath.SkipDir
		}
		if nil == info {
			return filepath.SkipDir
		}
		if !info.IsDir() {
			modTime := info.ModTime()
			if modTime.Before(lastMod) {
				lastMod = modTime
				fpath = file
			}
		}
		return nil
	})

	Log.Debugf("Deleting oldest chunk file %v", fpath)
	os.Remove(fpath)

	return err
}

// dirSize gets the total directory size
func dirSize(path string) (int64, error) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if nil != err {
			Log.Tracef("%v", err)
			return filepath.SkipDir
		}
		if nil == info {
			return filepath.SkipDir
		}
		if !info.IsDir() {
			size += info.Size()
		}
		return nil
	})
	return size, err
}
<commit_msg>cosmetics for debug logging<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strconv"

	"time"

	"strings"

	. "github.com/claudetech/loggo/default"
	"github.com/mxk/go-flowrate/flowrate"
	"github.com/orcaman/concurrent-map"
)

var instances cmap.ConcurrentMap
var chunkPath string
var chunkSize int64
var chunkDirMaxSize int64
var speedLimit int64

func init() {
	instances = cmap.New()
}

// Buffer is a buffered stream
type Buffer struct {
	numberOfInstances int
	client            *http.Client
	object            *APIObject
	tempDir           string
	preload           bool
	chunks            cmap.ConcurrentMap
}

// GetBufferInstance gets a singleton instance of buffer
func GetBufferInstance(client *http.Client, object *APIObject) (*Buffer, error) {
	if !instances.Has(object.ObjectID) {
		i, err := newBuffer(client, object)
		if nil != err {
			return nil, err
		}

		instances.Set(object.ObjectID, i)
	}

	instance, ok := instances.Get(object.ObjectID)
	// if buffer allocation failed due to race conditions it will try to fetch a new one
	if !ok {
		i, err := GetBufferInstance(client, object)
		if nil != err {
			return nil, err
		}
		instance = i
	}
	instance.(*Buffer).numberOfInstances++
	return instance.(*Buffer), nil
}

// SetChunkPath sets the global chunk path
func SetChunkPath(path string) {
	chunkPath = path
}

// SetChunkSize sets the global chunk size
func SetChunkSize(size int64) {
	chunkSize = size
}

// SetChunkDirMaxSize sets the maximum size of the chunk directory
func SetChunkDirMaxSize(size int64) {
	chunkDirMaxSize = size
}

// SetDownloadSpeedLimit sets the download speed limit per chunk
func SetDownloadSpeedLimit(downloadSpeedLimit int64) {
	speedLimit = downloadSpeedLimit
}

// NewBuffer creates a new buffer instance
func newBuffer(client *http.Client, object *APIObject) (*Buffer, error) {
	Log.Infof("Starting playback of %v (%v)", object.ObjectID, object.Name)
	Log.Debugf("Creating buffer for object %v (%v)", object.ObjectID, object.Name)

	tempDir := filepath.Join(chunkPath, object.ObjectID)
	if err := os.MkdirAll(tempDir, 0777); nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not create temp path for object %v (%v)", object.ObjectID, object.Name)
	}

	if 0 == chunkSize {
		Log.Debugf("ChunkSize was 0, setting to default (5 MB)")
		chunkSize = 5 * 1024 * 1024
	}

	buffer := Buffer{
		numberOfInstances: 0,
		client:            client,
		object:            object,
		tempDir:           tempDir,
		preload:           true,
		chunks:            cmap.New(),
	}

	return &buffer, nil
}

// Close all handles
func (b *Buffer) Close() error {
	b.numberOfInstances--
	if 0 == b.numberOfInstances {
		Log.Infof("Stopping playback of %v (%v)", b.object.ObjectID, b.object.Name)
		Log.Debugf("Stop buffering for object %v (%v)", b.object.ObjectID, b.object.Name)

		b.preload = false
		instances.Remove(b.object.ObjectID)
	}
	return nil
}

// ReadBytes on a specific location
func (b *Buffer) ReadBytes(start, size int64, preload bool, delay int32) ([]byte, error) {
	fOffset := start % chunkSize
	offset := start - fOffset
	offsetEnd := offset + chunkSize

	Log.Tracef("Getting object %v (%v) - chunk %v - offset %v for %v bytes",
		b.object.ObjectID, b.object.Name, strconv.Itoa(int(offset)), fOffset, size)

	if !preload && b.preload && uint64(offsetEnd) < b.object.Size {
		defer func() {
			go func() {
				preloadStart := strconv.Itoa(int(offsetEnd))
				if !b.chunks.Has(preloadStart) {
					b.chunks.Set(preloadStart, true)
					b.ReadBytes(offsetEnd, size, true, 0)
				}
			}()
		}()
	}

	filename := filepath.Join(b.tempDir, strconv.Itoa(int(offset)))
	if f, err := os.Open(filename); nil == err {
		defer f.Close()

		buf := make([]byte, size)
		if n, err := f.ReadAt(buf, fOffset); n > 0 && (nil == err || io.EOF == err || io.ErrUnexpectedEOF == err) {
			Log.Tracef("Found file %s bytes %v - %v in cache", filename, offset, offsetEnd)

			// update the last modified time for files that are often in use
			if err := os.Chtimes(filename, time.Now(), time.Now()); nil != err {
				Log.Warningf("Could not update last modified time for %v", filename)
			}

			eOffset := int64(math.Min(float64(size), float64(len(buf))))
			return buf[:eOffset], nil
		}

		Log.Debugf("%v", err)
		Log.Debugf("Could not read file %s at %v", filename, fOffset)
	}

	if chunkDirMaxSize > 0 {
		go func() {
			if err := cleanChunkDir(chunkPath); nil != err {
				Log.Debugf("%v", err)
				Log.Warningf("Could not delete oldest chunk")
			}
		}()
	}

	// sleep if request is throttled
	if delay > 0 {
		time.Sleep(time.Duration(delay) * time.Second)
	}

	Log.Debugf("Requesting object %v (%v) bytes %v - %v from API", b.object.ObjectID, b.object.Name, offset, offsetEnd)
	req, err := http.NewRequest("GET", b.object.DownloadURL, nil)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not create request object %v (%v) from API", b.object.ObjectID, b.object.Name)
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%v-%v", offset, offsetEnd))

	Log.Tracef("Sending HTTP Request %v", req)

	res, err := b.client.Do(req)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not request object %v (%v) from API", b.object.ObjectID, b.object.Name)
	}
	defer res.Body.Close()

	reader := res.Body
	if speedLimit > 0 {
		reader = flowrate.NewReader(res.Body, speedLimit)
	}

	if res.StatusCode != 206 {
		if res.StatusCode != 403 {
			Log.Debugf("Request\n----------\n%v\n----------\n", req)
			Log.Debugf("Response\n----------\n%v\n----------\n", res)
			return nil, fmt.Errorf("Wrong status code %v", res.StatusCode)
		}

		// throttle requests
		if delay > 8 {
			return nil, fmt.Errorf("Maximum throttle interval has been reached")
		}
		bytes, err := ioutil.ReadAll(reader)
		if nil != err {
			Log.Debugf("%v", err)
			return nil, fmt.Errorf("Could not read body of 403 error")
		}
		body := string(bytes)
		if strings.Contains(body, "dailyLimitExceeded") ||
			strings.Contains(body, "userRateLimitExceeded") ||
			strings.Contains(body, "rateLimitExceeded") ||
			strings.Contains(body, "backendError") {
			if 0 == delay {
				delay = 1
			} else {
				delay = delay * 2
			}
			return b.ReadBytes(start, size, true, delay)
		}

		// return an error if other 403 error occurred
		Log.Debugf("%v", body)
		return nil, fmt.Errorf("Could not read object %v (%v) / StatusCode: %v",
			b.object.ObjectID, b.object.Name, res.StatusCode)
	}

	bytes, err := ioutil.ReadAll(reader)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not read objects %v (%v) API response", b.object.ObjectID, b.object.Name)
	}

	if _, err := os.Stat(b.tempDir); os.IsNotExist(err) {
		if err := os.MkdirAll(b.tempDir, 0777); nil != err {
			Log.Debugf("%v", err)
			return nil, fmt.Errorf("Could not create chunk temp path for chunk %v", filename)
		}
	}

	if err := ioutil.WriteFile(filename, bytes, 0777); nil != err {
		Log.Debugf("%v", err)
		Log.Warningf("Could not write chunk temp file %v", filename)
	}

	sOffset := int64(math.Min(float64(fOffset), float64(len(bytes))))
	eOffset := int64(math.Min(float64(fOffset+size), float64(len(bytes))))
	return bytes[sOffset:eOffset], nil
}

// cleanChunkDir checks if the chunk folder is grown to big and clears the oldest file if necessary
func cleanChunkDir(chunkPath string) error {
	chunkDirSize, err := dirSize(chunkPath)
	if nil != err {
		return err
	}

	if chunkDirSize+chunkSize*2 > chunkDirMaxSize {
		if err := deleteOldestFile(chunkPath); nil != err {
			return err
		}
	}

	return nil
}

// deleteOldestFile deletes the oldest file in the directory
func deleteOldestFile(path string) error {
	var fpath string
	lastMod := time.Now()

	err := filepath.Walk(path, func(file string, info os.FileInfo, err error) error {
		if nil != err {
			Log.Tracef("%v", err)
			return filepath.SkipDir
		}
		if nil == info {
			return filepath.SkipDir
		}
		if !info.IsDir() {
			modTime := info.ModTime()
			if modTime.Before(lastMod) {
				lastMod = modTime
				fpath = file
			}
		}
		return nil
	})

	Log.Debugf("Deleting oldest chunk file %v", fpath)
	os.Remove(fpath)

	return err
}

// dirSize gets the total directory size
func dirSize(path string) (int64, error) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if nil != err {
			Log.Tracef("%v", err)
			return filepath.SkipDir
		}
		if nil == info {
			return filepath.SkipDir
		}
		if !info.IsDir() {
			size += info.Size()
		}
		return nil
	})
	return size, err
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"strings"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	pcidev "github.com/lxc/lxd/lxd/device/pci"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/resources"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type gpuMIG struct {
	deviceCommon
}

// GPUNvidiaDeviceKey is the key used for NVIDIA devices through libnvidia-container.
const GPUNvidiaDeviceKey = "nvidia.device"

// validateConfig checks the supplied config for correctness.
func (d *gpuMIG) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container, instancetype.VM) {
		return ErrUnsupportedDevType
	}

	requiredFields := []string{}

	optionalFields := []string{
		"vendorid",
		"productid",
		"id",
		"pci",
		"mig.gi",
		"mig.ci",
		"mig.uuid",
	}

	err := d.config.Validate(gpuValidationRules(requiredFields, optionalFields))
	if err != nil {
		return err
	}

	if d.config["pci"] != "" {
		for _, field := range []string{"id", "productid", "vendorid"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when when "pci" is set`, field)
			}
		}

		d.config["pci"] = pcidev.NormaliseAddress(d.config["pci"])
	}

	if d.config["id"] != "" {
		for _, field := range []string{"pci", "productid", "vendorid"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when when "id" is set`, field)
			}
		}
	}

	if d.config["mig.uuid"] != "" {
		for _, field := range []string{"mig.gi", "mig.ci"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when "mig.uuid" is set`, field)
			}
		}
	} else if d.config["mig.gi"] == "" || d.config["mig.ci"] == "" {
		return fmt.Errorf(`Either "mig.uuid" or both "mig.gi" and "mig.ci" must be set`)
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *gpuMIG) validateEnvironment() error {
	if !shared.IsTrue(d.inst.ExpandedConfig()["nvidia.runtime"]) {
		return fmt.Errorf("nvidia.runtime must be set to true for MIG GPUs to work")
	}

	return validatePCIDevice(d.config["pci"])
}

// buildMIGDeviceName builds the name of the MIG device based on old/new format
func (d *gpuMIG) buildMIGDeviceName(gpu api.ResourcesGPUCard) string {
	if d.config["mig.uuid"] != "" {
		if strings.HasPrefix(d.config["mig.uuid"], "MIG-") {
			return d.config["mig.uuid"]
		}
		return fmt.Sprintf("MIG-%s", d.config["mig.uuid"])
	}
	return fmt.Sprintf("MIG-%s/%s/%s", gpu.Nvidia.UUID, d.config["mig.gi"], d.config["mig.ci"])
}

// CanHotPlug returns whether the device can be managed whilst the instance is running,
func (d *gpuMIG) CanHotPlug() bool {
	return false
}

// Start is run when the device is added to the container.
func (d *gpuMIG) Start() (*deviceConfig.RunConfig, error) {
	// Check the basic config.
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}

	// Get all the GPUs.
	gpus, err := resources.GetGPU()
	if err != nil {
		return nil, err
	}

	var pciAddress string
	for _, gpu := range gpus.Cards {
		// Skip any cards that don't match the vendorid, pci or productid settings (if specified).
		if (d.config["vendorid"] != "" && gpu.VendorID != d.config["vendorid"]) ||
			(d.config["pci"] != "" && gpu.PCIAddress != d.config["pci"]) ||
			(d.config["productid"] != "" && gpu.ProductID != d.config["productid"]) {
			continue
		}

		// We found a match.
		if pciAddress != "" {
			return nil, fmt.Errorf("More than one GPU matched the MIG device")
		}
		pciAddress = gpu.PCIAddress

		// Validate the GPU.
		if gpu.Nvidia == nil {
			return nil, fmt.Errorf("Card isn't a NVIDIA GPU or driver isn't properly setup")
		}

		// Validate the MIG.
		fields := strings.SplitN(gpu.Nvidia.CardDevice, ":", 2)
		if len(fields) != 2 {
			return nil, fmt.Errorf("Bad NVIDIA GPU (couldn't find ID)")
		}
		gpuID := fields[1]

		if d.config["mig.uuid"] == "" {
			if !shared.PathExists(fmt.Sprintf("/proc/driver/nvidia/capabilities/gpu%s/mig/gi%s/ci%s/access", gpuID, d.config["mig.gi"], d.config["mig.ci"])) {
				return nil, fmt.Errorf("MIG device gi=%s ci=%s doesn't exist on GPU %s", d.config["mig.gi"], d.config["mig.ci"], gpuID)
			}
		}

		runConf.GPUDevice = append(runConf.GPUDevice, []deviceConfig.RunConfigItem{
			{Key: GPUNvidiaDeviceKey, Value: d.buildMIGDeviceName(gpu)},
		}...)
	}

	if pciAddress == "" {
		return nil, fmt.Errorf("Failed to detect requested GPU device")
	}

	return &runConf, nil
}

// Stop is run when the device is removed from the instance.
func (d *gpuMIG) Stop() (*deviceConfig.RunConfig, error) {
	return nil, nil
}
<commit_msg>lxd/device: Removes VM from list of supported instance types.<commit_after>package device

import (
	"fmt"
	"strings"

	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	pcidev "github.com/lxc/lxd/lxd/device/pci"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/resources"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

type gpuMIG struct {
	deviceCommon
}

// GPUNvidiaDeviceKey is the key used for NVIDIA devices through libnvidia-container.
const GPUNvidiaDeviceKey = "nvidia.device"

// validateConfig checks the supplied config for correctness.
func (d *gpuMIG) validateConfig(instConf instance.ConfigReader) error {
	if !instanceSupported(instConf.Type(), instancetype.Container) {
		return ErrUnsupportedDevType
	}

	requiredFields := []string{}

	optionalFields := []string{
		"vendorid",
		"productid",
		"id",
		"pci",
		"mig.gi",
		"mig.ci",
		"mig.uuid",
	}

	err := d.config.Validate(gpuValidationRules(requiredFields, optionalFields))
	if err != nil {
		return err
	}

	if d.config["pci"] != "" {
		for _, field := range []string{"id", "productid", "vendorid"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when when "pci" is set`, field)
			}
		}

		d.config["pci"] = pcidev.NormaliseAddress(d.config["pci"])
	}

	if d.config["id"] != "" {
		for _, field := range []string{"pci", "productid", "vendorid"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when when "id" is set`, field)
			}
		}
	}

	if d.config["mig.uuid"] != "" {
		for _, field := range []string{"mig.gi", "mig.ci"} {
			if d.config[field] != "" {
				return fmt.Errorf(`Cannot use %q when "mig.uuid" is set`, field)
			}
		}
	} else if d.config["mig.gi"] == "" || d.config["mig.ci"] == "" {
		return fmt.Errorf(`Either "mig.uuid" or both "mig.gi" and "mig.ci" must be set`)
	}

	return nil
}

// validateEnvironment checks the runtime environment for correctness.
func (d *gpuMIG) validateEnvironment() error {
	if !shared.IsTrue(d.inst.ExpandedConfig()["nvidia.runtime"]) {
		return fmt.Errorf("nvidia.runtime must be set to true for MIG GPUs to work")
	}

	return validatePCIDevice(d.config["pci"])
}

// buildMIGDeviceName builds the name of the MIG device based on old/new format
func (d *gpuMIG) buildMIGDeviceName(gpu api.ResourcesGPUCard) string {
	if d.config["mig.uuid"] != "" {
		if strings.HasPrefix(d.config["mig.uuid"], "MIG-") {
			return d.config["mig.uuid"]
		}
		return fmt.Sprintf("MIG-%s", d.config["mig.uuid"])
	}
	return fmt.Sprintf("MIG-%s/%s/%s", gpu.Nvidia.UUID, d.config["mig.gi"], d.config["mig.ci"])
}

// CanHotPlug returns whether the device can be managed whilst the instance is running,
func (d *gpuMIG) CanHotPlug() bool {
	return false
}

// Start is run when the device is added to the container.
func (d *gpuMIG) Start() (*deviceConfig.RunConfig, error) {
	// Check the basic config.
	err := d.validateEnvironment()
	if err != nil {
		return nil, err
	}

	runConf := deviceConfig.RunConfig{}

	// Get all the GPUs.
	gpus, err := resources.GetGPU()
	if err != nil {
		return nil, err
	}

	var pciAddress string
	for _, gpu := range gpus.Cards {
		// Skip any cards that don't match the vendorid, pci or productid settings (if specified).
		if (d.config["vendorid"] != "" && gpu.VendorID != d.config["vendorid"]) ||
			(d.config["pci"] != "" && gpu.PCIAddress != d.config["pci"]) ||
			(d.config["productid"] != "" && gpu.ProductID != d.config["productid"]) {
			continue
		}

		// We found a match.
		if pciAddress != "" {
			return nil, fmt.Errorf("More than one GPU matched the MIG device")
		}
		pciAddress = gpu.PCIAddress

		// Validate the GPU.
		if gpu.Nvidia == nil {
			return nil, fmt.Errorf("Card isn't a NVIDIA GPU or driver isn't properly setup")
		}

		// Validate the MIG.
		fields := strings.SplitN(gpu.Nvidia.CardDevice, ":", 2)
		if len(fields) != 2 {
			return nil, fmt.Errorf("Bad NVIDIA GPU (couldn't find ID)")
		}
		gpuID := fields[1]

		if d.config["mig.uuid"] == "" {
			if !shared.PathExists(fmt.Sprintf("/proc/driver/nvidia/capabilities/gpu%s/mig/gi%s/ci%s/access", gpuID, d.config["mig.gi"], d.config["mig.ci"])) {
				return nil, fmt.Errorf("MIG device gi=%s ci=%s doesn't exist on GPU %s", d.config["mig.gi"], d.config["mig.ci"], gpuID)
			}
		}

		runConf.GPUDevice = append(runConf.GPUDevice, []deviceConfig.RunConfigItem{
			{Key: GPUNvidiaDeviceKey, Value: d.buildMIGDeviceName(gpu)},
		}...)
	}

	if pciAddress == "" {
		return nil, fmt.Errorf("Failed to detect requested GPU device")
	}

	return &runConf, nil
}

// Stop is run when the device is removed from the instance.
func (d *gpuMIG) Stop() (*deviceConfig.RunConfig, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string `mapstructure:"inline"`

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string `mapstructure:"script"`

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// An array of multiple scripts to run.
	Scripts []string `mapstructure:"scripts"`

	TargetPath string `mapstructure:"target"`

	tpl *packer.ConfigTemplate
}

type ShellPostProcessor struct {
	cfg Config
}

type OutputPathTemplate struct {
	ArtifactId string
	BuildName  string
	Provider   string
}

func (p *ShellPostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.cfg, raws...)
	if err != nil {
		return err
	}

	errs := new(packer.MultiError)

	if p.cfg.InlineShebang == "" {
		p.cfg.InlineShebang = "/bin/sh"
	}

	if p.cfg.Scripts == nil {
		p.cfg.Scripts = make([]string, 0)
	}

	if p.cfg.Vars == nil {
		p.cfg.Vars = make([]string, 0)
	}

	if p.cfg.Script != "" && len(p.cfg.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.cfg.Script != "" {
		p.cfg.Scripts = []string{p.cfg.Script}
	}

	p.cfg.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.cfg.tpl.UserVars = p.cfg.PackerUserVars

	if p.cfg.TargetPath == "" {
		p.cfg.TargetPath = "packer_{{ .BuildName }}_{{.Provider}}"
	}

	if err = p.cfg.tpl.Validate(p.cfg.TargetPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"inline_shebang": &p.cfg.InlineShebang,
		"script":         &p.cfg.Script,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.cfg.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.cfg.Inline,
		"scripts":          p.cfg.Scripts,
		"environment_vars": p.cfg.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.cfg.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.cfg.Scripts) == 0 && p.cfg.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.cfg.Scripts) > 0 && p.cfg.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.cfg.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.cfg.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *ShellPostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	scripts := make([]string, len(p.cfg.Scripts))
	copy(scripts, p.cfg.Scripts)

	if p.cfg.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.cfg.InlineShebang))
		for _, command := range p.cfg.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	envVars := make([]string, len(p.cfg.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.cfg.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.cfg.PackerBuilderType
	copy(envVars[2:], p.cfg.Vars)

	for _, artifact := range artifact.Files() {

		for _, path := range scripts {
			ui.Say(fmt.Sprintf("Process with shell script: %s", path))

			log.Printf("Opening %s for reading", path)
			f, err := os.Open(path)
			if err != nil {
				return nil, false, fmt.Errorf("Error opening shell script: %s", err)
			}
			defer f.Close()

			ui.Message(fmt.Sprintf("Executing script with artifact: %s", artifact))
			args := []string{path, artifact}
			cmd := exec.Command("/bin/sh", args...)
			var buffer bytes.Buffer
			cmd.Stdout = &buffer
			cmd.Stderr = &buffer
			cmd.Env = envVars
			err = cmd.Run()
			if err != nil {
				return nil, false, fmt.Errorf("Unable to execute script: %s", buffer.String())
			}
			ui.Message(fmt.Sprintf("%s", buffer.String()))
		}
	}
	return artifact, false, nil
}
<commit_msg>not override standard variable<commit_after>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string `mapstructure:"inline"`

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string `mapstructure:"script"`

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// An array of multiple scripts to run.
	Scripts []string `mapstructure:"scripts"`

	TargetPath string `mapstructure:"target"`

	tpl *packer.ConfigTemplate
}

type ShellPostProcessor struct {
	cfg Config
}

type OutputPathTemplate struct {
	ArtifactId string
	BuildName  string
	Provider   string
}

func (p *ShellPostProcessor) Configure(raws ...interface{}) error {
	_, err := common.DecodeConfig(&p.cfg, raws...)
	if err != nil {
		return err
	}

	errs := new(packer.MultiError)

	if p.cfg.InlineShebang == "" {
		p.cfg.InlineShebang = "/bin/sh"
	}

	if p.cfg.Scripts == nil {
		p.cfg.Scripts = make([]string, 0)
	}

	if p.cfg.Vars == nil {
		p.cfg.Vars = make([]string, 0)
	}

	if p.cfg.Script != "" && len(p.cfg.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.cfg.Script != "" {
		p.cfg.Scripts = []string{p.cfg.Script}
	}

	p.cfg.tpl, err = packer.NewConfigTemplate()
	if err != nil {
		return err
	}
	p.cfg.tpl.UserVars = p.cfg.PackerUserVars

	if p.cfg.TargetPath == "" {
		p.cfg.TargetPath = "packer_{{ .BuildName }}_{{.Provider}}"
	}

	if err = p.cfg.tpl.Validate(p.cfg.TargetPath); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"inline_shebang": &p.cfg.InlineShebang,
		"script":         &p.cfg.Script,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.cfg.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.cfg.Inline,
		"scripts":          p.cfg.Scripts,
		"environment_vars": p.cfg.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.cfg.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.cfg.Scripts) == 0 && p.cfg.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.cfg.Scripts) > 0 && p.cfg.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.cfg.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.cfg.Vars {
		vs := strings.SplitN(kv, "=", 2)
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *ShellPostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {
	scripts := make([]string, len(p.cfg.Scripts))
	copy(scripts, p.cfg.Scripts)

	if p.cfg.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.cfg.InlineShebang))
		for _, command := range p.cfg.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return nil, false, fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	envVars := make([]string, len(p.cfg.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.cfg.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.cfg.PackerBuilderType
	copy(envVars[2:], p.cfg.Vars)

	for _, art := range artifact.Files() {

		for _, path := range scripts {
			ui.Say(fmt.Sprintf("Process with shell script: %s", path))

			log.Printf("Opening %s for reading", path)
			f, err := os.Open(path)
			if err != nil {
				return nil, false, fmt.Errorf("Error opening shell script: %s", err)
			}
			defer f.Close()

			ui.Message(fmt.Sprintf("Executing script with artifact: %s", art))
			args := []string{path, art}
			cmd := exec.Command("/bin/sh", args...)
			var buffer bytes.Buffer
			cmd.Stdout = &buffer
			cmd.Stderr = &buffer
			cmd.Env = envVars
			err = cmd.Run()
			if err != nil {
				return nil, false, fmt.Errorf("Unable to execute script: %s", buffer.String())
			}
			ui.Message(fmt.Sprintf("%s", buffer.String()))
		}
	}
	return artifact, false, nil
}
<|endoftext|>


<commit_before>package node

import (
	"errors"
	"github.com/MG-RAST/Shock/shock-server/conf"
	"github.com/MG-RAST/Shock/shock-server/db"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"io/ioutil"
	"path/filepath"
	"strings"
)

// mongodb has hard limit of 16 MB docuemnt size
var DocumentMaxByte = 16777216

// Initialize creates a copy of the mongodb connection and then uses that connection to
// create the Nodes collection in mongodb. Then, it ensures that there is a unique index
// on the id key in this collection, creating the index if necessary.
func Initialize() {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	c.EnsureIndex(mgo.Index{Key: []string{"acl.owner"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.read"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.write"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.delete"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"created_on"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"expiration"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"type"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"priority"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.path"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.virtual_parts"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.checksum.md5"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"id"}, Unique: true})
	if conf.MONGODB_ATTRIBUTE_INDEXES != "" {
		for _, v := range strings.Split(conf.MONGODB_ATTRIBUTE_INDEXES, ",") {
			v = "attributes." + strings.TrimSpace(v)
			c.EnsureIndex(mgo.Index{Key: []string{v}, Background: true})
		}
	}
}

func HasAttributeField(a string) bool {
	for _, b := range strings.Split(conf.MONGODB_ATTRIBUTE_INDEXES, ",") {
		if a == strings.TrimSpace(b) {
			return true
		}
	}
	return false
}

func dbDelete(q bson.M) (err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	_, err = c.RemoveAll(q)
	return
}

func dbUpsert(n *Node) (err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	_, err = c.Upsert(bson.M{"id": n.Id}, &n)
	return
}

func dbFind(q bson.M, results *Nodes, order string, options map[string]int) (count int, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	if order == "" {
		order = "created_on"
	}
	if limit, has := options["limit"]; has {
		if offset, has := options["offset"]; has {
			query := c.Find(q).Sort(order)
			count, err = query.Count()
			if err != nil {
				count = 0
				return
			}
			err = query.Limit(limit).Skip(offset).All(results)
		} else {
			err = errors.New("store.db.Find options limit and offset must be used together")
		}
	}
	if err != nil {
		return
	}
	err = c.Find(q).Sort(order).All(results)
	if err != nil {
		return
	}
	results.DBInit()
	return
}

func DbFindDistinct(q bson.M, d string) (results interface{}, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	err = c.Find(q).Distinct("attributes."+d, &results)
	return
}

func Load(id string) (n *Node, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	n = new(Node)
	err = c.Find(bson.M{"id": id}).One(&n)
	if err != nil {
		n = nil
		return
	}
	n.DBInit()
	return
}

func LoadNodes(ids []string) (n Nodes, err error) {
	if _, err = dbFind(bson.M{"id": bson.M{"$in": ids}}, &n, "", nil); err == nil {
		return n, err
	}
	return nil, err
}

func ReloadFromDisk(path string) (err error) {
	id := filepath.Base(path)
	nbson, err := ioutil.ReadFile(path + "/" + id + ".bson")
	if err != nil {
		return
	}
	node := new(Node)
	if err = bson.Unmarshal(nbson, &node); err == nil {
		if err = dbUpsert(node); err != nil {
			return err
		}
	}
	return
}
<commit_msg>bug fix - missing else clause in query<commit_after>package node

import (
	"errors"
	"github.com/MG-RAST/Shock/shock-server/conf"
	"github.com/MG-RAST/Shock/shock-server/db"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"io/ioutil"
	"path/filepath"
	"strings"
)

// mongodb has hard limit of 16 MB docuemnt size
var DocumentMaxByte = 16777216

// Initialize creates a copy of the mongodb connection and then uses that connection to
// create the Nodes collection in mongodb. Then, it ensures that there is a unique index
// on the id key in this collection, creating the index if necessary.
func Initialize() {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	c.EnsureIndex(mgo.Index{Key: []string{"acl.owner"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.read"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.write"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"acl.delete"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"created_on"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"expiration"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"type"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"priority"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.path"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.virtual_parts"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"file.checksum.md5"}, Background: true})
	c.EnsureIndex(mgo.Index{Key: []string{"id"}, Unique: true})
	if conf.MONGODB_ATTRIBUTE_INDEXES != "" {
		for _, v := range strings.Split(conf.MONGODB_ATTRIBUTE_INDEXES, ",") {
			v = "attributes." + strings.TrimSpace(v)
			c.EnsureIndex(mgo.Index{Key: []string{v}, Background: true})
		}
	}
}

func HasAttributeField(a string) bool {
	for _, b := range strings.Split(conf.MONGODB_ATTRIBUTE_INDEXES, ",") {
		if a == strings.TrimSpace(b) {
			return true
		}
	}
	return false
}

func dbDelete(q bson.M) (err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	_, err = c.RemoveAll(q)
	return
}

func dbUpsert(n *Node) (err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	_, err = c.Upsert(bson.M{"id": n.Id}, &n)
	return
}

func dbFind(q bson.M, results *Nodes, order string, options map[string]int) (count int, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	if order == "" {
		order = "created_on"
	}
	if limit, has := options["limit"]; has {
		if offset, has := options["offset"]; has {
			query := c.Find(q).Sort(order)
			count, err = query.Count()
			if err != nil {
				count = 0
				return
			}
			err = query.Limit(limit).Skip(offset).All(results)
			if err != nil {
				return
			}
		} else {
			err = errors.New("store.db.Find options limit and offset must be used together")
			return
		}
	} else {
		err = c.Find(q).Sort(order).All(results)
		if err != nil {
			return
		}
	}
	results.DBInit()
	return
}

func DbFindDistinct(q bson.M, d string) (results interface{}, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	err = c.Find(q).Distinct("attributes."+d, &results)
	return
}

func Load(id string) (n *Node, err error) {
	session := db.Connection.Session.Copy()
	defer session.Close()
	c := session.DB(conf.MONGODB_DATABASE).C("Nodes")
	n = new(Node)
	err = c.Find(bson.M{"id": id}).One(&n)
	if err != nil {
		n = nil
		return
	}
	n.DBInit()
	return
}

func LoadNodes(ids []string) (n Nodes, err error) {
	if _, err = dbFind(bson.M{"id": bson.M{"$in": ids}}, &n, "", nil); err == nil {
		return n, err
	}
	return nil, err
}

func ReloadFromDisk(path string) (err error) {
	id := filepath.Base(path)
	nbson, err := ioutil.ReadFile(path + "/" + id + ".bson")
	if err != nil {
		return
	}
	node := new(Node)
	if err = bson.Unmarshal(nbson, &node); err == nil {
		if err = dbUpsert(node); err != nil {
			return err
		}
	}
	return
}
<|endoftext|>


<commit_before>package regression

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"

	"github.com/golang/protobuf/proto"
	"github.com/taskgraph/taskgraph"
	pb "github.com/taskgraph/taskgraph/example/regression/proto"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// dummyMaster is prototype of parameter server, for now it does not
// carry out optimization yet. But it should be easy to add support when
// this full tests out.
// Note: in theory, since there should be no parent of this, so we should
// add error checking in the right places. We will skip these test for now.
type dummyMaster struct {
	dataChan           chan int32
	NodeProducer       chan bool
	framework          taskgraph.Framework
	epoch, taskID      uint64
	logger             *log.Logger
	config             map[string]string
	numberOfIterations uint64

	param        *pb.Parameter
	gradient     *pb.Gradient
	fromChildren map[uint64]*pb.Gradient

	epochChange    chan *event
	getP           chan *event
	getG           chan *event
	gradientReady  chan *event
	childDataReady chan *event
	getGReqs       []*event
	exitChan       chan struct{}
}

type event struct {
	ctx      context.Context
	epoch    uint64
	input    *pb.Input
	retP     chan *pb.Parameter
	retG     chan *pb.Gradient
	gradient *pb.Gradient
	fromID   uint64
	output   proto.Message
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummyMaster) Init(taskID uint64, framework taskgraph.Framework) {
	t.taskID = taskID
	t.framework = framework
	t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)

	t.epochChange = make(chan *event, 1)
	t.getP = make(chan *event, 1)
	t.getG = make(chan *event, 1)
	t.childDataReady = make(chan *event, 1)
	t.exitChan = make(chan struct{})
	go t.run()
}

func (t *dummyMaster) run() {
	for {
		select {
		case et := <-t.epochChange:
			for _, gG := range t.getGReqs {
				close(gG.retG)
			}
			t.getGReqs = nil
			t.enterEpoch(et.ctx, et.epoch)
		case gP := <-t.getP:
			// We have to check epoch here in user level because grpc doesn't
			// allow use to intercept messages. This should be fixed later.
			err := t.framework.CheckEpoch(gP.input.Epoch)
			if err != nil {
				close(gP.retP)
			}
			gP.retP <- t.param
		case gG := <-t.getG:
			err := t.framework.CheckEpoch(gG.input.Epoch)
			if err != nil {
				close(gG.retG)
			}
			if t.gradient != nil {
				gG.retG <- t.gradient
				break
			}
			// Waiting queue. Requests will get notified later. The number of request
			// won't be huge presumingly.
			t.getGReqs = append(t.getGReqs, gG)
		case <-t.gradientReady:
			for _, gG := range t.getGReqs {
				gG.retG <- t.gradient
			}
		case cr := <-t.childDataReady:
			t.ChildDataReady(cr.ctx, cr.fromID, cr.output)
		case <-t.exitChan:
			return
		}
	}
}

func (t *dummyMaster) Exit() {
	close(t.exitChan)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) EnterEpoch(ctx context.Context, epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.testablyFail("SetEpoch", strconv.FormatUint(epoch, 10)) {
		return
	}
	t.epochChange <- &event{ctx: ctx, epoch: epoch}
}

func (t *dummyMaster) enterEpoch(ctx context.Context, epoch uint64) {
	t.param = new(pb.Parameter)
	t.gradient = nil
	t.fromChildren = make(map[uint64]*pb.Gradient)
	t.gradientReady = make(chan *event, 1)

	t.epoch = epoch
	t.param.Value = int32(t.epoch)
	for _, c := range t.framework.GetTopology().GetNeighbors("Children", t.epoch) {
		t.framework.DataRequest(ctx, c, "/proto.Regression/GetGradient", &pb.Input{t.epoch})
	}
}

// These are payload rpc for application purpose.
func (t *dummyMaster) GetParameter(ctx context.Context, input *pb.Input) (*pb.Parameter, error) {
	retP := make(chan *pb.Parameter, 1)
	t.getP <- &event{ctx: ctx, input: input, retP: retP}
	p, ok := <-retP
	if !ok {
		return nil, fmt.Errorf("epoch changed")
	}
	return p, nil
}

func (t *dummyMaster) GetGradient(ctx context.Context, input *pb.Input) (*pb.Gradient, error) {
	retG := make(chan *pb.Gradient, 1)
	t.getP <- &event{ctx: ctx, input: input, retG: retG}
	g, ok := <-retG
	if !ok {
		return nil, fmt.Errorf("epoch changed")
	}
	return g, nil
}
func (t *dummyMaster) DataReady(ctx context.Context, fromID uint64, method string, output proto.Message) {
	if method == "/proto.Regression/GetGradient" {
		t.childDataReady <- &event{ctx: ctx, fromID: fromID, output: output}
		return
	}
	panic("")
}

func (t *dummyMaster) gradReady() {
	close(t.gradientReady)
	// In testing, we need to make sure dataChan has enough space and don't block.
	t.dataChan <- t.gradient.Value
	// In real ML, we modify the gradient first. But here it is noop.
	if t.epoch == t.numberOfIterations {
		if t.config["writefile"] != "" {
			data := []byte(fmt.Sprintf("Finished job. Gradient value: %v\n", t.gradient.Value))
			ioutil.WriteFile(t.config["writefile"], data, 0644)
		}
		t.framework.ShutdownJob()
	} else {
		t.logger.Printf("master finished current epoch, task: %d, epoch: %d", t.taskID, t.epoch)
		t.framework.IncEpoch(ctx)
	}
}

func (t *dummyMaster) ChildDataReady(ctx context.Context, childID uint64, output proto.Message) {
	d, ok := output.(*pb.Gradient)
	if !ok {
		t.logger.Fatalf("Can't convert proto message to Gradient: %v", output)
	}
	t.fromChildren[childID] = d

	t.logger.Printf("master ChildDataReady, task: %d, epoch: %d, child: %d, ready: %d\n",
		t.taskID, t.epoch, childID, len(t.fromChildren))

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetNeighbors("Children", t.epoch)) {
		t.gradient = new(pb.Gradient)
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}
		t.gradReady()
	}
}

func (t *dummyMaster) CreateOutputMessage(methodName string) proto.Message {
	switch methodName {
	case "/proto.Regression/GetParameter":
		return new(pb.Parameter)
	case "/proto.Regression/GetGradient":
		return new(pb.Gradient)
	default:
		t.logger.Panicf("Unknown method: %s", methodName)
		return nil
	}
}

func (t *dummyMaster) CreateServer() *grpc.Server {
	server := grpc.NewServer()
	pb.RegisterRegressionServer(server, t)
	return server
}

func (t *dummyMaster) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if len(args) >= 1 && t.config["failepoch"] != "" {
		// we need to care about fail at specific epoch
		if t.config["failepoch"] != args[0] {
			return false
		}
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("master task %d testably fail, method: %s\n", t.taskID, method)
	// Very hack. Need some internal knowledge. Don't change this.
	t.Exit()
	t.framework.Kill()
	t.NodeProducer <- true
	return true
}

func (t *dummyMaster) MetaReady(ctx context.Context, fromID uint64, linkType, meta string) {}
<commit_msg>fix master gradient ready<commit_after>package regression

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"

	"github.com/golang/protobuf/proto"
	"github.com/taskgraph/taskgraph"
	pb "github.com/taskgraph/taskgraph/example/regression/proto"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// dummyMaster is prototype of parameter server, for now it does not
// carry out optimization yet. But it should be easy to add support when
// this full tests out.
// Note: in theory, since there should be no parent of this, so we should
// add error checking in the right places. We will skip these test for now.
type dummyMaster struct {
	dataChan           chan int32
	NodeProducer       chan bool
	framework          taskgraph.Framework
	epoch, taskID      uint64
	logger             *log.Logger
	config             map[string]string
	numberOfIterations uint64

	param        *pb.Parameter
	gradient     *pb.Gradient
	fromChildren map[uint64]*pb.Gradient

	epochChange    chan *event
	getP           chan *event
	getG           chan *event
	childDataReady chan *event
	getGReqs       []*event
	exitChan       chan struct{}
}

type event struct {
	ctx      context.Context
	epoch    uint64
	input    *pb.Input
	retP     chan *pb.Parameter
	retG     chan *pb.Gradient
	gradient *pb.Gradient
	fromID   uint64
	output   proto.Message
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummyMaster) Init(taskID uint64, framework taskgraph.Framework) {
	t.taskID = taskID
	t.framework = framework
	t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)

	t.epochChange = make(chan *event, 1)
	t.getP = make(chan *event, 1)
	t.getG = make(chan *event, 1)
	t.childDataReady = make(chan *event, 1)
	t.exitChan = make(chan struct{})
	go t.run()
}

func (t *dummyMaster) run() {
	for {
		select {
		case et := <-t.epochChange:
			for _, gG := range t.getGReqs {
				close(gG.retG)
			}
			t.getGReqs = nil
			t.enterEpoch(et.ctx, et.epoch)
		case gP := <-t.getP:
			// We have to check epoch here in user level because grpc doesn't
			// allow use to intercept messages. This should be fixed later.
			err := t.framework.CheckEpoch(gP.input.Epoch)
			if err != nil {
				close(gP.retP)
			}
			gP.retP <- t.param
		case gG := <-t.getG:
			err := t.framework.CheckEpoch(gG.input.Epoch)
			if err != nil {
				close(gG.retG)
			}
			if t.gradient != nil {
				gG.retG <- t.gradient
				break
			}
			// Waiting queue. Requests will get notified later. The number of request
			// won't be huge presumingly.
			t.getGReqs = append(t.getGReqs, gG)
		case cr := <-t.childDataReady:
			t.ChildDataReady(cr.ctx, cr.fromID, cr.output)
		case <-t.exitChan:
			return
		}
	}
}

func (t *dummyMaster) Exit() {
	close(t.exitChan)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) EnterEpoch(ctx context.Context, epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.testablyFail("SetEpoch", strconv.FormatUint(epoch, 10)) {
		return
	}
	t.epochChange <- &event{ctx: ctx, epoch: epoch}
}

func (t *dummyMaster) enterEpoch(ctx context.Context, epoch uint64) {
	t.param = new(pb.Parameter)
	t.gradient = nil
	t.fromChildren = make(map[uint64]*pb.Gradient)

	t.epoch = epoch
	t.param.Value = int32(t.epoch)
	for _, c := range t.framework.GetTopology().GetNeighbors("Children", t.epoch) {
		t.framework.DataRequest(ctx, c, "/proto.Regression/GetGradient", &pb.Input{t.epoch})
	}
}

// These are payload rpc for application purpose.
func (t *dummyMaster) GetParameter(ctx context.Context, input *pb.Input) (*pb.Parameter, error) {
	retP := make(chan *pb.Parameter, 1)
	t.getP <- &event{ctx: ctx, input: input, retP: retP}
	p, ok := <-retP
	if !ok {
		return nil, fmt.Errorf("epoch changed")
	}
	return p, nil
}

func (t *dummyMaster) GetGradient(ctx context.Context, input *pb.Input) (*pb.Gradient, error) {
	retG := make(chan *pb.Gradient, 1)
	t.getP <- &event{ctx: ctx, input: input, retG: retG}
	g, ok := <-retG
	if !ok {
		return nil, fmt.Errorf("epoch changed")
	}
	return g, nil
}
func (t *dummyMaster) DataReady(ctx context.Context, fromID uint64, method string, output proto.Message) {
	if method == "/proto.Regression/GetGradient" {
		t.childDataReady <- &event{ctx: ctx, fromID: fromID, output: output}
		return
	}
	panic("")
}

func (t *dummyMaster) gradientReady(ctx context.Context) {
	for _, gG := range t.getGReqs {
		gG.retG <- t.gradient
	}
	t.getGReqs = nil
	// In testing, we need to make sure dataChan has enough space and don't block.
	t.dataChan <- t.gradient.Value
	// In real ML, we modify the gradient first. But here it is noop.
	if t.epoch == t.numberOfIterations {
		if t.config["writefile"] != "" {
			data := []byte(fmt.Sprintf("Finished job. Gradient value: %v\n", t.gradient.Value))
			ioutil.WriteFile(t.config["writefile"], data, 0644)
		}
		t.framework.ShutdownJob()
	} else {
		t.logger.Printf("master finished current epoch, task: %d, epoch: %d", t.taskID, t.epoch)
		t.framework.IncEpoch(ctx)
	}
}

func (t *dummyMaster) ChildDataReady(ctx context.Context, childID uint64, output proto.Message) {
	d, ok := output.(*pb.Gradient)
	if !ok {
		t.logger.Fatalf("Can't convert proto message to Gradient: %v", output)
	}
	t.fromChildren[childID] = d

	t.logger.Printf("master ChildDataReady, task: %d, epoch: %d, child: %d, ready: %d\n",
		t.taskID, t.epoch, childID, len(t.fromChildren))

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetNeighbors("Children", t.epoch)) {
		t.gradient = new(pb.Gradient)
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}
		t.gradientReady(ctx)
	}
}

func (t *dummyMaster) CreateOutputMessage(methodName string) proto.Message {
	switch methodName {
	case "/proto.Regression/GetParameter":
		return new(pb.Parameter)
	case "/proto.Regression/GetGradient":
		return new(pb.Gradient)
	default:
		t.logger.Panicf("Unknown method: %s", methodName)
		return nil
	}
}

func (t *dummyMaster) CreateServer() *grpc.Server {
	server := grpc.NewServer()
	pb.RegisterRegressionServer(server, t)
	return server
}

func (t *dummyMaster) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if len(args) >= 1 && t.config["failepoch"] != "" {
		// we need to care about fail at specific epoch
		if t.config["failepoch"] != args[0] {
			return false
		}
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("master task %d testably fail, method: %s\n", t.taskID, method)
	// Very hack. Need some internal knowledge. Don't change this.
	t.Exit()
	t.framework.Kill()
	t.NodeProducer <- true
	return true
}

func (t *dummyMaster) MetaReady(ctx context.Context, fromID uint64, linkType, meta string) {}
<|endoftext|>


<commit_before>package panos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/scottdware/go-rested"
	"strings"
)

// AddressObjects contains a slice of all address objects.
type AddressObjects struct {
	XMLName   xml.Name  `xml:"response"`
	Status    string    `xml:"status,attr"`
	Code      string    `xml:"code,attr"`
	Addresses []Address `xml:"result>address>entry"`
}

// Address contains information about each individual address object.
type Address struct {
	Name        string `xml:"name,attr"`
	IPAddress   string `xml:"ip-netmask,omitempty"`
	IPRange     string `xml:"ip-range,omitempty"`
	FQDN        string `xml:"fqdn,omitempty"`
	Description string `xml:"description,omitempty"`
}

// AddressGroups contains a slice of all address groups.
type AddressGroups struct {
	Groups []AddressGroup
}

// AddressGroup contains information about each individual address group.
type AddressGroup struct {
	Name          string
	Type          string
	Members       []string
	DynamicFilter string
	Description   string
}

// xmlAddressGroups is used for parsing of all address groups.
type xmlAddressGroups struct {
	XMLName xml.Name          `xml:"response"`
	Status  string            `xml:"status,attr"`
	Code    string            `xml:"code,attr"`
	Groups  []xmlAddressGroup `xml:"result>address-group>entry"`
}

// xmlAddressGroup is used for parsing each individual address group.
type xmlAddressGroup struct {
	Name          string   `xml:"name,attr"`
	Members       []string `xml:"static>member,omitempty"`
	DynamicFilter string   `xml:"dynamic>filter,omitempty"`
	Description   string   `xml:"description,omitempty"`
}

// Addresses returns information about all of the address objects. When run against a Panorama device,
// addresses from all device-groups are returned.
func (p *PaloAlto) Addresses() (*AddressObjects, error) {
	var addrs AddressObjects
	xpath := "/config/devices/entry//address"
	// xpath := "/config/devices/entry/vsys/entry/address"
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == true {
		xpath = "/config/panorama//address"
	}

	if p.DeviceType == "panorama" {
		// xpath = "/config/devices/entry/device-group/entry/address"
		xpath = "/config/devices/entry//address"
	}

	query := map[string]string{
		"type":   "config",
		"action": "get",
		"xpath":  xpath,
		"key":    p.Key,
	}
	addrData := r.Send("get", p.URI, nil, headers, query)

	if err := xml.Unmarshal(addrData.Body, &addrs); err != nil {
		return nil, err
	}

	if addrs.Status != "success" {
		return nil, fmt.Errorf("error code %s: %s", addrs.Code, errorCodes[addrs.Code])
	}

	return &addrs, nil
}

// AddressGroups returns information about all of the address groups. When run against a Panorama device,
// address groups from all device-groups are returned.
func (p *PaloAlto) AddressGroups() (*AddressGroups, error) {
	var parsedGroups xmlAddressGroups
	var groups AddressGroups
	xpath := "/config/devices/entry//address-group"
	// xpath := "/config/devices/entry/vsys/entry/address-group"
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == true {
		xpath = "/config/panorama//address-group"
	}

	if p.DeviceType == "panorama" {
		// xpath = "/config/devices/entry/device-group/entry/address-group"
		xpath = "/config/devices/entry//address-group"
	}

	query := map[string]string{
		"type":   "config",
		"action": "get",
		"xpath":  xpath,
		"key":    p.Key,
	}
	groupData := r.Send("get", p.URI, nil, headers, query)

	if err := xml.Unmarshal(groupData.Body, &parsedGroups); err != nil {
		return nil, err
	}

	if parsedGroups.Status != "success" {
		return nil, fmt.Errorf("error code %s: %s", parsedGroups.Code, errorCodes[parsedGroups.Code])
	}

	for _, g := range parsedGroups.Groups {
		gname := g.Name
		gtype := "Static"
		gmembers := g.Members
		gfilter := strings.TrimSpace(g.DynamicFilter)
		gdesc := g.Description

		if g.DynamicFilter != "" {
			gtype = "Dynamic"
		}

		groups.Groups = append(groups.Groups, AddressGroup{Name: gname, Type: gtype, Members: gmembers, DynamicFilter: gfilter, Description: gdesc})
	}

	return &groups, nil
}

// CreateAddress will add a new address object to the device. addrtype should be one of: ip, range, or fqdn. If creating
// an address object on a Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) CreateAddress(name, addrtype, address, description string, devicegroup ...string) error {
	var xmlBody string
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	switch addrtype {
	case "ip":
		xmlBody = fmt.Sprintf("<ip-netmask>%s</ip-netmask>", address)
	case "range":
		xmlBody = fmt.Sprintf("<ip-range>%s</ip-range>", address)
	case "fqdn":
		xmlBody = fmt.Sprintf("<fqdn>%s</fqdn>", address)
	}

	if description != "" {
		xmlBody += fmt.Sprintf("<description>%s</description>", description)
	}

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":    "config",
		"action":  "set",
		"xpath":   xpath,
		"element": xmlBody,
		"key":     p.Key,
	}

	resp := r.Send("post", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}

// DeleteAddress will remove an address object from the device. If deleting an address object on a
// Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) DeleteAddress(name string, devicegroup ...string) error {
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":   "config",
		"action": "delete",
		"xpath":  xpath,
		"key":    p.Key,
	}

	resp := r.Send("get", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}
<commit_msg>Added create/delete address group functions<commit_after>package panos

import (
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/scottdware/go-rested"
	"strings"
)

// AddressObjects contains a slice of all address objects.
type AddressObjects struct {
	XMLName   xml.Name  `xml:"response"`
	Status    string    `xml:"status,attr"`
	Code      string    `xml:"code,attr"`
	Addresses []Address `xml:"result>address>entry"`
}

// Address contains information about each individual address object.
type Address struct {
	Name        string `xml:"name,attr"`
	IPAddress   string `xml:"ip-netmask,omitempty"`
	IPRange     string `xml:"ip-range,omitempty"`
	FQDN        string `xml:"fqdn,omitempty"`
	Description string `xml:"description,omitempty"`
}

// AddressGroups contains a slice of all address groups.
type AddressGroups struct {
	Groups []AddressGroup
}

// AddressGroup contains information about each individual address group.
type AddressGroup struct {
	Name          string
	Type          string
	Members       []string
	DynamicFilter string
	Description   string
}

// xmlAddressGroups is used for parsing of all address groups.
type xmlAddressGroups struct {
	XMLName xml.Name          `xml:"response"`
	Status  string            `xml:"status,attr"`
	Code    string            `xml:"code,attr"`
	Groups  []xmlAddressGroup `xml:"result>address-group>entry"`
}

// xmlAddressGroup is used for parsing each individual address group.
type xmlAddressGroup struct {
	Name          string   `xml:"name,attr"`
	Members       []string `xml:"static>member,omitempty"`
	DynamicFilter string   `xml:"dynamic>filter,omitempty"`
	Description   string   `xml:"description,omitempty"`
}

// Addresses returns information about all of the address objects. When run against a Panorama device,
// addresses from all device-groups are returned.
func (p *PaloAlto) Addresses() (*AddressObjects, error) {
	var addrs AddressObjects
	xpath := "/config/devices/entry//address"
	// xpath := "/config/devices/entry/vsys/entry/address"
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == true {
		xpath = "/config/panorama//address"
	}

	if p.DeviceType == "panorama" {
		// xpath = "/config/devices/entry/device-group/entry/address"
		xpath = "/config/devices/entry//address"
	}

	query := map[string]string{
		"type":   "config",
		"action": "get",
		"xpath":  xpath,
		"key":    p.Key,
	}
	addrData := r.Send("get", p.URI, nil, headers, query)

	if err := xml.Unmarshal(addrData.Body, &addrs); err != nil {
		return nil, err
	}

	if addrs.Status != "success" {
		return nil, fmt.Errorf("error code %s: %s", addrs.Code, errorCodes[addrs.Code])
	}

	return &addrs, nil
}

// AddressGroups returns information about all of the address groups. When run against a Panorama device,
// address groups from all device-groups are returned.
func (p *PaloAlto) AddressGroups() (*AddressGroups, error) {
	var parsedGroups xmlAddressGroups
	var groups AddressGroups
	xpath := "/config/devices/entry//address-group"
	// xpath := "/config/devices/entry/vsys/entry/address-group"
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == true {
		xpath = "/config/panorama//address-group"
	}

	if p.DeviceType == "panorama" {
		// xpath = "/config/devices/entry/device-group/entry/address-group"
		xpath = "/config/devices/entry//address-group"
	}

	query := map[string]string{
		"type":   "config",
		"action": "get",
		"xpath":  xpath,
		"key":    p.Key,
	}
	groupData := r.Send("get", p.URI, nil, headers, query)

	if err := xml.Unmarshal(groupData.Body, &parsedGroups); err != nil {
		return nil, err
	}

	if parsedGroups.Status != "success" {
		return nil, fmt.Errorf("error code %s: %s", parsedGroups.Code, errorCodes[parsedGroups.Code])
	}

	for _, g := range parsedGroups.Groups {
		gname := g.Name
		gtype := "Static"
		gmembers := g.Members
		gfilter := strings.TrimSpace(g.DynamicFilter)
		gdesc := g.Description

		if g.DynamicFilter != "" {
			gtype = "Dynamic"
		}

		groups.Groups = append(groups.Groups, AddressGroup{Name: gname, Type: gtype, Members: gmembers, DynamicFilter: gfilter, Description: gdesc})
	}

	return &groups, nil
}

// CreateAddress will add a new address object to the device. addrtype should be one of: ip, range, or fqdn. If creating
// an address object on a Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) CreateAddress(name, addrtype, address, description string, devicegroup ...string) error {
	var xmlBody string
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	switch addrtype {
	case "ip":
		xmlBody = fmt.Sprintf("<ip-netmask>%s</ip-netmask>", address)
	case "range":
		xmlBody = fmt.Sprintf("<ip-range>%s</ip-range>", address)
	case "fqdn":
		xmlBody = fmt.Sprintf("<fqdn>%s</fqdn>", address)
	}

	if description != "" {
		xmlBody += fmt.Sprintf("<description>%s</description>", description)
	}

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":    "config",
		"action":  "set",
		"xpath":   xpath,
		"element": xmlBody,
		"key":     p.Key,
	}

	resp := r.Send("post", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}

// CreateStaticGroup will create a new static address group on the device. If creating an address group on
// a Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) CreateStaticGroup(name string, members []string, description string, devicegroup ...string) error {
	var xmlBody string
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	xmlBody = "<static>"

	for _, m := range members {
		xmlBody += fmt.Sprintf("<member>%s</member>", m)
	}

	xmlBody += "</static>"

	if description != "" {
		xmlBody += fmt.Sprintf("<description>%s</description>", description)
	}

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address-group/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address-group/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":    "config",
		"action":  "set",
		"xpath":   xpath,
		"element": xmlBody,
		"key":     p.Key,
	}

	resp := r.Send("post", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}

// CreateDynamicGroup will create a new dynamic address group on the device. The filter must be written like so:
// 'vm-servers' and 'some tag' or 'pcs' - using the tags as the match criteria. If creating an address group on a
// Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) CreateDynamicGroup(name, criteria, description string, devicegroup ...string) error {
	xmlBody := fmt.Sprintf("<dynamic><filter>%s</filter></dynamic>", criteria)
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	if description != "" {
		xmlBody += fmt.Sprintf("<description>%s</description>", description)
	}

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address-group/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address-group/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":    "config",
		"action":  "set",
		"xpath":   xpath,
		"element": xmlBody,
		"key":     p.Key,
	}

	resp := r.Send("post", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}

// DeleteAddress will remove an address object from the device. If deleting an address object on a
// Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) DeleteAddress(name string, devicegroup ...string) error {
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":   "config",
		"action": "delete",
		"xpath":  xpath,
		"key":    p.Key,
	}

	resp := r.Send("get", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}

// DeleteAddressGroup will remove an address group from the device. If deleting an address group on a
// Panorama device, then specify the given device-group name as the last parameter.
func (p *PaloAlto) DeleteAddressGroup(name string, devicegroup ...string) error {
	var xpath string
	var reqError requestError
	r := rested.NewRequest()

	if p.DeviceType == "panos" && p.Panorama == false {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/vsys/entry[@name='vsys1']/address-group/entry[@name='%s']", name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) > 0 {
		xpath = fmt.Sprintf("/config/devices/entry[@name='localhost.localdomain']/device-group/entry[@name='%s']/address-group/entry[@name='%s']", devicegroup[0], name)
	}

	if p.DeviceType == "panorama" && len(devicegroup) <= 0 {
		return errors.New("you must specify a device-group when connected to a Panorama device")
	}

	query := map[string]string{
		"type":   "config",
		"action": "delete",
		"xpath":  xpath,
		"key":    p.Key,
	}

	resp := r.Send("get", p.URI, nil, nil, query)
	if resp.Error != nil {
		return resp.Error
	}

	if err := xml.Unmarshal(resp.Body, &reqError); err != nil {
		return err
	}

	if reqError.Status != "success" {
		return fmt.Errorf("error code %s: %s", reqError.Code, errorCodes[reqError.Code])
	}

	return nil
}
<|endoftext|>


<commit_before>package signalfx

import (
	"math"
	"sync"
	"sync/atomic"

	"github.com/zvelo/go-signalfx/sfxproto"
)

var (
	// BucketMetricCount represents the count of datapoints seen
	BucketMetricCount = 1
	// BucketMetricMin represents the smallest datapoint seen
	BucketMetricMin = 2
	// BucketMetricMax represents the largest datapoint seen
	BucketMetricMax = 3
	// BucketMetricSum represents the sum of all datapoints seen
	BucketMetricSum = 4
	// BucketMetricSumOfSquares represents the sum of squares of all datapoints seen
	BucketMetricSumOfSquares = 5
)

// A Bucket trakcs groups of values, reporting metrics as gauges and
// resetting each time it reports. All operations on Buckets are goroutine safe.
type Bucket struct {
	metric       string
	dimensions   map[string]string
	count        uint64
	min          int64
	max          int64
	sum          int64
	sumOfSquares int64
	mu           sync.Mutex

	DisabledMetrics map[int]bool
}

func (b *Bucket) lock() {
	b.mu.Lock()
}

func (b *Bucket) unlock() {
	b.mu.Unlock()
}

// Clone makes a deep copy of a Bucket
func (b *Bucket) Clone() *Bucket {
	// use the mutex to keep all operations part of the same transaction
	b.lock()
	defer b.unlock()

	return &Bucket{
		metric:          b.metric,                                  // can't use Metric() since we already have a lock
		dimensions:      sfxproto.Dimensions(b.dimensions).Clone(), // can't use Dimensions() since we already have a lock
		count:           b.Count(),
		min:             b.Min(),
		max:             b.Max(),
		sum:             b.Sum(),
		sumOfSquares:    b.SumOfSquares(),
		DisabledMetrics: make(map[int]bool, 0),
	}
}

// Disable disables the given metrics for this bucket.
// They will be collected, but not reported.
func (b *Bucket) Disable(metrics ...int) {
	b.lock()
	defer b.unlock()
	for _, metric := range metrics {
		b.DisabledMetrics[metric] = true
	}
}

// Equal returns whether two buckets are exactly equal
func (b *Bucket) Equal(r *Bucket) bool {
	// lock the state of both buckets as this operation is effectively a
	// transaction on an exact state of a whole bucket

	b.lock()
	defer b.unlock()

	r.lock()
	defer r.unlock()

	if b.metric != r.metric {
		return false
	}

	if !sfxproto.Dimensions(b.dimensions).Equal(r.dimensions) {
		return false
	}

	if b.Count() != r.Count() {
		return false
	}

	if b.Min() != r.Min() {
		return false
	}

	if b.Max() != r.Max() {
		return false
	}

	if b.Sum() != r.Sum() {
		return false
	}

	if b.SumOfSquares() != r.SumOfSquares() {
		return false
	}

	return true
}

// Metric returns the metric name of the Bucket
func (b *Bucket) Metric() string {
	b.lock()
	defer b.unlock()

	return b.metric
}

// SetMetric sets the metric name of the Bucket
func (b *Bucket) SetMetric(name string) {
	b.lock()
	defer b.unlock()

	b.metric = name
}

// Dimensions returns a copy of the dimensions of the Bucket. Changes are not
// reflected inside the Bucket itself.
func (b *Bucket) Dimensions() map[string]string {
	b.lock()
	defer b.unlock()

	return sfxproto.Dimensions(b.dimensions).Clone()
}

// SetDimension adds or overwrites the dimension at key with value. If the key
// or value is empty, no changes are made
func (b *Bucket) SetDimension(key, value string) {
	if key == "" || value == "" {
		return
	}

	b.lock()
	defer b.unlock()

	b.dimensions[key] = value
}

// SetDimensions adds or overwrites multiple dimensions. Because the passed in
// dimensions can not be locked by this method, it is important that the caller
// ensures its state does not change for the duration of the operation.
func (b *Bucket) SetDimensions(dims map[string]string) {
	for key, value := range dims {
		b.SetDimension(key, value)
	}
}

// RemoveDimension removes one or more dimensions with the given keys
func (b *Bucket) RemoveDimension(keys ...string) {
	b.lock()
	defer b.unlock()

	for _, key := range keys {
		delete(b.dimensions, key)
	}
}

// Count returns the number of items added to the Bucket
func (b *Bucket) Count() uint64 {
	return atomic.LoadUint64(&b.count)
}

// Min returns the lowest item added to the Bucket
func (b *Bucket) Min() int64 {
	return atomic.LoadInt64(&b.min)
}

// Max returns the highest item added to the Bucket
func (b *Bucket) Max() int64 {
	return atomic.LoadInt64(&b.max)
}

// Sum returns the sum of all items added to the Bucket
func (b *Bucket) Sum() int64 {
	return atomic.LoadInt64(&b.sum)
}

// SumOfSquares returns the sum of the square of all items added to the Bucket
func (b *Bucket) SumOfSquares() int64 {
	return atomic.LoadInt64(&b.sumOfSquares)
}

// NewBucket creates a new Bucket. Because the passed in dimensions can not be
// locked by this method, it is important that the caller ensures its state does
// not change for the duration of the operation.
func NewBucket(metric string, dimensions map[string]string) *Bucket {
	return &Bucket{
		metric:     metric,
		dimensions: sfxproto.Dimensions(dimensions).Clone(),
		min:        math.MaxInt64,
		max:        math.MinInt64,
	}
}

// Add an item to the Bucket, later reporting the result in the next report
// cycle.
func (b *Bucket) Add(val int64) {
	// wrap all of these changes in the mutex lock since they need to occur as a
	// transaction, not a set of atomic operations
	b.lock()
	defer b.unlock()

	// still use atomic though, so that the atomic "getters" will never read
	// from an inconsistent state
	_ = atomic.AddUint64(&b.count, 1)
	atomic.AddInt64(&b.sum, val)
	atomic.AddInt64(&b.sumOfSquares, val*val)

	if b.Min() > val {
		atomic.StoreInt64(&b.min, val)
	}

	if b.Max() < val {
		atomic.StoreInt64(&b.max, val)
	}
}

func (b *Bucket) dimFor(defaultDims map[string]string, rollup string) map[string]string {
	b.lock()
	defer b.unlock()

	dims := sfxproto.Dimensions(defaultDims).Append(b.dimensions)
	dims["rollup"] = rollup

	return dims
}

// CountDataPoint returns a DataPoint representing the Bucket's Count. Because
// the passed in dimensions can not be locked by this method, it is important
// that the caller ensures its state does not change for the duration of the
// operation.
func (b *Bucket) CountDataPoint(defaultDims map[string]string) *DataPoint {
	cnt := atomic.SwapUint64(&b.count, 0)
	dp, _ := NewGauge(b.Metric(), cnt, b.dimFor(defaultDims, "count"))
	return dp
}

// SumDataPoint returns a DataPoint representing the Bucket's Sum. Because the
// passed in dimensions can not be locked by this method, it is important that
// the caller ensures its state does not change for the duration of the
// operation.
func (b *Bucket) SumDataPoint(defaultDims map[string]string) *DataPoint {
	sum := atomic.SwapInt64(&b.sum, 0)
	dp, _ := NewGauge(b.Metric(), sum, b.dimFor(defaultDims, "sum"))
	return dp
}

// SumOfSquaresDataPoint returns a DataPoint representing the Bucket's
// SumOfSquares. Because the passed in dimensions can not be locked by this
// method, it is important that the caller ensures its state does not change for
// the duration of the operation.
func (b *Bucket) SumOfSquaresDataPoint(defaultDims map[string]string) *DataPoint {
	sps := atomic.SwapInt64(&b.sumOfSquares, 0)
	dp, _ := NewGauge(b.Metric(), sps, b.dimFor(defaultDims, "sumsquare"))
	return dp
}

// MinDataPoint returns a DataPoint representing the Bucket's Min. Note that
// this resets the Min value. nil is returned if no items have been added to the
// bucket since it was created or last reset. Because the passed in dimensions
// can not be locked by this method, it is important that the caller ensures its
// state does not change for the duration of the operation.
func (b *Bucket) MinDataPoint(defaultDims map[string]string) *DataPoint {
	min := atomic.SwapInt64(&b.min, math.MaxInt64)

	if b.Count() != 0 && min != math.MaxInt64 {
		dp, _ := NewGauge(b.Metric(), min, b.dimFor(defaultDims, "min"))
		return dp
	}

	return nil
}

// MaxDataPoint returns a DataPoint representing the Bucket's Max. Note that
// this resets the Max value. nil is returned if no items have been added to the
// bucket since it was created or last reset. Because the passed in dimensions
// can not be locked by this method, it is important that the caller ensures its
// state does not change for the duration of the operation.
func (b *Bucket) MaxDataPoint(defaultDims map[string]string) *DataPoint {
	max := atomic.SwapInt64(&b.max, math.MinInt64)

	if b.Count() != 0 && max != math.MinInt64 {
		dp, _ := NewGauge(b.Metric(), max, b.dimFor(defaultDims, "max"))
		return dp
	}

	return nil
}

// DataPoints returns a DataPoints object with DataPoint values for Count, Sum,
// SumOfSquares, Min and Max (if set). Note that this resets all values.
// Because the passed in dimensions can not be locked by this method, it is
// important that the caller ensures its state does not change for the duration
// of the operation.
func (b *Bucket) DataPoints(defaultDims map[string]string) *DataPoints {
	dp := NewDataPoints(0)
	if disabled := b.DisabledMetrics[BucketMetricCount]; !disabled {
		dp = dp.Add(b.CountDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricSum]; !disabled {
		dp = dp.Add(b.SumDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricMin]; !disabled {
		dp = dp.Add(b.MinDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricMax]; !disabled {
		dp = dp.Add(b.MaxDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricSumOfSquares]; !disabled {
		dp = dp.Add(b.SumOfSquaresDataPoint(defaultDims))
	}

	return dp
}
<commit_msg>Fixed bucket nil map<commit_after>package signalfx

import (
	"math"
	"sync"
	"sync/atomic"

	"github.com/zvelo/go-signalfx/sfxproto"
)

var (
	// BucketMetricCount represents the count of datapoints seen
	BucketMetricCount = 1
	// BucketMetricMin represents the smallest datapoint seen
	BucketMetricMin = 2
	// BucketMetricMax represents the largest datapoint seen
	BucketMetricMax = 3
	// BucketMetricSum represents the sum of all datapoints seen
	BucketMetricSum = 4
	// BucketMetricSumOfSquares represents the sum of squares of all datapoints seen
	BucketMetricSumOfSquares = 5
)

// A Bucket trakcs groups of values, reporting metrics as gauges and
// resetting each time it reports. All operations on Buckets are goroutine safe.
type Bucket struct {
	metric       string
	dimensions   map[string]string
	count        uint64
	min          int64
	max          int64
	sum          int64
	sumOfSquares int64
	mu           sync.Mutex

	DisabledMetrics map[int]bool
}

func (b *Bucket) lock() {
	b.mu.Lock()
}

func (b *Bucket) unlock() {
	b.mu.Unlock()
}

// Clone makes a deep copy of a Bucket
func (b *Bucket) Clone() *Bucket {
	// use the mutex to keep all operations part of the same transaction
	b.lock()
	defer b.unlock()

	return &Bucket{
		metric:          b.metric,                                  // can't use Metric() since we already have a lock
		dimensions:      sfxproto.Dimensions(b.dimensions).Clone(), // can't use Dimensions() since we already have a lock
		count:           b.Count(),
		min:             b.Min(),
		max:             b.Max(),
		sum:             b.Sum(),
		sumOfSquares:    b.SumOfSquares(),
		DisabledMetrics: b.DisabledMetrics,
	}
}

// Disable disables the given metrics for this bucket.
// They will be collected, but not reported.
func (b *Bucket) Disable(metrics ...int) {
	b.lock()
	defer b.unlock()
	for _, metric := range metrics {
		b.DisabledMetrics[metric] = true
	}
}

// Equal returns whether two buckets are exactly equal
func (b *Bucket) Equal(r *Bucket) bool {
	// lock the state of both buckets as this operation is effectively a
	// transaction on an exact state of a whole bucket

	b.lock()
	defer b.unlock()

	r.lock()
	defer r.unlock()

	if b.metric != r.metric {
		return false
	}

	if !sfxproto.Dimensions(b.dimensions).Equal(r.dimensions) {
		return false
	}

	if b.Count() != r.Count() {
		return false
	}

	if b.Min() != r.Min() {
		return false
	}

	if b.Max() != r.Max() {
		return false
	}

	if b.Sum() != r.Sum() {
		return false
	}

	if b.SumOfSquares() != r.SumOfSquares() {
		return false
	}

	return true
}

// Metric returns the metric name of the Bucket
func (b *Bucket) Metric() string {
	b.lock()
	defer b.unlock()

	return b.metric
}

// SetMetric sets the metric name of the Bucket
func (b *Bucket) SetMetric(name string) {
	b.lock()
	defer b.unlock()

	b.metric = name
}

// Dimensions returns a copy of the dimensions of the Bucket. Changes are not
// reflected inside the Bucket itself.
func (b *Bucket) Dimensions() map[string]string {
	b.lock()
	defer b.unlock()

	return sfxproto.Dimensions(b.dimensions).Clone()
}

// SetDimension adds or overwrites the dimension at key with value. If the key
// or value is empty, no changes are made
func (b *Bucket) SetDimension(key, value string) {
	if key == "" || value == "" {
		return
	}

	b.lock()
	defer b.unlock()

	b.dimensions[key] = value
}

// SetDimensions adds or overwrites multiple dimensions. Because the passed in
// dimensions can not be locked by this method, it is important that the caller
// ensures its state does not change for the duration of the operation.
func (b *Bucket) SetDimensions(dims map[string]string) {
	for key, value := range dims {
		b.SetDimension(key, value)
	}
}

// RemoveDimension removes one or more dimensions with the given keys
func (b *Bucket) RemoveDimension(keys ...string) {
	b.lock()
	defer b.unlock()

	for _, key := range keys {
		delete(b.dimensions, key)
	}
}

// Count returns the number of items added to the Bucket
func (b *Bucket) Count() uint64 {
	return atomic.LoadUint64(&b.count)
}

// Min returns the lowest item added to the Bucket
func (b *Bucket) Min() int64 {
	return atomic.LoadInt64(&b.min)
}

// Max returns the highest item added to the Bucket
func (b *Bucket) Max() int64 {
	return atomic.LoadInt64(&b.max)
}

// Sum returns the sum of all items added to the Bucket
func (b *Bucket) Sum() int64 {
	return atomic.LoadInt64(&b.sum)
}

// SumOfSquares returns the sum of the square of all items added to the Bucket
func (b *Bucket) SumOfSquares() int64 {
	return atomic.LoadInt64(&b.sumOfSquares)
}

// NewBucket creates a new Bucket. Because the passed in dimensions can not be
// locked by this method, it is important that the caller ensures its state does
// not change for the duration of the operation.
func NewBucket(metric string, dimensions map[string]string) *Bucket {
	return &Bucket{
		metric:          metric,
		dimensions:      sfxproto.Dimensions(dimensions).Clone(),
		min:             math.MaxInt64,
		max:             math.MinInt64,
		DisabledMetrics: make(map[int]bool, 0),
	}
}

// Add an item to the Bucket, later reporting the result in the next report
// cycle.
func (b *Bucket) Add(val int64) {
	// wrap all of these changes in the mutex lock since they need to occur as a
	// transaction, not a set of atomic operations
	b.lock()
	defer b.unlock()

	// still use atomic though, so that the atomic "getters" will never read
	// from an inconsistent state
	_ = atomic.AddUint64(&b.count, 1)
	atomic.AddInt64(&b.sum, val)
	atomic.AddInt64(&b.sumOfSquares, val*val)

	if b.Min() > val {
		atomic.StoreInt64(&b.min, val)
	}

	if b.Max() < val {
		atomic.StoreInt64(&b.max, val)
	}
}

func (b *Bucket) dimFor(defaultDims map[string]string, rollup string) map[string]string {
	b.lock()
	defer b.unlock()

	dims := sfxproto.Dimensions(defaultDims).Append(b.dimensions)
	dims["rollup"] = rollup

	return dims
}

// CountDataPoint returns a DataPoint representing the Bucket's Count. Because
// the passed in dimensions can not be locked by this method, it is important
// that the caller ensures its state does not change for the duration of the
// operation.
func (b *Bucket) CountDataPoint(defaultDims map[string]string) *DataPoint {
	cnt := atomic.SwapUint64(&b.count, 0)
	dp, _ := NewGauge(b.Metric(), cnt, b.dimFor(defaultDims, "count"))
	return dp
}

// SumDataPoint returns a DataPoint representing the Bucket's Sum. Because the
// passed in dimensions can not be locked by this method, it is important that
// the caller ensures its state does not change for the duration of the
// operation.
func (b *Bucket) SumDataPoint(defaultDims map[string]string) *DataPoint {
	sum := atomic.SwapInt64(&b.sum, 0)
	dp, _ := NewGauge(b.Metric(), sum, b.dimFor(defaultDims, "sum"))
	return dp
}

// SumOfSquaresDataPoint returns a DataPoint representing the Bucket's
// SumOfSquares. Because the passed in dimensions can not be locked by this
// method, it is important that the caller ensures its state does not change for
// the duration of the operation.
func (b *Bucket) SumOfSquaresDataPoint(defaultDims map[string]string) *DataPoint {
	sps := atomic.SwapInt64(&b.sumOfSquares, 0)
	dp, _ := NewGauge(b.Metric(), sps, b.dimFor(defaultDims, "sumsquare"))
	return dp
}

// MinDataPoint returns a DataPoint representing the Bucket's Min. Note that
// this resets the Min value. nil is returned if no items have been added to the
// bucket since it was created or last reset. Because the passed in dimensions
// can not be locked by this method, it is important that the caller ensures its
// state does not change for the duration of the operation.
func (b *Bucket) MinDataPoint(defaultDims map[string]string) *DataPoint {
	min := atomic.SwapInt64(&b.min, math.MaxInt64)

	if b.Count() != 0 && min != math.MaxInt64 {
		dp, _ := NewGauge(b.Metric(), min, b.dimFor(defaultDims, "min"))
		return dp
	}

	return nil
}

// MaxDataPoint returns a DataPoint representing the Bucket's Max. Note that
// this resets the Max value. nil is returned if no items have been added to the
// bucket since it was created or last reset. Because the passed in dimensions
// can not be locked by this method, it is important that the caller ensures its
// state does not change for the duration of the operation.
func (b *Bucket) MaxDataPoint(defaultDims map[string]string) *DataPoint {
	max := atomic.SwapInt64(&b.max, math.MinInt64)

	if b.Count() != 0 && max != math.MinInt64 {
		dp, _ := NewGauge(b.Metric(), max, b.dimFor(defaultDims, "max"))
		return dp
	}

	return nil
}

// DataPoints returns a DataPoints object with DataPoint values for Count, Sum,
// SumOfSquares, Min and Max (if set). Note that this resets all values.
// Because the passed in dimensions can not be locked by this method, it is
// important that the caller ensures its state does not change for the duration
// of the operation.
func (b *Bucket) DataPoints(defaultDims map[string]string) *DataPoints {
	dp := NewDataPoints(0)
	if disabled := b.DisabledMetrics[BucketMetricCount]; !disabled {
		dp = dp.Add(b.CountDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricSum]; !disabled {
		dp = dp.Add(b.SumDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricMin]; !disabled {
		dp = dp.Add(b.MinDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricMax]; !disabled {
		dp = dp.Add(b.MaxDataPoint(defaultDims))
	}
	if disabled := b.DisabledMetrics[BucketMetricSumOfSquares]; !disabled {
		dp = dp.Add(b.SumOfSquaresDataPoint(defaultDims))
	}

	return dp
}
<|endoftext|>