<commit_before>package azurerm

import (
	"fmt"
	"log"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
)

func resourceArmImage() *schema.Resource {
	return &schema.Resource{
		Create: resourceArmImageCreateUpdate,
		Read:   resourceArmImageRead,
		Update: resourceArmImageCreateUpdate,
		Delete: resourceArmImageDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"location": locationSchema(),

			"resource_group_name": resourceGroupNameSchema(),

			"source_virtual_machine_id": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"os_disk": {
				Type:     schema.TypeSet,
				Optional: true,
				MaxItems: 1,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"os_type": {
							Type:     schema.TypeString,
							Optional: true,
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.Linux),
								string(compute.Windows),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"os_state": {
							Type:     schema.TypeString,
							Optional: true,
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.Generalized),
								string(compute.Specialized),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"managed_disk_id": {
							Type:     schema.TypeString,
							Computed: true,
							Optional: true,
						},

						"blob_uri": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: true,
						},

						"caching": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  string(compute.None),
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.None),
								string(compute.ReadOnly),
								string(compute.ReadWrite),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"size_gb": {
							Type:     schema.TypeInt,
							Computed: true,
							Optional: true,
						},
					},
				},
			},

			"data_disk": {
				Type:     schema.TypeList,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{

						"lun": {
							Type:     schema.TypeInt,
							Optional: true,
						},

						"managed_disk_id": {
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: true,
						},

						"blob_uri": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},

						"caching": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  string(compute.None),
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.None),
								string(compute.ReadOnly),
								string(compute.ReadWrite),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"size_gb": {
							Type:     schema.TypeInt,
							Optional: true,
							Computed: true,
						},
					},
				},
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceArmImageCreateUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	imageClient := client.imageClient

	log.Printf("[INFO] preparing arguments for AzureRM Image creation.")

	name := d.Get("name").(string)
	location := d.Get("location").(string)
	resGroup := d.Get("resource_group_name").(string)
	tags := d.Get("tags").(map[string]interface{})
	expandedTags := expandTags(tags)
	properties := compute.ImageProperties{}

	osDisk, err := expandAzureRmImageOsDisk(d)
	if err != nil {
		return err
	}

	dataDisks, err := expandAzureRmImageDataDisks(d)
	if err != nil {
		return err
	}

	storageProfile := compute.ImageStorageProfile{
		OsDisk:    osDisk,
		DataDisks: &dataDisks,
	}

	sourceVM := compute.SubResource{}
	if v, ok := d.GetOk("source_virtual_machine_id"); ok {
		vmID := v.(string)
		sourceVM = compute.SubResource{
			ID: &vmID,
		}
	}

	//either source VM or storage profile can be specified, but not both
	if sourceVM.ID == nil {
		//if both sourceVM and storageProfile are empty, return an error
		if storageProfile.OsDisk == nil && len(*storageProfile.DataDisks) == 0 {
			return fmt.Errorf("[ERROR] Cannot create image when both source VM and storage profile are empty")
		}

		properties = compute.ImageProperties{
			StorageProfile: &storageProfile,
		}
	} else {
		//creating an image from source VM
		properties = compute.ImageProperties{
			SourceVirtualMachine: &sourceVM,
		}
	}

	createImage := compute.Image{
		Name:            &name,
		Location:        &location,
		Tags:            expandedTags,
		ImageProperties: &properties,
	}

	_, imageErr := imageClient.CreateOrUpdate(resGroup, name, createImage, make(chan struct{}))
	err = <-imageErr
	if err != nil {
		return err
	}

	read, err := imageClient.Get(resGroup, name, "")
	if err != nil {
		return err
	}
	if read.ID == nil {
		return fmt.Errorf("[ERROR] Cannot read AzureRM Image %s (resource group %s) ID", name, resGroup)
	}

	d.SetId(*read.ID)

	return resourceArmImageRead(d, meta)
}

func resourceArmImageRead(d *schema.ResourceData, meta interface{}) error {
	imageClient := meta.(*ArmClient).imageClient

	id, err := parseAzureResourceID(d.Id())
	if err != nil {
		return err
	}
	resGroup := id.ResourceGroup
	name := id.Path["images"]

	resp, err := imageClient.Get(resGroup, name, "")
	if err != nil {
		if utils.ResponseWasNotFound(resp.Response) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("[ERROR] Error making Read request on AzureRM Image %q (resource group %q): %+v", name, resGroup, err)
	}

	d.Set("name", resp.Name)
	d.Set("resource_group_name", resGroup)
	d.Set("location", resp.Location)

	//either source VM or storage profile can be specified, but not both
	if resp.SourceVirtualMachine != nil {
		d.Set("source_virtual_machine_id", resp.SourceVirtualMachine.ID)
	} else if resp.StorageProfile != nil {
		if disk := resp.StorageProfile.OsDisk; disk != nil {
			if err := d.Set("os_disk", flattenAzureRmImageOSDisk(d, disk)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting AzureRM Image OS Disk error: %+v", err)
			}
		}

		if disks := resp.StorageProfile.DataDisks; disks != nil {
			if err := d.Set("data_disk", flattenAzureRmImageDataDisks(d, disks)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting AzureRM Image Data Disks error: %+v", err)
			}
		}
	}

	flattenAndSetTags(d, resp.Tags)

	return nil
}

func resourceArmImageDelete(d *schema.ResourceData, meta interface{}) error {
	imageClient := meta.(*ArmClient).imageClient

	id, err := parseAzureResourceID(d.Id())
	if err != nil {
		return err
	}
	resGroup := id.ResourceGroup
	name := id.Path["images"]

	_, deleteErr := imageClient.Delete(resGroup, name, make(chan struct{}))
	err = <-deleteErr
	if err != nil {
		return err
	}

	return nil
}

func flattenAzureRmImageOSDisk(d *schema.ResourceData, osDisk *compute.ImageOSDisk) []interface{} {
	result := make(map[string]interface{})

	if disk := osDisk; disk != nil {
		result["blob_uri"] = *osDisk.BlobURI
		result["caching"] = string(osDisk.Caching)
		if osDisk.DiskSizeGB != nil {
			result["size_gb"] = *osDisk.DiskSizeGB
		}
		if osDisk.ManagedDisk != nil {
			result["managed_disk_id"] = *osDisk.ManagedDisk.ID
		}
		result["os_type"] = osDisk.OsType
		result["os_state"] = osDisk.OsState
	}

	return []interface{}{result}
}

func flattenAzureRmImageDataDisks(d *schema.ResourceData, diskImages *[]compute.ImageDataDisk) []interface{} {
	result := make([]interface{}, 0)

	if images := diskImages; images != nil {
		for i, disk := range *images {
			l := make(map[string]interface{})
			l["blob_uri"] = disk.BlobURI
			l["caching"] = string(disk.Caching)
			if disk.DiskSizeGB != nil {
				l["size_gb"] = *disk.DiskSizeGB
			}
			l["lun"] = *disk.Lun
			if disk.ManagedDisk != nil {
				l["managed_disk_id"] = *disk.ManagedDisk.ID
			}

			result[i] = l
		}
	}

	return result
}

func expandAzureRmImageOsDisk(d *schema.ResourceData) (*compute.ImageOSDisk, error) {

	osDisk := &compute.ImageOSDisk{}
	disks := d.Get("os_disk").(*schema.Set).List()

	if len(disks) > 0 {
		config := disks[0].(map[string]interface{})

		if v := config["os_type"].(string); v != "" {
			osType := compute.OperatingSystemTypes(v)
			osDisk.OsType = osType
		}

		if v := config["os_state"].(string); v != "" {
			osState := compute.OperatingSystemStateTypes(v)
			osDisk.OsState = osState
		}

		managedDiskID := config["managed_disk_id"].(string)
		if managedDiskID != "" {
			managedDisk := &compute.SubResource{
				ID: &managedDiskID,
			}
			osDisk.ManagedDisk = managedDisk
		}

		blobURI := config["blob_uri"].(string)
		osDisk.BlobURI = &blobURI

		if v := config["caching"].(string); v != "" {
			caching := compute.CachingTypes(v)
			osDisk.Caching = caching
		}

		if size := config["size_gb"]; size != 0 {
			diskSize := int32(size.(int))
			osDisk.DiskSizeGB = &diskSize
		}
	}

	return osDisk, nil
}

func expandAzureRmImageDataDisks(d *schema.ResourceData) ([]compute.ImageDataDisk, error) {

	disks := d.Get("data_disk").([]interface{})

	dataDisks := make([]compute.ImageDataDisk, 0, len(disks))
	for _, diskConfig := range disks {
		config := diskConfig.(map[string]interface{})

		managedDiskID := d.Get("managed_disk_id").(string)
		blobURI := d.Get("blob_uri").(string)
		lun := int32(config["lun"].(int))

		dataDisk := compute.ImageDataDisk{
			Lun:     &lun,
			BlobURI: &blobURI,
		}

		if size := d.Get("size_gb"); size != 0 {
			diskSize := int32(size.(int))
			dataDisk.DiskSizeGB = &diskSize
		}

		if v := d.Get("caching").(string); v != "" {
			caching := compute.CachingTypes(v)
			dataDisk.Caching = caching
		}

		if managedDiskID != "" {
			managedDisk := &compute.SubResource{
				ID: &managedDiskID,
			}
			dataDisk.ManagedDisk = managedDisk
		}

		dataDisks = append(dataDisks, dataDisk)
	}

	return dataDisks, nil

}
<commit_msg>Fixing Image<commit_after>package azurerm

import (
	"fmt"
	"log"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
)

