<commit_before>package google

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
)

// Provider returns a terraform.ResourceProvider.
func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		Schema: map[string]*schema.Schema{
			"account_file": &schema.Schema{
				Type:        schema.TypeString,
				Optional:    true,
				DefaultFunc: schema.EnvDefaultFunc("GOOGLE_ACCOUNT_FILE", nil),
				Removed:     "Use the credentials field instead",
			},

			"credentials": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_CREDENTIALS",
					"GOOGLE_CLOUD_KEYFILE_JSON",
					"GCLOUD_KEYFILE_JSON",
				}, nil),
				ValidateFunc: validateCredentials,
			},

			"project": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_PROJECT",
					"GCLOUD_PROJECT",
					"CLOUDSDK_CORE_PROJECT",
				}, nil),
			},

			"region": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_REGION",
					"GCLOUD_REGION",
					"CLOUDSDK_COMPUTE_REGION",
				}, nil),
			},
		},

		DataSourcesMap: map[string]*schema.Resource{
			"google_iam_policy":    dataSourceGoogleIamPolicy(),
			"google_compute_zones": dataSourceGoogleComputeZones(),
		},

		ResourcesMap: map[string]*schema.Resource{
			"google_compute_autoscaler":             resourceComputeAutoscaler(),
			"google_compute_address":                resourceComputeAddress(),
			"google_compute_backend_service":        resourceComputeBackendService(),
			"google_compute_disk":                   resourceComputeDisk(),
			"google_compute_firewall":               resourceComputeFirewall(),
			"google_compute_forwarding_rule":        resourceComputeForwardingRule(),
			"google_compute_global_address":         resourceComputeGlobalAddress(),
			"google_compute_global_forwarding_rule": resourceComputeGlobalForwardingRule(),
			"google_compute_health_check":           resourceComputeHealthCheck(),
			"google_compute_http_health_check":      resourceComputeHttpHealthCheck(),
			"google_compute_https_health_check":     resourceComputeHttpsHealthCheck(),
			"google_compute_image":                  resourceComputeImage(),
			"google_compute_instance":               resourceComputeInstance(),
			"google_compute_instance_group":         resourceComputeInstanceGroup(),
			"google_compute_instance_group_manager": resourceComputeInstanceGroupManager(),
			"google_compute_instance_template":      resourceComputeInstanceTemplate(),
			"google_compute_network":                resourceComputeNetwork(),
			"google_compute_project_metadata":       resourceComputeProjectMetadata(),
			"google_compute_region_backend_service": resourceComputeRegionBackendService(),
			"google_compute_route":                  resourceComputeRoute(),
			"google_compute_ssl_certificate":        resourceComputeSslCertificate(),
			"google_compute_subnetwork":             resourceComputeSubnetwork(),
			"google_compute_target_http_proxy":      resourceComputeTargetHttpProxy(),
			"google_compute_target_https_proxy":     resourceComputeTargetHttpsProxy(),
			"google_compute_target_pool":            resourceComputeTargetPool(),
			"google_compute_url_map":                resourceComputeUrlMap(),
			"google_compute_vpn_gateway":            resourceComputeVpnGateway(),
			"google_compute_vpn_tunnel":             resourceComputeVpnTunnel(),
			"google_container_cluster":              resourceContainerCluster(),
			"google_container_node_pool":            resourceContainerNodePool(),
			"google_dns_managed_zone":               resourceDnsManagedZone(),
			"google_dns_record_set":                 resourceDnsRecordSet(),
			"google_sql_database":                   resourceSqlDatabase(),
			"google_sql_database_instance":          resourceSqlDatabaseInstance(),
			"google_sql_user":                       resourceSqlUser(),
			"google_project":                        resourceGoogleProject(),
			"google_project_iam_policy":             resourceGoogleProjectIamPolicy(),
			"google_project_services":               resourceGoogleProjectServices(),
			"google_pubsub_topic":                   resourcePubsubTopic(),
			"google_pubsub_subscription":            resourcePubsubSubscription(),
			"google_service_account":                resourceGoogleServiceAccount(),
			"google_storage_bucket":                 resourceStorageBucket(),
			"google_storage_bucket_acl":             resourceStorageBucketAcl(),
			"google_storage_bucket_object":          resourceStorageBucketObject(),
			"google_storage_object_acl":             resourceStorageObjectAcl(),
		},

		ConfigureFunc: providerConfigure,
	}
}

func providerConfigure(d *schema.ResourceData) (interface{}, error) {
	credentials := d.Get("credentials").(string)
	config := Config{
		Credentials: credentials,
		Project:     d.Get("project").(string),
		Region:      d.Get("region").(string),
	}

	if err := config.loadAndValidate(); err != nil {
		return nil, err
	}

	return &config, nil
}

func validateCredentials(v interface{}, k string) (warnings []string, errors []error) {
	if v == nil || v.(string) == "" {
		return
	}
	creds := v.(string)
	var account accountFile
	if err := json.Unmarshal([]byte(creds), &account); err != nil {
		errors = append(errors,
			fmt.Errorf("credentials are not valid JSON '%s': %s", creds, err))
	}

	return
}

// getRegionFromZone returns the region from a zone for Google cloud.
func getRegionFromZone(zone string) string {
	if zone != "" && len(zone) > 2 {
		region := zone[:len(zone)-2]
		return region
	}
	return ""
}

// getRegion reads the "region" field from the given resource data and falls
// back to the provider's value if not given. If the provider's value is not
// given, an error is returned.
func getRegion(d *schema.ResourceData, config *Config) (string, error) {
	res, ok := d.GetOk("region")
	if !ok {
		if config.Region != "" {
			return config.Region, nil
		}
		return "", fmt.Errorf("%q: required field is not set", "region")
	}
	return res.(string), nil
}

// getProject reads the "project" field from the given resource data and falls
// back to the provider's value if not given. If the provider's value is not
// given, an error is returned.
func getProject(d *schema.ResourceData, config *Config) (string, error) {
	res, ok := d.GetOk("project")
	if !ok {
		if config.Project != "" {
			return config.Project, nil
		}
		return "", fmt.Errorf("%q: required field is not set", "project")
	}
	return res.(string), nil
}

func getZonalResourceFromRegion(getResource func(string) (interface{}, error), region string, compute *compute.Service, project string) (interface{}, error) {
	zoneList, err := compute.Zones.List(project).Do()
	if err != nil {
		return nil, err
	}
	var resource interface{}
	for _, zone := range zoneList.Items {
		if strings.Contains(zone.Name, region) {
			resource, err = getResource(zone.Name)
			if err != nil {
				if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
					// Resource was not found in this zone
					continue
				}
				return nil, fmt.Errorf("Error reading Resource: %s", err)
			}
			// Resource was found
			return resource, nil
		}
	}
	// Resource does not exist in this region
	return nil, nil
}

// getNetworkLink reads the "network" field from the given resource data and if the value:
// - is a resource URL, returns the string unchanged
// - is the network name only, then looks up the resource URL using the google client
func getNetworkLink(d *schema.ResourceData, config *Config, field string) (string, error) {
	if v, ok := d.GetOk(field); ok {
		network := v.(string)

		project, err := getProject(d, config)
		if err != nil {
			return "", err
		}

		if !strings.HasPrefix(network, "https://www.googleapis.com/compute/") {
			// Network value provided is just the name, lookup the network SelfLink
			networkData, err := config.clientCompute.Networks.Get(
				project, network).Do()
			if err != nil {
				return "", fmt.Errorf("Error reading network: %s", err)
			}
			network = networkData.SelfLink
		}

		return network, nil

	} else {
		return "", nil
	}
}

// getNetworkName reads the "network" field from the given resource data and if the value:
// - is a resource URL, extracts the network name from the URL and returns it
// - is the network name only (i.e not prefixed with http://www.googleapis.com/compute/...), is returned unchanged
func getNetworkName(d *schema.ResourceData, field string) (string, error) {
	if v, ok := d.GetOk(field); ok {
		network := v.(string)
		return getNetworkNameFromSelfLink(network)
	}
	return "", nil
}

func getNetworkNameFromSelfLink(network string) (string, error) {
	if strings.HasPrefix(network, "https://www.googleapis.com/compute/") {
		// extract the network name from SelfLink URL
		networkName := network[strings.LastIndex(network, "/")+1:]
		if networkName == "" {
			return "", fmt.Errorf("network url not valid")
		}
		return networkName, nil
	}

	return network, nil
}
<commit_msg>provider/google: drop the account file.<commit_after>package google

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
)

// Provider returns a terraform.ResourceProvider.
func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		Schema: map[string]*schema.Schema{
			"credentials": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_CREDENTIALS",
					"GOOGLE_CLOUD_KEYFILE_JSON",
					"GCLOUD_KEYFILE_JSON",
				}, nil),
				ValidateFunc: validateCredentials,
			},

			"project": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_PROJECT",
					"GCLOUD_PROJECT",
					"CLOUDSDK_CORE_PROJECT",
				}, nil),
			},

			"region": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"GOOGLE_REGION",
					"GCLOUD_REGION",
					"CLOUDSDK_COMPUTE_REGION",
				}, nil),
			},
		},

		DataSourcesMap: map[string]*schema.Resource{
			"google_iam_policy":    dataSourceGoogleIamPolicy(),
			"google_compute_zones": dataSourceGoogleComputeZones(),
		},

		ResourcesMap: map[string]*schema.Resource{
			"google_compute_autoscaler":             resourceComputeAutoscaler(),
			"google_compute_address":                resourceComputeAddress(),
			"google_compute_backend_service":        resourceComputeBackendService(),
			"google_compute_disk":                   resourceComputeDisk(),
			"google_compute_firewall":               resourceComputeFirewall(),
			"google_compute_forwarding_rule":        resourceComputeForwardingRule(),
			"google_compute_global_address":         resourceComputeGlobalAddress(),
			"google_compute_global_forwarding_rule": resourceComputeGlobalForwardingRule(),
			"google_compute_health_check":           resourceComputeHealthCheck(),
			"google_compute_http_health_check":      resourceComputeHttpHealthCheck(),
			"google_compute_https_health_check":     resourceComputeHttpsHealthCheck(),
			"google_compute_image":                  resourceComputeImage(),
			"google_compute_instance":               resourceComputeInstance(),
			"google_compute_instance_group":         resourceComputeInstanceGroup(),
			"google_compute_instance_group_manager": resourceComputeInstanceGroupManager(),
			"google_compute_instance_template":      resourceComputeInstanceTemplate(),
			"google_compute_network":                resourceComputeNetwork(),
			"google_compute_project_metadata":       resourceComputeProjectMetadata(),
			"google_compute_region_backend_service": resourceComputeRegionBackendService(),
			"google_compute_route":                  resourceComputeRoute(),
			"google_compute_ssl_certificate":        resourceComputeSslCertificate(),
			"google_compute_subnetwork":             resourceComputeSubnetwork(),
			"google_compute_target_http_proxy":      resourceComputeTargetHttpProxy(),
			"google_compute_target_https_proxy":     resourceComputeTargetHttpsProxy(),
			"google_compute_target_pool":            resourceComputeTargetPool(),
			"google_compute_url_map":                resourceComputeUrlMap(),
			"google_compute_vpn_gateway":            resourceComputeVpnGateway(),
			"google_compute_vpn_tunnel":             resourceComputeVpnTunnel(),
			"google_container_cluster":              resourceContainerCluster(),
			"google_container_node_pool":            resourceContainerNodePool(),
			"google_dns_managed_zone":               resourceDnsManagedZone(),
			"google_dns_record_set":                 resourceDnsRecordSet(),
			"google_sql_database":                   resourceSqlDatabase(),
			"google_sql_database_instance":          resourceSqlDatabaseInstance(),
			"google_sql_user":                       resourceSqlUser(),
			"google_project":                        resourceGoogleProject(),
			"google_project_iam_policy":             resourceGoogleProjectIamPolicy(),
			"google_project_services":               resourceGoogleProjectServices(),
			"google_pubsub_topic":                   resourcePubsubTopic(),
			"google_pubsub_subscription":            resourcePubsubSubscription(),
			"google_service_account":                resourceGoogleServiceAccount(),
			"google_storage_bucket":                 resourceStorageBucket(),
			"google_storage_bucket_acl":             resourceStorageBucketAcl(),
			"google_storage_bucket_object":          resourceStorageBucketObject(),
			"google_storage_object_acl":             resourceStorageObjectAcl(),
		},

		ConfigureFunc: providerConfigure,
	}
}

