<commit_before>package drivers

import (
	"fmt"
	"os"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/storage/locking"
	"github.com/lxc/lxd/shared"
)

// tmpVolSuffix Suffix to use for any temporary volumes created by LXD.
const tmpVolSuffix = ".lxdtmp"

// defaultBlockSize Default size of block volumes.
const defaultBlockSize = "10GB"

// vmBlockFilesystemSize is the size of a VM block volume's associated filesystem volume.
const vmBlockFilesystemSize = "50MB"

// DefaultFilesystem filesytem to use for block devices by default.
const DefaultFilesystem = "ext4"

// volIDQuotaSkip is used to indicate to drivers that quotas should not be setup, used during backup import.
const volIDQuotaSkip = int64(-1)

// VolumeType represents a storage volume type.
type VolumeType string

// VolumeTypeImage represents an image storage volume.
const VolumeTypeImage = VolumeType("images")

// VolumeTypeCustom represents a custom storage volume.
const VolumeTypeCustom = VolumeType("custom")

// VolumeTypeContainer represents a container storage volume.
const VolumeTypeContainer = VolumeType("containers")

// VolumeTypeVM represents a virtual-machine storage volume.
const VolumeTypeVM = VolumeType("virtual-machines")

// ContentType indicates the format of the volume.
type ContentType string

// ContentTypeFS indicates the volume will be populated with a mountabble filesystem.
const ContentTypeFS = ContentType("fs")

// ContentTypeBlock indicates the volume will be a block device and its contents and we do not
// know which filesystem(s) (if any) are in use.
const ContentTypeBlock = ContentType("block")

// BaseDirectories maps volume types to the expected directories.
var BaseDirectories = map[VolumeType][]string{
	VolumeTypeContainer: {"containers", "containers-snapshots"},
	VolumeTypeCustom:    {"custom", "custom-snapshots"},
	VolumeTypeImage:     {"images"},
	VolumeTypeVM:        {"virtual-machines", "virtual-machines-snapshots"},
}

// Volume represents a storage volume, and provides functions to mount and unmount it.
type Volume struct {
	name            string
	pool            string
	poolConfig      map[string]string
	volType         VolumeType
	contentType     ContentType
	config          map[string]string
	driver          Driver
	keepDevice      bool
	customMountPath string
}

// NewVolume instantiates a new Volume struct.
func NewVolume(driver Driver, poolName string, volType VolumeType, contentType ContentType, volName string, volConfig, poolConfig map[string]string) Volume {
	return Volume{
		name:        volName,
		pool:        poolName,
		poolConfig:  poolConfig,
		volType:     volType,
		contentType: contentType,
		config:      volConfig,
		driver:      driver,
	}
}

// Name returns volume's name.
func (v Volume) Name() string {
	return v.name
}

// Config returns the volumes (unexpanded) config.
func (v Volume) Config() map[string]string {
	return v.config
}

// ExpandedConfig returns either the value of the volume's config key or the pool's config "volume.{key}" value.
func (v Volume) ExpandedConfig(key string) string {
	volVal, ok := v.config[key]
	if ok {
		return volVal
	}

	return v.poolConfig[fmt.Sprintf("volume.%s", key)]
}

// NewSnapshot instantiates a new Volume struct representing a snapshot of the parent volume.
func (v Volume) NewSnapshot(snapshotName string) (Volume, error) {
	if v.IsSnapshot() {
		return Volume{}, fmt.Errorf("Cannot create a snapshot volume from a snapshot")
	}

	fullSnapName := GetSnapshotVolumeName(v.name, snapshotName)
	return NewVolume(v.driver, v.pool, v.volType, v.contentType, fullSnapName, v.config, v.poolConfig), nil
}

// IsSnapshot indicates if volume is a snapshot.
func (v Volume) IsSnapshot() bool {
	return shared.IsSnapshot(v.name)
}

// MountPath returns the path where the volume will be mounted.
func (v Volume) MountPath() string {
	if v.customMountPath != "" {
		return v.customMountPath
	}

	return GetVolumeMountPath(v.pool, v.volType, v.name)
}

// EnsureMountPath creates the volume's mount path if missing, then sets the correct permission for the type.
func (v Volume) EnsureMountPath() error {
	volPath := v.MountPath()

	// Create volume's mount path, with any created directories set to 0711.
	err := os.Mkdir(volPath, 0711)
	if err != nil && !os.IsExist(err) {
		return errors.Wrapf(err, "Failed to create directory '%s'", volPath)
	}

	// Set very restrictive mode 0100 for non-custom and non-image volumes.
	mode := os.FileMode(0711)
	if v.volType != VolumeTypeCustom && v.volType != VolumeTypeImage {
		mode = os.FileMode(0100)
	}

	// Set mode of actual volume's mount path.
	err = os.Chmod(volPath, mode)
	if err != nil {
		return errors.Wrapf(err, "Failed to chmod '%s'", volPath)
	}

	return nil
}

// MountTask runs the supplied task after mounting the volume if needed. If the volume was mounted
// for this then it is unmounted when the task finishes.
func (v Volume) MountTask(task func(mountPath string, op *operations.Operation) error, op *operations.Operation) error {
	isSnap := v.IsSnapshot()

	// If the volume is a snapshot then call the snapshot specific mount/unmount functions as
	// these will mount the snapshot read only.
	if isSnap {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourMount, err := v.driver.MountVolumeSnapshot(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourMount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.UnmountVolumeSnapshot(v, op)
				unlock()
			}()
		}
	} else {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourMount, err := v.driver.MountVolume(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourMount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.UnmountVolume(v, op)
				unlock()
			}()
		}
	}

	return task(v.MountPath(), op)
}

// UnmountTask runs the supplied task after unmounting the volume if needed. If the volume was unmounted
// for this then it is mounted when the task finishes.
func (v Volume) UnmountTask(task func(op *operations.Operation) error, op *operations.Operation) error {
	isSnap := v.IsSnapshot()

	// If the volume is a snapshot then call the snapshot specific mount/unmount functions as
	// these will mount the snapshot read only.
	if isSnap {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourUnmount, err := v.driver.UnmountVolumeSnapshot(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourUnmount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.MountVolumeSnapshot(v, op)
				unlock()
			}()
		}
	} else {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourUnmount, err := v.driver.UnmountVolume(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourUnmount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.MountVolume(v, op)
				unlock()
			}()
		}
	}

	return task(op)
}

// Snapshots returns a list of snapshots for the volume.
func (v Volume) Snapshots(op *operations.Operation) ([]Volume, error) {
	if v.IsSnapshot() {
		return nil, fmt.Errorf("Volume is a snapshot")
	}

	snapshots, err := v.driver.VolumeSnapshots(v, op)
	if err != nil {
		return nil, err
	}

	snapVols := []Volume{}
	for _, snapName := range snapshots {
		snapshot, err := v.NewSnapshot(snapName)
		if err != nil {
			return nil, err
		}
		snapVols = append(snapVols, snapshot)
	}

	return snapVols, nil
}

// IsBlockBacked indicates whether storage device is block backed.
func (v Volume) IsBlockBacked() bool {
	return v.driver.Info().BlockBacking
}

// Type returns the volume type.
func (v Volume) Type() VolumeType {
	return v.volType
}

// ContentType returns the content type.
func (v Volume) ContentType() ContentType {
	return v.contentType
}

// IsVMBlock returns true if volume is a block volume for virtual machines or associated images.
func (v Volume) IsVMBlock() bool {
	return (v.volType == VolumeTypeVM || v.volType == VolumeTypeImage) && v.contentType == ContentTypeBlock
}

// NewVMBlockFilesystemVolume returns a copy of the volume with the content type set to ContentTypeFS and the
// config "size" property set to vmBlockFilesystemSize.
func (v Volume) NewVMBlockFilesystemVolume() Volume {
	// Copy volume config so modifications don't affect original volume.
	newConf := make(map[string]string, len(v.config))
	for k, v := range v.config {
		newConf[k] = v
	}

	// VM Block filesystems are a fixed size.
	newConf["size"] = vmBlockFilesystemSize

	return NewVolume(v.driver, v.pool, v.volType, ContentTypeFS, v.name, newConf, v.poolConfig)
}
<commit_msg>lxd/storage/drivers/volume: Only chmod if needed in EnsureMountPath<commit_after>package drivers

import (
	"fmt"
	"os"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/storage/locking"
	"github.com/lxc/lxd/shared"
)

// tmpVolSuffix Suffix to use for any temporary volumes created by LXD.
const tmpVolSuffix = ".lxdtmp"

// defaultBlockSize Default size of block volumes.
const defaultBlockSize = "10GB"

// vmBlockFilesystemSize is the size of a VM block volume's associated filesystem volume.
const vmBlockFilesystemSize = "50MB"

// DefaultFilesystem filesytem to use for block devices by default.
const DefaultFilesystem = "ext4"

// volIDQuotaSkip is used to indicate to drivers that quotas should not be setup, used during backup import.
const volIDQuotaSkip = int64(-1)

// VolumeType represents a storage volume type.
type VolumeType string

// VolumeTypeImage represents an image storage volume.
const VolumeTypeImage = VolumeType("images")

// VolumeTypeCustom represents a custom storage volume.
const VolumeTypeCustom = VolumeType("custom")

// VolumeTypeContainer represents a container storage volume.
const VolumeTypeContainer = VolumeType("containers")

// VolumeTypeVM represents a virtual-machine storage volume.
const VolumeTypeVM = VolumeType("virtual-machines")

// ContentType indicates the format of the volume.
type ContentType string

// ContentTypeFS indicates the volume will be populated with a mountabble filesystem.
const ContentTypeFS = ContentType("fs")

// ContentTypeBlock indicates the volume will be a block device and its contents and we do not
// know which filesystem(s) (if any) are in use.
const ContentTypeBlock = ContentType("block")

// BaseDirectories maps volume types to the expected directories.
var BaseDirectories = map[VolumeType][]string{
	VolumeTypeContainer: {"containers", "containers-snapshots"},
	VolumeTypeCustom:    {"custom", "custom-snapshots"},
	VolumeTypeImage:     {"images"},
	VolumeTypeVM:        {"virtual-machines", "virtual-machines-snapshots"},
}