func resourceArmImage() *schema.Resource {
	return &schema.Resource{
		Create: resourceArmImageCreateUpdate,
		Read:   resourceArmImageRead,
		Update: resourceArmImageCreateUpdate,
		Delete: resourceArmImageDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"location": locationSchema(),

			"resource_group_name": resourceGroupNameSchema(),

			"source_virtual_machine_id": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"os_disk": {
				Type:     schema.TypeSet,
				Optional: true,
				MaxItems: 1,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"os_type": {
							Type:     schema.TypeString,
							Optional: true,
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.Linux),
								string(compute.Windows),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"os_state": {
							Type:     schema.TypeString,
							Optional: true,
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.Generalized),
								string(compute.Specialized),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"managed_disk_id": {
							Type:     schema.TypeString,
							Computed: true,
							Optional: true,
						},

						"blob_uri": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
							ForceNew: true,
						},

						"caching": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  string(compute.None),
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.CachingTypesNone),
								string(compute.CachingTypesReadOnly),
								string(compute.CachingTypesReadWrite),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"size_gb": {
							Type:     schema.TypeInt,
							Computed: true,
							Optional: true,
						},
					},
				},
			},

			"data_disk": {
				Type:     schema.TypeList,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{

						"lun": {
							Type:     schema.TypeInt,
							Optional: true,
						},

						"managed_disk_id": {
							Type:     schema.TypeString,
							Optional: true,
							ForceNew: true,
						},

						"blob_uri": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},

						"caching": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  string(compute.None),
							ValidateFunc: validation.StringInSlice([]string{
								string(compute.CachingTypesNone),
								string(compute.CachingTypesReadOnly),
								string(compute.CachingTypesReadWrite),
							}, true),
							DiffSuppressFunc: ignoreCaseDiffSuppressFunc,
						},

						"size_gb": {
							Type:     schema.TypeInt,
							Optional: true,
							Computed: true,
						},
					},
				},
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceArmImageCreateUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient)
	imageClient := client.imageClient

	log.Printf("[INFO] preparing arguments for AzureRM Image creation.")

	name := d.Get("name").(string)
	location := d.Get("location").(string)
	resGroup := d.Get("resource_group_name").(string)
	tags := d.Get("tags").(map[string]interface{})
	expandedTags := expandTags(tags)
	properties := compute.ImageProperties{}

	osDisk, err := expandAzureRmImageOsDisk(d)
	if err != nil {
		return err
	}

	dataDisks, err := expandAzureRmImageDataDisks(d)
	if err != nil {
		return err
	}

	storageProfile := compute.ImageStorageProfile{
		OsDisk:    osDisk,
		DataDisks: &dataDisks,
	}

	sourceVM := compute.SubResource{}
	if v, ok := d.GetOk("source_virtual_machine_id"); ok {
		vmID := v.(string)
		sourceVM = compute.SubResource{
			ID: &vmID,
		}
	}

	//either source VM or storage profile can be specified, but not both
	if sourceVM.ID == nil {
		//if both sourceVM and storageProfile are empty, return an error
		if storageProfile.OsDisk == nil && len(*storageProfile.DataDisks) == 0 {
			return fmt.Errorf("[ERROR] Cannot create image when both source VM and storage profile are empty")
		}

		properties = compute.ImageProperties{
			StorageProfile: &storageProfile,
		}
	} else {
		//creating an image from source VM
		properties = compute.ImageProperties{
			SourceVirtualMachine: &sourceVM,
		}
	}

	createImage := compute.Image{
		Name:            &name,
		Location:        &location,
		Tags:            expandedTags,
		ImageProperties: &properties,
	}

	_, imageErr := imageClient.CreateOrUpdate(resGroup, name, createImage, make(chan struct{}))
	err = <-imageErr
	if err != nil {
		return err
	}

	read, err := imageClient.Get(resGroup, name, "")
	if err != nil {
		return err
	}
	if read.ID == nil {
		return fmt.Errorf("[ERROR] Cannot read AzureRM Image %s (resource group %s) ID", name, resGroup)
	}

	d.SetId(*read.ID)

	return resourceArmImageRead(d, meta)
}

func resourceArmImageRead(d *schema.ResourceData, meta interface{}) error {
	imageClient := meta.(*ArmClient).imageClient

	id, err := parseAzureResourceID(d.Id())
	if err != nil {
		return err
	}
	resGroup := id.ResourceGroup
	name := id.Path["images"]

	resp, err := imageClient.Get(resGroup, name, "")
	if err != nil {
		if utils.ResponseWasNotFound(resp.Response) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("[ERROR] Error making Read request on AzureRM Image %q (resource group %q): %+v", name, resGroup, err)
	}

	d.Set("name", resp.Name)
	d.Set("resource_group_name", resGroup)
	d.Set("location", azureRMNormalizeLocation(*resp.Location))

	//either source VM or storage profile can be specified, but not both
	if resp.SourceVirtualMachine != nil {
		d.Set("source_virtual_machine_id", resp.SourceVirtualMachine.ID)
	} else if resp.StorageProfile != nil {
		if disk := resp.StorageProfile.OsDisk; disk != nil {
			if err := d.Set("os_disk", flattenAzureRmImageOSDisk(d, disk)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting AzureRM Image OS Disk error: %+v", err)
			}
		}

		if disks := resp.StorageProfile.DataDisks; disks != nil {
			if err := d.Set("data_disk", flattenAzureRmImageDataDisks(d, disks)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting AzureRM Image Data Disks error: %+v", err)
			}
		}
	}

	flattenAndSetTags(d, resp.Tags)

	return nil
}

func resourceArmImageDelete(d *schema.ResourceData, meta interface{}) error {
	imageClient := meta.(*ArmClient).imageClient

	id, err := parseAzureResourceID(d.Id())
	if err != nil {
		return err
	}
	resGroup := id.ResourceGroup
	name := id.Path["images"]

	_, deleteErr := imageClient.Delete(resGroup, name, make(chan struct{}))
	err = <-deleteErr
	if err != nil {
		return err
	}

	return nil
}

func flattenAzureRmImageOSDisk(d *schema.ResourceData, osDisk *compute.ImageOSDisk) []interface{} {
	result := make(map[string]interface{})

	if disk := osDisk; disk != nil {
		result["blob_uri"] = *osDisk.BlobURI
		result["caching"] = string(osDisk.Caching)
		if osDisk.DiskSizeGB != nil {
			result["size_gb"] = *osDisk.DiskSizeGB
		}
		if osDisk.ManagedDisk != nil {
			result["managed_disk_id"] = *osDisk.ManagedDisk.ID
		}
		result["os_type"] = osDisk.OsType
		result["os_state"] = osDisk.OsState
	}

	return []interface{}{result}
}

func flattenAzureRmImageDataDisks(d *schema.ResourceData, diskImages *[]compute.ImageDataDisk) []interface{} {
	result := make([]interface{}, 0)

	if images := diskImages; images != nil {
		for i, disk := range *images {
			l := make(map[string]interface{})
			l["blob_uri"] = disk.BlobURI
			l["caching"] = string(disk.Caching)
			if disk.DiskSizeGB != nil {
				l["size_gb"] = *disk.DiskSizeGB
			}
			l["lun"] = *disk.Lun
			if disk.ManagedDisk != nil {
				l["managed_disk_id"] = *disk.ManagedDisk.ID
			}

			result[i] = l
		}
	}

	return result
}

func expandAzureRmImageOsDisk(d *schema.ResourceData) (*compute.ImageOSDisk, error) {

	osDisk := &compute.ImageOSDisk{}
	disks := d.Get("os_disk").(*schema.Set).List()

	if len(disks) > 0 {
		config := disks[0].(map[string]interface{})

		if v := config["os_type"].(string); v != "" {
			osType := compute.OperatingSystemTypes(v)
			osDisk.OsType = osType
		}

		if v := config["os_state"].(string); v != "" {
			osState := compute.OperatingSystemStateTypes(v)
			osDisk.OsState = osState
		}

		managedDiskID := config["managed_disk_id"].(string)
		if managedDiskID != "" {
			managedDisk := &compute.SubResource{
				ID: &managedDiskID,
			}
			osDisk.ManagedDisk = managedDisk
		}

		blobURI := config["blob_uri"].(string)
		osDisk.BlobURI = &blobURI

		if v := config["caching"].(string); v != "" {
			caching := compute.CachingTypes(v)
			osDisk.Caching = caching
		}

		if size := config["size_gb"]; size != 0 {
			diskSize := int32(size.(int))
			osDisk.DiskSizeGB = &diskSize
		}
	}

	return osDisk, nil
}

func expandAzureRmImageDataDisks(d *schema.ResourceData) ([]compute.ImageDataDisk, error) {

	disks := d.Get("data_disk").([]interface{})

	dataDisks := make([]compute.ImageDataDisk, 0, len(disks))
	for _, diskConfig := range disks {
		config := diskConfig.(map[string]interface{})

		managedDiskID := d.Get("managed_disk_id").(string)
		blobURI := d.Get("blob_uri").(string)
		lun := int32(config["lun"].(int))

		dataDisk := compute.ImageDataDisk{
			Lun:     &lun,
			BlobURI: &blobURI,
		}

		if size := d.Get("size_gb"); size != 0 {
			diskSize := int32(size.(int))
			dataDisk.DiskSizeGB = &diskSize
		}

		if v := d.Get("caching").(string); v != "" {
			caching := compute.CachingTypes(v)
			dataDisk.Caching = caching
		}

		if managedDiskID != "" {
			managedDisk := &compute.SubResource{
				ID: &managedDiskID,
			}
			dataDisk.ManagedDisk = managedDisk
		}

		dataDisks = append(dataDisks, dataDisk)
	}

	return dataDisks, nil

}
<|endoftext|>


<commit_before>package main

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

	lib "github.com/bazooka-ci/bazooka/commons"
)

type Generator struct {
	config       *lib.Config
	outputFolder string
	index        string
}