func providerConfigure(d *schema.ResourceData) (interface{}, error) {
	credentials := d.Get("credentials").(string)
	config := Config{
		Credentials: credentials,
		Project:     d.Get("project").(string),
		Region:      d.Get("region").(string),
	}

	if err := config.loadAndValidate(); err != nil {
		return nil, err
	}

	return &config, nil
}

func validateCredentials(v interface{}, k string) (warnings []string, errors []error) {
	if v == nil || v.(string) == "" {
		return
	}
	creds := v.(string)
	var account accountFile
	if err := json.Unmarshal([]byte(creds), &account); err != nil {
		errors = append(errors,
			fmt.Errorf("credentials are not valid JSON '%s': %s", creds, err))
	}

	return
}

// getRegionFromZone returns the region from a zone for Google cloud.
func getRegionFromZone(zone string) string {
	if zone != "" && len(zone) > 2 {
		region := zone[:len(zone)-2]
		return region
	}
	return ""
}

// getRegion reads the "region" field from the given resource data and falls
// back to the provider's value if not given. If the provider's value is not
// given, an error is returned.
func getRegion(d *schema.ResourceData, config *Config) (string, error) {
	res, ok := d.GetOk("region")
	if !ok {
		if config.Region != "" {
			return config.Region, nil
		}
		return "", fmt.Errorf("%q: required field is not set", "region")
	}
	return res.(string), nil
}

// getProject reads the "project" field from the given resource data and falls
// back to the provider's value if not given. If the provider's value is not
// given, an error is returned.
func getProject(d *schema.ResourceData, config *Config) (string, error) {
	res, ok := d.GetOk("project")
	if !ok {
		if config.Project != "" {
			return config.Project, nil
		}
		return "", fmt.Errorf("%q: required field is not set", "project")
	}
	return res.(string), nil
}

func getZonalResourceFromRegion(getResource func(string) (interface{}, error), region string, compute *compute.Service, project string) (interface{}, error) {
	zoneList, err := compute.Zones.List(project).Do()
	if err != nil {
		return nil, err
	}
	var resource interface{}
	for _, zone := range zoneList.Items {
		if strings.Contains(zone.Name, region) {
			resource, err = getResource(zone.Name)
			if err != nil {
				if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
					// Resource was not found in this zone
					continue
				}
				return nil, fmt.Errorf("Error reading Resource: %s", err)
			}
			// Resource was found
			return resource, nil
		}
	}
	// Resource does not exist in this region
	return nil, nil
}

// getNetworkLink reads the "network" field from the given resource data and if the value:
// - is a resource URL, returns the string unchanged
// - is the network name only, then looks up the resource URL using the google client
func getNetworkLink(d *schema.ResourceData, config *Config, field string) (string, error) {
	if v, ok := d.GetOk(field); ok {
		network := v.(string)

		project, err := getProject(d, config)
		if err != nil {
			return "", err
		}

		if !strings.HasPrefix(network, "https://www.googleapis.com/compute/") {
			// Network value provided is just the name, lookup the network SelfLink
			networkData, err := config.clientCompute.Networks.Get(
				project, network).Do()
			if err != nil {
				return "", fmt.Errorf("Error reading network: %s", err)
			}
			network = networkData.SelfLink
		}

		return network, nil

	} else {
		return "", nil
	}
}

// getNetworkName reads the "network" field from the given resource data and if the value:
// - is a resource URL, extracts the network name from the URL and returns it
// - is the network name only (i.e not prefixed with http://www.googleapis.com/compute/...), is returned unchanged
func getNetworkName(d *schema.ResourceData, field string) (string, error) {
	if v, ok := d.GetOk(field); ok {
		network := v.(string)
		return getNetworkNameFromSelfLink(network)
	}
	return "", nil
}

func getNetworkNameFromSelfLink(network string) (string, error) {
	if strings.HasPrefix(network, "https://www.googleapis.com/compute/") {
		// extract the network name from SelfLink URL
		networkName := network[strings.LastIndex(network, "/")+1:]
		if networkName == "" {
			return "", fmt.Errorf("network url not valid")
		}
		return networkName, nil
	}

	return network, nil
}
<|endoftext|>


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

package proxybundle

import (
	"archive/zip"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/google/go-github/github"
	genapi "github.com/srinandan/apigeecli/bundlegen"
	apiproxy "github.com/srinandan/apigeecli/bundlegen/apiproxydef"
	policies "github.com/srinandan/apigeecli/bundlegen/policies"
	proxies "github.com/srinandan/apigeecli/bundlegen/proxies"
	target "github.com/srinandan/apigeecli/bundlegen/targets"
	"github.com/srinandan/apigeecli/clilog"
	"golang.org/x/oauth2"
)

const rootDir = "apiproxy"

func GenerateAPIProxyBundleFromOAS(name string,
	content string,
	fileName string,
	skipPolicy bool,
	addCORS bool,
	oasGoogleAcessTokenScopeLiteral string,
	oasGoogleIdTokenAudLiteral string,
	oasGoogleIdTokenAudRef string,
	oasTargetUrlRef string) (err error) {

	var apiProxyData, proxyEndpointData, targetEndpointData string
	const resourceType = "oas"

	if err = os.Mkdir(rootDir, os.ModePerm); err != nil {
		return err
	}

	// write API Proxy file
	if apiProxyData, err = apiproxy.GetAPIProxy(); err != nil {
		return err
	}

	err = writeXMLData(rootDir+string(os.PathSeparator)+name+".xml", apiProxyData)
	if err != nil {
		return err
	}

	proxiesDirPath := rootDir + string(os.PathSeparator) + "proxies"
	policiesDirPath := rootDir + string(os.PathSeparator) + "policies"
	targetDirPath := rootDir + string(os.PathSeparator) + "targets"
	resDirPath := rootDir + string(os.PathSeparator) + "resources" + string(os.PathSeparator) + resourceType //"oas"

	if err = os.Mkdir(proxiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if proxyEndpointData, err = proxies.GetProxyEndpoint(); err != nil {
		return err
	}

	err = writeXMLData(proxiesDirPath+string(os.PathSeparator)+"default.xml", proxyEndpointData)
	if err != nil {
		return err
	}

	if err = os.Mkdir(targetDirPath, os.ModePerm); err != nil {
		return err
	}

	if targetEndpointData, err = target.GetTargetEndpoint(); err != nil {
		return err
	}

	if err = writeXMLData(targetDirPath+string(os.PathSeparator)+"default.xml", targetEndpointData); err != nil {
		return err
	}

	if !skipPolicy {
		if err = os.MkdirAll(resDirPath, os.ModePerm); err != nil {
			return err
		}
		if err = writeXMLData(resDirPath+string(os.PathSeparator)+fileName, content); err != nil {
			return err
		}
	}

	if err = os.Mkdir(policiesDirPath, os.ModePerm); err != nil {
		return err
	}

	//add set target url
	if genapi.GenerateSetTargetPolicy() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Set-Target-1.xml",
			policies.AddSetTargetEndpoint(oasTargetUrlRef)); err != nil {
			return err
		}
	}

	//add security policies
	for _, securityScheme := range genapi.GetSecuritySchemesList() {
		if securityScheme.APIKeyPolicy.APIKeyPolicyEnabled {
			if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Verify-API-Key-"+securityScheme.SchemeName+".xml",
				policies.AddVerifyApiKeyPolicy(securityScheme.APIKeyPolicy.APIKeyLocation,
					securityScheme.SchemeName,
					securityScheme.APIKeyPolicy.APIKeyName)); err != nil {
				return err
			}
		}
		if securityScheme.OAuthPolicy.OAuthPolicyEnabled {
			if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"OAuth-v20-1.xml",
				policies.AddOAuth2Policy(securityScheme.OAuthPolicy.Scope)); err != nil {
				return err
			}
		}
	}

	//add quota policies
	for quotaPolicyName, quotaPolicyContent := range genapi.GetQuotaPolicies() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Quota-"+quotaPolicyName+".xml", quotaPolicyContent); err != nil {
			return err
		}
	}

	//add spike arrest policies
	for spikeArrestPolicyName, spikeArrestPolicyContent := range genapi.GetSpikeArrestPolicies() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Spike-Arrest-"+spikeArrestPolicyName+".xml", spikeArrestPolicyContent); err != nil {
			return err
		}
	}

	if !skipPolicy {
		//add oas policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"OpenAPI-Spec-Validation-1.xml",
			policies.AddOpenAPIValidatePolicy(fileName)); err != nil {
			return err
		}
	}

	if addCORS {
		//add cors policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Add-CORS.xml", policies.AddCORSPolicy()); err != nil {
			return err
		}
	}

	if err = archiveBundle(rootDir, name+".zip"); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return nil
}

func GenerateAPIProxyBundleFromGQL(name string,
	content string,
	fileName string,
	action string,
	location string,
	keyName string,
	skipPolicy bool,
	addCORS bool,
	targetUrlRef string) (err error) {

	var apiProxyData, proxyEndpointData, targetEndpointData string
	const resourceType = "graphql"

	if err = os.Mkdir(rootDir, os.ModePerm); err != nil {
		return err
	}

	// write API Proxy file
	if apiProxyData, err = apiproxy.GetAPIProxy(); err != nil {
		return err
	}

	err = writeXMLData(rootDir+string(os.PathSeparator)+name+".xml", apiProxyData)
	if err != nil {
		return err
	}

	proxiesDirPath := rootDir + string(os.PathSeparator) + "proxies"
	policiesDirPath := rootDir + string(os.PathSeparator) + "policies"
	targetDirPath := rootDir + string(os.PathSeparator) + "targets"
	resDirPath := rootDir + string(os.PathSeparator) + "resources" + string(os.PathSeparator) + resourceType //"graphql"

	if err = os.Mkdir(proxiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if proxyEndpointData, err = proxies.GetProxyEndpoint(); err != nil {
		return err
	}

	err = writeXMLData(proxiesDirPath+string(os.PathSeparator)+"default.xml", proxyEndpointData)
	if err != nil {
		return err
	}

	if err = os.Mkdir(targetDirPath, os.ModePerm); err != nil {
		return err
	}

	if targetEndpointData, err = target.GetTargetEndpoint(); err != nil {
		return err
	}

	if err = writeXMLData(targetDirPath+string(os.PathSeparator)+"default.xml", targetEndpointData); err != nil {
		return err
	}

	if !skipPolicy {
		if err = os.MkdirAll(resDirPath, os.ModePerm); err != nil {
			return err
		}
		if err = writeXMLData(resDirPath+string(os.PathSeparator)+fileName, content); err != nil {
			return err
		}
	}

	if err = os.Mkdir(policiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Set-Target-1.xml",
		policies.AddSetTargetEndpoint(targetUrlRef)); err != nil {
		return err
	}

	if !skipPolicy {
		//add gql policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Validate-"+name+"-Schema.xml",
			policies.AddGraphQLPolicy(name, action, fileName)); err != nil {
			return err
		}
	}

	if keyName != "" {
		//add verifyapi key policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Verify-API-Key-"+name+".xml",
			policies.AddVerifyApiKeyPolicy(location, name, keyName)); err != nil {
			return err
		}
	}

	if addCORS {
		//add cors policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Add-CORS.xml", policies.AddCORSPolicy()); err != nil {
			return err
		}
	}

	if err = archiveBundle(rootDir, name+".zip"); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return nil
}

func writeXMLData(fileName string, data string) error {
	fileWriter, err := os.Create(fileName)
	if err != nil {
		return err
	}
	_, err = fileWriter.WriteString(data)
	if err != nil {
		return err
	}

	fileWriter.Close()
	return nil
}

func GenerateArchiveBundle(pathToZip, destinationPath string) error {
	return archiveBundle(pathToZip, destinationPath)
}