// Volume represents a storage volume, and provides functions to mount and unmount it.
type Volume struct {
	name            string
	pool            string
	poolConfig      map[string]string
	volType         VolumeType
	contentType     ContentType
	config          map[string]string
	driver          Driver
	keepDevice      bool
	customMountPath string
}

// NewVolume instantiates a new Volume struct.
func NewVolume(driver Driver, poolName string, volType VolumeType, contentType ContentType, volName string, volConfig, poolConfig map[string]string) Volume {
	return Volume{
		name:        volName,
		pool:        poolName,
		poolConfig:  poolConfig,
		volType:     volType,
		contentType: contentType,
		config:      volConfig,
		driver:      driver,
	}
}

// Name returns volume's name.
func (v Volume) Name() string {
	return v.name
}

// Config returns the volumes (unexpanded) config.
func (v Volume) Config() map[string]string {
	return v.config
}

// ExpandedConfig returns either the value of the volume's config key or the pool's config "volume.{key}" value.
func (v Volume) ExpandedConfig(key string) string {
	volVal, ok := v.config[key]
	if ok {
		return volVal
	}

	return v.poolConfig[fmt.Sprintf("volume.%s", key)]
}

// NewSnapshot instantiates a new Volume struct representing a snapshot of the parent volume.
func (v Volume) NewSnapshot(snapshotName string) (Volume, error) {
	if v.IsSnapshot() {
		return Volume{}, fmt.Errorf("Cannot create a snapshot volume from a snapshot")
	}

	fullSnapName := GetSnapshotVolumeName(v.name, snapshotName)
	return NewVolume(v.driver, v.pool, v.volType, v.contentType, fullSnapName, v.config, v.poolConfig), nil
}

// IsSnapshot indicates if volume is a snapshot.
func (v Volume) IsSnapshot() bool {
	return shared.IsSnapshot(v.name)
}

// MountPath returns the path where the volume will be mounted.
func (v Volume) MountPath() string {
	if v.customMountPath != "" {
		return v.customMountPath
	}

	return GetVolumeMountPath(v.pool, v.volType, v.name)
}

// EnsureMountPath creates the volume's mount path if missing, then sets the correct permission for the type.
// If permission setting fails and the volume is a snapshot then the error is ignored as snapshots are read only.
func (v Volume) EnsureMountPath() error {
	volPath := v.MountPath()

	revert := revert.New()
	defer revert.Fail()

	// Create volume's mount path if missing, with any created directories set to 0711.
	if !shared.PathExists(volPath) {
		err := os.Mkdir(volPath, 0711)
		if err != nil {
			return errors.Wrapf(err, "Failed to create mount directory %q", volPath)
		}
		revert.Add(func() { os.Remove(volPath) })
	}

	// Set very restrictive mode 0100 for non-custom and non-image volumes.
	mode := os.FileMode(0711)
	if v.volType != VolumeTypeCustom && v.volType != VolumeTypeImage {
		mode = os.FileMode(0100)
	}

	fInfo, err := os.Lstat(volPath)
	if err != nil {
		return errors.Wrapf(err, "Error getting mount directory info %q", volPath)
	}

	// We expect the mount path to be a directory, so use this for comparison.
	compareMode := os.ModeDir | mode

	// Set mode of actual volume's mount path if needed.
	if fInfo.Mode() != compareMode {
		err = os.Chmod(volPath, mode)

		// If the chmod failed, return the error as long as the volume is not a snapshot.
		// If the volume is a snapshot, we must ignore the error as snapshots are readonly and cannot be
		// modified after they are taken, such that any permission error is not fixable at mount time.
		if err != nil && !v.IsSnapshot() {
			return errors.Wrapf(err, "Failed to chmod mount directory %q (%04o)", volPath, mode)
		}
	}

	revert.Success()
	return nil
}

// MountTask runs the supplied task after mounting the volume if needed. If the volume was mounted
// for this then it is unmounted when the task finishes.
func (v Volume) MountTask(task func(mountPath string, op *operations.Operation) error, op *operations.Operation) error {
	isSnap := v.IsSnapshot()

	// If the volume is a snapshot then call the snapshot specific mount/unmount functions as
	// these will mount the snapshot read only.
	if isSnap {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourMount, err := v.driver.MountVolumeSnapshot(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourMount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.UnmountVolumeSnapshot(v, op)
				unlock()
			}()
		}
	} else {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourMount, err := v.driver.MountVolume(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourMount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.UnmountVolume(v, op)
				unlock()
			}()
		}
	}

	return task(v.MountPath(), op)
}

// UnmountTask runs the supplied task after unmounting the volume if needed. If the volume was unmounted
// for this then it is mounted when the task finishes.
func (v Volume) UnmountTask(task func(op *operations.Operation) error, op *operations.Operation) error {
	isSnap := v.IsSnapshot()

	// If the volume is a snapshot then call the snapshot specific mount/unmount functions as
	// these will mount the snapshot read only.
	if isSnap {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourUnmount, err := v.driver.UnmountVolumeSnapshot(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourUnmount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.MountVolumeSnapshot(v, op)
				unlock()
			}()
		}
	} else {
		unlock := locking.Lock(v.pool, string(v.volType), v.name)

		ourUnmount, err := v.driver.UnmountVolume(v, op)
		if err != nil {
			unlock()
			return err
		}

		unlock()

		if ourUnmount {
			defer func() {
				unlock := locking.Lock(v.pool, string(v.volType), v.name)
				v.driver.MountVolume(v, op)
				unlock()
			}()
		}
	}

	return task(op)
}

// Snapshots returns a list of snapshots for the volume.
func (v Volume) Snapshots(op *operations.Operation) ([]Volume, error) {
	if v.IsSnapshot() {
		return nil, fmt.Errorf("Volume is a snapshot")
	}

	snapshots, err := v.driver.VolumeSnapshots(v, op)
	if err != nil {
		return nil, err
	}

	snapVols := []Volume{}
	for _, snapName := range snapshots {
		snapshot, err := v.NewSnapshot(snapName)
		if err != nil {
			return nil, err
		}
		snapVols = append(snapVols, snapshot)
	}

	return snapVols, nil
}

// IsBlockBacked indicates whether storage device is block backed.
func (v Volume) IsBlockBacked() bool {
	return v.driver.Info().BlockBacking
}

// Type returns the volume type.
func (v Volume) Type() VolumeType {
	return v.volType
}

// ContentType returns the content type.
func (v Volume) ContentType() ContentType {
	return v.contentType
}

// IsVMBlock returns true if volume is a block volume for virtual machines or associated images.
func (v Volume) IsVMBlock() bool {
	return (v.volType == VolumeTypeVM || v.volType == VolumeTypeImage) && v.contentType == ContentTypeBlock
}

// NewVMBlockFilesystemVolume returns a copy of the volume with the content type set to ContentTypeFS and the
// config "size" property set to vmBlockFilesystemSize.
func (v Volume) NewVMBlockFilesystemVolume() Volume {
	// Copy volume config so modifications don't affect original volume.
	newConf := make(map[string]string, len(v.config))
	for k, v := range v.config {
		newConf[k] = v
	}

	// VM Block filesystems are a fixed size.
	newConf["size"] = vmBlockFilesystemSize

	return NewVolume(v.driver, v.pool, v.volType, ContentTypeFS, v.name, newConf, v.poolConfig)
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/docker/docker-registry/common/testutil"
	"github.com/docker/docker-registry/digest"
	"github.com/docker/docker-registry/storage"
)

type testBlob struct {
	digest   digest.Digest
	contents []byte
}

func TestPush(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	uploadLocations := make([]string, len(testBlobs))
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		// TODO(bbland): this is returning the same location for all uploads,
		// because we can't know which blob will get which location.
		// It's sort of okay because we're using unique digests, but this needs
		// to change at some point.
		uploadLocations[i] = fmt.Sprintf("/v2/%s/blob/test-uuid", name)
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "POST",
				Route:  "/v2/" + name + "/blob/upload/",
			},
			Response: testutil.Response{
				StatusCode: http.StatusAccepted,
				Headers: http.Header(map[string][]string{
					"Location": {uploadLocations[i]},
				}),
			},
		}
		blobRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  uploadLocations[i],
				QueryParams: map[string][]string{
					"length": {fmt.Sprint(len(blob.contents))},
					"digest": {blob.digest.String()},
				},
				Body: blob.contents,
			},
			Response: testutil.Response{
				StatusCode: http.StatusCreated,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  "/v2/" + name + "/manifest/" + tag,
				Body:   manifestBytes,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		writer, err := l.Writer()
		if err != nil {
			t.Fatal(err)
		}

		writer.SetSize(len(blob.contents))
		writer.Write(blob.contents)
		writer.Close()
	}

	objectStore.WriteManifest(name, tag, manifest)

	err = Push(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPull(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/manifest/" + tag,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       manifestBytes,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	err = Pull(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		blobBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(blobBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}

func TestPullResume(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	layers := make([]registry.FSLayer, len(testBlobs))
	history := make([]registry.ManifestHistory, len(testBlobs))

	for i, layer := range testBlobs {
		layers[i] = registry.FSLayer{BlobSum: layer.digest}
		history[i] = registry.ManifestHistory{V1Compatibility: layer.digest.String()}
	}

	manifest := &registry.ImageManifest{
		Name:          name,
		Tag:           tag,
		Architecture:  "x86",
		FSLayers:      layers,
		History:       history,
		SchemaVersion: 1,
	}
	manifestBytes, err := json.Marshal(manifest)

	layerRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		layerRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[:len(blob.contents)/2],
				Headers: http.Header(map[string][]string{
					"Content-Length": {fmt.Sprint(len(blob.contents))},
				}),
			},
		}
		layerRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[len(blob.contents)/2:],
			},
		}
	}

	for i := 0; i < 3; i++ {
		layerRequestResponseMappings = append(layerRequestResponseMappings, testutil.RequestResponseMap{
			testutil.RequestResponseMapping{
				Request: testutil.Request{
					Method: "GET",
					Route:  "/v2/" + name + "/manifest/" + tag,
				},
				Response: testutil.Response{
					StatusCode: http.StatusOK,
					Body:       manifestBytes,
				},
			},
		}...)
	}

	handler := testutil.NewHandler(layerRequestResponseMappings)
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*registry.ImageManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for attempts := 0; attempts < 3; attempts++ {
		err = Pull(client, objectStore, name, tag)
		if err == nil {
			break
		}
	}

	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		layerBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(layerBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}