func (g *Generator) GenerateDockerfile() error {
	err := os.MkdirAll(fmt.Sprintf("%s/%s", g.outputFolder, g.index), 0755)
	if err != nil {
		return err
	}

	var dockerBuffer bytes.Buffer

	dockerBuffer.WriteString(fmt.Sprintf("FROM %s\n\n", g.config.FromImage))

	envMap := lib.GetEnvMap(g.config.Env)

	bzkBuildDir := envMap["BZK_BUILD_DIR"][0].Value

	dockerBuffer.WriteString(fmt.Sprintf("COPY source %s/\n\n", bzkBuildDir))

	dockerBuffer.WriteString(fmt.Sprintf("COPY work/%s/bazooka_run.sh %s/\n", g.index, bzkBuildDir))
	dockerBuffer.WriteString(fmt.Sprintf("RUN  chmod +x %s/bazooka_run.sh\n\n", bzkBuildDir))

	type buildPhase struct {
		name      string
		commands  []string
		beforeCmd []string
		runCmd    []string
		always    bool
	}

	phases := []*buildPhase{
		&buildPhase{
			name:      "setup",
			commands:  g.config.Setup,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"./bazooka_setup.sh",
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "before_install",
			commands:  g.config.BeforeInstall,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"./bazooka_before_install.sh",
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "install",
			commands:  g.config.Install,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"./bazooka_install.sh",
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "before_script",
			commands:  g.config.BeforeScript,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"./bazooka_before_script.sh",
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "script",
			commands:  g.config.Script,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"./bazooka_script.sh",
				"exitCode=$?",
			},
		},
		&buildPhase{
			name:     "archive",
			commands: archiveCommands(g.config.Archive),
		},
		&buildPhase{
			name:     "archive_success",
			commands: archiveCommands(g.config.ArchiveSuccess),
			runCmd: []string{
				"if [[ $exitCode == 0 ]] ; then",
				"  ./bazooka_archive_success.sh",
				"fi",
			},
		},
		&buildPhase{
			name:     "archive_failure",
			commands: archiveCommands(g.config.ArchiveFailure),
			runCmd: []string{
				"if [[ $exitCode != 0 ]] ; then",
				"  ./bazooka_archive_failure.sh",
				"fi",
			},
		},
		&buildPhase{
			name:      "after_success",
			commands:  g.config.AfterSuccess,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"if [[ $exitCode == 0 ]] ; then",
				"  ./bazooka_after_success.sh",
				"fi",
			},
		},
		&buildPhase{
			name:      "after_failure",
			commands:  g.config.AfterFailure,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"if [[ $exitCode != 0 ]] ; then",
				"  ./bazooka_after_failure.sh",
				"fi"},
		},
		&buildPhase{
			name:      "after_script",
			commands:  g.config.AfterScript,
			beforeCmd: []string{"set -e"},
			runCmd:    []string{},
		},
		&buildPhase{
			name:   "end",
			always: true,
			runCmd: []string{
				"exit $exitCode",
			},
		},
	}

	var bufferRun bytes.Buffer
	bufferRun.WriteString("#!/bin/bash\n")

	for _, phase := range phases {
		if len(phase.commands) > 0 {
			var phaseBuffer bytes.Buffer
			phaseBuffer.WriteString("#!/bin/bash\n\n")
			phaseBuffer.WriteString(fmt.Sprintf("echo %s\n", strconv.Quote(fmt.Sprintf("<PHASE:%s>", phase.name))))
			for _, action := range phase.beforeCmd {
				phaseBuffer.WriteString(fmt.Sprintf("%s\n", action))
			}
			for _, action := range phase.commands {
				phaseBuffer.WriteString(fmt.Sprintf("echo %s\n", strconv.Quote(fmt.Sprintf("<CMD:%s>", action))))
				phaseBuffer.WriteString(fmt.Sprintf("%s\n", action))
			}
			err = ioutil.WriteFile(fmt.Sprintf("%s/%s/bazooka_%s.sh", g.outputFolder, g.index, phase.name), phaseBuffer.Bytes(), 0644)
			if err != nil {
				return fmt.Errorf("Phase [%s/%s]: writing file failed: %v", g.index, phase.name, err)
			}

			dockerBuffer.WriteString(fmt.Sprintf("COPY work/%s/bazooka_%s.sh %s/\n", g.index, phase.name, bzkBuildDir))
			dockerBuffer.WriteString(fmt.Sprintf("RUN  chmod +x %s/bazooka_%s.sh\n\n", bzkBuildDir, phase.name))

			if len(phase.runCmd) == 0 {
				bufferRun.WriteString(fmt.Sprintf("./bazooka_%s.sh\n", phase.name))
			} else {
				for _, action := range phase.runCmd {
					bufferRun.WriteString(fmt.Sprintf("%s\n", action))
				}
			}
		} else if phase.always {
			for _, action := range phase.runCmd {
				bufferRun.WriteString(fmt.Sprintf("%s\n", action))
			}
		}
	}

	err = ioutil.WriteFile(fmt.Sprintf("%s/%s/bazooka_run.sh", g.outputFolder, g.index), bufferRun.Bytes(), 0644)
	if err != nil {
		return fmt.Errorf("Phase [%s/run]: writing file failed: %v", g.index, err)
	}

	for _, env := range g.config.Env {
		dockerBuffer.WriteString(fmt.Sprintf("ENV  %s %s\n", env.Name, env.Value))
	}

	dockerBuffer.WriteString(fmt.Sprintf("WORKDIR %s\n\n", bzkBuildDir))

	dockerBuffer.WriteString("CMD  ./bazooka_run.sh\n")

	err = ioutil.WriteFile(fmt.Sprintf("%s/%s/Dockerfile", g.outputFolder, g.index), dockerBuffer.Bytes(), 0644)
	if err != nil {
		return fmt.Errorf("Phase [%s/docker]: writing file failed: %v", g.index, err)
	}

	if len(g.config.Services) > 0 {
		err = lib.Flush(g.config.Services, fmt.Sprintf("%s/%s/services", g.outputFolder, g.index))
		if err != nil {
			return fmt.Errorf("Phase [%s/services]: writing file failed: %v", g.index, err)
		}
	}
	return nil
}

func archiveCommands(globs lib.Globs) []string {
	res := make([]string, len(globs))
	for i, pat := range globs {
		res[i] = fmt.Sprintf("cp -R %s /artifacts/", pat)
	}
	return res
}
<commit_msg>Update generator to manage changes in working directory<commit_after>package main

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

	lib "github.com/bazooka-ci/bazooka/commons"
)

type Generator struct {
	config       *lib.Config
	outputFolder string
	index        string
}

func (g *Generator) GenerateDockerfile() error {
	err := os.MkdirAll(fmt.Sprintf("%s/%s", g.outputFolder, g.index), 0755)
	if err != nil {
		return err
	}

	var dockerBuffer bytes.Buffer

	dockerBuffer.WriteString(fmt.Sprintf("FROM %s\n\n", g.config.FromImage))

	envMap := lib.GetEnvMap(g.config.Env)

	bzkBuildDir := envMap["BZK_BUILD_DIR"][0].Value

	dockerBuffer.WriteString(fmt.Sprintf("COPY source %s/\n\n", bzkBuildDir))

	dockerBuffer.WriteString(fmt.Sprintf("COPY work/%s/bazooka_run.sh %s/\n", g.index, bzkBuildDir))
	dockerBuffer.WriteString(fmt.Sprintf("RUN  chmod +x %s/bazooka_run.sh\n\n", bzkBuildDir))

	type buildPhase struct {
		name      string
		commands  []string
		beforeCmd []string
		runCmd    []string
		always    bool
	}

	phases := []*buildPhase{
		&buildPhase{
			name:      "setup",
			commands:  g.config.Setup,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				fmt.Sprintf("%s/bazooka_setup.sh", bzkBuildDir),
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "before_install",
			commands:  g.config.BeforeInstall,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				fmt.Sprintf("%s/bazooka_before_install.sh", bzkBuildDir),
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "install",
			commands:  g.config.Install,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				fmt.Sprintf("%s/bazooka_install.sh", bzkBuildDir),
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "before_script",
			commands:  g.config.BeforeScript,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				fmt.Sprintf("%s/bazooka_before_script.sh", bzkBuildDir),
				"rc=$?",
				"if [[ $rc != 0 ]] ; then",
				"    exit 42",
				"fi",
			},
		},
		&buildPhase{
			name:      "script",
			commands:  g.config.Script,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				fmt.Sprintf("%s/bazooka_script.sh", bzkBuildDir),
				"exitCode=$?",
			},
		},
		&buildPhase{
			name:     "archive",
			commands: archiveCommands(g.config.Archive),
		},
		&buildPhase{
			name:     "archive_success",
			commands: archiveCommands(g.config.ArchiveSuccess),
			runCmd: []string{
				"if [[ $exitCode == 0 ]] ; then",
				fmt.Sprintf("  %s/bazooka_archive_success.sh", bzkBuildDir),
				"fi",
			},
		},
		&buildPhase{
			name:     "archive_failure",
			commands: archiveCommands(g.config.ArchiveFailure),
			runCmd: []string{
				"if [[ $exitCode != 0 ]] ; then",
				fmt.Sprintf("  %s/bazooka_archive_failure.sh", bzkBuildDir),
				"fi",
			},
		},
		&buildPhase{
			name:      "after_success",
			commands:  g.config.AfterSuccess,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"if [[ $exitCode == 0 ]] ; then",
				fmt.Sprintf("  %s/bazooka_after_success.sh", bzkBuildDir),
				"fi",
			},
		},
		&buildPhase{
			name:      "after_failure",
			commands:  g.config.AfterFailure,
			beforeCmd: []string{"set -e"},
			runCmd: []string{
				"if [[ $exitCode != 0 ]] ; then",
				fmt.Sprintf("  %s/bazooka_after_failure.sh", bzkBuildDir),
				"fi"},
		},
		&buildPhase{
			name:      "after_script",
			commands:  g.config.AfterScript,
			beforeCmd: []string{"set -e"},
			runCmd:    []string{},
		},
		&buildPhase{
			name:   "end",
			always: true,
			runCmd: []string{
				"exit $exitCode",
			},
		},
	}

	var bufferRun bytes.Buffer
	bufferRun.WriteString("#!/bin/bash\n")

	for _, phase := range phases {
		if len(phase.commands) > 0 {
			bufferRun.WriteString("source /.bzkenv\n")
			var phaseBuffer bytes.Buffer
			phaseBuffer.WriteString("#!/bin/bash\n\n")
			phaseBuffer.WriteString(fmt.Sprintf("echo %s\n", strconv.Quote(fmt.Sprintf("<PHASE:%s>", phase.name))))
			for _, action := range phase.beforeCmd {
				phaseBuffer.WriteString(fmt.Sprintf("%s\n", action))
			}
			for _, action := range phase.commands {
				phaseBuffer.WriteString(fmt.Sprintf("echo %s\n", strconv.Quote(fmt.Sprintf("<CMD:%s>", action))))
				phaseBuffer.WriteString(fmt.Sprintf("%s\n", action))
			}
			phaseBuffer.WriteString(`echo "cd \"$(pwd)\"" > /.bzkenv`)
			phaseBuffer.WriteString("\n")
			err = ioutil.WriteFile(fmt.Sprintf("%s/%s/bazooka_%s.sh", g.outputFolder, g.index, phase.name), phaseBuffer.Bytes(), 0644)
			if err != nil {
				return fmt.Errorf("Phase [%s/%s]: writing file failed: %v", g.index, phase.name, err)
			}

			dockerBuffer.WriteString(fmt.Sprintf("COPY work/%s/bazooka_%s.sh %s/\n", g.index, phase.name, bzkBuildDir))
			dockerBuffer.WriteString(fmt.Sprintf("RUN  chmod +x %s/bazooka_%s.sh\n\n", bzkBuildDir, phase.name))

			if len(phase.runCmd) == 0 {
				bufferRun.WriteString(fmt.Sprintf("%s/bazooka_%s.sh\n", bzkBuildDir, phase.name))
			} else {
				for _, action := range phase.runCmd {
					bufferRun.WriteString(fmt.Sprintf("%s\n", action))
				}
			}
		} else if phase.always {
			bufferRun.WriteString("source /.bzkenv\n")
			for _, action := range phase.runCmd {
				bufferRun.WriteString(fmt.Sprintf("%s\n", action))
			}
		}
	}

	err = ioutil.WriteFile(fmt.Sprintf("%s/%s/bazooka_run.sh", g.outputFolder, g.index), bufferRun.Bytes(), 0644)
	if err != nil {
		return fmt.Errorf("Phase [%s/run]: writing file failed: %v", g.index, err)
	}

	for _, env := range g.config.Env {
		dockerBuffer.WriteString(fmt.Sprintf("ENV  %s %s\n", env.Name, env.Value))
	}

	dockerBuffer.WriteString(fmt.Sprintf("WORKDIR %s\n\n", bzkBuildDir))

	dockerBuffer.WriteString(fmt.Sprintf("CMD  %s/bazooka_run.sh\n", bzkBuildDir))

	err = ioutil.WriteFile(fmt.Sprintf("%s/%s/Dockerfile", g.outputFolder, g.index), dockerBuffer.Bytes(), 0644)
	if err != nil {
		return fmt.Errorf("Phase [%s/docker]: writing file failed: %v", g.index, err)
	}

	if len(g.config.Services) > 0 {
		err = lib.Flush(g.config.Services, fmt.Sprintf("%s/%s/services", g.outputFolder, g.index))
		if err != nil {
			return fmt.Errorf("Phase [%s/services]: writing file failed: %v", g.index, err)
		}
	}
	return nil
}