func archiveBundle(pathToZip, destinationPath string) (err error) {

	var destinationFile *os.File

	destinationFile, err = os.Create(destinationPath)
	if err != nil {
		return err
	}

	myZip := zip.NewWriter(destinationFile)
	err = filepath.Walk(pathToZip, func(filePath string, info os.FileInfo, err error) error {
		if info.IsDir() {
			relPath := strings.TrimPrefix(filePath, filepath.Dir(pathToZip))
			_, err = myZip.Create(strings.TrimPrefix(relPath, string(os.PathSeparator)) + string(os.PathSeparator))
			return err
		}
		if err != nil {
			return err
		}
		relPath := strings.TrimPrefix(filePath, filepath.Dir(pathToZip))
		zipFile, err := myZip.Create(strings.TrimPrefix(relPath, string(os.PathSeparator)))
		if err != nil {
			return err
		}
		fsFile, err := os.Open(filePath)
		if err != nil {
			return err
		}
		_, err = io.Copy(zipFile, fsFile)
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return err
	}
	err = myZip.Close()
	if err != nil {
		return err
	}
	return nil
}

func GitHubImportBundle(owner string, repo string, repopath string) (err error) {

	//clean up any files or folders
	CleanUp()
	os.RemoveAll(rootDir)

	//
	token := os.Getenv("GITHUB_TOKEN")
	ctx := context.Background()
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})
	tc := oauth2.NewClient(ctx, ts)
	client := github.NewClient(tc)

	//1. download the proxy
	if err := downloadProxyFromRepo(client, ctx, owner, repo, repopath); err != nil {
		return err
	}

	if client != nil {
		fmt.Println("")
	}

	//2. compress the proxy folder
	curDir, _ := os.Getwd()
	if err := archiveBundle(path.Join(curDir, rootDir), path.Join(curDir, rootDir+".zip")); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return err
}

func CleanUp() {
	if _, err := os.Stat(rootDir + ".zip"); err == nil {
		_ = os.Remove(rootDir + ".zip")
	}
}

func downloadProxyFromRepo(client *github.Client, ctx context.Context, owner string, repo string, repopath string) (err error) {

	var fileContent *github.RepositoryContent
	var directoryContents []*github.RepositoryContent

	if fileContent, directoryContents, _, err = client.Repositories.GetContents(ctx, owner, repo, repopath, nil); err != nil {
		return err
	}

	if fileContent != nil {
		if err = downloadResource(*fileContent.Path, *fileContent.DownloadURL); err != nil {
			return err
		}
	}

	if len(directoryContents) > 0 {
		for _, directoryContent := range directoryContents {
			if *directoryContent.Type == "dir" {
				if err = downloadProxyFromRepo(client, ctx, owner, repo, path.Join(repopath, *directoryContent.Name)); err != nil {
					return err
				}
			} else if *directoryContent.Type == "file" {
				if err = downloadResource(*directoryContent.Path, *directoryContent.DownloadURL); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func getApiProxyFolder(repoPath string) (apiProxyFolder string, apiProxyFile string) {
	re := regexp.MustCompile(`(\S*)?(\/?)apiproxy`)

	apiProxyFileBytes := re.ReplaceAll([]byte(repoPath), []byte(rootDir))
	apiProxyFile = string(apiProxyFileBytes)

	apiProxyFolder = filepath.Dir(apiProxyFile)
	return apiProxyFolder, apiProxyFile
}

//downloadResource method is used to download resources, proxy bundles, sharedflows
func downloadResource(repoPath string, url string) (err error) {

	var apiproxyFolder, apiproxyFile string

	if apiproxyFolder, apiproxyFile = getApiProxyFolder(repoPath); err != nil {
		return err
	}

	_ = os.MkdirAll(apiproxyFolder, 0755)

	out, err := os.Create(apiproxyFile)
	if err != nil {
		clilog.Info.Println("error creating file: ", err)
		return err
	}
	defer out.Close()

	client := &http.Client{}

	clilog.Info.Println("Connecting to : ", url)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		clilog.Info.Println("error in client: ", err)
		return err
	}

	resp, err := client.Do(req)

	if err != nil {
		clilog.Info.Println("error connecting: ", err)
		return err
	} else if resp.StatusCode > 299 {
		clilog.Info.Println("error in response: ", resp.Body)
		return errors.New("error in response")
	}

	if resp != nil {
		defer resp.Body.Close()
	}

	if resp == nil {
		clilog.Info.Println("error in response: Response was null")
		return fmt.Errorf("error in response: Response was null")
	}

	_, err = io.Copy(out, resp.Body)
	if err != nil {
		clilog.Info.Println("error writing response to file: ", err)
		return err
	}

	clilog.Info.Println("Resource " + apiproxyFolder + " completed")
	return nil
}
<commit_msg>fix to issue #19, remove duplicate path sep<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package proxybundle

import (
	"archive/zip"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/google/go-github/github"
	genapi "github.com/srinandan/apigeecli/bundlegen"
	apiproxy "github.com/srinandan/apigeecli/bundlegen/apiproxydef"
	policies "github.com/srinandan/apigeecli/bundlegen/policies"
	proxies "github.com/srinandan/apigeecli/bundlegen/proxies"
	target "github.com/srinandan/apigeecli/bundlegen/targets"
	"github.com/srinandan/apigeecli/clilog"
	"golang.org/x/oauth2"
)

const rootDir = "apiproxy"

func GenerateAPIProxyBundleFromOAS(name string,
	content string,
	fileName string,
	skipPolicy bool,
	addCORS bool,
	oasGoogleAcessTokenScopeLiteral string,
	oasGoogleIdTokenAudLiteral string,
	oasGoogleIdTokenAudRef string,
	oasTargetUrlRef string) (err error) {

	var apiProxyData, proxyEndpointData, targetEndpointData string
	const resourceType = "oas"

	if err = os.Mkdir(rootDir, os.ModePerm); err != nil {
		return err
	}

	// write API Proxy file
	if apiProxyData, err = apiproxy.GetAPIProxy(); err != nil {
		return err
	}

	err = writeXMLData(rootDir+string(os.PathSeparator)+name+".xml", apiProxyData)
	if err != nil {
		return err
	}

	proxiesDirPath := rootDir + string(os.PathSeparator) + "proxies"
	policiesDirPath := rootDir + string(os.PathSeparator) + "policies"
	targetDirPath := rootDir + string(os.PathSeparator) + "targets"
	resDirPath := rootDir + string(os.PathSeparator) + "resources" + string(os.PathSeparator) + resourceType //"oas"

	if err = os.Mkdir(proxiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if proxyEndpointData, err = proxies.GetProxyEndpoint(); err != nil {
		return err
	}

	err = writeXMLData(proxiesDirPath+string(os.PathSeparator)+"default.xml", proxyEndpointData)
	if err != nil {
		return err
	}

	if err = os.Mkdir(targetDirPath, os.ModePerm); err != nil {
		return err
	}

	if targetEndpointData, err = target.GetTargetEndpoint(); err != nil {
		return err
	}

	if err = writeXMLData(targetDirPath+string(os.PathSeparator)+"default.xml", targetEndpointData); err != nil {
		return err
	}

	if !skipPolicy {
		if err = os.MkdirAll(resDirPath, os.ModePerm); err != nil {
			return err
		}
		if err = writeXMLData(resDirPath+string(os.PathSeparator)+fileName, content); err != nil {
			return err
		}
	}

	if err = os.Mkdir(policiesDirPath, os.ModePerm); err != nil {
		return err
	}

	//add set target url
	if genapi.GenerateSetTargetPolicy() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Set-Target-1.xml",
			policies.AddSetTargetEndpoint(oasTargetUrlRef)); err != nil {
			return err
		}
	}

	//add security policies
	for _, securityScheme := range genapi.GetSecuritySchemesList() {
		if securityScheme.APIKeyPolicy.APIKeyPolicyEnabled {
			if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Verify-API-Key-"+securityScheme.SchemeName+".xml",
				policies.AddVerifyApiKeyPolicy(securityScheme.APIKeyPolicy.APIKeyLocation,
					securityScheme.SchemeName,
					securityScheme.APIKeyPolicy.APIKeyName)); err != nil {
				return err
			}
		}
		if securityScheme.OAuthPolicy.OAuthPolicyEnabled {
			if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"OAuth-v20-1.xml",
				policies.AddOAuth2Policy(securityScheme.OAuthPolicy.Scope)); err != nil {
				return err
			}
		}
	}

	//add quota policies
	for quotaPolicyName, quotaPolicyContent := range genapi.GetQuotaPolicies() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Quota-"+quotaPolicyName+".xml", quotaPolicyContent); err != nil {
			return err
		}
	}

	//add spike arrest policies
	for spikeArrestPolicyName, spikeArrestPolicyContent := range genapi.GetSpikeArrestPolicies() {
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Spike-Arrest-"+spikeArrestPolicyName+".xml", spikeArrestPolicyContent); err != nil {
			return err
		}
	}

	if !skipPolicy {
		//add oas policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"OpenAPI-Spec-Validation-1.xml",
			policies.AddOpenAPIValidatePolicy(fileName)); err != nil {
			return err
		}
	}

	if addCORS {
		//add cors policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Add-CORS.xml", policies.AddCORSPolicy()); err != nil {
			return err
		}
	}

	if err = archiveBundle(rootDir, name+".zip"); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return nil
}

func GenerateAPIProxyBundleFromGQL(name string,
	content string,
	fileName string,
	action string,
	location string,
	keyName string,
	skipPolicy bool,
	addCORS bool,
	targetUrlRef string) (err error) {

	var apiProxyData, proxyEndpointData, targetEndpointData string
	const resourceType = "graphql"

	if err = os.Mkdir(rootDir, os.ModePerm); err != nil {
		return err
	}

	// write API Proxy file
	if apiProxyData, err = apiproxy.GetAPIProxy(); err != nil {
		return err
	}

	err = writeXMLData(rootDir+string(os.PathSeparator)+name+".xml", apiProxyData)
	if err != nil {
		return err
	}

	proxiesDirPath := rootDir + string(os.PathSeparator) + "proxies"
	policiesDirPath := rootDir + string(os.PathSeparator) + "policies"
	targetDirPath := rootDir + string(os.PathSeparator) + "targets"
	resDirPath := rootDir + string(os.PathSeparator) + "resources" + string(os.PathSeparator) + resourceType //"graphql"

	if err = os.Mkdir(proxiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if proxyEndpointData, err = proxies.GetProxyEndpoint(); err != nil {
		return err
	}

	err = writeXMLData(proxiesDirPath+string(os.PathSeparator)+"default.xml", proxyEndpointData)
	if err != nil {
		return err
	}

	if err = os.Mkdir(targetDirPath, os.ModePerm); err != nil {
		return err
	}

	if targetEndpointData, err = target.GetTargetEndpoint(); err != nil {
		return err
	}

	if err = writeXMLData(targetDirPath+string(os.PathSeparator)+"default.xml", targetEndpointData); err != nil {
		return err
	}

	if !skipPolicy {
		if err = os.MkdirAll(resDirPath, os.ModePerm); err != nil {
			return err
		}
		if err = writeXMLData(resDirPath+string(os.PathSeparator)+fileName, content); err != nil {
			return err
		}
	}

	if err = os.Mkdir(policiesDirPath, os.ModePerm); err != nil {
		return err
	}

	if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Set-Target-1.xml",
		policies.AddSetTargetEndpoint(targetUrlRef)); err != nil {
		return err
	}

	if !skipPolicy {
		//add gql policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Validate-"+name+"-Schema.xml",
			policies.AddGraphQLPolicy(name, action, fileName)); err != nil {
			return err
		}
	}

	if keyName != "" {
		//add verifyapi key policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Verify-API-Key-"+name+".xml",
			policies.AddVerifyApiKeyPolicy(location, name, keyName)); err != nil {
			return err
		}
	}

	if addCORS {
		//add cors policy
		if err = writeXMLData(policiesDirPath+string(os.PathSeparator)+"Add-CORS.xml", policies.AddCORSPolicy()); err != nil {
			return err
		}
	}

	if err = archiveBundle(rootDir, name+".zip"); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return nil
}

func writeXMLData(fileName string, data string) error {
	fileWriter, err := os.Create(fileName)
	if err != nil {
		return err
	}
	_, err = fileWriter.WriteString(data)
	if err != nil {
		return err
	}

	fileWriter.Close()
	return nil
}