<commit_msg>Correct references after move of Manifest into storage package<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/docker/docker-registry/common/testutil"
	"github.com/docker/docker-registry/digest"
	"github.com/docker/docker-registry/storage"
)

type testBlob struct {
	digest   digest.Digest
	contents []byte
}

func TestPush(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	uploadLocations := make([]string, len(testBlobs))
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		// TODO(bbland): this is returning the same location for all uploads,
		// because we can't know which blob will get which location.
		// It's sort of okay because we're using unique digests, but this needs
		// to change at some point.
		uploadLocations[i] = fmt.Sprintf("/v2/%s/blob/test-uuid", name)
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "POST",
				Route:  "/v2/" + name + "/blob/upload/",
			},
			Response: testutil.Response{
				StatusCode: http.StatusAccepted,
				Headers: http.Header(map[string][]string{
					"Location": {uploadLocations[i]},
				}),
			},
		}
		blobRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  uploadLocations[i],
				QueryParams: map[string][]string{
					"length": {fmt.Sprint(len(blob.contents))},
					"digest": {blob.digest.String()},
				},
				Body: blob.contents,
			},
			Response: testutil.Response{
				StatusCode: http.StatusCreated,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "PUT",
				Route:  "/v2/" + name + "/manifest/" + tag,
				Body:   manifestBytes,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		writer, err := l.Writer()
		if err != nil {
			t.Fatal(err)
		}

		writer.SetSize(len(blob.contents))
		writer.Write(blob.contents)
		writer.Close()
	}

	objectStore.WriteManifest(name, tag, manifest)

	err = Push(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPull(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	blobs := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, blob := range testBlobs {
		blobs[i] = storage.FSLayer{BlobSum: blob.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: blob.digest.String()}
	}

	manifest := &storage.SignedManifest{
		Manifest: storage.Manifest{
			Name:         name,
			Tag:          tag,
			Architecture: "x86",
			FSLayers:     blobs,
			History:      history,
			Versioned: storage.Versioned{
				SchemaVersion: 1,
			},
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	blobRequestResponseMappings := make([]testutil.RequestResponseMapping, len(testBlobs))
	for i, blob := range testBlobs {
		blobRequestResponseMappings[i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents,
			},
		}
	}

	handler := testutil.NewHandler(append(blobRequestResponseMappings, testutil.RequestResponseMap{
		testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/manifest/" + tag,
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       manifestBytes,
			},
		},
	}...))
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	err = Pull(client, objectStore, name, tag)
	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		blobBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(blobBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}

func TestPullResume(t *testing.T) {
	name := "hello/world"
	tag := "sometag"
	testBlobs := []testBlob{
		{
			digest:   "12345",
			contents: []byte("some contents"),
		},
		{
			digest:   "98765",
			contents: []byte("some other contents"),
		},
	}
	layers := make([]storage.FSLayer, len(testBlobs))
	history := make([]storage.ManifestHistory, len(testBlobs))

	for i, layer := range testBlobs {
		layers[i] = storage.FSLayer{BlobSum: layer.digest}
		history[i] = storage.ManifestHistory{V1Compatibility: layer.digest.String()}
	}

	manifest := &storage.Manifest{
		Name:         name,
		Tag:          tag,
		Architecture: "x86",
		FSLayers:     layers,
		History:      history,
		Versioned: storage.Versioned{
			SchemaVersion: 1,
		},
	}
	manifestBytes, err := json.Marshal(manifest)

	layerRequestResponseMappings := make([]testutil.RequestResponseMapping, 2*len(testBlobs))
	for i, blob := range testBlobs {
		layerRequestResponseMappings[2*i] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[:len(blob.contents)/2],
				Headers: http.Header(map[string][]string{
					"Content-Length": {fmt.Sprint(len(blob.contents))},
				}),
			},
		}
		layerRequestResponseMappings[2*i+1] = testutil.RequestResponseMapping{
			Request: testutil.Request{
				Method: "GET",
				Route:  "/v2/" + name + "/blob/" + blob.digest.String(),
			},
			Response: testutil.Response{
				StatusCode: http.StatusOK,
				Body:       blob.contents[len(blob.contents)/2:],
			},
		}
	}

	for i := 0; i < 3; i++ {
		layerRequestResponseMappings = append(layerRequestResponseMappings, testutil.RequestResponseMap{
			testutil.RequestResponseMapping{
				Request: testutil.Request{
					Method: "GET",
					Route:  "/v2/" + name + "/manifest/" + tag,
				},
				Response: testutil.Response{
					StatusCode: http.StatusOK,
					Body:       manifestBytes,
				},
			},
		}...)
	}

	handler := testutil.NewHandler(layerRequestResponseMappings)
	server := httptest.NewServer(handler)
	client := New(server.URL)
	objectStore := &memoryObjectStore{
		mutex:           new(sync.Mutex),
		manifestStorage: make(map[string]*storage.SignedManifest),
		layerStorage:    make(map[digest.Digest]Layer),
	}

	for attempts := 0; attempts < 3; attempts++ {
		err = Pull(client, objectStore, name, tag)
		if err == nil {
			break
		}
	}

	if err != nil {
		t.Fatal(err)
	}

	m, err := objectStore.Manifest(name, tag)
	if err != nil {
		t.Fatal(err)
	}

	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}

	if string(mBytes) != string(manifestBytes) {
		t.Fatal("Incorrect manifest")
	}

	for _, blob := range testBlobs {
		l, err := objectStore.Layer(blob.digest)
		if err != nil {
			t.Fatal(err)
		}

		reader, err := l.Reader()
		if err != nil {
			t.Fatal(err)
		}
		defer reader.Close()

		layerBytes, err := ioutil.ReadAll(reader)
		if err != nil {
			t.Fatal(err)
		}

		if string(layerBytes) != string(blob.contents) {
			t.Fatal("Incorrect blob")
		}
	}
}
<|endoftext|>


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

// program which takes
// - a rule file
// - a sample list of metrics
// and sees how well the rule performs against the metrics.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"runtime"
	"sort"
	"sync"

	"github.com/square/metrics/api"
	"github.com/square/metrics/main/common"
	"github.com/square/metrics/metric_metadata/cassandra"
	"github.com/square/metrics/util"
)

var (
	metricsFile      = flag.String("metrics-file", "", "Location of YAML configuration file.")
	unmatchedFile    = flag.String("unmatched-file", "", "location of metrics list to output unmatched transformations.")
	insertToDatabase = flag.Bool("insert-to-db", false, "If true, insert rows to database.")
	reverse          = flag.Bool("reverse", false, "If true, then attempt the reverse-rule lookup also.")
)

// Statistics represents the aggregated result of rules
// after running through the test file.
type Statistics struct {
	perMetric map[api.MetricKey]PerMetricStatistics
	matched   int // number of matched rows
	unmatched int // number of unmatched rows
}

// PerMetricStatistics represents per-metric result of rules
// after running through the test file.
type PerMetricStatistics struct {
	matched          int // number of matched rows
	reverseSuccess   int // number of reversed entries
	reverseError     int // number of incorrectly reversed entries.
	reverseIncorrect int // number of incorrectly reversed entries.
}

func main() {
	flag.Parse()
	common.SetupLogger()

	config := common.LoadConfig()

	// graphiteConfig := util.GraphiteConverterConfig{ConversionRulesPath: config.MetricMetadataAPI.ConversionRulesPath}
	//TODO(cchandler): Make a constructor for a graphite converter so we don't
	//have to stich everything together outside of the package.

	ruleset, err := util.LoadRules(config.MetricMetadataConfig.ConversionRulesPath)

	if err != nil {
		common.ExitWithMessage(fmt.Sprintf("Error while reading rules: %s", err.Error()))
	}
	_ = util.RuleBasedGraphiteConverter{Ruleset: ruleset}

	metricFile, err := os.Open(*metricsFile)
	if err != nil {
		common.ExitWithMessage("No metric file.")
	}
	scanner := bufio.NewScanner(metricFile)
	cassandraConfig := cassandra.CassandraMetricMetadataConfig{
		Hosts:    config.MetricMetadataConfig.Hosts,
		Keyspace: config.MetricMetadataConfig.Keyspace,
	}
	apiInstance := common.NewMetricMetadataAPI(cassandraConfig)

	var output *os.File
	if *unmatchedFile != "" {
		output, err = os.Create(*unmatchedFile)
		if err != nil {
			common.ExitWithMessage(fmt.Sprintf("Error creating the output file: %s", err.Error()))
		}
	}
	stat := run(ruleset, scanner, apiInstance, output)
	report(stat)
}

func run(ruleset util.RuleSet, scanner *bufio.Scanner, apiInstance api.MetricMetadataAPI, unmatched *os.File) Statistics {
	var wg sync.WaitGroup
	stat := Statistics{
		perMetric: make(map[api.MetricKey]PerMetricStatistics),
	}
	type result struct {
		input   string
		result  api.TaggedMetric
		success bool
	}
	inputBuffer := make(chan string, 10)
	outputBuffer := make(chan result, 10)
	done := make(chan struct{})
	for id := 0; id < runtime.NumCPU(); id++ {
		go func() {
			for {
				select {
				case <-done:
					return
				case input := <-inputBuffer:
					metric, matched := ruleset.MatchRule(input)
					outputBuffer <- result{
						input,
						metric,
						matched,
					}
				}
			}
		}()
	}
	go func() {
		// aggregate function.
		for {
			select {
			case <-done:
				return
			case output := <-outputBuffer:
				converted, matched := output.result, output.success
				if matched {
					stat.matched++
					perMetric := stat.perMetric[converted.MetricKey]
					perMetric.matched++
					if *insertToDatabase {
						apiInstance.AddMetric(converted)
					}
					if *reverse {
						reversed, err := ruleset.ToGraphiteName(converted)
						if err != nil {
							perMetric.reverseError++
						} else if string(reversed) != output.input {
							perMetric.reverseIncorrect++
						} else {
							perMetric.reverseSuccess++
						}
					}
					stat.perMetric[converted.MetricKey] = perMetric
				} else {
					stat.unmatched++
					if unmatched != nil {
						unmatched.WriteString(output.input)
						unmatched.WriteString("\n")
					}
				}
				wg.Done()
			}
		}
	}()

	for scanner.Scan() {
		wg.Add(1)
		input := scanner.Text()
		inputBuffer <- input
	}
	wg.Wait()
	close(done) // broadcast to shutdown all goroutines.
	return stat
}