func archiveCommands(globs lib.Globs) []string {
	res := make([]string, len(globs))
	for i, pat := range globs {
		res[i] = fmt.Sprintf("cp -R %s /artifacts/", pat)
	}
	return res
}
<|endoftext|>


<commit_before>/*
    PulseHA - HA Cluster Daemon
    Copyright (C) 2017  Andrew Zak <andrew@pulseha.com>

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

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

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

import (
	"reflect"
	"errors"
)

/**
 * Memberlist struct type
 */
type Memberlist struct {
	Members []*Member
	Config *Config
}

/**
 * Member struct type
 */
type Member struct {
	Hostname   string
	Client Client
}

/**
 * Add a member to the client list
 */
func (m *Memberlist) MemberAdd(hostname string, client Client) {
	newMember := &Member{
		Hostname: hostname,
		Client:client,
	}

	m.Members = append(m.Members, newMember)
}

/**
 * Remove a member from the client list by hostname
 */
func (m *Memberlist) MemberRemoveByName(hostname string) () {
	for i, member := range m.Members {
		if member.Hostname == hostname {
			m.Members = append(m.Members[:i], m.Members[i+1:]...)
		}
	}
}

/**
 * Return Member by hostname
 */
func (m *Memberlist) GetMemberByHostname(hostname string) (*Member) {
	for _, member := range m.Members {
		if member.Hostname == hostname {
			return member
		}
	}
	return nil
}

/**
 * Return true/false whether a member exists or not.
 */
func (m *Memberlist) MemberExists(hostname string) (bool) {
	for _, member := range m.Members {
		if member.Hostname == hostname {
			return true
		}
	}
	return false
}

/**
 * Attempt to broadcast a client function to other nodes (clients) within the memberlist
 */
func (m *Memberlist) Broadcast(funcName string, params ... interface{}) (interface{}, error) {
	for _, member := range m.Members {
		funcList := member.Client.GetFuncBroadcastList()
		f := reflect.ValueOf(funcList[funcName])
		if len(params) != f.Type().NumIn() {
			return nil, errors.New("the number of passed parameters do not match the function")
		}
		vals := make([]reflect.Value, len(params))
		for k, param := range params {
			vals[k] = reflect.ValueOf(param)
		}
		f.Call(vals)
	}
	return nil, nil
}

/**
 * check how many are in the cluster
 * if more than one, request who's active.
 * if no one responds assume active.
 */
func (m *Memberlist) Setup() {
	if clusterCheck(m.Config) {
		// we are in a cluster
	}
}

<commit_msg>continued development<commit_after>/*
    PulseHA - HA Cluster Daemon
    Copyright (C) 2017  Andrew Zak <andrew@pulseha.com>

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

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

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

import (
	"reflect"
	"errors"
)

/**
 * Memberlist struct type
 */
type Memberlist struct {
	Members []*Member
	Config *Config
}

/**
 * Member struct type
 */
type Member struct {
	Hostname   string
	Client Client
}

/**
 * Add a member to the client list
 */
func (m *Memberlist) MemberAdd(hostname string, client Client) {
	newMember := &Member{
		Hostname: hostname,
		Client:client,
	}

	m.Members = append(m.Members, newMember)
}

/**
 * Remove a member from the client list by hostname
 */
func (m *Memberlist) MemberRemoveByName(hostname string) () {
	for i, member := range m.Members {
		if member.Hostname == hostname {
			m.Members = append(m.Members[:i], m.Members[i+1:]...)
		}
	}
}

/**
 * Return Member by hostname
 */
func (m *Memberlist) GetMemberByHostname(hostname string) (*Member) {
	for _, member := range m.Members {
		if member.Hostname == hostname {
			return member
		}
	}
	return nil
}

/**
 * Return true/false whether a member exists or not.
 */
func (m *Memberlist) MemberExists(hostname string) (bool) {
	for _, member := range m.Members {
		if member.Hostname == hostname {
			return true
		}
	}
	return false
}

/**
 * Attempt to broadcast a client function to other nodes (clients) within the memberlist
 */
func (m *Memberlist) Broadcast(funcName string, params ... interface{}) (interface{}, error) {
	for _, member := range m.Members {
		funcList := member.Client.GetFuncBroadcastList()
		f := reflect.ValueOf(funcList[funcName])
		if len(params) != f.Type().NumIn() {
			return nil, errors.New("the number of passed parameters do not match the function")
		}
		vals := make([]reflect.Value, len(params))
		for k, param := range params {
			vals[k] = reflect.ValueOf(param)
		}
		f.Call(vals)
	}
	return nil, nil
}

/**
 * check how many are in the cluster
 * if more than one, request who's active.
 * if no one responds assume active.
 * This should probably populate the memberlist
 */
func (m *Memberlist) Setup() {
	// Check to see if we are in a cluster
	if clusterCheck(m.Config) {
		// Are we the only member in the cluster?
		if clusterTotal(m.Config) == 1 {
			// We are the only member in the cluster so
			// we are assume that we are now the active appliance.
		} else {
			// Contact a member in the list to see who is the "active" node.
			// Iterate through the memberlist until a response is receive.
			// If no response has been made then assume we should become the active appliance.
		}
	}
}

/**
	Promote a member within the memberlist to become the active
	node
 */
func (m *Member) PromoteMember() {
	// Inform everyone in the cluster that a specific node is now the new active
	// Demote if old active is no longer active. promote if the passive is the new active.
}

<|endoftext|>


<commit_before>package main

import "os"
import "github.com/bgentry/heroku-go"
import hk "github.com/heroku/hk/hkclient"

const (
	PLUGIN_NAME    = "deploy"
	PLUGIN_VERSION = 1
	// PLUGIN_USER_AGENT = "hk-" + PLUGIN_NAME "/1"
)

var client *heroku.Client
var nrc *hk.NetRc

func help() {}

func init() {
	nrc, err := hk.LoadNetRc()
	if err != nil && os.IsNotExist(err) {
		nrc = &hk.NetRc{}
	}

	clients, err := hk.New(nrc, "TODO user agent")

	if err == nil {
		client = clients.Client
	} else {
		// TODO
	}
}

func main() {
	if os.Getenv("HKPLUGINMODE") == "info" {
		help()
		os.Exit(0)
	}

	if len(os.Args) < 2 {
		help()
		os.Exit(1)
	}

	/*
		TODO:
			* Check that we have an APP context or set it
			* tarzip current directory or git root
			* ideally with gitignoring, if possible
			* upload to S3 with an object expiry of ~5min
			* hit build API with that link
			* tail output (if build api has implemented that)
	*/
}
<commit_msg>Add a rough pass at creating a .tar.gz from provided directory<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"github.com/bgentry/heroku-go"
	hk "github.com/heroku/hk/hkclient"
	// "io"
	"io/ioutil"
	"os"
	"path/filepath"
)

const (
	PLUGIN_NAME    = "deploy"
	PLUGIN_VERSION = 1
	// PLUGIN_USER_AGENT = "hk-" + PLUGIN_NAME "/1"
)

var client *heroku.Client
var nrc *hk.NetRc

func help() {}

func init() {
	nrc, err := hk.LoadNetRc()
	if err != nil && os.IsNotExist(err) {
		nrc = &hk.NetRc{}
	}

	clients, err := hk.New(nrc, "TODO user agent")

	if err == nil {
		client = clients.Client
	} else {
		// TODO
	}
}

func shouldIgnore(path string) bool {
	// TODO: gitignore-ish rules, if a .gitignore exists?
	return path == ".git"
}

// FIXME: This is currently building a tgz that fails to extract without error:
//	x deploy.go
//	x web/README.md
//	x web/main.go: Truncated input file (needed 418 bytes, only 0 available)
//	tar: Error exit delayed from previous errors.
//
// But the files seem intact?
func buildTgz(root string) bytes.Buffer {
	buf := new(bytes.Buffer)
	gz := gzip.NewWriter(buf)
	tw := tar.NewWriter(gz)

	filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		// TODO: handle incoming err

		if shouldIgnore(path) {
			return filepath.SkipDir
		}

		if info.IsDir() {
			return nil
		}

		fmt.Printf("Adding %s (size: %d).\n", path, info.Size())
		hdr, err := tar.FileInfoHeader(info, path)
		if err != nil {
			return err
		}
		hdr.Name = path
		tw.WriteHeader(hdr)

		body, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		tw.Write(body)

		return nil
	})

	gz.Close() // `buf` now contains the tgz

	return *buf
}