func GenerateArchiveBundle(pathToZip, destinationPath string) error {
	return archiveBundle(pathToZip, destinationPath)
}

func archiveBundle(pathToZip, destinationPath string) (err error) {

	var destinationFile *os.File

	destinationFile, err = os.Create(destinationPath)
	if err != nil {
		return err
	}

	myZip := zip.NewWriter(destinationFile)
	err = filepath.Walk(pathToZip, func(filePath string, info os.FileInfo, err error) error {
		if info.IsDir() {
			relPath := strings.TrimPrefix(filePath, filepath.Dir(pathToZip))
			zipEntry := strings.ReplaceAll(strings.TrimPrefix(relPath, string(os.PathSeparator))+string(os.PathSeparator), string(os.PathSeparator)+string(os.PathSeparator), string(os.PathSeparator))
			_, err = myZip.Create(zipEntry)
			return err
		}
		if err != nil {
			return err
		}
		relPath := strings.TrimPrefix(filePath, filepath.Dir(pathToZip))
		zipEntry := strings.ReplaceAll(strings.TrimPrefix(relPath, string(os.PathSeparator)), string(os.PathSeparator)+string(os.PathSeparator), string(os.PathSeparator))
		zipFile, err := myZip.Create(zipEntry)
		if err != nil {
			return err
		}
		fsFile, err := os.Open(filePath)
		if err != nil {
			return err
		}
		_, err = io.Copy(zipFile, fsFile)
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		return err
	}
	err = myZip.Close()
	if err != nil {
		return err
	}
	return nil
}

func GitHubImportBundle(owner string, repo string, repopath string) (err error) {

	//clean up any files or folders
	CleanUp()
	os.RemoveAll(rootDir)

	//
	token := os.Getenv("GITHUB_TOKEN")
	ctx := context.Background()
	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})
	tc := oauth2.NewClient(ctx, ts)
	client := github.NewClient(tc)

	//1. download the proxy
	if err := downloadProxyFromRepo(client, ctx, owner, repo, repopath); err != nil {
		return err
	}

	if client != nil {
		fmt.Println("")
	}

	//2. compress the proxy folder
	curDir, _ := os.Getwd()
	if err := archiveBundle(path.Join(curDir, rootDir), path.Join(curDir, rootDir+".zip")); err != nil {
		return err
	}

	defer os.RemoveAll(rootDir) // clean up
	return err
}

func CleanUp() {
	if _, err := os.Stat(rootDir + ".zip"); err == nil {
		_ = os.Remove(rootDir + ".zip")
	}
}

func downloadProxyFromRepo(client *github.Client, ctx context.Context, owner string, repo string, repopath string) (err error) {

	var fileContent *github.RepositoryContent
	var directoryContents []*github.RepositoryContent

	if fileContent, directoryContents, _, err = client.Repositories.GetContents(ctx, owner, repo, repopath, nil); err != nil {
		return err
	}

	if fileContent != nil {
		if err = downloadResource(*fileContent.Path, *fileContent.DownloadURL); err != nil {
			return err
		}
	}

	if len(directoryContents) > 0 {
		for _, directoryContent := range directoryContents {
			if *directoryContent.Type == "dir" {
				if err = downloadProxyFromRepo(client, ctx, owner, repo, path.Join(repopath, *directoryContent.Name)); err != nil {
					return err
				}
			} else if *directoryContent.Type == "file" {
				if err = downloadResource(*directoryContent.Path, *directoryContent.DownloadURL); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func getApiProxyFolder(repoPath string) (apiProxyFolder string, apiProxyFile string) {
	re := regexp.MustCompile(`(\S*)?(\/?)apiproxy`)

	apiProxyFileBytes := re.ReplaceAll([]byte(repoPath), []byte(rootDir))
	apiProxyFile = string(apiProxyFileBytes)

	apiProxyFolder = filepath.Dir(apiProxyFile)
	return apiProxyFolder, apiProxyFile
}

//downloadResource method is used to download resources, proxy bundles, sharedflows
func downloadResource(repoPath string, url string) (err error) {

	var apiproxyFolder, apiproxyFile string

	if apiproxyFolder, apiproxyFile = getApiProxyFolder(repoPath); err != nil {
		return err
	}

	_ = os.MkdirAll(apiproxyFolder, 0755)

	out, err := os.Create(apiproxyFile)
	if err != nil {
		clilog.Info.Println("error creating file: ", err)
		return err
	}
	defer out.Close()

	client := &http.Client{}

	clilog.Info.Println("Connecting to : ", url)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		clilog.Info.Println("error in client: ", err)
		return err
	}

	resp, err := client.Do(req)

	if err != nil {
		clilog.Info.Println("error connecting: ", err)
		return err
	} else if resp.StatusCode > 299 {
		clilog.Info.Println("error in response: ", resp.Body)
		return errors.New("error in response")
	}

	if resp != nil {
		defer resp.Body.Close()
	}

	if resp == nil {
		clilog.Info.Println("error in response: Response was null")
		return fmt.Errorf("error in response: Response was null")
	}

	_, err = io.Copy(out, resp.Body)
	if err != nil {
		clilog.Info.Println("error writing response to file: ", err)
		return err
	}

	clilog.Info.Println("Resource " + apiproxyFolder + " completed")
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	c "github.com/couchbase/indexing/secondary/common"
	qclient "github.com/couchbase/indexing/secondary/queryport/client"
	"math"
	"sync"
	"sync/atomic"
	"time"
)

var (
	defaultLatencyBuckets = []int64{
		500, 1000, 10000, 200000, 500000, 800000,
	}
)

type Job struct {
	spec   *ScanConfig
	result *ScanResult
}

func RunScan(client *qclient.GsiClient,
	spec *ScanConfig, result *ScanResult) {
	var err error
	var rows int64

	result.Id = spec.Id

	errFn := func(e string) {
		fmt.Printf("REQ:%d scan error occured: %s\n", spec.Id, e)
		atomic.AddUint64(&result.ErrorCount, 1)
	}

	callb := func(res qclient.ResponseReader) bool {
		if res.Error() != nil {
			errFn(res.Error().Error())
			return false
		} else {
			_, pkeys, err := res.GetEntries()
			if err != nil {
				errFn(err.Error())
				return false
			}

			rows += int64(len(pkeys))
		}

		return true
	}

	startTime := time.Now()
	switch spec.Type {
	case "All":
		err = client.ScanAll(spec.DefnId, spec.Limit, c.AnyConsistency, nil, callb)
	case "Range":
		err = client.Range(spec.DefnId, spec.Low, spec.High,
			qclient.Inclusion(spec.Inclusion), false, spec.Limit, c.AnyConsistency, nil, callb)
	case "Lookup":
		err = client.Lookup(spec.DefnId, spec.Lookups, false,
			spec.Limit, c.AnyConsistency, nil, callb)
	}

	if err != nil {
		errFn(err.Error())
	}

	var lat int64
	dur := time.Now().Sub(startTime)
	atomic.AddUint64(&result.Rows, uint64(rows))
	if rows > 0 {
		lat = dur.Nanoseconds() / rows
	}
	result.LatencyHisto.Add(lat)
	atomic.AddInt64(&result.Duration, dur.Nanoseconds())
}

func Worker(jobQ chan Job, clientQ chan *qclient.GsiClient, wg *sync.WaitGroup) {
	defer wg.Done()

	for job := range jobQ {
		c := <-clientQ
		RunScan(c, job.spec, job.result)
		clientQ <- c
	}
}

func RunCommands(cluster string, cfg *Config) (*Result, error) {
	var result Result

	var clientQ chan *qclient.GsiClient
	var jobQ chan Job
	var wg sync.WaitGroup

	if len(cfg.LatencyBuckets) == 0 {
		cfg.LatencyBuckets = defaultLatencyBuckets
	}

	config := c.SystemConfig.SectionConfig("queryport.client.", true)
	client, err := qclient.NewGsiClient(cluster, config)
	if err != nil {
		return nil, err
	}
	defer client.Close()

	indexes, err := client.Refresh()
	if err != nil {
		return nil, err
	}

	clientQ = make(chan *qclient.GsiClient, cfg.Clients)
	for i := 0; i < cfg.Clients; i++ {
		c, err := qclient.NewGsiClient(cluster, config)
		if err != nil {
			return nil, err
		}

		defer c.Close()
		clientQ <- c
	}

	jobQ = make(chan Job, cfg.Concurrency)
	for i := 0; i < cfg.Concurrency; i++ {
		wg.Add(1)
		go Worker(jobQ, clientQ, &wg)
	}

	for i, spec := range cfg.ScanSpecs {
		if spec.Id == 0 {
			spec.Id = uint64(i)
		}

		for _, index := range indexes {
			if index.Definition.Bucket == spec.Bucket &&
				index.Definition.Name == spec.Index {
				spec.DefnId = uint64(index.Definition.DefnId)
			}
		}

		hFn := func(v int64) string {
			if v == math.MinInt64 {
				return "0"
			} else if v == math.MaxInt64 {
				return "inf"
			}
			return fmt.Sprint(time.Nanosecond * time.Duration(v))
		}

		res := new(ScanResult)
		res.LatencyHisto.Init(cfg.LatencyBuckets, hFn)
		res.Id = spec.Id
		for i := 0; i < spec.Repeat+1; i++ {
			j := Job{
				spec:   spec,
				result: res,
			}

			jobQ <- j
		}

		result.ScanResults = append(result.ScanResults, res)
	}

	close(jobQ)
	wg.Wait()

	return &result, err
}
<commit_msg>MB-14641 cbindexperf: Fix client distribution bug<commit_after>package main

import (
	"fmt"
	c "github.com/couchbase/indexing/secondary/common"
	qclient "github.com/couchbase/indexing/secondary/queryport/client"
	"math"
	"sync"
	"sync/atomic"
	"time"
)

var (
	defaultLatencyBuckets = []int64{
		500, 1000, 10000, 200000, 500000, 800000,
	}
)

type Job struct {
	spec   *ScanConfig
	result *ScanResult
}

func RunScan(client *qclient.GsiClient,
	spec *ScanConfig, result *ScanResult) {
	var err error
	var rows int64

	result.Id = spec.Id

	errFn := func(e string) {
		fmt.Printf("REQ:%d scan error occured: %s\n", spec.Id, e)
		atomic.AddUint64(&result.ErrorCount, 1)
	}

	callb := func(res qclient.ResponseReader) bool {
		if res.Error() != nil {
			errFn(res.Error().Error())
			return false
		} else {
			_, pkeys, err := res.GetEntries()
			if err != nil {
				errFn(err.Error())
				return false
			}

			rows += int64(len(pkeys))
		}

		return true
	}

	startTime := time.Now()
	switch spec.Type {
	case "All":
		err = client.ScanAll(spec.DefnId, spec.Limit, c.AnyConsistency, nil, callb)
	case "Range":
		err = client.Range(spec.DefnId, spec.Low, spec.High,
			qclient.Inclusion(spec.Inclusion), false, spec.Limit, c.AnyConsistency, nil, callb)
	case "Lookup":
		err = client.Lookup(spec.DefnId, spec.Lookups, false,
			spec.Limit, c.AnyConsistency, nil, callb)
	}

	if err != nil {
		errFn(err.Error())
	}

	var lat int64
	dur := time.Now().Sub(startTime)
	atomic.AddUint64(&result.Rows, uint64(rows))
	if rows > 0 {
		lat = dur.Nanoseconds() / rows
	}
	result.LatencyHisto.Add(lat)
	atomic.AddInt64(&result.Duration, dur.Nanoseconds())
}

func Worker(jobQ chan Job, c *qclient.GsiClient, wg *sync.WaitGroup) {
	defer wg.Done()

	for job := range jobQ {
		RunScan(c, job.spec, job.result)
	}
}

func RunCommands(cluster string, cfg *Config) (*Result, error) {
	var result Result

	var clients []*qclient.GsiClient
	var jobQ chan Job
	var wg sync.WaitGroup

	if len(cfg.LatencyBuckets) == 0 {
		cfg.LatencyBuckets = defaultLatencyBuckets
	}

	config := c.SystemConfig.SectionConfig("queryport.client.", true)
	client, err := qclient.NewGsiClient(cluster, config)
	if err != nil {
		return nil, err
	}
	defer client.Close()

	indexes, err := client.Refresh()
	if err != nil {
		return nil, err
	}

	clients = make([]*qclient.GsiClient, cfg.Clients)
	for i := 0; i < cfg.Clients; i++ {
		c, err := qclient.NewGsiClient(cluster, config)
		if err != nil {
			return nil, err
		}

		defer c.Close()
		clients[i] = c
	}

	jobQ = make(chan Job, cfg.Concurrency*1000)
	for i := 0; i < cfg.Concurrency; i++ {
		wg.Add(1)
		go Worker(jobQ, clients[i%cfg.Clients], &wg)
	}

	for i, spec := range cfg.ScanSpecs {
		if spec.Id == 0 {
			spec.Id = uint64(i)
		}

		for _, index := range indexes {
			if index.Definition.Bucket == spec.Bucket &&
				index.Definition.Name == spec.Index {
				spec.DefnId = uint64(index.Definition.DefnId)
			}
		}

		hFn := func(v int64) string {
			if v == math.MinInt64 {
				return "0"
			} else if v == math.MaxInt64 {
				return "inf"
			}
			return fmt.Sprint(time.Nanosecond * time.Duration(v))
		}

		res := new(ScanResult)
		res.LatencyHisto.Init(cfg.LatencyBuckets, hFn)
		res.Id = spec.Id
		for i := 0; i < spec.Repeat+1; i++ {
			j := Job{
				spec:   spec,
				result: res,
			}

			jobQ <- j
		}

		result.ScanResults = append(result.ScanResults, res)
	}

	close(jobQ)
	wg.Wait()

	return &result, err
}
<|endoftext|>


<commit_before>// +build local

/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

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

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

package local

import (
	"context"
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/docker/docker/errdefs"

	"github.com/compose-spec/compose-go/types"
	"github.com/docker/buildx/build"
	"github.com/docker/buildx/driver"
	_ "github.com/docker/buildx/driver/docker" // required to get default driver registered
	"github.com/docker/buildx/util/progress"
)

func (s *composeService) ensureImagesExists(ctx context.Context, project *types.Project) error {
	opts := map[string]build.Options{}
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}

		// TODO build vs pull should be controlled by pull policy, see https://github.com/compose-spec/compose-spec/issues/26
		if service.Image != "" {
			needPull, err := s.needPull(ctx, service)
			if err != nil {
				return err
			}
			if !needPull {
				continue
			}
		}
		if service.Build != nil {
			imageName := service.Image
			if imageName == "" {
				imageName = project.Name + "_" + service.Name
			}
			opts[imageName] = s.toBuildOptions(service, project.WorkingDir)
			continue
		}

		// Buildx has no command to "just pull", see
		// so we bake a temporary dockerfile that will just pull and export pulled image
		opts[service.Name] = build.Options{
			Inputs: build.Inputs{
				ContextPath:    ".",
				DockerfilePath: "-",
				InStream:       strings.NewReader("FROM " + service.Image),
			},
			Tags: []string{service.Image},
			Pull: true,
		}

	}

	return s.build(ctx, project, opts)
}

func (s *composeService) needPull(ctx context.Context, service types.ServiceConfig) (bool, error) {
	_, _, err := s.apiClient.ImageInspectWithRaw(ctx, service.Image)
	if err != nil {
		if errdefs.IsNotFound(err) {
			return true, nil
		}
		return false, err
	}
	return false, nil
}

func (s *composeService) build(ctx context.Context, project *types.Project, opts map[string]build.Options) error {
	if len(opts) == 0 {
		return nil
	}
	const drivername = "default"
	d, err := driver.GetDriver(ctx, drivername, nil, s.apiClient, nil, nil, "", nil, project.WorkingDir)
	if err != nil {
		return err
	}
	driverInfo := []build.DriverInfo{
		{
			Name:   "default",
			Driver: d,
		},
	}
	// We rely on buildx "docker" builder integrated in docker engine, so don't need a DockerAPI here
	w := progress.NewPrinter(ctx, os.Stdout, "auto")
	_, err = build.Build(ctx, driverInfo, opts, nil, nil, w)
	return err
}

func (s *composeService) toBuildOptions(service types.ServiceConfig, contextPath string) build.Options {
	var tags []string
	if service.Image != "" {
		tags = append(tags, service.Image)
	}

	if service.Build.Dockerfile == "" {
		service.Build.Dockerfile = "Dockerfile"
	}
	var buildArgs map[string]string

	return build.Options{
		Inputs: build.Inputs{
			ContextPath:    path.Join(contextPath, service.Build.Context),
			DockerfilePath: path.Join(contextPath, service.Build.Context, service.Build.Dockerfile),
		},
		BuildArgs: flatten(mergeArgs(service.Build.Args, buildArgs)),
		Tags:      tags,
	}
}

func flatten(in types.MappingWithEquals) map[string]string {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]string)
	for k, v := range in {
		if v == nil {
			continue
		}
		out[k] = *v
	}
	return out
}