func report(stat Statistics) {
	total := stat.matched + stat.unmatched
	fmt.Printf("Processed %d entries\n", total)
	fmt.Printf("Matched:   %d\n", stat.matched)
	fmt.Printf("Unmatched: %d\n", stat.unmatched)
	fmt.Printf("Per-rule statistics\n")
	rowformat := "%-60s %7d %7d %7d %7d\n"
	headformat := "%-60s %7s %7s %7s %7s\n"
	fmt.Printf(headformat, "name", "match", "rev-suc", "rev-err", "rev-fail")
	sortedKeys := make([]string, len(stat.perMetric))
	index := 0
	for key := range stat.perMetric {
		sortedKeys[index] = string(key)
		index++
	}
	sort.Strings(sortedKeys)
	for _, key := range sortedKeys {
		perMetric := stat.perMetric[api.MetricKey(key)]
		fmt.Printf(rowformat,
			string(key),
			perMetric.matched,
			perMetric.reverseSuccess,
			perMetric.reverseError,
			perMetric.reverseIncorrect,
		)
	}
}
<commit_msg>Update ruletester to read metrics-indexer cache files<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// program which takes
// - a rule file
// - a sample list of metrics
// and sees how well the rule performs against the metrics.
package main

import (
	"bytes"
	"compress/zlib"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	// "runtime"
	// "sort"
	"sync"

	"github.com/square/metrics/api"
	"github.com/square/metrics/main/common"
	"github.com/square/metrics/metric_metadata/cassandra"
	"github.com/square/metrics/util"
)

var (
	metricsFile   = flag.String("metrics-file", "", "Location of zlib compressed gob string file.")
	unmatchedFile = flag.String("unmatched-file", "", "location of metrics list to output unmatched transformations.")
	reverse       = flag.Bool("reverse", false, "If true, then attempt the reverse-rule lookup also.")
)

// Statistics represents the aggregated result of rules
// after running through the test file.
type Statistics struct {
	perMetric map[api.MetricKey]PerMetricStatistics
	matched   int // number of matched rows
	unmatched int // number of unmatched rows
}

// PerMetricStatistics represents per-metric result of rules
// after running through the test file.
type PerMetricStatistics struct {
	matched          int // number of matched rows
	reverseSuccess   int // number of reversed entries
	reverseError     int // number of incorrectly reversed entries.
	reverseIncorrect int // number of incorrectly reversed entries.
}

func ReadMetricsFile(file string) ([]string, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	result := make(map[string]int)
	r, err := zlib.NewReader(bytes.NewBuffer(data))
	b := new(bytes.Buffer)
	io.Copy(b, r)
	d := gob.NewDecoder(b)

	err = d.Decode(&result)
	if err != nil {
		return nil, err
	}

	strings := []string{}
	for k, _ := range result {
		strings = append(strings, k)
	}
	return strings, nil
}

func main() {
	flag.Parse()
	common.SetupLogger()

	if *metricsFile == "" {
		common.ExitWithMessage("No metric file.")
		fmt.Printf("You must specify a metrics file.\n")
		os.Exit(1)
	}

	config := common.LoadConfig()

	// graphiteConfig := util.GraphiteConverterConfig{ConversionRulesPath: config.MetricMetadataAPI.ConversionRulesPath}
	//TODO(cchandler): Make a constructor for a graphite converter so we don't
	//have to stich everything together outside of the package.

	ruleset, err := util.LoadRules(config.MetricMetadataConfig.ConversionRulesPath)

	if err != nil {
		common.ExitWithMessage(fmt.Sprintf("Error while reading rules: %s", err.Error()))
	}
	graphiteConverter := util.RuleBasedGraphiteConverter{Ruleset: ruleset}

	// metricFile, err := os.Open(*metricsFile)
	metrics, err := ReadMetricsFile(*metricsFile)
	if err != nil {
		fmt.Printf("Fatal error reading metrics file %s\n", err)
		os.Exit(1)
	}

	DoAnalysis(metrics, graphiteConverter)

	fmt.Printf("Total metric count %d\n", len(metrics))

	// matched := 0
	// unmatched := 0
	// reverse_convert_failed := 0

	// fmt.Printf("Matched: %d\n", matched)
	// fmt.Printf("Unmatched: %d\n", unmatched)
	// fmt.Printf("Reverse convert failed: %d\n", reverse_convert_failed)

	if err != nil {
		common.ExitWithMessage("No metric file.")
	}
	// scanner := bufio.NewScanner(metricFile)
	cassandraConfig := cassandra.CassandraMetricMetadataConfig{
		Hosts:    config.MetricMetadataConfig.Hosts,
		Keyspace: config.MetricMetadataConfig.Keyspace,
	}
	_ = common.NewMetricMetadataAPI(cassandraConfig)

	// var output *os.File
	// if *unmatchedFile != "" {
	// 	output, err = os.Create(*unmatchedFile)
	// 	if err != nil {
	// 		common.ExitWithMessage(fmt.Sprintf("Error creating the output file: %s", err.Error()))
	// 	}
	// }
	// stat := run(ruleset, scanner, apiInstance, output)
	// report(stat)
}

type ChunkResult struct {
	matched                int
	unmatched              int
	reverse_convert_failed int
}

func DoAnalysis(metrics []string, graphiteConverter util.RuleBasedGraphiteConverter) {
	workQueue := make(chan []string, len(metrics)/25+1)
	resultQueue := make(chan ChunkResult, len(metrics))
	var wg sync.WaitGroup

	i := 0
	for i = 0; i < len(metrics); i = i + 25 {
		workQueue <- metrics[i : i+25]
	}
	if i <= len(metrics) {
		workQueue <- metrics[i:]
	}

	close(workQueue)

	fmt.Printf("Starting work...\n")
	for j := 0; j < 10; j++ {
		wg.Add(1)
		go func() {
			counter := 0
			defer wg.Done()
			for {
				select {
				case metrics, more := <-workQueue:
					counter++
					if counter%100 == 0 && counter != 0 {
						fmt.Printf(".")
					}
					chunk_result := ChunkResult{}
					for _, metric := range metrics {
						graphiteMetric := util.GraphiteMetric(metric)
						taggedMetric, err := graphiteConverter.ToTaggedName(graphiteMetric)
						if err != nil {
							_, err := graphiteConverter.ToGraphiteName(taggedMetric)
							if err != nil {
								chunk_result.reverse_convert_failed++
							} else {

							}
							chunk_result.matched++
						} else {
							chunk_result.unmatched++
						}
					}
					resultQueue <- chunk_result

					if !more {
						return
					}
				default:
					fmt.Printf("Nothing more!\n")
				}
			}

		}()
	}
	wg.Wait()
	close(resultQueue)

	fmt.Printf("\n")
	fmt.Printf("Processing results!")
	totalResults := ChunkResult{}
	//Merge chunks
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case result, more := <-resultQueue:
				totalResults.matched += result.matched
				totalResults.unmatched += result.unmatched
				totalResults.reverse_convert_failed += result.reverse_convert_failed

				if !more {
					return
				}
			}
		}
	}()

	wg.Wait()
	fmt.Printf("\n")
	fmt.Printf("Matched: %d\n", totalResults.matched)
	fmt.Printf("Unmatched: %d\n", totalResults.unmatched)
	fmt.Printf("Reverse convert failed: %d\n", totalResults.reverse_convert_failed)

}

func ProcessWork(workQueue chan []string) {

}

// func run(ruleset util.RuleSet, scanner *bufio.Scanner, apiInstance api.MetricMetadataAPI, unmatched *os.File) Statistics {
// 	var wg sync.WaitGroup
// 	stat := Statistics{
// 		perMetric: make(map[api.MetricKey]PerMetricStatistics),
// 	}
// 	type result struct {
// 		input   string
// 		result  api.TaggedMetric
// 		success bool
// 	}
// 	inputBuffer := make(chan string, 10)
// 	outputBuffer := make(chan result, 10)
// 	done := make(chan struct{})
// 	for id := 0; id < runtime.NumCPU(); id++ {
// 		go func() {
// 			for {
// 				select {
// 				case <-done:
// 					return
// 				case input := <-inputBuffer:
// 					metric, matched := ruleset.MatchRule(input)
// 					outputBuffer <- result{
// 						input,
// 						metric,
// 						matched,
// 					}
// 				}
// 			}
// 		}()
// 	}
// 	go func() {
// 		// aggregate function.
// 		for {
// 			select {
// 			case <-done:
// 				return
// 			case output := <-outputBuffer:
// 				converted, matched := output.result, output.success
// 				if matched {
// 					stat.matched++
// 					perMetric := stat.perMetric[converted.MetricKey]
// 					perMetric.matched++
// 					if *insertToDatabase {
// 						apiInstance.AddMetric(converted)
// 					}
// 					if *reverse {
// 						reversed, err := ruleset.ToGraphiteName(converted)
// 						if err != nil {
// 							perMetric.reverseError++
// 						} else if string(reversed) != output.input {
// 							perMetric.reverseIncorrect++
// 						} else {
// 							perMetric.reverseSuccess++
// 						}
// 					}
// 					stat.perMetric[converted.MetricKey] = perMetric
// 				} else {
// 					stat.unmatched++
// 					if unmatched != nil {
// 						unmatched.WriteString(output.input)
// 						unmatched.WriteString("\n")
// 					}
// 				}
// 				wg.Done()
// 			}
// 		}
// 	}()