func main() {
	if os.Getenv("HKPLUGINMODE") == "info" {
		help()
		os.Exit(0)
	}

	if len(os.Args) < 2 {
		help()
		os.Exit(1)
	}

	dir := os.Args[1] // TODO: Maybe fallback to CWD or Git root?
	tgz := buildTgz(dir)
	fmt.Printf("%v %d\n", tgz.Bytes(), tgz.Len())
	// fmt.Println(string(tgz.Bytes()))

	/*
		TODO:
			* Check that we have an APP context or set it
			* upload tgz to S3 with an object expiry of ~5min
			* hit build API with that link
			* tail output (if build api has implemented that)
	*/
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"strings"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

//  Will return the latest commit hash, waiting on https://github.com/google/go-github/pull/49
func latestGitHubCommit(c *github.Client, repoName string) *github.Repository {
	repo, _, err := c.Repositories.Get("gengo", repoName)
	if err != nil {
		log.Panic(err)
	}
	return repo
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, _ := template.ParseFiles("templates/index.html")
	t.Execute(w, map[string]interface{}{"Projects": projects})
}

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name     string
	RepoPath string
	Hosts    []Host
	Branch   string
}

type Project struct {
	Name         string
	GitHubURL    string
	Environments []Environment
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = v.(yaml.Map)["branch"].(yaml.Scalar).String()
		e.RepoPath = v.(yaml.Map)["repo_path"].(yaml.Scalar).String()
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for k, v := range p.(yaml.Map) {
			proj := Project{Name: k, GitHubURL: v.(yaml.Map)["github_url"].(yaml.Scalar).String()}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, environment))
				host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
				projects[i].Environments[j].Hosts[k] = host
			}
			projects[i].Environments[j] = environment
		}
	}
	return projects
}

func main() {
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	fmt.Println(client)

	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<commit_msg>reorganize (move structs to top, HomeHandler to bottom)<commit_after>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"strings"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name     string
	RepoPath string
	Hosts    []Host
	Branch   string
}

type Project struct {
	Name         string
	GitHubURL    string
	Environments []Environment
}

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

//  Will return the latest commit hash, waiting on https://github.com/google/go-github/pull/49
func latestGitHubCommit(c *github.Client, repoName string) *github.Repository {
	repo, _, err := c.Repositories.Get("gengo", repoName)
	if err != nil {
		log.Panic(err)
	}
	return repo
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = v.(yaml.Map)["branch"].(yaml.Scalar).String()
		e.RepoPath = v.(yaml.Map)["repo_path"].(yaml.Scalar).String()
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for k, v := range p.(yaml.Map) {
			proj := Project{Name: k, GitHubURL: v.(yaml.Map)["github_url"].(yaml.Scalar).String()}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, environment))
				host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
				projects[i].Environments[j].Hosts[k] = host
			}
			projects[i].Environments[j] = environment
		}
	}
	return projects
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, _ := template.ParseFiles("templates/index.html")
	t.Execute(w, map[string]interface{}{"Projects": projects})
}

func main() {
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	fmt.Println(client)

	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<|endoftext|>


<commit_before>package server
import (
	"testing"
	"gopkg.in/DATA-DOG/go-sqlmock.v1"
	"database/sql"
	"fmt"
)

func TestReplaceSqlParams(t *testing.T) {
	p1 := map[string][]string{ "a":[]string{"1"} }
	p2 := map[string][]string{ "a":[]string{"1"},"b":[]string{"2"} }

	s, p := replaceSqlParams("select 1", p1)
	if s != "select 1" || len(p) != 0 {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}", p1)
	if s != "select ?" || len(p) != 1 || p[0] != "1" {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}", p2)
	if s != "select ?" || len(p) != 1 || p[0] != "1" {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}, ${b}", p2)
	if s != "select ?, ?" || len(p) != 2 || p[0] != "1" || p[1] != "2" {
		t.Fail()
	}
}

func mockRowsToSqlRows(mockRows sqlmock.Rows) *sql.Rows {
	db, mock, _ := sqlmock.New()
	mock.ExpectQuery("select").WillReturnRows(mockRows)
	rows, _ := db.Query("select")
	return rows
}

func TestRowToResult(t *testing.T) {
	mockRows := sqlmock.NewRows([]string{"a","b","c"})
	mockRows.AddRow(1, 2, 3)
	result, err := rowsToResult(mockRowsToSqlRows(mockRows))
	if err != nil {
		t.Error(err)
	}
	if fmt.Sprintf("%v", result) != "[map[a:1 b:2 c:3]]" {
		t.Error(result)
	}
}

<commit_msg>fix test case<commit_after>package server
import (
	"testing"
	"gopkg.in/DATA-DOG/go-sqlmock.v1"
	"database/sql"
)

func TestReplaceSqlParams(t *testing.T) {
	p1 := map[string][]string{ "a":[]string{"1"} }
	p2 := map[string][]string{ "a":[]string{"1"},"b":[]string{"2"} }

	s, p := replaceSqlParams("select 1", p1)
	if s != "select 1" || len(p) != 0 {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}", p1)
	if s != "select ?" || len(p) != 1 || p[0] != "1" {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}", p2)
	if s != "select ?" || len(p) != 1 || p[0] != "1" {
		t.Fail()
	}

	s, p = replaceSqlParams("select ${a}, ${b}", p2)
	if s != "select ?, ?" || len(p) != 2 || p[0] != "1" || p[1] != "2" {
		t.Fail()
	}
}

func mockRowsToSqlRows(mockRows sqlmock.Rows) *sql.Rows {
	db, mock, _ := sqlmock.New()
	mock.ExpectQuery("select").WillReturnRows(mockRows)
	rows, _ := db.Query("select")
	return rows
}

func TestRowToResult(t *testing.T) {
	mockRows := sqlmock.NewRows([]string{"a","b","c"})
	mockRows.AddRow(1, 2, 3)
	result, err := rowsToResult(mockRowsToSqlRows(mockRows))
	if err != nil {
		t.Error(err)
	}
	if len(result) != 1 || len(result[0]) != 3 {
		t.Fail()
	}
	if result[0]["a"] != "1" || result[0]["b"] != "2" || result[0]["c"] != "3" {
		t.Fail()
	}

}

<|endoftext|>


<commit_before>// Package handler is the handler for the `micro debug stats` service
package handler

import (
	"context"
	"errors"
	"sync"
	"time"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/cmd"
	debug "github.com/micro/go-micro/debug/service/proto"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/registry/cache"
	"github.com/micro/go-micro/util/log"
	stats "github.com/micro/micro/debug/stats/proto"
)

// New initialises and returns a new Stats service handler
func New(done <-chan bool) (*Stats, error) {
	s := &Stats{
		registry: cache.New(*cmd.DefaultOptions().Registry),
		client:   *cmd.DefaultOptions().Client,
	}

	if err := s.scan(); err != nil {
		return nil, err
	}
	s.Start(done)
	return s, nil
}

// Stats is the Debug.Stats handler
type Stats struct {
	registry  registry.Registry
	client    client.Client
	snapshots []*stats.Snapshot
	next      []*stats.Snapshot
	cached    []*registry.Service
	sync.RWMutex
}

// Read returns gets a snapshot of all current stats
func (s *Stats) Read(ctx context.Context, req *stats.ReadRequest, rsp *stats.ReadResponse) error {
	s.RLock()
	rsp.Stats = s.snapshots
	s.RUnlock()
	return nil
}

// Stream starts streaming stats
func (s *Stats) Stream(ctx context.Context, req *stats.StreamRequest, rsp stats.Stats_StreamStream) error {
	return errors.New("Not Implemented")
}

// Start Starts scraping other services
// close the returned channel to stop scraping
func (s *Stats) Start(done <-chan bool) {
	go func(s *Stats) {
		for {
			select {
			case <-done:
				return
			default:
				s.scrape()
				time.Sleep(time.Second)
			}
		}
	}(s)
	go func(s *Stats) {
		rescan := time.NewTicker(10 * time.Second)
		for {
			select {
			case <-done:
				rescan.Stop()
				return
			case <-rescan.C:
				s.scan()
			}
		}
	}(s)
}

func (s *Stats) scan() error {
	services, err := s.registry.ListServices()
	if err != nil {
		return err
	}

	if err != nil {
		return err
	}

	serviceMap := make(map[string]*registry.Service)

	// check each service has nodes
	for _, service := range services {
		if len(service.Nodes) > 0 {
			serviceMap[service.Name+service.Version] = service
			continue
		}

		// get nodes that does not exist
		newServices, err := s.registry.GetService(service.Name)
		if err != nil {
			continue
		}

		// store service by version
		for _, service := range newServices {
			serviceMap[service.Name+service.Version] = service
		}
	}

	// flatten the map
	var serviceList []*registry.Service

	for _, service := range serviceMap {
		serviceList = append(serviceList, service)
	}

	// save the list
	s.Lock()
	s.cached = serviceList
	s.Unlock()
	return nil
}

func (s *Stats) scrape() {
	s.RLock()
	// Create a local copy of cached services
	services := make([]*registry.Service, len(s.cached))
	copy(services, s.cached)
	s.RUnlock()

	// Start building the next list of snapshots
	s.Lock()
	s.next = make([]*stats.Snapshot, 0)
	s.Unlock()

	// Call each node of each service in goroutines
	var wg sync.WaitGroup
	for _, svc := range services {

		// Ignore nodeless and non mucp services
		if len(svc.Nodes) == 0 {
			continue
		}
		if svc.Nodes[0].Metadata["protocol"] != "mucp" {
			continue
		}

		// Call every node
		for _, node := range svc.Nodes {
			go func(st *Stats, service *registry.Service, node *registry.Node) {
				wg.Add(1)
				defer wg.Done()

				// create new context to cancel within a few seconds
				ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
				defer cancel()

				req := s.client.NewRequest(svc.Name, "Debug.Stats", &debug.StatsRequest{})
				rsp := new(debug.StatsResponse)
				if err := s.client.Call(ctx, req, rsp, client.WithAddress(node.Address)); err != nil {
					log.Errorf("Error calling %s@%s (%s)", svc.Name, node.Address, err.Error())
					return
				}

				// Append the new snapshot
				snap := stats.Snapshot{
					Service: &stats.Service{
						Name:    service.Name,
						Version: service.Version,
						Node: &stats.Node{
							Id: node.Address,
						},
					},
					Started: int64(rsp.Started),
					Uptime:  rsp.Uptime,
					Memory:  rsp.Memory,
					Threads: rsp.Threads,
					Gc:      rsp.Gc,
				}
				s.Lock()
				s.next = append(s.next, &snap)
				s.Unlock()
			}(s, svc, node)
		}
	}
	wg.Wait()

	// Swap in the snapshots
	s.Lock()
	s.snapshots = s.next
	s.Unlock()
}
<commit_msg>Fix handler comment<commit_after>// Package handler is the handler for the `micro debug stats` service
package handler

import (
	"context"
	"errors"
	"sync"
	"time"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/cmd"
	debug "github.com/micro/go-micro/debug/service/proto"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/registry/cache"
	"github.com/micro/go-micro/util/log"
	stats "github.com/micro/micro/debug/stats/proto"
)

// New initialises and returns a new Stats service handler
func New(done <-chan bool) (*Stats, error) {
	s := &Stats{
		registry: cache.New(*cmd.DefaultOptions().Registry),
		client:   *cmd.DefaultOptions().Client,
	}

	if err := s.scan(); err != nil {
		return nil, err
	}
	s.Start(done)
	return s, nil
}

// Stats is the Debug.Stats handler
type Stats struct {
	registry  registry.Registry
	client    client.Client
	snapshots []*stats.Snapshot
	next      []*stats.Snapshot
	cached    []*registry.Service
	sync.RWMutex
}

// Read returns gets a snapshot of all current stats
func (s *Stats) Read(ctx context.Context, req *stats.ReadRequest, rsp *stats.ReadResponse) error {
	s.RLock()
	rsp.Stats = s.snapshots
	s.RUnlock()
	return nil
}

// Stream starts streaming stats
func (s *Stats) Stream(ctx context.Context, req *stats.StreamRequest, rsp stats.Stats_StreamStream) error {
	return errors.New("Not Implemented")
}

// Start Starts scraping other services until the provided channel is closed
func (s *Stats) Start(done <-chan bool) {
	go func(s *Stats) {
		for {
			select {
			case <-done:
				return
			default:
				s.scrape()
				time.Sleep(time.Second)
			}
		}
	}(s)
	go func(s *Stats) {
		rescan := time.NewTicker(10 * time.Second)
		for {
			select {
			case <-done:
				rescan.Stop()
				return
			case <-rescan.C:
				s.scan()
			}
		}
	}(s)
}

func (s *Stats) scan() error {
	services, err := s.registry.ListServices()
	if err != nil {
		return err
	}

	if err != nil {
		return err
	}

	serviceMap := make(map[string]*registry.Service)

	// check each service has nodes
	for _, service := range services {
		if len(service.Nodes) > 0 {
			serviceMap[service.Name+service.Version] = service
			continue
		}

		// get nodes that does not exist
		newServices, err := s.registry.GetService(service.Name)
		if err != nil {
			continue
		}

		// store service by version
		for _, service := range newServices {
			serviceMap[service.Name+service.Version] = service
		}
	}

	// flatten the map
	var serviceList []*registry.Service

	for _, service := range serviceMap {
		serviceList = append(serviceList, service)
	}

	// save the list
	s.Lock()
	s.cached = serviceList
	s.Unlock()
	return nil
}

func (s *Stats) scrape() {
	s.RLock()
	// Create a local copy of cached services
	services := make([]*registry.Service, len(s.cached))
	copy(services, s.cached)
	s.RUnlock()

	// Start building the next list of snapshots
	s.Lock()
	s.next = make([]*stats.Snapshot, 0)
	s.Unlock()

	// Call each node of each service in goroutines
	var wg sync.WaitGroup
	for _, svc := range services {

		// Ignore nodeless and non mucp services
		if len(svc.Nodes) == 0 {
			continue
		}
		if svc.Nodes[0].Metadata["protocol"] != "mucp" {
			continue
		}

		// Call every node
		for _, node := range svc.Nodes {
			go func(st *Stats, service *registry.Service, node *registry.Node) {
				wg.Add(1)
				defer wg.Done()

				// create new context to cancel within a few seconds
				ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
				defer cancel()

				req := s.client.NewRequest(svc.Name, "Debug.Stats", &debug.StatsRequest{})
				rsp := new(debug.StatsResponse)
				if err := s.client.Call(ctx, req, rsp, client.WithAddress(node.Address)); err != nil {
					log.Errorf("Error calling %s@%s (%s)", svc.Name, node.Address, err.Error())
					return
				}

				// Append the new snapshot
				snap := stats.Snapshot{
					Service: &stats.Service{
						Name:    service.Name,
						Version: service.Version,
						Node: &stats.Node{
							Id: node.Address,
						},
					},
					Started: int64(rsp.Started),
					Uptime:  rsp.Uptime,
					Memory:  rsp.Memory,
					Threads: rsp.Threads,
					Gc:      rsp.Gc,
				}
				s.Lock()
				s.next = append(s.next, &snap)
				s.Unlock()
			}(s, svc, node)
		}
	}
	wg.Wait()

	// Swap in the snapshots
	s.Lock()
	s.snapshots = s.next
	s.Unlock()
}
<|endoftext|>


<commit_before>package statsd

import (
	"net"
	"os"
	"strconv"
	"time"
)

// Address of the StatsD server.
var Address = "localhost:8125"

// AlsoAppendHost appends hostname along with any metric sent.
var AlsoAppendHost = true

// Cache the conn for perf.
var conn net.Conn
var host string

func init() {
	var err error

	if host, err = os.Hostname(); err != nil {
		panic(err)
	}
}

func Timer() timer {
	return timer{time.Now()}
}

// Timer
type timer struct {
	start time.Time
}

func (t *timer) Reset() {
	t.start = time.Now()
}

func (t *timer) Send(names ...interface{}) (took time.Duration) {
	took = time.Since(t.start)

	if err := getConnection(); err != nil {
		return
	}

	value := ":" + strconv.FormatUint(uint64(took.Nanoseconds()/1e6), 10) + "|ms"

	for _, name := range names {
		conn.Write([]byte(name.(string) + value))

		// Send a host suffixed stat too.
		if AlsoAppendHost {
			conn.Write([]byte(name.(string) + "." + host + value))
		}
	}

	return
}

// Gauge sets arbitrary numeric value for a given metric.
func Gauge(name string, value int64) {
	if err := getConnection(); err != nil {
		return
	}

	suffix := ":" + strconv.FormatInt(value, 10) + "|g"

	conn.Write([]byte(name + suffix))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + suffix))
	}

	return
}