func mergeArgs(src types.MappingWithEquals, values map[string]string) types.MappingWithEquals {
	for key := range src {
		if val, ok := values[key]; ok {
			if val == "" {
				src[key] = nil
			} else {
				src[key] = &val
			}
		}
	}
	return src
}
<commit_msg>Give progress its own context<commit_after>// +build local

/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

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

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

package local

import (
	"context"
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/docker/docker/errdefs"

	"github.com/compose-spec/compose-go/types"
	"github.com/docker/buildx/build"
	"github.com/docker/buildx/driver"
	_ "github.com/docker/buildx/driver/docker" // required to get default driver registered
	"github.com/docker/buildx/util/progress"
)

func (s *composeService) ensureImagesExists(ctx context.Context, project *types.Project) error {
	opts := map[string]build.Options{}
	for _, service := range project.Services {
		if service.Image == "" && service.Build == nil {
			return fmt.Errorf("invalid service %q. Must specify either image or build", service.Name)
		}

		// TODO build vs pull should be controlled by pull policy, see https://github.com/compose-spec/compose-spec/issues/26
		if service.Image != "" {
			needPull, err := s.needPull(ctx, service)
			if err != nil {
				return err
			}
			if !needPull {
				continue
			}
		}
		if service.Build != nil {
			imageName := service.Image
			if imageName == "" {
				imageName = project.Name + "_" + service.Name
			}
			opts[imageName] = s.toBuildOptions(service, project.WorkingDir)
			continue
		}

		// Buildx has no command to "just pull", see
		// so we bake a temporary dockerfile that will just pull and export pulled image
		opts[service.Name] = build.Options{
			Inputs: build.Inputs{
				ContextPath:    ".",
				DockerfilePath: "-",
				InStream:       strings.NewReader("FROM " + service.Image),
			},
			Tags: []string{service.Image},
			Pull: true,
		}

	}

	return s.build(ctx, project, opts)
}

func (s *composeService) needPull(ctx context.Context, service types.ServiceConfig) (bool, error) {
	_, _, err := s.apiClient.ImageInspectWithRaw(ctx, service.Image)
	if err != nil {
		if errdefs.IsNotFound(err) {
			return true, nil
		}
		return false, err
	}
	return false, nil
}

func (s *composeService) build(ctx context.Context, project *types.Project, opts map[string]build.Options) error {
	if len(opts) == 0 {
		return nil
	}
	const drivername = "default"
	d, err := driver.GetDriver(ctx, drivername, nil, s.apiClient, nil, nil, "", nil, project.WorkingDir)
	if err != nil {
		return err
	}
	driverInfo := []build.DriverInfo{
		{
			Name:   "default",
			Driver: d,
		},
	}

	// Progress needs its own context that lives longer than the
	// build one otherwise it won't read all the messages from
	// build and will lock
	progressCtx, cancel := context.WithCancel(context.Background())
	defer cancel()
	w := progress.NewPrinter(progressCtx, os.Stdout, "auto")

	// We rely on buildx "docker" builder integrated in docker engine, so don't need a DockerAPI here
	_, err = build.Build(ctx, driverInfo, opts, nil, nil, w)
	return err
}

func (s *composeService) toBuildOptions(service types.ServiceConfig, contextPath string) build.Options {
	var tags []string
	if service.Image != "" {
		tags = append(tags, service.Image)
	}

	if service.Build.Dockerfile == "" {
		service.Build.Dockerfile = "Dockerfile"
	}
	var buildArgs map[string]string

	return build.Options{
		Inputs: build.Inputs{
			ContextPath:    path.Join(contextPath, service.Build.Context),
			DockerfilePath: path.Join(contextPath, service.Build.Context, service.Build.Dockerfile),
		},
		BuildArgs: flatten(mergeArgs(service.Build.Args, buildArgs)),
		Tags:      tags,
	}
}

func flatten(in types.MappingWithEquals) map[string]string {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]string)
	for k, v := range in {
		if v == nil {
			continue
		}
		out[k] = *v
	}
	return out
}

func mergeArgs(src types.MappingWithEquals, values map[string]string) types.MappingWithEquals {
	for key := range src {
		if val, ok := values[key]; ok {
			if val == "" {
				src[key] = nil
			} else {
				src[key] = &val
			}
		}
	}
	return src
}
<|endoftext|>


<commit_before>package ddtxn

import (
	"ddtxn/dlog"
	"flag"
	"log"
	"runtime/debug"
	"time"
)

// Local per-worker store. Specific types to more quickly apply local
// changes

// Phases
const (
	SPLIT = iota
	MERGE
	JOIN
)

var SampleRate = flag.Int64("sr", 100, "Sample every sr nanoseconds\n")

type LocalStore struct {
	sums       map[Key]int32
	max        map[Key]int32
	bw         map[Key]Value
	lists      map[Key][]Entry
	s          *Store
	phase      uint32
	Ncopy      int64
	candidates *Candidates
	start      time.Time
	count      bool
}

func NewLocalStore(s *Store) *LocalStore {
	x := make([]*OneStat, 0)
	sh := StatsHeap(x)
	ls := &LocalStore{
		sums:       make(map[Key]int32),
		max:        make(map[Key]int32),
		bw:         make(map[Key]Value),
		lists:      make(map[Key][]Entry),
		s:          s,
		candidates: &Candidates{make(map[Key]*OneStat), &sh},
		start:      time.Now(),
	}
	return ls
}

func (ls *LocalStore) Apply(key Key, key_type KeyType, v Value, op KeyType) {
	if op != key_type {
		// Perhaps do something.  When is this set?
		dlog.Printf("Different op types %v %v\n", key_type, op)
	}
	switch op {
	case SUM:
		ls.sums[key] += v.(int32)
	case MAX:
		delta := v.(int32)
		if ls.max[key] < delta {
			ls.max[key] = delta
		}
	case WRITE:
		ls.bw[key] = v
	case LIST:
		entry := v.(Entry)
		l, ok := ls.lists[key]
		if !ok {
			l = make([]Entry, 0)
			ls.lists[key] = l
		}
		// TODO: Use listApply or add one to list to keep them sorted
		ls.lists[key] = append(l, entry)
	}
}

func (ls *LocalStore) Merge() {
	for k, v := range ls.sums {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}
		if v == 0 {
			continue
		}
		if v == 0 {
			continue
		}
		d := ls.s.getOrCreateTypedKey(k, int32(0), SUM)
		d.Apply(v)
		ls.sums[k] = 0
		ls.Ncopy++
	}

	for k, v := range ls.max {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}

		if v == 0 {
			continue
		}
		d := ls.s.getOrCreateTypedKey(k, int32(0), MAX)
		d.Apply(v)
		ls.Ncopy++
	}

	for k, v := range ls.bw {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}

		d := ls.s.getOrCreateTypedKey(k, "", WRITE)
		d.Apply(v)
		ls.Ncopy++
	}

	for k, v := range ls.lists {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}
		if len(v) == 0 {
			continue
		}

		d := ls.s.getOrCreateTypedKey(k, nil, LIST)
		d.Apply(v)
		delete(ls.lists, k)
		ls.Ncopy++
	}
}

// TODO: 2PL

type Write struct {
	key    Key
	br     *BRecord
	v      Value
	op     KeyType
	create bool
	locked bool
	vint32 int32
	dd     bool
}