// 	for scanner.Scan() {
// 		wg.Add(1)
// 		input := scanner.Text()
// 		inputBuffer <- input
// 	}
// 	wg.Wait()
// 	close(done) // broadcast to shutdown all goroutines.
// 	return stat
// }

// func report(stat Statistics) {
// 	total := stat.matched + stat.unmatched
// 	fmt.Printf("Processed %d entries\n", total)
// 	fmt.Printf("Matched:   %d\n", stat.matched)
// 	fmt.Printf("Unmatched: %d\n", stat.unmatched)
// 	fmt.Printf("Per-rule statistics\n")
// 	rowformat := "%-60s %7d %7d %7d %7d\n"
// 	headformat := "%-60s %7s %7s %7s %7s\n"
// 	fmt.Printf(headformat, "name", "match", "rev-suc", "rev-err", "rev-fail")
// 	sortedKeys := make([]string, len(stat.perMetric))
// 	index := 0
// 	for key := range stat.perMetric {
// 		sortedKeys[index] = string(key)
// 		index++
// 	}
// 	sort.Strings(sortedKeys)
// 	for _, key := range sortedKeys {
// 		perMetric := stat.perMetric[api.MetricKey(key)]
// 		fmt.Printf(rowformat,
// 			string(key),
// 			perMetric.matched,
// 			perMetric.reverseSuccess,
// 			perMetric.reverseError,
// 			perMetric.reverseIncorrect,
// 		)
// 	}
// }
<|endoftext|>


<commit_before>package editor

import (
	"bytes"
	"fmt"
	"reflect"
)

type element struct {
	TagName string
	Attrs   map[string]string
	Name    string
	label   string
	data    string
	viewBuf *bytes.Buffer
}

// Input returns the []byte of an <input> HTML element with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Input(fieldName string, p interface{}, attrs map[string]string) []byte {
	e := newElement("input", attrs["label"], fieldName, p, attrs)

	return domElementSelfClose(e)
}

// Textarea returns the []byte of a <textarea> HTML element with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Textarea(fieldName string, p interface{}, attrs map[string]string) []byte {
	e := newElement("textarea", attrs["label"], fieldName, p, attrs)

	return domElement(e)
}

// Select returns the []byte of a <select> HTML element plus internal <options> with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Select(fieldName string, p interface{}, attrs, options map[string]string) []byte {
	// options are the value attr and the display value, i.e.
	// <option value="{map key}">{map value}</option>

	// find the field value in p to determine if an option is pre-selected
	fieldVal := valueFromStructField(fieldName, p).String()

	// may need to alloc a buffer, as we will probably loop through options
	// and append the []byte from domElement() called for each option
	sel := newElement("select", attrs["label"], fieldName, p, attrs)
	var opts []*element

	// provide a call to action for the select element
	cta := &element{
		TagName: "option",
		Attrs:   map[string]string{"disabled": "true", "selected": "true"},
		data:    "Select an option...",
		viewBuf: &bytes.Buffer{},
	}

	// provide a selection reset (will store empty string in db)
	reset := &element{
		TagName: "option",
		Attrs:   map[string]string{"value": ""},
		data:    "None",
		viewBuf: &bytes.Buffer{},
	}

	opts = append(opts, cta, reset)

	var val string
	for k, v := range options {
		if k == fieldVal {
			val = "true"
		} else {
			val = "false"
		}
		opt := &element{
			TagName: "option",
			Attrs:   map[string]string{"value": k, "selected": val},
			data:    v,
			viewBuf: &bytes.Buffer{},
		}

		// if val is false (unselected option), delete the attr for clarity
		if val == "false" {
			delete(opt.Attrs, "selected")
		}

		opts = append(opts, opt)
	}

	return domElementWithChildren(sel, opts)
}

// Checkbox returns the []byte of a set of <input type="checkbox"> HTML elements
// wrapped in a <div> with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Checkbox(fieldName string, p interface{}, attrs, options map[string]string) []byte {
	// options are the value attr and the display value, i.e.
	/*
		<label>
			{map value}
			<input type="checkbox" name="{fieldName}" value="{map key}"/>
		</label>
	*/

	div := newElement("div", attrs["label"], "", p, attrs)
	var opts []*element

	// get the pre-checked options if this is already an existing post
	checkedVals := valueFromStructField(fieldName, p)                         // returns refelct.Value
	checked := checkedVals.Slice(0, checkedVals.Len()).Interface().([]string) // casts reflect.Value to []string

	i := 0
	for k, v := range options {
		// check if k is in the pre-checked values and set to checked
		var val string
		for _, x := range checked {
			if k == x {
				val = "true"
			}
		}

		// create a *element manually using the maodified tagNameFromStructFieldMulti
		// func since this is for a multi-value name
		input := &element{
			TagName: "input",
			Attrs: map[string]string{
				"type":    "checkbox",
				"checked": val,
				"value":   k,
			},
			Name:    tagNameFromStructFieldMulti(fieldName, i, p),
			label:   v,
			data:    "",
			viewBuf: &bytes.Buffer{},
		}

		// if checked == false, delete it from input.Attrs for clarity
		if input.Attrs["checked"] == "" {
			delete(input.Attrs, "checked")
		}

		opts = append(opts, input)
		i++
	}

	return domElementWithChildrenCheckbox(div, opts)
}