// Inc is a simple counter adding one to a given metric.
func Inc(name string) {
	if err := getConnection(); err != nil {
		return
	}

	conn.Write([]byte(name + ":1|c"))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + ":1|c"))
	}

	return
}

// IncSampled increments a counter with the given sample rate (between 0.0 - 1.0).
// Example: Passing 0.1 tells statsd that this metric has been sample 1/10 times.
// IE Reported metric will be 10x the number of times called.
func IncSampled(name string, rate float64){
	if err := getConnection(); err != nil {
		return
	}

	rateString := "|@" + strconv.FormatFloat(rate, 'f', -1, 64)

	conn.Write([]byte(name + ":1|c" + rateString))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + ":1|c" + rateString))
	}

	return
}

// Time sends duration in ms for a given metric.
func Time(name string, took time.Duration) {
	if err := getConnection(); err != nil {
		return
	}

	value := ":" + strconv.FormatUint(uint64(took.Nanoseconds()/1e6), 10) + "|ms"

	conn.Write([]byte(name + value))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + value))
	}

	return
}

func getConnection() (err error) {
	// If we don't have a conn, make one.
	if conn == nil {
		if conn, err = net.Dial("udp", Address); err != nil {
			conn = nil
			return
		}
	}

	return
}
<commit_msg>Simplify increment code<commit_after>package statsd

import (
	"net"
	"os"
	"strconv"
	"time"
)

// Address of the StatsD server.
var Address = "localhost:8125"

// AlsoAppendHost appends hostname along with any metric sent.
var AlsoAppendHost = true

// Cache the conn for perf.
var conn net.Conn
var host string

func init() {
	var err error

	if host, err = os.Hostname(); err != nil {
		panic(err)
	}
}

func Timer() timer {
	return timer{time.Now()}
}

// Timer
type timer struct {
	start time.Time
}

func (t *timer) Reset() {
	t.start = time.Now()
}

func (t *timer) Send(names ...interface{}) (took time.Duration) {
	took = time.Since(t.start)

	if err := getConnection(); err != nil {
		return
	}

	value := ":" + strconv.FormatUint(uint64(took.Nanoseconds()/1e6), 10) + "|ms"

	for _, name := range names {
		conn.Write([]byte(name.(string) + value))

		// Send a host suffixed stat too.
		if AlsoAppendHost {
			conn.Write([]byte(name.(string) + "." + host + value))
		}
	}

	return
}

// Gauge sets arbitrary numeric value for a given metric.
func Gauge(name string, value int64) {
	if err := getConnection(); err != nil {
		return
	}

	suffix := ":" + strconv.FormatInt(value, 10) + "|g"

	conn.Write([]byte(name + suffix))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + suffix))
	}

	return
}

// Inc is a simple counter adding one to a given metric.
func Inc(name string) {
	return IncSampled(name, 1.0)
}

// IncSampled increments a counter with the given sample rate (between 0.0 - 1.0).
// Example: Passing 0.1 tells statsd that this metric has been sample 1/10 times.
// IE Reported metric will be 10x the number of times called.
func IncSampled(name string, rate float64) {
	if err := getConnection(); err != nil {
		return
	}

	message := ":1|c"
	if rate < 1.0 {
		message = message + "|@" + strconv.FormatFloat(rate, 'f', -1, 64)
	}

	conn.Write([]byte(name + message))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + message))
	}

	return
}

// Time sends duration in ms for a given metric.
func Time(name string, took time.Duration) {
	if err := getConnection(); err != nil {
		return
	}

	value := ":" + strconv.FormatUint(uint64(took.Nanoseconds()/1e6), 10) + "|ms"

	conn.Write([]byte(name + value))

	// Send a host suffixed stat too.
	if AlsoAppendHost {
		conn.Write([]byte(name + "." + host + value))
	}

	return
}