type ETransaction struct {
	read   []*BRecord
	lasts  []uint64
	w      *Worker
	s      *Store
	ls     *LocalStore
	writes []Write
}

// Re-use this?
func StartTransaction(w *Worker) *ETransaction {
	tx := &ETransaction{
		read:   make([]*BRecord, 0, 30),
		lasts:  make([]uint64, 0, 30),
		writes: make([]Write, 0, 60),
		w:      w,
		s:      w.store,
		ls:     w.local_store,
	}
	return tx
}

func (tx *ETransaction) Reset() {
	tx.lasts = tx.lasts[:0]
	tx.read = tx.read[:0]
	tx.writes = tx.writes[:0]
	d := time.Since(tx.ls.start)
	tx.ls.count = (*SysType == DOPPEL && d.Nanoseconds()%*SampleRate == 0)
	if tx.ls.count {
		tx.w.Nsamples++
	}
}

func (tx *ETransaction) Read(k Key) (*BRecord, error) {
	if *SysType == DOPPEL {
		if tx.ls.phase == SPLIT {
			if tx.s.IsDD(k) {
				if tx.ls.count {
					tx.ls.candidates.Stash(k)
				}
				return nil, ESTASH
			}
		}
	}
	// TODO: If I wrote the key, return that value instead
	br, err := tx.s.getKey(k)
	tx.w.NGetKeyCalls++
	if *CountKeys {
		p, r := UndoCKey(k)
		if r == 117 {
			tx.w.NKeyAccesses[p]++
		}
	}
	if err != nil {
		return nil, err
	}
	ok, last := br.IsUnlocked()
	// if locked and not by me, abort
	// else note the last timestamp, save it, return value
	if !ok {
		if tx.ls.count {
			tx.ls.candidates.Conflict(k)
		}
		tx.Abort()
		return nil, EABORT
	}
	n := len(tx.read)
	tx.read = tx.read[0 : n+1]
	tx.read[n] = br
	tx.lasts = tx.lasts[0 : n+1]
	tx.lasts[n] = last
	return br, nil
}

func (tx *ETransaction) add(k Key, v Value, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].v = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addInt32(k Key, v int32, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].vint32 = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) WriteInt32(k Key, a int32, op KeyType) {
	tx.addInt32(k, a, op, false)
}

func (tx *ETransaction) Write(k Key, v Value, kt KeyType) {
	if kt == SUM || kt == MAX {
		tx.addInt32(k, v.(int32), kt, true)
		return
	}
	tx.add(k, v, kt, true)
}

func (tx *ETransaction) Abort() TID {
	for i, _ := range tx.writes {
		if tx.writes[i].locked {
			tx.writes[i].br.Unlock(0)
		}
	}
	return 0
}

func (tx *ETransaction) Commit() TID {
	// for each write key
	//  if global get from global store and lock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if *SysType == DOPPEL && tx.ls.phase == SPLIT && tx.s.IsDD(w.key) {
			w.dd = true
			continue
		}
		if w.br == nil {
			br, err := tx.s.getKey(w.key)
			tx.w.NGetKeyCalls++
			if *CountKeys {
				p, r := UndoCKey(w.key)
				if r == 117 {
					tx.w.NKeyAccesses[p]++
				}
			}
			if br == nil || err != nil {
				switch w.op {
				case SUM:
					br = tx.s.CreateInt32Key(w.key, w.vint32, w.op)
				case MAX:
					br = tx.s.CreateInt32Key(w.key, w.vint32, w.op)
				default:
					br = tx.s.CreateKey(w.key, "", WRITE)
				}
			}
			w.br = br
		}
		if !w.br.Lock() {
			if tx.ls.count {
				tx.ls.candidates.Conflict(w.key)
			}
			return tx.Abort()
		}
		w.locked = true
	}
	// TODO: acquire timestamp higher than anything i've read or am
	// writing
	tid := tx.w.commitTID()

	// for each read key
	//  verify
	if len(tx.read) != len(tx.lasts) {
		debug.PrintStack()
		log.Fatalf("Mismatch in lengths reads: %v, lasts: %v\n", tx.read, tx.lasts)
	}
	for i, _ := range tx.read {
		if tx.ls.count {
			tx.ls.candidates.Read(tx.read[i].key)
		}
		rd := false
		if !tx.read[i].Verify(tx.lasts[i]) {
			for j, _ := range tx.writes {
				if tx.writes[j].key == tx.read[i].key {
					// We would have aborted if we did not successfully
					// lock this earlier
					rd = true
					break
				}
			}
			if rd {
				continue
			}
			if tx.ls.count {
				tx.ls.candidates.Conflict(tx.read[i].key)
			}
			return tx.Abort()
		}
	}
	// for each write key
	//  if dd and split phase, apply locally
	//  else apply globally and unlock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if tx.ls.phase == SPLIT && w.dd {
			if tx.ls.count {
				tx.ls.candidates.Write(w.key)
			}
			tx.w.NDDWrites++
			switch w.op {
			case SUM:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case MAX:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			default:
				tx.ls.Apply(w.key, w.op, w.v, w.op)
			}
		} else {
			switch w.op {
			case SUM:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case MAX:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			default:
				tx.s.Set(w.br, w.v, w.op)
			}
			w.br.Unlock(tid)
		}
	}
	return tid
}
<commit_msg>don't use time when a rough count will do<commit_after>package ddtxn

import (
	"ddtxn/dlog"
	"flag"
	"log"
	"runtime/debug"
)

// Local per-worker store. Specific types to more quickly apply local
// changes

// Phases
const (
	SPLIT = iota
	MERGE
	JOIN
)

var SampleRate = flag.Int("sr", 1000, "Sample every sr nanoseconds\n")

type LocalStore struct {
	sums       map[Key]int32
	max        map[Key]int32
	bw         map[Key]Value
	lists      map[Key][]Entry
	s          *Store
	phase      uint32
	Ncopy      int64
	candidates *Candidates
	count      bool
}

func NewLocalStore(s *Store) *LocalStore {
	x := make([]*OneStat, 0)
	sh := StatsHeap(x)
	ls := &LocalStore{
		sums:       make(map[Key]int32),
		max:        make(map[Key]int32),
		bw:         make(map[Key]Value),
		lists:      make(map[Key][]Entry),
		s:          s,
		candidates: &Candidates{make(map[Key]*OneStat), &sh},
	}
	return ls
}

func (ls *LocalStore) Apply(key Key, key_type KeyType, v Value, op KeyType) {
	if op != key_type {
		// Perhaps do something.  When is this set?
		dlog.Printf("Different op types %v %v\n", key_type, op)
	}
	switch op {
	case SUM:
		ls.sums[key] += v.(int32)
	case MAX:
		delta := v.(int32)
		if ls.max[key] < delta {
			ls.max[key] = delta
		}
	case WRITE:
		ls.bw[key] = v
	case LIST:
		entry := v.(Entry)
		l, ok := ls.lists[key]
		if !ok {
			l = make([]Entry, 0)
			ls.lists[key] = l
		}
		// TODO: Use listApply or add one to list to keep them sorted
		ls.lists[key] = append(l, entry)
	}
}

func (ls *LocalStore) Merge() {
	for k, v := range ls.sums {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}
		if v == 0 {
			continue
		}
		if v == 0 {
			continue
		}
		d := ls.s.getOrCreateTypedKey(k, int32(0), SUM)
		d.Apply(v)
		ls.sums[k] = 0
		ls.Ncopy++
	}

	for k, v := range ls.max {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}

		if v == 0 {
			continue
		}
		d := ls.s.getOrCreateTypedKey(k, int32(0), MAX)
		d.Apply(v)
		ls.Ncopy++
	}

	for k, v := range ls.bw {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}

		d := ls.s.getOrCreateTypedKey(k, "", WRITE)
		d.Apply(v)
		ls.Ncopy++
	}

	for k, v := range ls.lists {
		if *SysType == OCC {
			debug.PrintStack()
			log.Fatalf("Why is there derived data %v %v\n", k, v)
		}
		if len(v) == 0 {
			continue
		}

		d := ls.s.getOrCreateTypedKey(k, nil, LIST)
		d.Apply(v)
		delete(ls.lists, k)
		ls.Ncopy++
	}
}

// TODO: 2PL

type Write struct {
	key    Key
	br     *BRecord
	v      Value
	op     KeyType
	create bool
	locked bool
	vint32 int32
	dd     bool
}

type ETransaction struct {
	read   []*BRecord
	lasts  []uint64
	w      *Worker
	s      *Store
	ls     *LocalStore
	writes []Write
	t      int // Used just as a rough count
}

// Re-use this?
func StartTransaction(w *Worker) *ETransaction {
	tx := &ETransaction{
		read:   make([]*BRecord, 0, 30),
		lasts:  make([]uint64, 0, 30),
		writes: make([]Write, 0, 60),
		w:      w,
		s:      w.store,
		ls:     w.local_store,
	}
	return tx
}

func (tx *ETransaction) Reset() {
	tx.lasts = tx.lasts[:0]
	tx.read = tx.read[:0]
	tx.writes = tx.writes[:0]
	tx.ls.count = (*SysType == DOPPEL && tx.t%*SampleRate == 0)
	if tx.ls.count {
		tx.w.Nsamples++
	}
	tx.t++
}

func (tx *ETransaction) Read(k Key) (*BRecord, error) {
	if *SysType == DOPPEL {
		if tx.ls.phase == SPLIT {
			if tx.s.IsDD(k) {
				if tx.ls.count {
					tx.ls.candidates.Stash(k)
				}
				return nil, ESTASH
			}
		}
	}
	// TODO: If I wrote the key, return that value instead
	br, err := tx.s.getKey(k)
	tx.w.NGetKeyCalls++
	if *CountKeys {
		p, r := UndoCKey(k)
		if r == 117 {
			tx.w.NKeyAccesses[p]++
		}
	}
	if err != nil {
		return nil, err
	}
	ok, last := br.IsUnlocked()
	// if locked and not by me, abort
	// else note the last timestamp, save it, return value
	if !ok {
		if tx.ls.count {
			tx.ls.candidates.Conflict(k)
		}
		tx.Abort()
		return nil, EABORT
	}
	n := len(tx.read)
	tx.read = tx.read[0 : n+1]
	tx.read[n] = br
	tx.lasts = tx.lasts[0 : n+1]
	tx.lasts[n] = last
	return br, nil
}

func (tx *ETransaction) add(k Key, v Value, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].v = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) addInt32(k Key, v int32, op KeyType, create bool) {
	if len(tx.writes) == cap(tx.writes) {
		// TODO: extend
		log.Fatalf("Ran out of room\n")
	}
	n := len(tx.writes)
	tx.writes = tx.writes[0 : n+1]
	tx.writes[n].key = k
	tx.writes[n].br = nil
	tx.writes[n].vint32 = v
	tx.writes[n].op = op
	tx.writes[n].create = create
	tx.writes[n].locked = false
}

func (tx *ETransaction) WriteInt32(k Key, a int32, op KeyType) {
	tx.addInt32(k, a, op, false)
}

func (tx *ETransaction) Write(k Key, v Value, kt KeyType) {
	if kt == SUM || kt == MAX {
		tx.addInt32(k, v.(int32), kt, true)
		return
	}
	tx.add(k, v, kt, true)
}

func (tx *ETransaction) Abort() TID {
	for i, _ := range tx.writes {
		if tx.writes[i].locked {
			tx.writes[i].br.Unlock(0)
		}
	}
	return 0
}