// domElementSelfClose is a special DOM element which is parsed as a
// self-closing tag and thus needs to be created differently
func domElementSelfClose(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` value="`))
	e.viewBuf.Write([]byte(e.data + `" `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` />`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

// domElementCheckbox is a special DOM element which is parsed as a
// checkbox input tag and thus needs to be created differently
func domElementCheckbox(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>`))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` /> `))

	if e.label != "" {
		e.viewBuf.Write([]byte(e.label + `</label>`))
	}

	return e.viewBuf.Bytes()
}

// domElement creates a DOM element
func domElement(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + string(value) + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	e.viewBuf.Write([]byte(e.data))
	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func domElementWithChildren(e *element, children []*element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + string(value) + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	// loop over children and create domElement for each child
	for _, child := range children {
		e.viewBuf.Write(domElement(child))
	}

	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func domElementWithChildrenCheckbox(e *element, children []*element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	// loop over children and create domElement for each child
	for _, child := range children {
		e.viewBuf.Write(domElementCheckbox(child))
	}

	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func tagNameFromStructField(name string, post interface{}) string {
	// sometimes elements in these environments will not have a name,
	// and thus no tag name in the struct which correlates to it.
	if name == "" {
		return name
	}

	field, ok := reflect.TypeOf(post).Elem().FieldByName(name)
	if !ok {
		panic("Couldn't get struct field for: " + name + ". Make sure you pass the right field name to editor field elements.")
	}

	tag, ok := field.Tag.Lookup("json")
	if !ok {
		panic("Couldn't get json struct tag for: " + name + ". Struct fields for content types must have 'json' tags.")
	}

	return tag
}

// due to the format in which gorilla/schema expects form names to be when
// one is associated with multiple values, we need to output the name as such.
// Ex. 'category.0', 'category.1', 'category.2' and so on.
func tagNameFromStructFieldMulti(name string, i int, post interface{}) string {
	tag := tagNameFromStructField(name, post)

	return fmt.Sprintf("%s.%d", tag, i)
}

func valueFromStructField(name string, post interface{}) reflect.Value {
	field := reflect.Indirect(reflect.ValueOf(post)).FieldByName(name)

	return field
}

func newElement(tagName, label, fieldName string, p interface{}, attrs map[string]string) *element {
	return &element{
		TagName: tagName,
		Attrs:   attrs,
		Name:    tagNameFromStructField(fieldName, p),
		label:   label,
		data:    valueFromStructField(fieldName, p).String(),
		viewBuf: &bytes.Buffer{},
	}
}
<commit_msg>refactor input attrs on Checkbox to be more logical<commit_after>package editor

import (
	"bytes"
	"fmt"
	"reflect"
)

type element struct {
	TagName string
	Attrs   map[string]string
	Name    string
	label   string
	data    string
	viewBuf *bytes.Buffer
}

// Input returns the []byte of an <input> HTML element with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Input(fieldName string, p interface{}, attrs map[string]string) []byte {
	e := newElement("input", attrs["label"], fieldName, p, attrs)

	return domElementSelfClose(e)
}

// Textarea returns the []byte of a <textarea> HTML element with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Textarea(fieldName string, p interface{}, attrs map[string]string) []byte {
	e := newElement("textarea", attrs["label"], fieldName, p, attrs)

	return domElement(e)
}

// Select returns the []byte of a <select> HTML element plus internal <options> with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Select(fieldName string, p interface{}, attrs, options map[string]string) []byte {
	// options are the value attr and the display value, i.e.
	// <option value="{map key}">{map value}</option>

	// find the field value in p to determine if an option is pre-selected
	fieldVal := valueFromStructField(fieldName, p).String()

	// may need to alloc a buffer, as we will probably loop through options
	// and append the []byte from domElement() called for each option
	sel := newElement("select", attrs["label"], fieldName, p, attrs)
	var opts []*element

	// provide a call to action for the select element
	cta := &element{
		TagName: "option",
		Attrs:   map[string]string{"disabled": "true", "selected": "true"},
		data:    "Select an option...",
		viewBuf: &bytes.Buffer{},
	}

	// provide a selection reset (will store empty string in db)
	reset := &element{
		TagName: "option",
		Attrs:   map[string]string{"value": ""},
		data:    "None",
		viewBuf: &bytes.Buffer{},
	}

	opts = append(opts, cta, reset)

	var val string
	for k, v := range options {
		if k == fieldVal {
			val = "true"
		} else {
			val = "false"
		}
		opt := &element{
			TagName: "option",
			Attrs:   map[string]string{"value": k, "selected": val},
			data:    v,
			viewBuf: &bytes.Buffer{},
		}

		// if val is false (unselected option), delete the attr for clarity
		if val == "false" {
			delete(opt.Attrs, "selected")
		}

		opts = append(opts, opt)
	}

	return domElementWithChildren(sel, opts)
}

// Checkbox returns the []byte of a set of <input type="checkbox"> HTML elements
// wrapped in a <div> with a label.
// IMPORTANT:
// The `fieldName` argument will cause a panic if it is not exactly the string
// form of the struct field that this editor input is representing
func Checkbox(fieldName string, p interface{}, attrs, options map[string]string) []byte {
	// options are the value attr and the display value, i.e.
	/*
		<label>
			{map value}
			<input type="checkbox" name="{fieldName}" value="{map key}"/>
		</label>
	*/

	div := newElement("div", attrs["label"], "", p, attrs)
	var opts []*element

	// get the pre-checked options if this is already an existing post
	checkedVals := valueFromStructField(fieldName, p)                         // returns refelct.Value
	checked := checkedVals.Slice(0, checkedVals.Len()).Interface().([]string) // casts reflect.Value to []string

	i := 0
	for k, v := range options {
		inputAttrs := map[string]string{
			"type":  "checkbox",
			"value": k,
		}

		// check if k is in the pre-checked values and set to checked
		for _, x := range checked {
			if k == x {
				inputAttrs["checked"] = "true"
			}
		}

		// create a *element manually using the maodified tagNameFromStructFieldMulti
		// func since this is for a multi-value name
		input := &element{
			TagName: "input",
			Attrs:   inputAttrs,
			Name:    tagNameFromStructFieldMulti(fieldName, i, p),
			label:   v,
			data:    "",
			viewBuf: &bytes.Buffer{},
		}

		opts = append(opts, input)
		i++
	}

	return domElementWithChildrenCheckbox(div, opts)
}

// domElementSelfClose is a special DOM element which is parsed as a
// self-closing tag and thus needs to be created differently
func domElementSelfClose(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` value="`))
	e.viewBuf.Write([]byte(e.data + `" `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` />`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

// domElementCheckbox is a special DOM element which is parsed as a
// checkbox input tag and thus needs to be created differently
func domElementCheckbox(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>`))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` /> `))

	if e.label != "" {
		e.viewBuf.Write([]byte(e.label + `</label>`))
	}

	return e.viewBuf.Bytes()
}

// domElement creates a DOM element
func domElement(e *element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + string(value) + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	e.viewBuf.Write([]byte(e.data))
	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func domElementWithChildren(e *element, children []*element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + string(value) + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	// loop over children and create domElement for each child
	for _, child := range children {
		e.viewBuf.Write(domElement(child))
	}

	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func domElementWithChildrenCheckbox(e *element, children []*element) []byte {
	if e.label != "" {
		e.viewBuf.Write([]byte(`<label>` + e.label))
	}
	e.viewBuf.Write([]byte(`<` + e.TagName + ` `))

	for attr, value := range e.Attrs {
		e.viewBuf.Write([]byte(attr + `="` + value + `" `))
	}
	e.viewBuf.Write([]byte(` name="` + e.Name + `"`))
	e.viewBuf.Write([]byte(` >`))

	// loop over children and create domElement for each child
	for _, child := range children {
		e.viewBuf.Write(domElementCheckbox(child))
	}

	e.viewBuf.Write([]byte(`</` + e.TagName + `>`))

	if e.label != "" {
		e.viewBuf.Write([]byte(`</label>`))
	}

	return e.viewBuf.Bytes()
}

func tagNameFromStructField(name string, post interface{}) string {
	// sometimes elements in these environments will not have a name,
	// and thus no tag name in the struct which correlates to it.
	if name == "" {
		return name
	}

	field, ok := reflect.TypeOf(post).Elem().FieldByName(name)
	if !ok {
		panic("Couldn't get struct field for: " + name + ". Make sure you pass the right field name to editor field elements.")
	}

	tag, ok := field.Tag.Lookup("json")
	if !ok {
		panic("Couldn't get json struct tag for: " + name + ". Struct fields for content types must have 'json' tags.")
	}

	return tag
}

// due to the format in which gorilla/schema expects form names to be when
// one is associated with multiple values, we need to output the name as such.
// Ex. 'category.0', 'category.1', 'category.2' and so on.
func tagNameFromStructFieldMulti(name string, i int, post interface{}) string {
	tag := tagNameFromStructField(name, post)

	return fmt.Sprintf("%s.%d", tag, i)
}

func valueFromStructField(name string, post interface{}) reflect.Value {
	field := reflect.Indirect(reflect.ValueOf(post)).FieldByName(name)

	return field
}

func newElement(tagName, label, fieldName string, p interface{}, attrs map[string]string) *element {
	return &element{
		TagName: tagName,
		Attrs:   attrs,
		Name:    tagNameFromStructField(fieldName, p),
		label:   label,
		data:    valueFromStructField(fieldName, p).String(),
		viewBuf: &bytes.Buffer{},
	}
}
<|endoftext|>


<commit_before>package staticdump

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/evepraisal/go-evepraisal"

	"gopkg.in/yaml.v2"
)

type TypeDB struct {
	staticDumpURL string
	dir           string

	typeMap map[string]evepraisal.EveType
}

func NewTypeDB(dir string, staticDumpURL string) (evepraisal.TypeDB, error) {

	typeDB := &TypeDB{
		typeMap:       make(map[string]evepraisal.EveType),
		staticDumpURL: staticDumpURL,
		dir:           dir,
	}

	if _, err := os.Stat(typeDB.staticDumpPath()); os.IsNotExist(err) {
		log.Printf("Downloading static dump to %s", typeDB.staticDumpPath())
		err := typeDB.downloadStaticDump()
		if err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	log.Println("Load type data")
	err := typeDB.loadTypeData()
	if err != nil {
		return nil, err
	}

	log.Println("Done loading type data")

	return typeDB, nil
}

func (db *TypeDB) staticDumpPath() string {
	return filepath.Join(db.dir, filepath.Base(db.staticDumpURL))
}

func (db *TypeDB) HasType(typeName string) bool {
	_, ok := db.GetType(typeName)
	return ok
}

func (db *TypeDB) GetType(typeName string) (evepraisal.EveType, bool) {
	t, ok := db.typeMap[strings.ToLower(typeName)]
	return t, ok
}

func (db *TypeDB) Close() error {
	return nil
}

func (db *TypeDB) downloadStaticDump() error {
	out, err := os.Create(db.staticDumpPath())
	defer out.Close()
	if err != nil {
		return err
	}

	req, err := http.NewRequest("GET", db.staticDumpURL, nil)
	if err != nil {
		return err
	}
	req.Header.Add("User-Agent", "go-evepraisal")
	resp, err := http.DefaultClient.Do(req)
	defer resp.Body.Close()
	n, err := io.Copy(out, resp.Body)
	if err != nil {
		return err
	}

	log.Println("Successfully wrote %d bytes to %s", n, db.staticDumpPath())
	return nil
}

func (db *TypeDB) loadTypeData() error {
	r, err := zip.OpenReader(db.staticDumpPath())
	if err != nil {
		return err
	}
	defer r.Close()

	f, err := findZipFile(r.File, "sde/fsd/typeIDs.yaml")
	if err != nil {
		return err
	}

	fr, err := f.Open()
	if err != nil {
		return err
	}

	typeIDContents, err := ioutil.ReadAll(fr)
	if err != nil {
		return err
	}

	var allTypes map[int64]Type

	err = yaml.Unmarshal(typeIDContents, &allTypes)
	if err != nil {
		return err
	}

	typeMap := make(map[string]evepraisal.EveType)
	for typeID, t := range allTypes {
		typeMap[strings.ToLower(t.Name.En)] = evepraisal.EveType{
			ID:   typeID,
			Name: t.Name.En,
		}
	}

	db.typeMap = typeMap

	return nil
}

type Type struct {
	Name struct {
		En string
	}
}

func findZipFile(files []*zip.File, filename string) (*zip.File, error) {
	for _, f := range files {
		if filename == f.Name {
			return f, nil
		}
	}
	return nil, fmt.Errorf("Could not locate %s in archive", filename)
}
<commit_msg>Fix lint issues<commit_after>package staticdump

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/evepraisal/go-evepraisal"

	"gopkg.in/yaml.v2"
)

type TypeDB struct {
	staticDumpURL string
	dir           string

	typeMap map[string]evepraisal.EveType
}

func NewTypeDB(dir string, staticDumpURL string) (evepraisal.TypeDB, error) {

	typeDB := &TypeDB{
		typeMap:       make(map[string]evepraisal.EveType),
		staticDumpURL: staticDumpURL,
		dir:           dir,
	}

	if _, err := os.Stat(typeDB.staticDumpPath()); os.IsNotExist(err) {
		log.Printf("Downloading static dump to %s", typeDB.staticDumpPath())
		err := typeDB.downloadStaticDump()
		if err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	log.Println("Load type data")
	err := typeDB.loadTypeData()
	if err != nil {
		return nil, err
	}

	log.Println("Done loading type data")

	return typeDB, nil
}

func (db *TypeDB) staticDumpPath() string {
	return filepath.Join(db.dir, filepath.Base(db.staticDumpURL))
}

func (db *TypeDB) HasType(typeName string) bool {
	_, ok := db.GetType(typeName)
	return ok
}

func (db *TypeDB) GetType(typeName string) (evepraisal.EveType, bool) {
	t, ok := db.typeMap[strings.ToLower(typeName)]
	return t, ok
}

func (db *TypeDB) Close() error {
	return nil
}

func (db *TypeDB) downloadStaticDump() error {
	out, err := os.Create(db.staticDumpPath())
	defer out.Close()
	if err != nil {
		return err
	}

	req, err := http.NewRequest("GET", db.staticDumpURL, nil)
	if err != nil {
		return err
	}
	req.Header.Add("User-Agent", "go-evepraisal")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	n, err := io.Copy(out, resp.Body)
	if err != nil {
		return err
	}

	log.Printf("Successfully wrote %d bytes to %s", n, db.staticDumpPath())
	return nil
}

func (db *TypeDB) loadTypeData() error {
	r, err := zip.OpenReader(db.staticDumpPath())
	if err != nil {
		return err
	}
	defer r.Close()

	f, err := findZipFile(r.File, "sde/fsd/typeIDs.yaml")
	if err != nil {
		return err
	}

	fr, err := f.Open()
	if err != nil {
		return err
	}

	typeIDContents, err := ioutil.ReadAll(fr)
	if err != nil {
		return err
	}

	var allTypes map[int64]Type

	err = yaml.Unmarshal(typeIDContents, &allTypes)
	if err != nil {
		return err
	}

	typeMap := make(map[string]evepraisal.EveType)
	for typeID, t := range allTypes {
		typeMap[strings.ToLower(t.Name.En)] = evepraisal.EveType{
			ID:   typeID,
			Name: t.Name.En,
		}
	}

	db.typeMap = typeMap

	return nil
}

type Type struct {
	Name struct {
		En string
	}
}

func findZipFile(files []*zip.File, filename string) (*zip.File, error) {
	for _, f := range files {
		if filename == f.Name {
			return f, nil
		}
	}
	return nil, fmt.Errorf("Could not locate %s in archive", filename)
}
<|endoftext|>


<commit_before>package driver

import (
	"encoding/json"
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/executor"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/client/getter"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reQemuVersion = regexp.MustCompile(`version (\d[\.\d+]+)`)
)

// QemuDriver is a driver for running images via Qemu
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type QemuDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type QemuDriverConfig struct {
	ArtifactSource string           `mapstructure:"artifact_source"`
	Checksum       string           `mapstructure:"checksum"`
	Accelerator    string           `mapstructure:"accelerator"`
	PortMap        []map[string]int `mapstructure:"port_map"` // A map of host port labels and to guest ports.
}

// qemuHandle is returned from Start/Open as a handle to the PID
type qemuHandle struct {
	cmd         executor.Executor
	killTimeout time.Duration
	logger      *log.Logger
	waitCh      chan *cstructs.WaitResult
	doneCh      chan struct{}
}

// NewQemuDriver is used to create a new exec driver
func NewQemuDriver(ctx *DriverContext) Driver {
	return &QemuDriver{DriverContext: *ctx}
}

func (d *QemuDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	bin := "qemu-system-x86_64"
	if runtime.GOOS == "windows" {
		// On windows, the "qemu-system-x86_64" command does not respond to the
		// version flag.
		bin = "qemu-img"
	}
	outBytes, err := exec.Command(bin, "--version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	matches := reQemuVersion.FindStringSubmatch(out)
	if len(matches) != 2 {
		return false, fmt.Errorf("Unable to parse Qemu version string: %#v", matches)
	}

	node.Attributes["driver.qemu"] = "1"
	node.Attributes["driver.qemu.version"] = matches[1]

	return true, nil
}

// Run an existing Qemu image. Start() will pull down an existing, valid Qemu
// image and save it to the Drivers Allocation Dir
func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig QemuDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}

	if len(driverConfig.PortMap) > 1 {
		return nil, fmt.Errorf("Only one port_map block is allowed in the qemu driver config")
	}

	// Get the image source
	source, ok := task.Config["artifact_source"]
	if !ok || source == "" {
		return nil, fmt.Errorf("Missing source image Qemu driver")
	}

	// Qemu defaults to 128M of RAM for a given VM. Instead, we force users to
	// supply a memory size in the tasks resources
	if task.Resources == nil || task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Missing required Task Resource: Memory")
	}

	// Get the tasks local directory.
	taskDir, ok := ctx.AllocDir.TaskDirs[d.DriverContext.taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}

	// Proceed to download an artifact to be executed.
	vmPath, err := getter.GetArtifact(
		filepath.Join(taskDir, allocdir.TaskLocal),
		driverConfig.ArtifactSource,
		driverConfig.Checksum,
		d.logger,
	)
	if err != nil {
		return nil, err
	}

	vmID := filepath.Base(vmPath)

	// Parse configuration arguments
	// Create the base arguments
	accelerator := "tcg"
	if driverConfig.Accelerator != "" {
		accelerator = driverConfig.Accelerator
	}
	// TODO: Check a lower bounds, e.g. the default 128 of Qemu
	mem := fmt.Sprintf("%dM", task.Resources.MemoryMB)

	args := []string{
		"qemu-system-x86_64",
		"-machine", "type=pc,accel=" + accelerator,
		"-name", vmID,
		"-m", mem,
		"-drive", "file=" + vmPath,
		"-nodefconfig",
		"-nodefaults",
		"-nographic",
	}

	// Check the Resources required Networks to add port mappings. If no resources
	// are required, we assume the VM is a purely compute job and does not require
	// the outside world to be able to reach it. VMs ran without port mappings can
	// still reach out to the world, but without port mappings it is effectively
	// firewalled
	protocols := []string{"udp", "tcp"}
	if len(task.Resources.Networks) > 0 && len(driverConfig.PortMap) == 1 {
		// Loop through the port map and construct the hostfwd string, to map
		// reserved ports to the ports listenting in the VM
		// Ex: hostfwd=tcp::22000-:22,hostfwd=tcp::80-:8080
		var forwarding []string
		taskPorts := task.Resources.Networks[0].MapLabelToValues(nil)
		for label, guest := range driverConfig.PortMap[0] {
			host, ok := taskPorts[label]
			if !ok {
				return nil, fmt.Errorf("Unknown port label %q", label)
			}

			for _, p := range protocols {
				forwarding = append(forwarding, fmt.Sprintf("hostfwd=%s::%d-:%d", p, host, guest))
			}
		}

		if len(forwarding) != 0 {
			args = append(args,
				"-netdev",
				fmt.Sprintf("user,id=user.0,%s", strings.Join(forwarding, ",")),
				"-device", "virtio-net,netdev=user.0",
			)
		}
	}

	// If using KVM, add optimization args
	if accelerator == "kvm" {
		args = append(args,
			"-enable-kvm",
			"-cpu", "host",
			// Do we have cores information available to the Driver?
			// "-smp", fmt.Sprintf("%d", cores),
		)
	}

	// Setup the command
	cmd := executor.Command(args[0], args[1:]...)
	if err := cmd.Limit(task.Resources); err != nil {
		return nil, fmt.Errorf("failed to constrain resources: %s", err)
	}

	if err := cmd.ConfigureTaskDir(d.taskName, ctx.AllocDir); err != nil {
		return nil, fmt.Errorf("failed to configure task directory: %v", err)
	}

	d.logger.Printf("[DEBUG] Starting QemuVM command: %q", strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("failed to start command: %v", err)
	}
	d.logger.Printf("[INFO] Started new QemuVM: %s", vmID)

	// Create and Return Handle
	h := &qemuHandle{
		cmd:         cmd,
		killTimeout: d.DriverContext.KillTimeout(task),
		logger:      d.logger,
		doneCh:      make(chan struct{}),
		waitCh:      make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

type qemuId struct {
	ExecutorId  string
	KillTimeout time.Duration
}

func (d *QemuDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	id := &qemuId{}
	if err := json.Unmarshal([]byte(handleID), id); err != nil {
		return nil, fmt.Errorf("Failed to parse handle '%s': %v", handleID, err)
	}

	// Find the process
	cmd, err := executor.OpenId(id.ExecutorId)
	if err != nil {
		return nil, fmt.Errorf("failed to open ID %v: %v", id.ExecutorId, err)
	}

	// Return a driver handle
	h := &execHandle{
		cmd:         cmd,
		logger:      d.logger,
		killTimeout: id.KillTimeout,
		doneCh:      make(chan struct{}),
		waitCh:      make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (h *qemuHandle) ID() string {
	executorId, _ := h.cmd.ID()
	id := qemuId{
		ExecutorId:  executorId,
		KillTimeout: h.killTimeout,
	}

	data, err := json.Marshal(id)
	if err != nil {
		h.logger.Printf("[ERR] driver.qemu: failed to marshal ID to JSON: %s", err)
	}
	return string(data)
}

func (h *qemuHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *qemuHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *qemuHandle) Kill() error {
	h.cmd.Shutdown()
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(h.killTimeout):
		return h.cmd.ForceStop()
	}
}

func (h *qemuHandle) run() {
	res := h.cmd.Wait()
	close(h.doneCh)
	h.waitCh <- res
	close(h.waitCh)
}
<commit_msg>Use basic executor<commit_after>package driver

import (
	"encoding/json"
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/executor"
	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/client/getter"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/mitchellh/mapstructure"
)

var (
	reQemuVersion = regexp.MustCompile(`version (\d[\.\d+]+)`)
)

// QemuDriver is a driver for running images via Qemu
// We attempt to chose sane defaults for now, with more configuration available
// planned in the future
type QemuDriver struct {
	DriverContext
	fingerprint.StaticFingerprinter
}

type QemuDriverConfig struct {
	ArtifactSource string           `mapstructure:"artifact_source"`
	Checksum       string           `mapstructure:"checksum"`
	Accelerator    string           `mapstructure:"accelerator"`
	PortMap        []map[string]int `mapstructure:"port_map"` // A map of host port labels and to guest ports.
}

// qemuHandle is returned from Start/Open as a handle to the PID
type qemuHandle struct {
	cmd         executor.Executor
	killTimeout time.Duration
	logger      *log.Logger
	waitCh      chan *cstructs.WaitResult
	doneCh      chan struct{}
}

// NewQemuDriver is used to create a new exec driver
func NewQemuDriver(ctx *DriverContext) Driver {
	return &QemuDriver{DriverContext: *ctx}
}

func (d *QemuDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool, error) {
	bin := "qemu-system-x86_64"
	if runtime.GOOS == "windows" {
		// On windows, the "qemu-system-x86_64" command does not respond to the
		// version flag.
		bin = "qemu-img"
	}
	outBytes, err := exec.Command(bin, "--version").Output()
	if err != nil {
		return false, nil
	}
	out := strings.TrimSpace(string(outBytes))

	matches := reQemuVersion.FindStringSubmatch(out)
	if len(matches) != 2 {
		return false, fmt.Errorf("Unable to parse Qemu version string: %#v", matches)
	}

	node.Attributes["driver.qemu"] = "1"
	node.Attributes["driver.qemu.version"] = matches[1]

	return true, nil
}

// Run an existing Qemu image. Start() will pull down an existing, valid Qemu
// image and save it to the Drivers Allocation Dir
func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
	var driverConfig QemuDriverConfig
	if err := mapstructure.WeakDecode(task.Config, &driverConfig); err != nil {
		return nil, err
	}

	if len(driverConfig.PortMap) > 1 {
		return nil, fmt.Errorf("Only one port_map block is allowed in the qemu driver config")
	}

	// Get the image source
	source, ok := task.Config["artifact_source"]
	if !ok || source == "" {
		return nil, fmt.Errorf("Missing source image Qemu driver")
	}

	// Qemu defaults to 128M of RAM for a given VM. Instead, we force users to
	// supply a memory size in the tasks resources
	if task.Resources == nil || task.Resources.MemoryMB == 0 {
		return nil, fmt.Errorf("Missing required Task Resource: Memory")
	}

	// Get the tasks local directory.
	taskDir, ok := ctx.AllocDir.TaskDirs[d.DriverContext.taskName]
	if !ok {
		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
	}

	// Proceed to download an artifact to be executed.
	vmPath, err := getter.GetArtifact(
		filepath.Join(taskDir, allocdir.TaskLocal),
		driverConfig.ArtifactSource,
		driverConfig.Checksum,
		d.logger,
	)
	if err != nil {
		return nil, err
	}

	vmID := filepath.Base(vmPath)

	// Parse configuration arguments
	// Create the base arguments
	accelerator := "tcg"
	if driverConfig.Accelerator != "" {
		accelerator = driverConfig.Accelerator
	}
	// TODO: Check a lower bounds, e.g. the default 128 of Qemu
	mem := fmt.Sprintf("%dM", task.Resources.MemoryMB)

	args := []string{
		"qemu-system-x86_64",
		"-machine", "type=pc,accel=" + accelerator,
		"-name", vmID,
		"-m", mem,
		"-drive", "file=" + vmPath,
		"-nodefconfig",
		"-nodefaults",
		"-nographic",
	}

	// Check the Resources required Networks to add port mappings. If no resources
	// are required, we assume the VM is a purely compute job and does not require
	// the outside world to be able to reach it. VMs ran without port mappings can
	// still reach out to the world, but without port mappings it is effectively
	// firewalled
	protocols := []string{"udp", "tcp"}
	if len(task.Resources.Networks) > 0 && len(driverConfig.PortMap) == 1 {
		// Loop through the port map and construct the hostfwd string, to map
		// reserved ports to the ports listenting in the VM
		// Ex: hostfwd=tcp::22000-:22,hostfwd=tcp::80-:8080
		var forwarding []string
		taskPorts := task.Resources.Networks[0].MapLabelToValues(nil)
		for label, guest := range driverConfig.PortMap[0] {
			host, ok := taskPorts[label]
			if !ok {
				return nil, fmt.Errorf("Unknown port label %q", label)
			}

			for _, p := range protocols {
				forwarding = append(forwarding, fmt.Sprintf("hostfwd=%s::%d-:%d", p, host, guest))
			}
		}

		if len(forwarding) != 0 {
			args = append(args,
				"-netdev",
				fmt.Sprintf("user,id=user.0,%s", strings.Join(forwarding, ",")),
				"-device", "virtio-net,netdev=user.0",
			)
		}
	}

	// If using KVM, add optimization args
	if accelerator == "kvm" {
		args = append(args,
			"-enable-kvm",
			"-cpu", "host",
			// Do we have cores information available to the Driver?
			// "-smp", fmt.Sprintf("%d", cores),
		)
	}

	// Setup the command
	cmd := executor.NewBasicExecutor()
	executor.SetCommand(cmd, args[0], args[1:])
	if err := cmd.Limit(task.Resources); err != nil {
		return nil, fmt.Errorf("failed to constrain resources: %s", err)
	}

	if err := cmd.ConfigureTaskDir(d.taskName, ctx.AllocDir); err != nil {
		return nil, fmt.Errorf("failed to configure task directory: %v", err)
	}

	d.logger.Printf("[DEBUG] Starting QemuVM command: %q", strings.Join(args, " "))
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("failed to start command: %v", err)
	}
	d.logger.Printf("[INFO] Started new QemuVM: %s", vmID)

	// Create and Return Handle
	h := &qemuHandle{
		cmd:         cmd,
		killTimeout: d.DriverContext.KillTimeout(task),
		logger:      d.logger,
		doneCh:      make(chan struct{}),
		waitCh:      make(chan *cstructs.WaitResult, 1),
	}

	go h.run()
	return h, nil
}

type qemuId struct {
	ExecutorId  string
	KillTimeout time.Duration
}

func (d *QemuDriver) Open(ctx *ExecContext, handleID string) (DriverHandle, error) {
	id := &qemuId{}
	if err := json.Unmarshal([]byte(handleID), id); err != nil {
		return nil, fmt.Errorf("Failed to parse handle '%s': %v", handleID, err)
	}

	// Find the process
	cmd := executor.NewBasicExecutor()
	if err := cmd.Open(id.ExecutorId); err != nil {
		return nil, fmt.Errorf("failed to open ID %v: %v", id.ExecutorId, err)
	}

	// Return a driver handle
	h := &execHandle{
		cmd:         cmd,
		logger:      d.logger,
		killTimeout: id.KillTimeout,
		doneCh:      make(chan struct{}),
		waitCh:      make(chan *cstructs.WaitResult, 1),
	}
	go h.run()
	return h, nil
}

func (h *qemuHandle) ID() string {
	executorId, _ := h.cmd.ID()
	id := qemuId{
		ExecutorId:  executorId,
		KillTimeout: h.killTimeout,
	}

	data, err := json.Marshal(id)
	if err != nil {
		h.logger.Printf("[ERR] driver.qemu: failed to marshal ID to JSON: %s", err)
	}
	return string(data)
}

func (h *qemuHandle) WaitCh() chan *cstructs.WaitResult {
	return h.waitCh
}

func (h *qemuHandle) Update(task *structs.Task) error {
	// Update is not possible
	return nil
}

// TODO: allow a 'shutdown_command' that can be executed over a ssh connection
// to the VM
func (h *qemuHandle) Kill() error {
	h.cmd.Shutdown()
	select {
	case <-h.doneCh:
		return nil
	case <-time.After(h.killTimeout):
		return h.cmd.ForceStop()
	}
}

func (h *qemuHandle) run() {
	res := h.cmd.Wait()
	close(h.doneCh)
	h.waitCh <- res
	close(h.waitCh)
}
<|endoftext|>


<commit_before>package grpc

import (
	"context"
	"io"
	"sync"

	"github.com/micro/go-micro/client"
	"google.golang.org/grpc"
)

// Implements the streamer interface
type grpcStream struct {
	sync.RWMutex
	err      error
	conn     *grpc.ClientConn
	stream   grpc.ClientStream
	request  client.Request
	response client.Response
	context  context.Context
}

func (g *grpcStream) Context() context.Context {
	return g.context
}

func (g *grpcStream) Request() client.Request {
	return g.request
}

func (g *grpcStream) Response() client.Response {
	return g.response
}

func (g *grpcStream) Send(msg interface{}) error {
	if err := g.stream.SendMsg(msg); err != nil {
		g.setError(err)
		return err
	}
	return nil
}

func (g *grpcStream) Recv(msg interface{}) (err error) {
	defer g.setError(err)
	if err = g.stream.RecvMsg(msg); err != nil {
		if err == io.EOF {
			// #202 - inconsistent gRPC stream behavior
			// the only way to tell if the stream is done is when we get a EOF on the Recv
			// here we should close the underlying gRPC ClientConn
			closeErr := g.conn.Close()
			if closeErr != nil {
				err = closeErr
			}
		}
	}
	return
}

func (g *grpcStream) Error() error {
	g.RLock()
	defer g.RUnlock()
	return g.err
}

func (g *grpcStream) setError(e error) {
	g.Lock()
	g.err = e
	g.Unlock()
}

// Close the gRPC send stream
// #202 - inconsistent gRPC stream behavior
// The underlying gRPC stream should not be closed here since the
// stream should still be able to receive after this function call
// TODO: should the conn be closed in another way?
func (g *grpcStream) Close() error {
	return g.stream.CloseSend()
}
<commit_msg>grpc client goroutine leak fix<commit_after>package grpc

import (
	"context"
	"io"
	"sync"

	"github.com/micro/go-micro/client"
	"google.golang.org/grpc"
)

// Implements the streamer interface
type grpcStream struct {
	sync.RWMutex
	err      error
	conn     *grpc.ClientConn
	stream   grpc.ClientStream
	request  client.Request
	response client.Response
	context  context.Context
}

func (g *grpcStream) Context() context.Context {
	return g.context
}

func (g *grpcStream) Request() client.Request {
	return g.request
}

func (g *grpcStream) Response() client.Response {
	return g.response
}

func (g *grpcStream) Send(msg interface{}) error {
	if err := g.stream.SendMsg(msg); err != nil {
		g.setError(err)
		return err
	}
	return nil
}

func (g *grpcStream) Recv(msg interface{}) (err error) {
	defer g.setError(err)
	if err = g.stream.RecvMsg(msg); err != nil {
		// #202 - inconsistent gRPC stream behavior
		// the only way to tell if the stream is done is when we get a EOF on the Recv
		// here we should close the underlying gRPC ClientConn
		closeErr := g.conn.Close()
		if err == io.EOF && closeErr != nil {
			err = closeErr
		}
	}
	return
}

func (g *grpcStream) Error() error {
	g.RLock()
	defer g.RUnlock()
	return g.err
}

func (g *grpcStream) setError(e error) {
	g.Lock()
	g.err = e
	g.Unlock()
}

// Close the gRPC send stream
// #202 - inconsistent gRPC stream behavior
// The underlying gRPC stream should not be closed here since the
// stream should still be able to receive after this function call
// TODO: should the conn be closed in another way?
func (g *grpcStream) Close() error {
	return g.stream.CloseSend()
}
<|endoftext|>