func getConnection() (err error) {
	// If we don't have a conn, make one.
	if conn == nil {
		if conn, err = net.Dial("udp", Address); err != nil {
			conn = nil
			return
		}
	}

	return
}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

package chaincode

import (
	"errors"
	"fmt"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"github.com/openblockchain/obc-peer/openchain/ledger"
	pb "github.com/openblockchain/obc-peer/protos"
)

//Execute - execute transaction or a query
func Execute(ctxt context.Context, chain *ChaincodeSupport, t *pb.Transaction) ([]byte, error) {
	var err error

	// get a handle to ledger to mark the begin/finish of a tx
	ledger, ledgerErr := ledger.GetLedger()
	if ledgerErr != nil {
		return nil, fmt.Errorf("Failed to get handle to ledger (%s)", ledgerErr)
	}

	if secHelper := chain.getSecHelper(); nil != secHelper {
		var err error
		t, err = secHelper.TransactionPreExecution(t)
		// Note that t is now decrypted and is a deep clone of the original input t
		if nil != err {
			return nil, err
		}
	}

	if t.Type == pb.Transaction_CHAINCODE_NEW {
		_, err := chain.DeployChaincode(ctxt, t)
		if err != nil {
			return nil, fmt.Errorf("Failed to deploy chaincode spec(%s)", err)
		}

		//launch and wait for ready
		markTxBegin(ledger, t)
		_, _, err = chain.LaunchChaincode(ctxt, t)
		if err != nil {
			//TODO rollback transaction as init might have set state
			markTxFinish(ledger, t, false)
			return nil, fmt.Errorf("%s", err)
		}
		markTxFinish(ledger, t, true)
	} else if t.Type == pb.Transaction_CHAINCODE_EXECUTE || t.Type == pb.Transaction_CHAINCODE_QUERY {
		//will launch if necessary (and wait for ready)
		cID, cMsg, err := chain.LaunchChaincode(ctxt, t)
		if err != nil {
			return nil, fmt.Errorf("Failed to launch chaincode spec(%s)", err)
		}

		//this should work because it worked above...
		chaincode := cID.Name

		if err != nil {
			return nil, fmt.Errorf("Failed to stablish stream to container %s", chaincode)
		}

		// TODO: Need to comment next line and uncomment call to getTimeout, when transaction blocks are being created
		timeout := time.Duration(30000) * time.Millisecond
		//timeout, err := getTimeout(cID)

		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve chaincode spec(%s)", err)
		}

		var ccMsg *pb.ChaincodeMessage
		if t.Type == pb.Transaction_CHAINCODE_EXECUTE {
			ccMsg, err = createTransactionMessage(t.Uuid, cMsg)
			if err != nil {
				return nil, fmt.Errorf("Failed to transaction message(%s)", err)
			}
		} else {
			ccMsg, err = createQueryMessage(t.Uuid, cMsg)
			if err != nil {
				return nil, fmt.Errorf("Failed to query message(%s)", err)
			}
		}

		markTxBegin(ledger, t)
		resp, err := chain.Execute(ctxt, chaincode, ccMsg, timeout, t)
		if err != nil {
			// Rollback transaction
			markTxFinish(ledger, t, false)
			fmt.Printf("Got ERROR inside execute %s\n", err)
			return nil, fmt.Errorf("Failed to execute transaction or query(%s)", err)
		} else if resp == nil {
			// Rollback transaction
			markTxFinish(ledger, t, false)
			return nil, fmt.Errorf("Failed to receive a response for (%s)", t.Uuid)
		} else {
			if resp.Type == pb.ChaincodeMessage_COMPLETED || resp.Type == pb.ChaincodeMessage_QUERY_COMPLETED {
				// Success
				markTxFinish(ledger, t, true)
				return resp.Payload, nil
			} else if resp.Type == pb.ChaincodeMessage_ERROR || resp.Type == pb.ChaincodeMessage_QUERY_ERROR {
				// Rollback transaction
				markTxFinish(ledger, t, false)
				return nil, fmt.Errorf("Transaction or query returned with failure: %s", string(resp.Payload))
			}
			markTxFinish(ledger, t, false)
			return resp.Payload, fmt.Errorf("receive a response for (%s) but in invalid state(%d)", t.Uuid, resp.Type)
		}

	} else {
		err = fmt.Errorf("Invalid transaction type %s", t.Type.String())
	}
	return nil, err
}

//ExecuteTransactions - will execute transactions on the array one by one
//will return an array of errors one for each transaction. If the execution
//succeeded, array element will be nil. returns state hash
func ExecuteTransactions(ctxt context.Context, cname ChainName, xacts []*pb.Transaction) ([]byte, []error) {
	var chain = GetChain(cname)
	if chain == nil {
		panic(fmt.Sprintf("[ExecuteTransactions]Chain %s not found\n", cname))
	}
	errs := make([]error, len(xacts)+1)
	for i, t := range xacts {
		_, errs[i] = Execute(ctxt, chain, t)
	}
	ledger, hasherr := ledger.GetLedger()
	var statehash []byte
	if hasherr == nil {
		statehash, hasherr = ledger.GetTempStateHash()
	}
	errs[len(errs)-1] = hasherr
	return statehash, errs
}

// GetSecureContext returns the security context from the context object or error
// Security context is nil if security is off from openchain.yaml file
// func GetSecureContext(ctxt context.Context) (crypto.Peer, error) {
// 	var err error
// 	temp := ctxt.Value("security")
// 	if nil != temp {
// 		if secCxt, ok := temp.(crypto.Peer); ok {
// 			return secCxt, nil
// 		}
// 		err = errors.New("Failed to convert security context type")
// 	}
// 	return nil, err
// }

var errFailedToGetChainCodeSpecForTransaction = errors.New("Failed to get ChainCodeSpec from Transaction")

func getTimeout(cID *pb.ChaincodeID) (time.Duration, error) {
	ledger, err := ledger.GetLedger()
	if err == nil {
		chaincodeID := cID.Name
		txUUID, err := ledger.GetState(chaincodeID, "github.com_openblockchain_obc-peer_chaincode_id", true)
		if err == nil {
			tx, err := ledger.GetTransactionByUUID(string(txUUID))
			if err == nil {
				chaincodeDeploymentSpec := &pb.ChaincodeDeploymentSpec{}
				proto.Unmarshal(tx.Payload, chaincodeDeploymentSpec)
				chaincodeSpec := chaincodeDeploymentSpec.GetChaincodeSpec()
				timeout := time.Duration(time.Duration(chaincodeSpec.Timeout) * time.Millisecond)
				return timeout, nil
			}
		}
	}

	return -1, errFailedToGetChainCodeSpecForTransaction
}

func markTxBegin(ledger *ledger.Ledger, t *pb.Transaction) {
	if t.Type == pb.Transaction_CHAINCODE_QUERY {
		return
	}
	ledger.TxBegin(t.Uuid)
}

func markTxFinish(ledger *ledger.Ledger, t *pb.Transaction, successful bool) {
	if t.Type == pb.Transaction_CHAINCODE_QUERY {
		return
	}
	ledger.TxFinished(t.Uuid, successful)
}
<commit_msg>clean up debug log and print<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

package chaincode

import (
	"errors"
	"fmt"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"github.com/openblockchain/obc-peer/openchain/ledger"
	pb "github.com/openblockchain/obc-peer/protos"
)

//Execute - execute transaction or a query
func Execute(ctxt context.Context, chain *ChaincodeSupport, t *pb.Transaction) ([]byte, error) {
	var err error

	// get a handle to ledger to mark the begin/finish of a tx
	ledger, ledgerErr := ledger.GetLedger()
	if ledgerErr != nil {
		return nil, fmt.Errorf("Failed to get handle to ledger (%s)", ledgerErr)
	}

	if secHelper := chain.getSecHelper(); nil != secHelper {
		var err error
		t, err = secHelper.TransactionPreExecution(t)
		// Note that t is now decrypted and is a deep clone of the original input t
		if nil != err {
			return nil, err
		}
	}

	if t.Type == pb.Transaction_CHAINCODE_NEW {
		_, err := chain.DeployChaincode(ctxt, t)
		if err != nil {
			return nil, fmt.Errorf("Failed to deploy chaincode spec(%s)", err)
		}

		//launch and wait for ready
		markTxBegin(ledger, t)
		_, _, err = chain.LaunchChaincode(ctxt, t)
		if err != nil {
			//TODO rollback transaction as init might have set state
			markTxFinish(ledger, t, false)
			return nil, fmt.Errorf("%s", err)
		}
		markTxFinish(ledger, t, true)
	} else if t.Type == pb.Transaction_CHAINCODE_EXECUTE || t.Type == pb.Transaction_CHAINCODE_QUERY {
		//will launch if necessary (and wait for ready)
		cID, cMsg, err := chain.LaunchChaincode(ctxt, t)
		if err != nil {
			return nil, fmt.Errorf("Failed to launch chaincode spec(%s)", err)
		}

		//this should work because it worked above...
		chaincode := cID.Name

		if err != nil {
			return nil, fmt.Errorf("Failed to stablish stream to container %s", chaincode)
		}

		// TODO: Need to comment next line and uncomment call to getTimeout, when transaction blocks are being created
		timeout := time.Duration(30000) * time.Millisecond
		//timeout, err := getTimeout(cID)

		if err != nil {
			return nil, fmt.Errorf("Failed to retrieve chaincode spec(%s)", err)
		}

		var ccMsg *pb.ChaincodeMessage
		if t.Type == pb.Transaction_CHAINCODE_EXECUTE {
			ccMsg, err = createTransactionMessage(t.Uuid, cMsg)
			if err != nil {
				return nil, fmt.Errorf("Failed to transaction message(%s)", err)
			}
		} else {
			ccMsg, err = createQueryMessage(t.Uuid, cMsg)
			if err != nil {
				return nil, fmt.Errorf("Failed to query message(%s)", err)
			}
		}

		markTxBegin(ledger, t)
		resp, err := chain.Execute(ctxt, chaincode, ccMsg, timeout, t)
		if err != nil {
			// Rollback transaction
			markTxFinish(ledger, t, false)
			//fmt.Printf("Got ERROR inside execute %s\n", err)
			return nil, fmt.Errorf("Failed to execute transaction or query(%s)", err)
		} else if resp == nil {
			// Rollback transaction
			markTxFinish(ledger, t, false)
			return nil, fmt.Errorf("Failed to receive a response for (%s)", t.Uuid)
		} else {
			if resp.Type == pb.ChaincodeMessage_COMPLETED || resp.Type == pb.ChaincodeMessage_QUERY_COMPLETED {
				// Success
				markTxFinish(ledger, t, true)
				return resp.Payload, nil
			} else if resp.Type == pb.ChaincodeMessage_ERROR || resp.Type == pb.ChaincodeMessage_QUERY_ERROR {
				// Rollback transaction
				markTxFinish(ledger, t, false)
				return nil, fmt.Errorf("Transaction or query returned with failure: %s", string(resp.Payload))
			}
			markTxFinish(ledger, t, false)
			return resp.Payload, fmt.Errorf("receive a response for (%s) but in invalid state(%d)", t.Uuid, resp.Type)
		}

	} else {
		err = fmt.Errorf("Invalid transaction type %s", t.Type.String())
	}
	return nil, err
}