func (tx *ETransaction) Commit() TID {
	// for each write key
	//  if global get from global store and lock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if *SysType == DOPPEL && tx.ls.phase == SPLIT && tx.s.IsDD(w.key) {
			w.dd = true
			continue
		}
		if w.br == nil {
			br, err := tx.s.getKey(w.key)
			tx.w.NGetKeyCalls++
			if *CountKeys {
				p, r := UndoCKey(w.key)
				if r == 117 {
					tx.w.NKeyAccesses[p]++
				}
			}
			if br == nil || err != nil {
				switch w.op {
				case SUM:
					br = tx.s.CreateInt32Key(w.key, w.vint32, w.op)
				case MAX:
					br = tx.s.CreateInt32Key(w.key, w.vint32, w.op)
				default:
					br = tx.s.CreateKey(w.key, "", WRITE)
				}
			}
			w.br = br
		}
		if !w.br.Lock() {
			if tx.ls.count {
				tx.ls.candidates.Conflict(w.key)
			}
			return tx.Abort()
		}
		w.locked = true
	}
	// TODO: acquire timestamp higher than anything i've read or am
	// writing
	tid := tx.w.commitTID()

	// for each read key
	//  verify
	if len(tx.read) != len(tx.lasts) {
		debug.PrintStack()
		log.Fatalf("Mismatch in lengths reads: %v, lasts: %v\n", tx.read, tx.lasts)
	}
	for i, _ := range tx.read {
		if tx.ls.count {
			tx.ls.candidates.Read(tx.read[i].key)
		}
		rd := false
		if !tx.read[i].Verify(tx.lasts[i]) {
			for j, _ := range tx.writes {
				if tx.writes[j].key == tx.read[i].key {
					// We would have aborted if we did not successfully
					// lock this earlier
					rd = true
					break
				}
			}
			if rd {
				continue
			}
			if tx.ls.count {
				tx.ls.candidates.Conflict(tx.read[i].key)
			}
			return tx.Abort()
		}
	}
	// for each write key
	//  if dd and split phase, apply locally
	//  else apply globally and unlock
	for i, _ := range tx.writes {
		w := &tx.writes[i]
		if tx.ls.phase == SPLIT && w.dd {
			if tx.ls.count {
				tx.ls.candidates.Write(w.key)
			}
			tx.w.NDDWrites++
			switch w.op {
			case SUM:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			case MAX:
				tx.ls.Apply(w.key, w.op, w.vint32, w.op)
			default:
				tx.ls.Apply(w.key, w.op, w.v, w.op)
			}
		} else {
			switch w.op {
			case SUM:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			case MAX:
				tx.s.SetInt32(w.br, w.vint32, w.op)
			default:
				tx.s.Set(w.br, w.v, w.op)
			}
			w.br.Unlock(tid)
		}
	}
	return tid
}
<|endoftext|>


<commit_before>package logger

import (
	"bytes"
	"testing"
)

const (
	namet = name + ".Test"
)

func TestGetLevel(t *testing.T) {
	n := New("logger.Test.GetLevel")

	n.Info(n, "Starting")
	m := make(map[Logger]Priority)
	m[""] = DefaultPriority
	m["."] = DefaultPriority
	m["Test"] = DefaultPriority
	m[".Test"] = DefaultPriority

	SetLevel("Test2", Emergency)
	m["Test2"] = Emergency
	m["Test2.Test"] = Emergency
	m["Test2.Test.Test"] = Emergency
	m["Test2.Test.Test.Test"] = Emergency
	m["Test2.Test.Test.Test.Test"] = Emergency
	m["Test2.Test.Test.Test.Test.Test"] = Emergency

	for k, v := range m {
		o := GetLevel(k)
		if o != v {
			n.Error(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
		n.Debug(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
	}
	n.Info(n, "Finished")
}

func TestGetParent(t *testing.T) {
	n := New("logger.Test.getParent")

	n.Info(n, "Starting")
	m := [][]Logger{
		{"", "."},
		{".Test", "."},
		{".", "."},
		{"Test", "."},
		{"Test.Test", "Test"},
		{"Test.Test.Test", "Test.Test"},
		{"Test.Test.Test.Test", "Test.Test.Test"},
	}

	for i := range m {
		a := m[i]

		k := a[0]
		v := a[1]

		o := getParent(k)
		if o != v {
			n.Error(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
		n.Debug(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
	}
	n.Info(n, "Finished")
}

func TestGetParentOutputSame(t *testing.T) {
	l := New(namet + ".GetParent.Output.Same")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	p.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "Test Parent,Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestGetParentOutputDifferent(t *testing.T) {
	l := New(namet + ".GetParent.Output.Different")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	c.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestGetParentOutputInheritance(t *testing.T) {
	l := New(namet + ".GetParent.Output.Inheritance")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	c.SetLevel(Debug)
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	p.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "TestParent,Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestPrintMessage(t *testing.T) {
	l := New(namet + ".PrintMessage")

	p := "\033[0m"
	b := "Test - " + p + p + "Debug" + p + " - "

	m := [][]string{
		{"", b},
		{"Test", b + "Test"},
		{"Test.Test", b + "Test.Test"},
		{"Test.Test.Test", b + "Test.Test.Test"},
	}

	r := list.GetLogger("Test")
	r.Format = "{{.Logger}} - {{.Priority}} - {{.Message}}"

	for _, d := range m {
		l.Info("Checking: ", d)

		k := d[0]
		v := d[1]

		var b bytes.Buffer
		r.Output = &b

		printMessage(r, Debug, k)
		o := b.String()

		l.Debug("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
		if o != v {
			l.Critical("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
	}
}

func TestPrintMessageNoColor(t *testing.T) {
	l := New(namet + ".PrintMessage")

	m := [][]string{
		{"", "Test - Debug - "},
		{"Test", "Test - Debug - Test"},
		{"Test.Test", "Test - Debug - Test.Test"},
		{"Test.Test.Test", "Test - Debug - Test.Test.Test"},
	}

	r := list.GetLogger("Test")
	r.Format = "{{.Logger}} - {{.Priority}} - {{.Message}}"
	r.NoColor = true

	for _, d := range m {
		l.Info("Checking: ", d)

		k := d[0]
		v := d[1]

		var b bytes.Buffer
		r.Output = &b

		printMessage(r, Debug, k)
		o := b.String()

		l.Debug("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
		if o != v {
			l.Critical("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
	}
}

func TestPrintColors(t *testing.T) {
	l := New("logger.Test.PrintColors")
	SetLevel("logger.Test.PrintColors", Disable)

	//TODO: Compare strings instead of printing.

	l.Debug("Debug")
	l.Info("Info")
	l.Notice("Notice")
	l.Warning("Warning")
	l.Error("Error")
	l.Critical("Critical")
	l.Alert("Alert")
	l.Emergency("Emergency")

	SetNoColor("logger.Test.PrintColors", true)
	l.Debug("NoColorDebug")
	l.Info("NoColorInfo")
	l.Notice("NoColorNotice")
	l.Warning("NoColorWarning")
	l.Error("NoColorError")
	l.Critical("NoColorCritical")
	l.Alert("NoColorAlert")
	l.Emergency("NoColorEmergency")
}

func TestCheckPriorityOK(t *testing.T) {
	l := New(namet + ".CheckPriority.OK")

	for k := range priorities {
		l.Info("Checking: ", k)

		e := checkPriority(k)
		l.Debug("Return of ", k, ": ", e)
		if e != nil {
			l.Critical(e)
			t.Fail()
		}
	}
}

func TestCheckPriorityFail(t *testing.T) {
	l := New(namet + ".CheckPriority.FAIL")

	k := Disable + 1

	l.Info("Checking: ", k)

	e := checkPriority(k)
	l.Debug("Return of ", k, ": ", e)
	if e == nil {
		l.Critical("Should not have succeeded")
		t.Fail()
		return
	}
}

func TestCheckPriorityFailDoesNotExist(t *testing.T) {
	l := New(namet + ".CheckPriority.FAIL.DoesNotExist")

	k := Disable + 1
	x := "priority does not exist"

	l.Info("Checking: ", k)

	e := checkPriority(k)
	l.Debug("Return of ", k, ": ", e)
	if e != nil {

		if e.Error() != x {
			l.Critical("Wrong error, EXPECTED: ", x, ", GOT: ", e.Error())
			t.Fail()
		}
	}
}

func TestGetPriorityFormat(t *testing.T) {
	l := New(namet + ".GetPriorityFormat")

	m := [][]int{
		{int(Debug), colornone, textnormal},
		{int(Notice), colorgreen, textnormal},
		{int(Info), colorblue, textnormal},
		{int(Warning), coloryellow, textnormal},
		{int(Error), coloryellow, textbold},
		{int(Critical), colorred, textnormal},
		{int(Alert), colorred, textbold},
		{int(Emergency), colorred, textblink},
	}

	for _, d := range m {
		p := Priority(d[0])
		n, e := NamePriority(p)
		if e != nil {
			l.Alert("Can not name priority: ", e)
			t.Fail()
		}

		c := d[1]
		f := d[2]

		a, b := getPriorityFormat(p)

		if c != a {
			l.Critical("Wrong color for ", n, ", EXPECTED: ", c, ", GOT: ", a)
			t.Fail()
		}

		if f != b {
			l.Critical("Wrong format for ", n, ", EXPECTED: ", c, ", GOT: ", b)
			t.Fail()
		}
	}
}

func BenchmarkLogRootEmergency(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage(".", Emergency, "Test")
	}
}

func BenchmarkLogRootEmergencyNoColor(b *testing.B) {
	SetNoColor(".", true)

	for i := 0; i < b.N; i++ {
		logMessage(".", Emergency, "Test")
	}
}

func BenchmarkLogRoot(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage(".", Debug, "Test")
	}
}

func BenchmarkLogChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChild", Debug, "Test")
	}
}

func BenchmarkLogChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChild.Test", Debug, "Test")
	}
}

func BenchmarkLogChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChildChild.Test.Test", Debug, "Test")
	}
}

func BenchmarkLogChildAllocated(b *testing.B) {
	SetLevel("BenchLogChildAllocated", Emergency)
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildAllocated", Debug, "Test")
	}
}

func BenchmarkLogChildChildAllocated(b *testing.B) {
	SetLevel("BenchLogChildChildAllocated.Test", Emergency)
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChildAllocated.Test", Debug, "Test")
	}
}

func BenchmarkGetParentRoot(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent(".")
	}
}

func BenchmarkGetParentChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChild")
	}
}

func BenchmarkGetParentChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChild.Test")
	}
}

func BenchmarkGetParentChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChild.Test.Test")
	}
}

func BenchmarkGetParentChildChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChildChild.Test.Test")
	}
}

func BenchmarkGetParentChildChildChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChildChildChild.Test.Test.Test")
	}
}

func BenchmarkPrintMessage(b *testing.B) {
	var a bytes.Buffer
	l := list.GetLogger("BenchprintMessage")
	l.Output = &a

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		printMessage(l, Debug, "Message")
	}
}

func BenchmarkFormatMessage(b *testing.B) {
	l := list.GetLogger("BenchformatMessage")

	m := new(message)
	m.Time = "Mo 30 Sep 2013 20:29:19 CEST"
	m.Logger = l.Logger
	m.Priority = "Debug"
	m.Message = "Test"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		formatMessage(m, l.Format)
	}
}
<commit_msg>Added TestSetLevelFail which tests that the function SetLevel fails correctly.<commit_after>package logger

import (
	"bytes"
	"testing"
)

const (
	namet = name + ".Test"
)

func TestGetLevel(t *testing.T) {
	n := New("logger.Test.GetLevel")

	n.Info(n, "Starting")
	m := make(map[Logger]Priority)
	m[""] = DefaultPriority
	m["."] = DefaultPriority
	m["Test"] = DefaultPriority
	m[".Test"] = DefaultPriority

	SetLevel("Test2", Emergency)
	m["Test2"] = Emergency
	m["Test2.Test"] = Emergency
	m["Test2.Test.Test"] = Emergency
	m["Test2.Test.Test.Test"] = Emergency
	m["Test2.Test.Test.Test.Test"] = Emergency
	m["Test2.Test.Test.Test.Test.Test"] = Emergency

	for k, v := range m {
		o := GetLevel(k)
		if o != v {
			n.Error(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
		n.Debug(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
	}
	n.Info(n, "Finished")
}

func TestSetLevelFail(t *testing.T) {
	l := New(namet + ".SetLevel.Fail")

	m := Disable + 1
	v := "priority does not exist"

	n := New("Test")
	o := n.SetLevel(m)

	if v != o.Error() {
		l.Critical("GOT: '", o, "', EXPECED: '", v, "'")
		t.Fail()
	}
}

func TestGetParent(t *testing.T) {
	n := New("logger.Test.getParent")

	n.Info(n, "Starting")
	m := [][]Logger{
		{"", "."},
		{".Test", "."},
		{".", "."},
		{"Test", "."},
		{"Test.Test", "Test"},
		{"Test.Test.Test", "Test.Test"},
		{"Test.Test.Test.Test", "Test.Test.Test"},
	}

	for i := range m {
		a := m[i]

		k := a[0]
		v := a[1]

		o := getParent(k)
		if o != v {
			n.Error(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
		n.Debug(n, "GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
	}
	n.Info(n, "Finished")
}

func TestGetParentOutputSame(t *testing.T) {
	l := New(namet + ".GetParent.Output.Same")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	p.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "Test Parent,Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestGetParentOutputDifferent(t *testing.T) {
	l := New(namet + ".GetParent.Output.Different")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	c.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestGetParentOutputInheritance(t *testing.T) {
	l := New(namet + ".GetParent.Output.Inheritance")

	p := Logger("Test")
	p.SetFormat("{{.Message}}")

	c := Logger("Test.Test")
	c.SetLevel(Debug)
	l.Info("Parent: '", getParent(c), "'")

	var b bytes.Buffer
	p.SetOutput(&b)

	p.Notice("Test Parent,")
	c.Notice("Test Child")

	o := b.String()
	v := "TestParent,Test Child"

	l.Debug("GOT: ", o, ", EXPECTED: ", v)
	if o != v {
		l.Critical("GOT: ", o, ", EXPECTED: ", v)
		t.Fail()
	}
}

func TestPrintMessage(t *testing.T) {
	l := New(namet + ".PrintMessage")

	p := "\033[0m"
	b := "Test - " + p + p + "Debug" + p + " - "

	m := [][]string{
		{"", b},
		{"Test", b + "Test"},
		{"Test.Test", b + "Test.Test"},
		{"Test.Test.Test", b + "Test.Test.Test"},
	}

	r := list.GetLogger("Test")
	r.Format = "{{.Logger}} - {{.Priority}} - {{.Message}}"

	for _, d := range m {
		l.Info("Checking: ", d)

		k := d[0]
		v := d[1]

		var b bytes.Buffer
		r.Output = &b

		printMessage(r, Debug, k)
		o := b.String()

		l.Debug("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
		if o != v {
			l.Critical("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
	}
}

func TestPrintMessageNoColor(t *testing.T) {
	l := New(namet + ".PrintMessage")

	m := [][]string{
		{"", "Test - Debug - "},
		{"Test", "Test - Debug - Test"},
		{"Test.Test", "Test - Debug - Test.Test"},
		{"Test.Test.Test", "Test - Debug - Test.Test.Test"},
	}

	r := list.GetLogger("Test")
	r.Format = "{{.Logger}} - {{.Priority}} - {{.Message}}"
	r.NoColor = true

	for _, d := range m {
		l.Info("Checking: ", d)

		k := d[0]
		v := d[1]

		var b bytes.Buffer
		r.Output = &b

		printMessage(r, Debug, k)
		o := b.String()

		l.Debug("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
		if o != v {
			l.Critical("GOT: '", o, "', EXPECED: '", v, "'", ", KEY: '", k, "'")
			t.Fail()
		}
	}
}

func TestPrintColors(t *testing.T) {
	l := New("logger.Test.PrintColors")
	SetLevel("logger.Test.PrintColors", Disable)

	//TODO: Compare strings instead of printing.

	l.Debug("Debug")
	l.Info("Info")
	l.Notice("Notice")
	l.Warning("Warning")
	l.Error("Error")
	l.Critical("Critical")
	l.Alert("Alert")
	l.Emergency("Emergency")

	SetNoColor("logger.Test.PrintColors", true)
	l.Debug("NoColorDebug")
	l.Info("NoColorInfo")
	l.Notice("NoColorNotice")
	l.Warning("NoColorWarning")
	l.Error("NoColorError")
	l.Critical("NoColorCritical")
	l.Alert("NoColorAlert")
	l.Emergency("NoColorEmergency")
}

func TestCheckPriorityOK(t *testing.T) {
	l := New(namet + ".CheckPriority.OK")

	for k := range priorities {
		l.Info("Checking: ", k)

		e := checkPriority(k)
		l.Debug("Return of ", k, ": ", e)
		if e != nil {
			l.Critical(e)
			t.Fail()
		}
	}
}

func TestCheckPriorityFail(t *testing.T) {
	l := New(namet + ".CheckPriority.FAIL")

	k := Disable + 1

	l.Info("Checking: ", k)

	e := checkPriority(k)
	l.Debug("Return of ", k, ": ", e)
	if e == nil {
		l.Critical("Should not have succeeded")
		t.Fail()
		return
	}
}

func TestCheckPriorityFailDoesNotExist(t *testing.T) {
	l := New(namet + ".CheckPriority.FAIL.DoesNotExist")

	k := Disable + 1
	x := "priority does not exist"

	l.Info("Checking: ", k)

	e := checkPriority(k)
	l.Debug("Return of ", k, ": ", e)
	if e != nil {

		if e.Error() != x {
			l.Critical("Wrong error, EXPECTED: ", x, ", GOT: ", e.Error())
			t.Fail()
		}
	}
}

func TestGetPriorityFormat(t *testing.T) {
	l := New(namet + ".GetPriorityFormat")

	m := [][]int{
		{int(Debug), colornone, textnormal},
		{int(Notice), colorgreen, textnormal},
		{int(Info), colorblue, textnormal},
		{int(Warning), coloryellow, textnormal},
		{int(Error), coloryellow, textbold},
		{int(Critical), colorred, textnormal},
		{int(Alert), colorred, textbold},
		{int(Emergency), colorred, textblink},
	}

	for _, d := range m {
		p := Priority(d[0])
		n, e := NamePriority(p)
		if e != nil {
			l.Alert("Can not name priority: ", e)
			t.Fail()
		}

		c := d[1]
		f := d[2]

		a, b := getPriorityFormat(p)

		if c != a {
			l.Critical("Wrong color for ", n, ", EXPECTED: ", c, ", GOT: ", a)
			t.Fail()
		}

		if f != b {
			l.Critical("Wrong format for ", n, ", EXPECTED: ", c, ", GOT: ", b)
			t.Fail()
		}
	}
}

func BenchmarkLogRootEmergency(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage(".", Emergency, "Test")
	}
}

func BenchmarkLogRootEmergencyNoColor(b *testing.B) {
	SetNoColor(".", true)

	for i := 0; i < b.N; i++ {
		logMessage(".", Emergency, "Test")
	}
}

func BenchmarkLogRoot(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage(".", Debug, "Test")
	}
}

func BenchmarkLogChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChild", Debug, "Test")
	}
}

func BenchmarkLogChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChild.Test", Debug, "Test")
	}
}

func BenchmarkLogChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChildChild.Test.Test", Debug, "Test")
	}
}

func BenchmarkLogChildAllocated(b *testing.B) {
	SetLevel("BenchLogChildAllocated", Emergency)
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildAllocated", Debug, "Test")
	}
}

func BenchmarkLogChildChildAllocated(b *testing.B) {
	SetLevel("BenchLogChildChildAllocated.Test", Emergency)
	for i := 0; i < b.N; i++ {
		logMessage("BenchLogChildChildAllocated.Test", Debug, "Test")
	}
}

func BenchmarkGetParentRoot(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent(".")
	}
}

func BenchmarkGetParentChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChild")
	}
}

func BenchmarkGetParentChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChild.Test")
	}
}

func BenchmarkGetParentChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChild.Test.Test")
	}
}

func BenchmarkGetParentChildChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChildChild.Test.Test")
	}
}

func BenchmarkGetParentChildChildChildChildChild(b *testing.B) {
	for i := 0; i < b.N; i++ {
		getParent("BenchgetParentChildChildChildChild.Test.Test.Test")
	}
}

func BenchmarkPrintMessage(b *testing.B) {
	var a bytes.Buffer
	l := list.GetLogger("BenchprintMessage")
	l.Output = &a

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		printMessage(l, Debug, "Message")
	}
}

func BenchmarkFormatMessage(b *testing.B) {
	l := list.GetLogger("BenchformatMessage")

	m := new(message)
	m.Time = "Mo 30 Sep 2013 20:29:19 CEST"
	m.Logger = l.Logger
	m.Priority = "Debug"
	m.Message = "Test"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		formatMessage(m, l.Format)
	}
}
<|endoftext|>


<commit_before>package logical

// Auth is the resulting authentication information that is part of
// Response for credential backends.
type Auth struct {
	// ClientToken is the token that is generated for the authentication.
	// This will be filled in by Vault core when an auth structure is
	// returned. Setting this manually will have no effect.
	ClientToken string

	// Policies is the list of policies that the authenticated user
	// is associated with.
	Policies []string

	// Metadata is used to attach arbitrary string-type metadata to
	// an authenticated user. This metadata will be outputted into the
	// audit log.
	Metadata map[string]string
}
<commit_msg>logical: GoStringer for Auth<commit_after>package logical

import (
	"fmt"
)

// Auth is the resulting authentication information that is part of
// Response for credential backends.
type Auth struct {
	// ClientToken is the token that is generated for the authentication.
	// This will be filled in by Vault core when an auth structure is
	// returned. Setting this manually will have no effect.
	ClientToken string

	// Policies is the list of policies that the authenticated user
	// is associated with.
	Policies []string

	// Metadata is used to attach arbitrary string-type metadata to
	// an authenticated user. This metadata will be outputted into the
	// audit log.
	Metadata map[string]string
}

func (a *Auth) GoString() string {
	return fmt.Sprintf("*%#v", *a)
}
<|endoftext|>


<commit_before>package login

import (
	"fmt"

	"github.com/ghodss/yaml"
)

// Config is a login configuration. It is similar to standard oidc.Config, except bind field that sets the address
// callback server can listen on.
type Config struct {
	// The base path of local server for OIDC callback. This is the base for redirectURL that all clients MUST register
	// first on the OIDC server. It can point to localhost. E.g http://127.0.0.1 -> redirectURL: http://127.0.0.1/callback
	// This is also and address that HTTP's callback server will listen on. Bind Address must include port. You can specify 0 if your
	// OIDC provider support wildcard on port (almost all server does NOT).
	BindAddress string `yaml:"address"`

	// Canonical URL for Provider that will be the target issuer that this server authenticate End Uusers against.
	Provider string `yaml:"provider"`

	ClientID     string `yaml:"client_id"`
	ClientSecret string `yaml:"secret"`

	Scopes     []string `yaml:"scopes"`
	NonceCheck bool     `yaml:"include_nonce"`
}

// ConfigFromYaml parses config from yaml file.
func ConfigFromYaml(yamlContent []byte) (Config, error) {
	var c Config
	if err := yaml.Unmarshal(yamlContent, &c); err != nil {
		return Config{}, fmt.Errorf("Config: Failed to parse config file: %v", err)
	}

	// TODO(bplotka) validate cfg.
	return c, nil
}
<commit_msg>Fixed parsing config.<commit_after>package login

import (
	"fmt"

	"github.com/ghodss/yaml"
)

// Config is a login configuration. It is similar to standard oidc.Config, except bind field that sets the address
// callback server can listen on.
type Config struct {
	// The base path of local server for OIDC callback. This is the base for redirectURL that all clients MUST register
	// first on the OIDC server. It can point to localhost. E.g http://127.0.0.1 -> redirectURL: http://127.0.0.1/callback
	// This is also and address that HTTP's callback server will listen on. Bind Address must include port. You can specify 0 if your
	// OIDC provider support wildcard on port (almost all server does NOT).
	BindAddress string `json:"address"`

	// Canonical URL for Provider that will be the target issuer that this server authenticate End Uusers against.
	Provider string `json:"provider"`

	ClientID     string `json:"client_id"`
	ClientSecret string `json:"secret"`

	Scopes     []string `json:"scopes"`
	NonceCheck bool     `json:"include_nonce"`
}

// ConfigFromYaml parses config from yaml file.
func ConfigFromYaml(yamlContent []byte) (Config, error) {
	var c Config
	if err := yaml.Unmarshal(yamlContent, &c); err != nil {
		return Config{}, fmt.Errorf("Config: Failed to parse config file: %v", err)
	}

	// TODO(bplotka) validate cfg.
	return c, nil
}
<|endoftext|>