//ExecuteTransactions - will execute transactions on the array one by one
//will return an array of errors one for each transaction. If the execution
//succeeded, array element will be nil. returns state hash
func ExecuteTransactions(ctxt context.Context, cname ChainName, xacts []*pb.Transaction) ([]byte, []error) {
	var chain = GetChain(cname)
	if chain == nil {
		// TODO: We should never get here, but otherwise a good reminder to better handle
		panic(fmt.Sprintf("[ExecuteTransactions]Chain %s not found\n", cname))
	}
	errs := make([]error, len(xacts)+1)
	for i, t := range xacts {
		_, errs[i] = Execute(ctxt, chain, t)
	}
	ledger, hasherr := ledger.GetLedger()
	var statehash []byte
	if hasherr == nil {
		statehash, hasherr = ledger.GetTempStateHash()
	}
	errs[len(errs)-1] = hasherr
	return statehash, errs
}

// GetSecureContext returns the security context from the context object or error
// Security context is nil if security is off from openchain.yaml file
// func GetSecureContext(ctxt context.Context) (crypto.Peer, error) {
// 	var err error
// 	temp := ctxt.Value("security")
// 	if nil != temp {
// 		if secCxt, ok := temp.(crypto.Peer); ok {
// 			return secCxt, nil
// 		}
// 		err = errors.New("Failed to convert security context type")
// 	}
// 	return nil, err
// }

var errFailedToGetChainCodeSpecForTransaction = errors.New("Failed to get ChainCodeSpec from Transaction")

func getTimeout(cID *pb.ChaincodeID) (time.Duration, error) {
	ledger, err := ledger.GetLedger()
	if err == nil {
		chaincodeID := cID.Name
		txUUID, err := ledger.GetState(chaincodeID, "github.com_openblockchain_obc-peer_chaincode_id", true)
		if err == nil {
			tx, err := ledger.GetTransactionByUUID(string(txUUID))
			if err == nil {
				chaincodeDeploymentSpec := &pb.ChaincodeDeploymentSpec{}
				proto.Unmarshal(tx.Payload, chaincodeDeploymentSpec)
				chaincodeSpec := chaincodeDeploymentSpec.GetChaincodeSpec()
				timeout := time.Duration(time.Duration(chaincodeSpec.Timeout) * time.Millisecond)
				return timeout, nil
			}
		}
	}

	return -1, errFailedToGetChainCodeSpecForTransaction
}

func markTxBegin(ledger *ledger.Ledger, t *pb.Transaction) {
	if t.Type == pb.Transaction_CHAINCODE_QUERY {
		return
	}
	ledger.TxBegin(t.Uuid)
}

func markTxFinish(ledger *ledger.Ledger, t *pb.Transaction, successful bool) {
	if t.Type == pb.Transaction_CHAINCODE_QUERY {
		return
	}
	ledger.TxFinished(t.Uuid, successful)
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// These tests exercise quota reclamation.

package test

import (
	"testing"
	"time"
)

// Check that simple quota reclamation works
func TestQRSimple(t *testing.T) {
	test(t,
		users("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", "hello"),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works, eventually, after enough iterations.
func TestQRLargePointerSet(t *testing.T) {
	var busyWork []fileOp
	iters := 100
	for i := 0; i < iters; i++ {
		busyWork = append(busyWork, mkfile("a", "hello"), rm("a"))
	}
	// 5 unreferenced pointers per iteration -- 3 updates to the root
	// block, one empty file written to, and one non-empty file
	// deleted.
	ptrsPerIter := 5
	var qrOps []optionOp
	// Each reclamation needs a sync after it (e.g., a new "as"
	// clause) to ensure it completes before the next force
	// reclamation.
	for i := 0; i < ptrsPerIter*iters/100; i++ {
		qrOps = append(qrOps, as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		))
	}
	totalOps := []optionOp{users("alice"), as(alice, busyWork...)}
	totalOps = append(totalOps, qrOps...)
	test(t, totalOps...)
}

// Test that quota reclamation handles conflict resolution correctly.
func TestQRAfterCR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			write("a/c", "world"),
		),
		as(bob, noSync(),
			rm("a/b"),
			reenableUpdates(),
		),
		as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works on multi-block files
func TestQRWithMultiBlockFiles(t *testing.T) {
	test(t,
		blockSize(20), users("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", ntimesString(15, "0123456789")),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Test that conflict resolution fails gracefully after quota
// reclamation deletes some necessary blocks.
func TestCRAfterQR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			rm("a/b"),
		),
		as(bob, noSync(),
			setex("a/b", true),
		),
		as(alice,
			addTime(2*time.Minute),
			// Force rmd.data.Dir.MTime to something recent. TODO: remove me.
			mkfile("c", "test"),
			forceQuotaReclamation(),
		),
		as(bob, noSync(),
			reenableUpdates(),
			notExists("a/b"),
		),
		as(alice,
			notExists("a/b"),
		),
	)
}

// Test that conflict resolution works gracefully after quota
// reclamation deletes a modified+deleted directory from the merged
// branch.  Regression for KBFS-1202.
func TestCRAfterRmdirAndQR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkdir("a/b"),
			mkfile("a/b/c", "world"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			mkfile("a/b/d", "world2"),
			rm("a/b/c"),
			rm("a/b/d"),
			rmdir("a/b"),
		),
		as(bob, noSync(),
			setmtime("a/b/c", time.Now().Add(1*time.Minute)),
			rm("a/b/c"),
			rmdir("a/b"),
		),
		as(alice,
			addTime(2*time.Minute),
			// Force rmd.data.Dir.MTime to something recent. TODO: remove me.
			mkfile("c", "test"),
			forceQuotaReclamation(),
		),
		as(bob, noSync(),
			reenableUpdates(),
			lsdir("a/", m{}),
		),
		as(alice,
			lsdir("a/", m{}),
		),
	)
}
<commit_msg>qr_test: add workaround to TestQRLargePointerSet for osxfuse bug (#171)<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// These tests exercise quota reclamation.

package test

import (
	"fmt"
	"testing"
	"time"
)

// Check that simple quota reclamation works
func TestQRSimple(t *testing.T) {
	test(t,
		users("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", "hello"),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works, eventually, after enough iterations.
func TestQRLargePointerSet(t *testing.T) {
	var busyWork []fileOp
	iters := 100
	for i := 0; i < iters; i++ {
		name := fmt.Sprintf("a%d", i)
		busyWork = append(busyWork, mkfile(name, "hello"), rm(name))
	}
	// 5 unreferenced pointers per iteration -- 3 updates to the root
	// block, one empty file written to, and one non-empty file
	// deleted.
	ptrsPerIter := 5
	var qrOps []optionOp
	// Each reclamation needs a sync after it (e.g., a new "as"
	// clause) to ensure it completes before the next force
	// reclamation.
	for i := 0; i < ptrsPerIter*iters/100; i++ {
		qrOps = append(qrOps, as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		))
	}
	totalOps := []optionOp{users("alice"), as(alice, busyWork...)}
	totalOps = append(totalOps, qrOps...)
	test(t, totalOps...)
}

// Test that quota reclamation handles conflict resolution correctly.
func TestQRAfterCR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			write("a/c", "world"),
		),
		as(bob, noSync(),
			rm("a/b"),
			reenableUpdates(),
		),
		as(alice,
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Check that quota reclamation works on multi-block files
func TestQRWithMultiBlockFiles(t *testing.T) {
	test(t,
		blockSize(20), users("alice"),
		as(alice,
			addTime(1*time.Minute),
			mkfile("a", ntimesString(15, "0123456789")),
			rm("a"),
			addTime(2*time.Minute),
			forceQuotaReclamation(),
		),
	)
}

// Test that conflict resolution fails gracefully after quota
// reclamation deletes some necessary blocks.
func TestCRAfterQR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkfile("a/b", "hello"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			rm("a/b"),
		),
		as(bob, noSync(),
			setex("a/b", true),
		),
		as(alice,
			addTime(2*time.Minute),
			// Force rmd.data.Dir.MTime to something recent. TODO: remove me.
			mkfile("c", "test"),
			forceQuotaReclamation(),
		),
		as(bob, noSync(),
			reenableUpdates(),
			notExists("a/b"),
		),
		as(alice,
			notExists("a/b"),
		),
	)
}

// Test that conflict resolution works gracefully after quota
// reclamation deletes a modified+deleted directory from the merged
// branch.  Regression for KBFS-1202.
func TestCRAfterRmdirAndQR(t *testing.T) {
	test(t,
		users("alice", "bob"),
		as(alice,
			mkdir("a/b"),
			mkfile("a/b/c", "world"),
		),
		as(bob,
			disableUpdates(),
		),
		as(alice,
			mkfile("a/b/d", "world2"),
			rm("a/b/c"),
			rm("a/b/d"),
			rmdir("a/b"),
		),
		as(bob, noSync(),
			setmtime("a/b/c", time.Now().Add(1*time.Minute)),
			rm("a/b/c"),
			rmdir("a/b"),
		),
		as(alice,
			addTime(2*time.Minute),
			// Force rmd.data.Dir.MTime to something recent. TODO: remove me.
			mkfile("c", "test"),
			forceQuotaReclamation(),
		),
		as(bob, noSync(),
			reenableUpdates(),
			lsdir("a/", m{}),
		),
		as(alice,
			lsdir("a/", m{}),
		),
	)
}
<|endoftext|>