<commit_before>// Copyright 2016-2019 terraform-provider-sakuracloud 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 sakuracloud

import (
	"context"
	"fmt"
	"time"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
)

func resourceSakuraCloudProxyLB() *schema.Resource {
	return &schema.Resource{
		Create: resourceSakuraCloudProxyLBCreate,
		Read:   resourceSakuraCloudProxyLBRead,
		Update: resourceSakuraCloudProxyLBUpdate,
		Delete: resourceSakuraCloudProxyLBDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(5 * time.Minute),
			Read:   schema.DefaultTimeout(5 * time.Minute),
			Update: schema.DefaultTimeout(5 * time.Minute),
			Delete: schema.DefaultTimeout(5 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},
			"plan": {
				Type:         schema.TypeInt,
				Optional:     true,
				Default:      types.ProxyLBPlans.CPS100.Int(),
				ValidateFunc: validation.IntInSlice(types.ProxyLBPlanValues),
			},
			"vip_failover": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
			},
			"sticky_session": {
				Type:     schema.TypeBool,
				Optional: true,
			},
			"timeout": {
				Type:     schema.TypeInt,
				Default:  10,
				Optional: true,
			},
			"region": {
				Type:         schema.TypeString,
				Optional:     true,
				Default:      types.ProxyLBRegions.IS1.String(),
				ValidateFunc: validation.StringInSlice(types.ProxyLBRegionStrings, false),
				ForceNew:     true,
			},
			"bind_port": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 2,
				MinItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"proxy_mode": {
							Type:         schema.TypeString,
							Required:     true,
							ValidateFunc: validation.StringInSlice(types.ProxyLBProxyModeStrings, false),
						},
						"port": {
							Type:     schema.TypeInt,
							Optional: true,
						},
						"redirect_to_https": {
							Type:     schema.TypeBool,
							Optional: true,
						},
						"support_http2": {
							Type:     schema.TypeBool,
							Optional: true,
						},
						"response_header": {
							Type:     schema.TypeList,
							Optional: true,
							MaxItems: 10,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"header": {
										Type:     schema.TypeString,
										Required: true,
									},
									"value": {
										Type:     schema.TypeString,
										Required: true,
									},
								},
							},
						},
					},
				},
			},
			"health_check": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 1,
				MinItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"protocol": {
							Type:         schema.TypeString,
							Required:     true,
							ValidateFunc: validation.StringInSlice(types.ProxyLBProtocolStrings, false),
						},
						"delay_loop": {
							Type:         schema.TypeInt,
							Optional:     true,
							ValidateFunc: validation.IntBetween(10, 60),
							Default:      10,
						},
						"host_header": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"path": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"sorry_server": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip_address": {
							Type:     schema.TypeString,
							Required: true,
						},
						"port": {
							Type:     schema.TypeInt,
							Optional: true,
						},
					},
				},
			},
			"certificate": {
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"server_cert": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"intermediate_cert": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"private_key": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"additional_certificate": {
							Type:     schema.TypeList,
							Optional: true,
							Computed: true,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"server_cert": {
										Type:     schema.TypeString,
										Required: true,
									},
									"intermediate_cert": {
										Type:     schema.TypeString,
										Optional: true,
									},
									"private_key": {
										Type:     schema.TypeString,
										Required: true,
									},
								},
							},
						},
					},
				},
			},
			"server": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 40,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip_address": {
							Type:     schema.TypeString,
							Required: true,
						},
						"port": {
							Type:         schema.TypeInt,
							Required:     true,
							ValidateFunc: validation.IntBetween(1, 65535),
						},
						"enabled": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
					},
				},
			},
			"icon_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateSakuracloudIDType,
			},
			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": {
				Type:     schema.TypeSet,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},
			"fqdn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"vip": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"proxy_networks": {
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Computed: true,
			},
		},
	}
}

func resourceSakuraCloudProxyLBCreate(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutCreate)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	proxyLB, err := proxyLBOp.Create(ctx, expandProxyLBCreateRequest(d))
	if err != nil {
		return fmt.Errorf("creating SakuraCloud ProxyLB is failed: %s", err)
	}

	certs := expandProxyLBCerts(d)
	if certs != nil {
		_, err := proxyLBOp.SetCertificates(ctx, proxyLB.ID, &sacloud.ProxyLBSetCertificatesRequest{
			PrimaryCerts:    certs.PrimaryCert,
			AdditionalCerts: certs.AdditionalCerts,
		})
		if err != nil {
			return fmt.Errorf("setting Certificates to ProxyLB[%s] is failed: %s", proxyLB.ID, err)
		}
	}

	d.SetId(proxyLB.ID.String())
	return resourceSakuraCloudProxyLBRead(d, meta)
}

func resourceSakuraCloudProxyLBRead(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutRead)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		if sacloud.IsNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	return setProxyLBResourceData(ctx, d, client, proxyLB)
}

func resourceSakuraCloudProxyLBUpdate(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutUpdate)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	sakuraMutexKV.Lock(d.Id())
	defer sakuraMutexKV.Unlock(d.Id())

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	proxyLB, err = proxyLBOp.Update(ctx, proxyLB.ID, expandProxyLBUpdateRequest(d))
	if err != nil {
		return fmt.Errorf("updating SakuraCloud ProxyLB[%s] is failed: %s", d.Id(), err)
	}

	if d.HasChange("plan") {
		newPlan := types.EProxyLBPlan(d.Get("plan").(int))
		upd, err := proxyLBOp.ChangePlan(ctx, proxyLB.ID, &sacloud.ProxyLBChangePlanRequest{Plan: newPlan})
		if err != nil {
			return fmt.Errorf("changing ProxyLB[%s] plan is failed: %s", d.Id(), err)
		}

		// update ID
		proxyLB = upd
		d.SetId(proxyLB.ID.String())
	}

	if proxyLB.LetsEncrypt == nil && d.HasChange("certificate") {
		certs := expandProxyLBCerts(d)
		if certs == nil {
			if err := proxyLBOp.DeleteCertificates(ctx, proxyLB.ID); err != nil {
				return fmt.Errorf("deleting Certificates of ProxyLB[%s] is failed: %s", d.Id(), err)
			}
		} else {
			if _, err := proxyLBOp.SetCertificates(ctx, proxyLB.ID, &sacloud.ProxyLBSetCertificatesRequest{
				PrimaryCerts:    certs.PrimaryCert,
				AdditionalCerts: certs.AdditionalCerts,
			}); err != nil {
				return fmt.Errorf("setting Certificates to ProxyLB[%s] is failed: %s", d.Id(), err)
			}
		}
	}
	return resourceSakuraCloudProxyLBRead(d, meta)
}

func resourceSakuraCloudProxyLBDelete(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutDelete)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	sakuraMutexKV.Lock(d.Id())
	defer sakuraMutexKV.Unlock(d.Id())

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		if sacloud.IsNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	if err := proxyLBOp.Delete(ctx, proxyLB.ID); err != nil {
		return fmt.Errorf("deleting ProxyLB[%s] is failed: %s", d.Id(), err)
	}
	return nil
}

func setProxyLBResourceData(ctx context.Context, d *schema.ResourceData, client *APIClient, data *sacloud.ProxyLB) error {
	// certificates
	proxyLBOp := sacloud.NewProxyLBOp(client)

	certs, err := proxyLBOp.GetCertificates(ctx, data.ID)
	if err != nil {
		// even if certificate is deleted, it will not result in an error
		return err
	}

	d.Set("name", data.Name)                                   // nolint
	d.Set("plan", data.Plan.Int())                             // nolint
	d.Set("vip_failover", data.UseVIPFailover)                 // nolint
	d.Set("sticky_session", flattenProxyLBStickySession(data)) // nolint
	d.Set("timeout", flattenProxyLBTimeout(data))              // nolint
	d.Set("region", data.Region.String())                      // nolint
	d.Set("fqdn", data.FQDN)                                   // nolint
	d.Set("vip", data.VirtualIPAddress)                        // nolint
	d.Set("proxy_networks", data.ProxyNetworks)                // nolint
	d.Set("icon_id", data.IconID.String())                     // nolint
	d.Set("description", data.Description)                     // nolint
	if err := d.Set("bind_port", flattenProxyLBBindPorts(data)); err != nil {
		return err
	}
	if err := d.Set("health_check", flattenProxyLBHealthCheck(data)); err != nil {
		return err
	}
	if err := d.Set("sorry_server", flattenProxyLBSorryServer(data)); err != nil {
		return err
	}
	if err := d.Set("server", flattenProxyLBServers(data)); err != nil {
		return err
	}
	if err := d.Set("certificate", flattenProxyLBCerts(certs)); err != nil {
		return err
	}
	return d.Set("tags", flattenTags(data.Tags))
}
<commit_msg>Set MaxItems:19 to proxylb<commit_after>// Copyright 2016-2019 terraform-provider-sakuracloud 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 sakuracloud

import (
	"context"
	"fmt"
	"time"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
)

func resourceSakuraCloudProxyLB() *schema.Resource {
	return &schema.Resource{
		Create: resourceSakuraCloudProxyLBCreate,
		Read:   resourceSakuraCloudProxyLBRead,
		Update: resourceSakuraCloudProxyLBUpdate,
		Delete: resourceSakuraCloudProxyLBDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Timeouts: &schema.ResourceTimeout{
			Create: schema.DefaultTimeout(5 * time.Minute),
			Read:   schema.DefaultTimeout(5 * time.Minute),
			Update: schema.DefaultTimeout(5 * time.Minute),
			Delete: schema.DefaultTimeout(5 * time.Minute),
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},
			"plan": {
				Type:         schema.TypeInt,
				Optional:     true,
				Default:      types.ProxyLBPlans.CPS100.Int(),
				ValidateFunc: validation.IntInSlice(types.ProxyLBPlanValues),
			},
			"vip_failover": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
			},
			"sticky_session": {
				Type:     schema.TypeBool,
				Optional: true,
			},
			"timeout": {
				Type:     schema.TypeInt,
				Default:  10,
				Optional: true,
			},
			"region": {
				Type:         schema.TypeString,
				Optional:     true,
				Default:      types.ProxyLBRegions.IS1.String(),
				ValidateFunc: validation.StringInSlice(types.ProxyLBRegionStrings, false),
				ForceNew:     true,
			},
			"bind_port": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 2,
				MinItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"proxy_mode": {
							Type:         schema.TypeString,
							Required:     true,
							ValidateFunc: validation.StringInSlice(types.ProxyLBProxyModeStrings, false),
						},
						"port": {
							Type:     schema.TypeInt,
							Optional: true,
						},
						"redirect_to_https": {
							Type:     schema.TypeBool,
							Optional: true,
						},
						"support_http2": {
							Type:     schema.TypeBool,
							Optional: true,
						},
						"response_header": {
							Type:     schema.TypeList,
							Optional: true,
							MaxItems: 10,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"header": {
										Type:     schema.TypeString,
										Required: true,
									},
									"value": {
										Type:     schema.TypeString,
										Required: true,
									},
								},
							},
						},
					},
				},
			},
			"health_check": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 1,
				MinItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"protocol": {
							Type:         schema.TypeString,
							Required:     true,
							ValidateFunc: validation.StringInSlice(types.ProxyLBProtocolStrings, false),
						},
						"delay_loop": {
							Type:         schema.TypeInt,
							Optional:     true,
							ValidateFunc: validation.IntBetween(10, 60),
							Default:      10,
						},
						"host_header": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"path": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"sorry_server": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip_address": {
							Type:     schema.TypeString,
							Required: true,
						},
						"port": {
							Type:     schema.TypeInt,
							Optional: true,
						},
					},
				},
			},
			"certificate": {
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"server_cert": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"intermediate_cert": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"private_key": {
							Type:     schema.TypeString,
							Optional: true,
							Computed: true,
						},
						"additional_certificate": {
							Type:     schema.TypeList,
							Optional: true,
							MaxItems: 19,
							Elem: &schema.Resource{
								Schema: map[string]*schema.Schema{
									"server_cert": {
										Type:     schema.TypeString,
										Required: true,
									},
									"intermediate_cert": {
										Type:     schema.TypeString,
										Optional: true,
									},
									"private_key": {
										Type:     schema.TypeString,
										Required: true,
									},
								},
							},
						},
					},
				},
			},
			"server": {
				Type:     schema.TypeList,
				Optional: true,
				MaxItems: 40,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ip_address": {
							Type:     schema.TypeString,
							Required: true,
						},
						"port": {
							Type:         schema.TypeInt,
							Required:     true,
							ValidateFunc: validation.IntBetween(1, 65535),
						},
						"enabled": {
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
					},
				},
			},
			"icon_id": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateSakuracloudIDType,
			},
			"description": {
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": {
				Type:     schema.TypeSet,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},
			"fqdn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"vip": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"proxy_networks": {
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Computed: true,
			},
		},
	}
}

func resourceSakuraCloudProxyLBCreate(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutCreate)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	proxyLB, err := proxyLBOp.Create(ctx, expandProxyLBCreateRequest(d))
	if err != nil {
		return fmt.Errorf("creating SakuraCloud ProxyLB is failed: %s", err)
	}

	certs := expandProxyLBCerts(d)
	if certs != nil {
		_, err := proxyLBOp.SetCertificates(ctx, proxyLB.ID, &sacloud.ProxyLBSetCertificatesRequest{
			PrimaryCerts:    certs.PrimaryCert,
			AdditionalCerts: certs.AdditionalCerts,
		})
		if err != nil {
			return fmt.Errorf("setting Certificates to ProxyLB[%s] is failed: %s", proxyLB.ID, err)
		}
	}

	d.SetId(proxyLB.ID.String())
	return resourceSakuraCloudProxyLBRead(d, meta)
}

func resourceSakuraCloudProxyLBRead(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutRead)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		if sacloud.IsNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	return setProxyLBResourceData(ctx, d, client, proxyLB)
}

func resourceSakuraCloudProxyLBUpdate(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutUpdate)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	sakuraMutexKV.Lock(d.Id())
	defer sakuraMutexKV.Unlock(d.Id())

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	proxyLB, err = proxyLBOp.Update(ctx, proxyLB.ID, expandProxyLBUpdateRequest(d))
	if err != nil {
		return fmt.Errorf("updating SakuraCloud ProxyLB[%s] is failed: %s", d.Id(), err)
	}

	if d.HasChange("plan") {
		newPlan := types.EProxyLBPlan(d.Get("plan").(int))
		upd, err := proxyLBOp.ChangePlan(ctx, proxyLB.ID, &sacloud.ProxyLBChangePlanRequest{Plan: newPlan})
		if err != nil {
			return fmt.Errorf("changing ProxyLB[%s] plan is failed: %s", d.Id(), err)
		}

		// update ID
		proxyLB = upd
		d.SetId(proxyLB.ID.String())
	}

	if proxyLB.LetsEncrypt == nil && d.HasChange("certificate") {
		certs := expandProxyLBCerts(d)
		if certs == nil {
			if err := proxyLBOp.DeleteCertificates(ctx, proxyLB.ID); err != nil {
				return fmt.Errorf("deleting Certificates of ProxyLB[%s] is failed: %s", d.Id(), err)
			}
		} else {
			if _, err := proxyLBOp.SetCertificates(ctx, proxyLB.ID, &sacloud.ProxyLBSetCertificatesRequest{
				PrimaryCerts:    certs.PrimaryCert,
				AdditionalCerts: certs.AdditionalCerts,
			}); err != nil {
				return fmt.Errorf("setting Certificates to ProxyLB[%s] is failed: %s", d.Id(), err)
			}
		}
	}
	return resourceSakuraCloudProxyLBRead(d, meta)
}

func resourceSakuraCloudProxyLBDelete(d *schema.ResourceData, meta interface{}) error {
	client, _, err := sakuraCloudClient(d, meta)
	if err != nil {
		return err
	}
	ctx, cancel := operationContext(d, schema.TimeoutDelete)
	defer cancel()

	proxyLBOp := sacloud.NewProxyLBOp(client)

	sakuraMutexKV.Lock(d.Id())
	defer sakuraMutexKV.Unlock(d.Id())

	proxyLB, err := proxyLBOp.Read(ctx, sakuraCloudID(d.Id()))
	if err != nil {
		if sacloud.IsNotFoundError(err) {
			d.SetId("")
			return nil
		}
		return fmt.Errorf("could not read SakuraCloud ProxyLB[%s]: %s", d.Id(), err)
	}

	if err := proxyLBOp.Delete(ctx, proxyLB.ID); err != nil {
		return fmt.Errorf("deleting ProxyLB[%s] is failed: %s", d.Id(), err)
	}
	return nil
}

func setProxyLBResourceData(ctx context.Context, d *schema.ResourceData, client *APIClient, data *sacloud.ProxyLB) error {
	// certificates
	proxyLBOp := sacloud.NewProxyLBOp(client)

	certs, err := proxyLBOp.GetCertificates(ctx, data.ID)
	if err != nil {
		// even if certificate is deleted, it will not result in an error
		return err
	}

	d.Set("name", data.Name)                                   // nolint
	d.Set("plan", data.Plan.Int())                             // nolint
	d.Set("vip_failover", data.UseVIPFailover)                 // nolint
	d.Set("sticky_session", flattenProxyLBStickySession(data)) // nolint
	d.Set("timeout", flattenProxyLBTimeout(data))              // nolint
	d.Set("region", data.Region.String())                      // nolint
	d.Set("fqdn", data.FQDN)                                   // nolint
	d.Set("vip", data.VirtualIPAddress)                        // nolint
	d.Set("proxy_networks", data.ProxyNetworks)                // nolint
	d.Set("icon_id", data.IconID.String())                     // nolint
	d.Set("description", data.Description)                     // nolint
	if err := d.Set("bind_port", flattenProxyLBBindPorts(data)); err != nil {
		return err
	}
	if err := d.Set("health_check", flattenProxyLBHealthCheck(data)); err != nil {
		return err
	}
	if err := d.Set("sorry_server", flattenProxyLBSorryServer(data)); err != nil {
		return err
	}
	if err := d.Set("server", flattenProxyLBServers(data)); err != nil {
		return err
	}
	if err := d.Set("certificate", flattenProxyLBCerts(certs)); err != nil {
		return err
	}
	return d.Set("tags", flattenTags(data.Tags))
}
<|endoftext|>


<commit_before>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Name:           name,
		Tags:           t.Tags,
		Service:        Service,
	}

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Name = name

	return sendSample(span)
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err, ok := r.(error)
			if !ok {
				err = fmt.Errorf("encountered panic while sending sample %#v", err)
			}
			logrus.WithError(err).Error("Panicked while dialing UDP connection - traces will not be sent")
		}
	}()

	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)

	// at this point, the connection should already be initialized

	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<commit_msg>Fix test<commit_after>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

var udpConn *net.UDPConn

// used to initialize udpConn inside sendSample
var udpInitOnce sync.Once
var udpInitFunc = func() { initUDPConn(localVeneurAddress) }
var udpConnUninitializedErr = errors.New("UDP connection is not yet initialized")

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	if t.End.IsZero() {
		t.End = time.Now()
	}
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Name:           name,
		Tags:           t.Tags,
		Service:        Service,
	}

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Name = name

	return sendSample(span)
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	s.Name = name
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err, ok := r.(error)
			if !ok {
				err = fmt.Errorf("encountered panic while sending sample %#v", err)
			}
			logrus.WithError(err).Error("Panicked while dialing UDP connection - traces will not be sent")
		}
	}()

	if Disabled() {
		return nil
	}

	udpInitOnce.Do(udpInitFunc)

	// at this point, the connection should already be initialized

	if udpConn == nil {
		return udpConnUninitializedErr
	}

	data, err := proto.Marshal(sample)
	if err != nil {
		return
	}

	_, err = udpConn.Write(data)
	if err != nil {
		return
	}

	return nil
}

// initUDPConn will initialize the global UDP connection.
// It will panic if it encounters an error.
// It should only be called via the corresponding sync.Once
func initUDPConn(address string) {
	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		panic(err)
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		panic(err)
	}
	udpConn = conn
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<|endoftext|>


<commit_before>package tracker

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/anacrolix/missinggo"

	"github.com/anacrolix/torrent/util"
)

type Action int32

const (
	Connect Action = iota
	Announce
	Scrape
	Error

	connectRequestConnectionId = 0x41727101980

	// BEP 41
	optionTypeEndOfOptions = 0
	optionTypeNOP          = 1
	optionTypeURLData      = 2
)

type ConnectionRequest struct {
	ConnectionId int64
	Action       int32
	TransctionId int32
}

type ConnectionResponse struct {
	ConnectionId int64
}

type ResponseHeader struct {
	Action        Action
	TransactionId int32
}

type RequestHeader struct {
	ConnectionId  int64
	Action        Action
	TransactionId int32
} // 16 bytes

type AnnounceResponseHeader struct {
	Interval int32
	Leechers int32
	Seeders  int32
}

func init() {
	RegisterClientScheme("udp", newClient)
}

func newClient(url *url.URL) Client {
	return &udpClient{
		url: *url,
	}
}

func newTransactionId() int32 {
	return int32(rand.Uint32())
}

func timeout(contiguousTimeouts int) (d time.Duration) {
	if contiguousTimeouts > 8 {
		contiguousTimeouts = 8
	}
	d = 15 * time.Second
	for ; contiguousTimeouts > 0; contiguousTimeouts-- {
		d *= 2
	}
	return
}

type udpClient struct {
	contiguousTimeouts   int
	connectionIdReceived time.Time
	connectionId         int64
	socket               net.Conn
	url                  url.URL
}

func (c *udpClient) URL() string {
	return c.url.String()
}

func (c *udpClient) String() string {
	return c.URL()
}

func (c *udpClient) Announce(req *AnnounceRequest) (res AnnounceResponse, err error) {
	if !c.connected() {
		err = ErrNotConnected
		return
	}
	reqURI := c.url.RequestURI()
	// Clearly this limits the request URI to 255 bytes. BEP 41 supports
	// longer but I'm not fussed.
	options := append([]byte{optionTypeURLData, byte(len(reqURI))}, []byte(reqURI)...)
	b, err := c.request(Announce, req, options)
	if err != nil {
		return
	}
	var h AnnounceResponseHeader
	err = readBody(b, &h)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		err = fmt.Errorf("error parsing announce response: %s", err)
		return
	}
	res.Interval = h.Interval
	res.Leechers = h.Leechers
	res.Seeders = h.Seeders
	cps, err := util.UnmarshalIPv4CompactPeers(b.Bytes())
	if err != nil {
		return
	}
	for _, cp := range cps {
		res.Peers = append(res.Peers, Peer{
			IP:   cp.IP[:],
			Port: int(cp.Port),
		})
	}
	return
}

// body is the binary serializable request body. trailer is optional data
// following it, such as for BEP 41.
func (c *udpClient) write(h *RequestHeader, body interface{}, trailer []byte) (err error) {
	var buf bytes.Buffer
	err = binary.Write(&buf, binary.BigEndian, h)
	if err != nil {
		panic(err)
	}
	if body != nil {
		err = binary.Write(&buf, binary.BigEndian, body)
		if err != nil {
			panic(err)
		}
	}
	_, err = buf.Write(trailer)
	if err != nil {
		return
	}
	n, err := c.socket.Write(buf.Bytes())
	if err != nil {
		return
	}
	if n != buf.Len() {
		panic("write should send all or error")
	}
	return
}

func read(r io.Reader, data interface{}) error {
	return binary.Read(r, binary.BigEndian, data)
}

func write(w io.Writer, data interface{}) error {
	return binary.Write(w, binary.BigEndian, data)
}

// args is the binary serializable request body. trailer is optional data
// following it, such as for BEP 41.
func (c *udpClient) request(action Action, args interface{}, options []byte) (responseBody *bytes.Buffer, err error) {
	tid := newTransactionId()
	err = c.write(&RequestHeader{
		ConnectionId:  c.connectionId,
		Action:        action,
		TransactionId: tid,
	}, args, options)
	if err != nil {
		return
	}
	c.socket.SetReadDeadline(time.Now().Add(timeout(c.contiguousTimeouts)))
	b := make([]byte, 0x800) // 2KiB
	for {
		var n int
		n, err = c.socket.Read(b)
		if opE, ok := err.(*net.OpError); ok {
			if opE.Timeout() {
				c.contiguousTimeouts++
				return
			}
		}
		if err != nil {
			return
		}
		buf := bytes.NewBuffer(b[:n])
		var h ResponseHeader
		err = binary.Read(buf, binary.BigEndian, &h)
		switch err {
		case io.ErrUnexpectedEOF:
			continue
		case nil:
		default:
			return
		}
		if h.TransactionId != tid {
			continue
		}
		c.contiguousTimeouts = 0
		if h.Action == Error {
			err = errors.New(buf.String())
		}
		responseBody = buf
		return
	}
}

func readBody(r io.Reader, data ...interface{}) (err error) {
	for _, datum := range data {
		err = binary.Read(r, binary.BigEndian, datum)
		if err != nil {
			break
		}
	}
	return
}

func (c *udpClient) connected() bool {
	return !c.connectionIdReceived.IsZero() && time.Now().Before(c.connectionIdReceived.Add(time.Minute))
}

func (c *udpClient) Connect() (err error) {
	if c.connected() {
		return nil
	}
	c.connectionId = connectRequestConnectionId
	if c.socket == nil {
		hmp := missinggo.SplitHostPort(c.url.Host)
		if hmp.NoPort {
			hmp.NoPort = false
			hmp.Port = 80
		}
		c.socket, err = net.Dial("udp", hmp.String())
		if err != nil {
			return
		}
	}
	b, err := c.request(Connect, nil, nil)
	if err != nil {
		return
	}
	var res ConnectionResponse
	err = readBody(b, &res)
	if err != nil {
		return
	}
	c.connectionId = res.ConnectionId
	c.connectionIdReceived = time.Now()
	return
}
<commit_msg>Begin tracking use of file descriptors in UDP tracker<commit_after>package tracker

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/missinggo/pproffd"

	"github.com/anacrolix/torrent/util"
)

type Action int32

const (
	Connect Action = iota
	Announce
	Scrape
	Error

	connectRequestConnectionId = 0x41727101980

	// BEP 41
	optionTypeEndOfOptions = 0
	optionTypeNOP          = 1
	optionTypeURLData      = 2
)

type ConnectionRequest struct {
	ConnectionId int64
	Action       int32
	TransctionId int32
}

type ConnectionResponse struct {
	ConnectionId int64
}

type ResponseHeader struct {
	Action        Action
	TransactionId int32
}

type RequestHeader struct {
	ConnectionId  int64
	Action        Action
	TransactionId int32
} // 16 bytes

type AnnounceResponseHeader struct {
	Interval int32
	Leechers int32
	Seeders  int32
}

func init() {
	RegisterClientScheme("udp", newClient)
}

func newClient(url *url.URL) Client {
	return &udpClient{
		url: *url,
	}
}

func newTransactionId() int32 {
	return int32(rand.Uint32())
}

func timeout(contiguousTimeouts int) (d time.Duration) {
	if contiguousTimeouts > 8 {
		contiguousTimeouts = 8
	}
	d = 15 * time.Second
	for ; contiguousTimeouts > 0; contiguousTimeouts-- {
		d *= 2
	}
	return
}

type udpClient struct {
	contiguousTimeouts   int
	connectionIdReceived time.Time
	connectionId         int64
	socket               net.Conn
	url                  url.URL
}

func (c *udpClient) URL() string {
	return c.url.String()
}

func (c *udpClient) String() string {
	return c.URL()
}

func (c *udpClient) Announce(req *AnnounceRequest) (res AnnounceResponse, err error) {
	if !c.connected() {
		err = ErrNotConnected
		return
	}
	reqURI := c.url.RequestURI()
	// Clearly this limits the request URI to 255 bytes. BEP 41 supports
	// longer but I'm not fussed.
	options := append([]byte{optionTypeURLData, byte(len(reqURI))}, []byte(reqURI)...)
	b, err := c.request(Announce, req, options)
	if err != nil {
		return
	}
	var h AnnounceResponseHeader
	err = readBody(b, &h)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		err = fmt.Errorf("error parsing announce response: %s", err)
		return
	}
	res.Interval = h.Interval
	res.Leechers = h.Leechers
	res.Seeders = h.Seeders
	cps, err := util.UnmarshalIPv4CompactPeers(b.Bytes())
	if err != nil {
		return
	}
	for _, cp := range cps {
		res.Peers = append(res.Peers, Peer{
			IP:   cp.IP[:],
			Port: int(cp.Port),
		})
	}
	return
}

// body is the binary serializable request body. trailer is optional data
// following it, such as for BEP 41.
func (c *udpClient) write(h *RequestHeader, body interface{}, trailer []byte) (err error) {
	var buf bytes.Buffer
	err = binary.Write(&buf, binary.BigEndian, h)
	if err != nil {
		panic(err)
	}
	if body != nil {
		err = binary.Write(&buf, binary.BigEndian, body)
		if err != nil {
			panic(err)
		}
	}
	_, err = buf.Write(trailer)
	if err != nil {
		return
	}
	n, err := c.socket.Write(buf.Bytes())
	if err != nil {
		return
	}
	if n != buf.Len() {
		panic("write should send all or error")
	}
	return
}

func read(r io.Reader, data interface{}) error {
	return binary.Read(r, binary.BigEndian, data)
}

func write(w io.Writer, data interface{}) error {
	return binary.Write(w, binary.BigEndian, data)
}

// args is the binary serializable request body. trailer is optional data
// following it, such as for BEP 41.
func (c *udpClient) request(action Action, args interface{}, options []byte) (responseBody *bytes.Buffer, err error) {
	tid := newTransactionId()
	err = c.write(&RequestHeader{
		ConnectionId:  c.connectionId,
		Action:        action,
		TransactionId: tid,
	}, args, options)
	if err != nil {
		return
	}
	c.socket.SetReadDeadline(time.Now().Add(timeout(c.contiguousTimeouts)))
	b := make([]byte, 0x800) // 2KiB
	for {
		var n int
		n, err = c.socket.Read(b)
		if opE, ok := err.(*net.OpError); ok {
			if opE.Timeout() {
				c.contiguousTimeouts++
				return
			}
		}
		if err != nil {
			return
		}
		buf := bytes.NewBuffer(b[:n])
		var h ResponseHeader
		err = binary.Read(buf, binary.BigEndian, &h)
		switch err {
		case io.ErrUnexpectedEOF:
			continue
		case nil:
		default:
			return
		}
		if h.TransactionId != tid {
			continue
		}
		c.contiguousTimeouts = 0
		if h.Action == Error {
			err = errors.New(buf.String())
		}
		responseBody = buf
		return
	}
}

func readBody(r io.Reader, data ...interface{}) (err error) {
	for _, datum := range data {
		err = binary.Read(r, binary.BigEndian, datum)
		if err != nil {
			break
		}
	}
	return
}

func (c *udpClient) connected() bool {
	return !c.connectionIdReceived.IsZero() && time.Now().Before(c.connectionIdReceived.Add(time.Minute))
}

func (c *udpClient) Connect() (err error) {
	if c.connected() {
		return nil
	}
	c.connectionId = connectRequestConnectionId
	if c.socket == nil {
		hmp := missinggo.SplitHostPort(c.url.Host)
		if hmp.NoPort {
			hmp.NoPort = false
			hmp.Port = 80
		}
		c.socket, err = net.Dial("udp", hmp.String())
		if err != nil {
			return
		}
		c.socket = pproffd.WrapNetConn(c.socket)
	}
	b, err := c.request(Connect, nil, nil)
	if err != nil {
		return
	}
	var res ConnectionResponse
	err = readBody(b, &res)
	if err != nil {
		return
	}
	c.connectionId = res.ConnectionId
	c.connectionIdReceived = time.Now()
	return
}
<|endoftext|>


<commit_before>package poke

import (
	"fmt"
	"github.com/raphamorim/go-rainbow"
	"strconv"
)

const LONGEST_TYPE_NAME_LEN = 8
const WHITE = "#FFFFFF"
const FAV_BLUE = "#42b3f4"

func strRepeat(amount int, str string) string {
	var toReturn string
	for i := 0; i < amount; i++ {
		toReturn += str
	}
	return toReturn
}

func typeLabel(t Type) string {
	n := LONGEST_TYPE_NAME_LEN - len(t.Name)
	txt := strRepeat(n, " ") + t.Name + ":"
	return rainbow.Bold(rainbow.Hex(WHITE, txt))
}

func combiLabel(c TypeCombination) string {
	n := LONGEST_TYPE_NAME_LEN*2 + 1 - len(c.FirstSlot.Name) - len(c.SecondSlot.Name)
	var txt string
	if c.SecondSlot.Name != "" {
		txt = strRepeat(n, " ") + c.FirstSlot.Name + "-" + c.SecondSlot.Name + ":"
	} else {
		txt = strRepeat(n+1, " ") + c.FirstSlot.Name + ":"
	}
	return rainbow.Bold(rainbow.Hex(WHITE, txt))
}

func PrintHisto(histo map[Type]int, sorted [18]Type) {
	for _, t := range sorted {
		bar := strRepeat(histo[t], "#")
		fmt.Printf("%s %s (%d)\n",
			typeLabel(t),
			rainbow.Hex(t.HexColor, bar),
			histo[t])
	}
}

func PrintRatios(ratios map[Type][2]int, sorted [18]Type) {
	for _, t := range sorted {
		fmt.Printf("%s %s / %s (%f)\n",
			typeLabel(t),
			rainbow.Hex(t.HexColor, strconv.Itoa(ratios[t][0])),
			rainbow.Hex(t.HexColor, strconv.Itoa(ratios[t][1])),
			float64(ratios[t][0])/float64(ratios[t][1]))
	}
}

func PrintCombiHisto(histo map[TypeCombination]int, sorted []TypeCombination) {
	for _, c := range sorted {
		bar := strRepeat(histo[c], "#")
		fmt.Printf("%s %s (%d)\n",
			combiLabel(c),
			rainbow.Hex(FAV_BLUE, bar),
			histo[c])
	}
}

func PrintBattlePokemon(r BestMoveSetResult) {
	s := rainbow.Hex("#ffffff", r.PokemonName+": [")
	for i, m := range r.MoveSet {
		if i != 0 {
			s += ", "
		}
		s += rainbow.Hex(m.Type.HexColor, m.Name)
	}
	fmt.Println(s + rainbow.Hex("#ffffff", "]"))
}

func PrintBestPokemonResults(results []BestMoveSetResult) {
	for _, r := range results {
		PrintBattlePokemon(r)
	}
}
<commit_msg>Enhance printing of moveSet and align BestPokemon command<commit_after>package poke

import (
	"fmt"
	"github.com/raphamorim/go-rainbow"
	"strconv"
)

const LONGEST_POKEMON_NAME = 14
const LONGEST_TYPE_NAME = 8
const WHITE = "#FFFFFF"
const FAV_BLUE = "#42b3f4"

func strRepeat(amount int, str string) string {
	var toReturn string
	for i := 0; i < amount; i++ {
		toReturn += str
	}
	return toReturn
}

func typeLabel(t Type) string {
	n := LONGEST_TYPE_NAME - len(t.Name)
	txt := strRepeat(n, " ") + t.Name + ":"
	return rainbow.Bold(rainbow.Hex(WHITE, txt))
}

func combiLabel(c TypeCombination) string {
	n := LONGEST_TYPE_NAME*2 + 1 - len(c.FirstSlot.Name) - len(c.SecondSlot.Name)
	var txt string
	if c.SecondSlot.Name != "" {
		txt = strRepeat(n, " ") + c.FirstSlot.Name + "-" + c.SecondSlot.Name + ":"
	} else {
		txt = strRepeat(n+1, " ") + c.FirstSlot.Name + ":"
	}
	return rainbow.Bold(rainbow.Hex(WHITE, txt))
}

func PrintHisto(histo map[Type]int, sorted [18]Type) {
	for _, t := range sorted {
		bar := strRepeat(histo[t], "#")
		fmt.Printf("%s %s (%d)\n",
			typeLabel(t),
			rainbow.Hex(t.HexColor, bar),
			histo[t])
	}
}

func PrintRatios(ratios map[Type][2]int, sorted [18]Type) {
	for _, t := range sorted {
		fmt.Printf("%s %s / %s (%f)\n",
			typeLabel(t),
			rainbow.Hex(t.HexColor, strconv.Itoa(ratios[t][0])),
			rainbow.Hex(t.HexColor, strconv.Itoa(ratios[t][1])),
			float64(ratios[t][0])/float64(ratios[t][1]))
	}
}

func PrintCombiHisto(histo map[TypeCombination]int, sorted []TypeCombination) {
	for _, c := range sorted {
		bar := strRepeat(histo[c], "#")
		fmt.Printf("%s %s (%d)\n",
			combiLabel(c),
			rainbow.Hex(FAV_BLUE, bar),
			histo[c])
	}
}

func PrintBattlePokemon(r BestMoveSetResult) {
	paddedName := strRepeat(LONGEST_POKEMON_NAME - len(r.PokemonName), " ") + r.PokemonName
	s := rainbow.Hex("#ffffff", paddedName+": [")
	s += MoveSetToString(r.MoveSet)
	fmt.Println(s + rainbow.Hex("#ffffff", "] ") + strconv.Itoa(r.TotalKt))
}

func MoveSetToString(moveSet [4]Move) string {
	s := ""
	for i, m := range moveSet {
		if i != 0 {
			s += ", "
		}
		var sOrP string
		if m.isPhysical {
			sOrP = "P"
		} else {
			sOrP = "S"
		}
		s += rainbow.Hex(m.Type.HexColor, fmt.Sprintf("%s(%s%d|%d)", m.Name, sOrP, m.Power, m.Accuracy))
	}
	return s
}

func PrintBestPokemonResults(results []BestMoveSetResult) {
	for _, r := range results {
		PrintBattlePokemon(r)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"fmt"
	"testing"

	"github.com/hailocab/go-hostpool"
)

// Tests of the round-robin host selection policy implementation
func TestRoundRobinHostPolicy(t *testing.T) {
	policy := RoundRobinHostPolicy()

	hosts := [...]*HostInfo{
		{hostId: "0"},
		{hostId: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// interleaved iteration should always increment the host
	iterA := policy.Pick(nil)
	if actual := iterA(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	iterB := policy.Pick(nil)
	if actual := iterB(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterB(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterA(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}

	iterC := policy.Pick(nil)
	if actual := iterC(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterC(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
}

// Tests of the token-aware host selection policy implementation with a
// round-robin host selection policy fallback.
func TestTokenAwareHostPolicy(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	query := &Query{}

	iter := policy.Pick(nil)
	if iter == nil {
		t.Fatal("host iterator was nil")
	}
	actual := iter()
	if actual != nil {
		t.Fatalf("expected nil from iterator, but was %v", actual)
	}

	// set the hosts
	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the token ring is not setup without the partitioner, but the fallback
	// should work
	if actual := policy.Pick(nil)(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}

	query.RoutingKey([]byte("30"))
	if actual := policy.Pick(query)(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}

	policy.SetPartitioner("OrderedPartitioner")

	// now the token ring is configured
	query.RoutingKey([]byte("20"))
	iter = policy.Pick(query)
	if actual := iter(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}
	// rest are round robin
	if actual := iter(); actual.Info().Peer() != "2" {
		t.Errorf("Expected peer 2 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "3" {
		t.Errorf("Expected peer 3 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}
}

// Tests of the host pool host selection policy implementation
func TestHostPoolHostPolicy(t *testing.T) {
	policy := HostPoolHostPolicy(hostpool.New(nil))

	hosts := [...]*HostInfo{
		{hostId: "0", peer: "0"},
		{hostId: "1", peer: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the first host selected is actually at [1], but this is ok for RR
	// interleaved iteration should always increment the host
	iter := policy.Pick(nil)
	actualA := iter()
	if actualA.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualA.Info().HostID())
	}
	actualA.Mark(nil)

	actualB := iter()
	if actualB.Info().HostID() != "1" {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actualB.Info().HostID())
	}
	actualB.Mark(fmt.Errorf("error"))

	actualC := iter()
	if actualC.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualC.Info().HostID())
	}
	actualC.Mark(nil)

	actualD := iter()
	if actualD.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualD.Info().HostID())
	}
	actualD.Mark(nil)
}

func TestRoundRobinNilHostInfo(t *testing.T) {
	policy := RoundRobinHostPolicy()

	host := &HostInfo{hostId: "host-1"}
	policy.AddHost(host)

	iter := policy.Pick(nil)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.HostID() != host.HostID() {
		t.Fatalf("expected host %v got %v", host, *v)
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestTokenAwareNilHostInfo(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}
	policy.SetPartitioner("OrderedPartitioner")

	query := &Query{}
	query.RoutingKey([]byte("20"))

	iter := policy.Pick(query)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.Peer() != "1" {
		t.Fatalf("expected peer 1 got %v", v.Peer())
	}

	// Empty the hosts to trigger the panic when using the fallback.
	for _, host := range hosts {
		policy.RemoveHost(host.Peer())
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestCOWList_Add(t *testing.T) {
	var cow cowHostList

	toAdd := [...]string{"peer1", "peer2", "peer3"}

	for _, addr := range toAdd {
		if !cow.add(&HostInfo{peer: addr}) {
			t.Fatal("did not add peer which was not in the set")
		}
	}

	hosts := cow.get()
	if len(hosts) != len(toAdd) {
		t.Fatalf("expected to have %d hosts got %d", len(toAdd), len(hosts))
	}

	set := make(map[string]bool)
	for _, host := range hosts {
		set[host.Peer()] = true
	}

	for _, addr := range toAdd {
		if !set[addr] {
			t.Errorf("addr was not in the host list: %q", addr)
		}
	}
}

func TestSimpleRetryPolicy(t *testing.T) {
	q := &Query{}
	rt := &SimpleRetryPolicy{NumRetries: 2}
	if !rt.Attempt(q) {
		t.Fatal("should allow retry after 0 attempts")
	}
	q.attempts = 5
	if rt.Attempt(q) {
		t.Fatal("should not allow retry after passing threshold")
	}
}
<commit_msg>Fixed TestHostPoolHostPolicy that was intermitently failing<commit_after>// Copyright (c) 2015 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"fmt"
	"testing"

	"github.com/hailocab/go-hostpool"
)

// Tests of the round-robin host selection policy implementation
func TestRoundRobinHostPolicy(t *testing.T) {
	policy := RoundRobinHostPolicy()

	hosts := [...]*HostInfo{
		{hostId: "0"},
		{hostId: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// interleaved iteration should always increment the host
	iterA := policy.Pick(nil)
	if actual := iterA(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	iterB := policy.Pick(nil)
	if actual := iterB(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterB(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterA(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}

	iterC := policy.Pick(nil)
	if actual := iterC(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterC(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
}

// Tests of the token-aware host selection policy implementation with a
// round-robin host selection policy fallback.
func TestTokenAwareHostPolicy(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	query := &Query{}

	iter := policy.Pick(nil)
	if iter == nil {
		t.Fatal("host iterator was nil")
	}
	actual := iter()
	if actual != nil {
		t.Fatalf("expected nil from iterator, but was %v", actual)
	}

	// set the hosts
	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the token ring is not setup without the partitioner, but the fallback
	// should work
	if actual := policy.Pick(nil)(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}

	query.RoutingKey([]byte("30"))
	if actual := policy.Pick(query)(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}

	policy.SetPartitioner("OrderedPartitioner")

	// now the token ring is configured
	query.RoutingKey([]byte("20"))
	iter = policy.Pick(query)
	if actual := iter(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}
	// rest are round robin
	if actual := iter(); actual.Info().Peer() != "2" {
		t.Errorf("Expected peer 2 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "3" {
		t.Errorf("Expected peer 3 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}
}

// Tests of the host pool host selection policy implementation
func TestHostPoolHostPolicy(t *testing.T) {
	policy := HostPoolHostPolicy(hostpool.New(nil))

	hosts := []*HostInfo{
		{hostId: "0", peer: "0"},
		{hostId: "1", peer: "1"},
	}

	// Using set host to control the ordering of the hosts as calling "AddHost" iterates the map
	// which will result in an unpredictable ordering
	policy.(*hostPoolHostPolicy).SetHosts(hosts)

	// the first host selected is actually at [1], but this is ok for RR
	// interleaved iteration should always increment the host
	iter := policy.Pick(nil)
	actualA := iter()
	if actualA.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualA.Info().HostID())
	}
	actualA.Mark(nil)

	actualB := iter()
	if actualB.Info().HostID() != "1" {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actualB.Info().HostID())
	}
	actualB.Mark(fmt.Errorf("error"))

	actualC := iter()
	if actualC.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualC.Info().HostID())
	}
	actualC.Mark(nil)

	actualD := iter()
	if actualD.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualD.Info().HostID())
	}
	actualD.Mark(nil)
}

func TestRoundRobinNilHostInfo(t *testing.T) {
	policy := RoundRobinHostPolicy()

	host := &HostInfo{hostId: "host-1"}
	policy.AddHost(host)

	iter := policy.Pick(nil)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.HostID() != host.HostID() {
		t.Fatalf("expected host %v got %v", host, *v)
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestTokenAwareNilHostInfo(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}
	policy.SetPartitioner("OrderedPartitioner")

	query := &Query{}
	query.RoutingKey([]byte("20"))

	iter := policy.Pick(query)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.Peer() != "1" {
		t.Fatalf("expected peer 1 got %v", v.Peer())
	}

	// Empty the hosts to trigger the panic when using the fallback.
	for _, host := range hosts {
		policy.RemoveHost(host.Peer())
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestCOWList_Add(t *testing.T) {
	var cow cowHostList

	toAdd := [...]string{"peer1", "peer2", "peer3"}

	for _, addr := range toAdd {
		if !cow.add(&HostInfo{peer: addr}) {
			t.Fatal("did not add peer which was not in the set")
		}
	}

	hosts := cow.get()
	if len(hosts) != len(toAdd) {
		t.Fatalf("expected to have %d hosts got %d", len(toAdd), len(hosts))
	}

	set := make(map[string]bool)
	for _, host := range hosts {
		set[host.Peer()] = true
	}

	for _, addr := range toAdd {
		if !set[addr] {
			t.Errorf("addr was not in the host list: %q", addr)
		}
	}
}

func TestSimpleRetryPolicy(t *testing.T) {
	q := &Query{}
	rt := &SimpleRetryPolicy{NumRetries: 2}
	if !rt.Attempt(q) {
		t.Fatal("should allow retry after 0 attempts")
	}
	q.attempts = 5
	if rt.Attempt(q) {
		t.Fatal("should not allow retry after passing threshold")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/garyburd/redigo/redis"
	"log"
	"strings"
)

var (
	protected = []string{"list", "add", "api", "counter", "css", "img", "js"}
)

type Url struct {
	id     string
	Link   string
	Short  string
	Clicks int64
}

type SiteStats struct {
	Clicks       int
	Links        int
	ClicksPerUrl float64
}

func GetUrlById(id string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	cr := UrlCache.Get(id)
	if cr != nil {
		log.Print("UrlCache: Cache HIT!")
		log.Print("Updating click count in goroutine")
		go DB.Do("INCR", "url:clicks:"+id)
		return cr.(*Url), nil
	}
	log.Print("UrlCache: Cache Miss, retrieving from DB")
	id = strings.Split(id, ":")[0]
	k, err := DB.Do("GET", "url:link:"+id)
	if err != nil {
		return nil, err
	}
	switch k.(type) {
	case nil:
		return nil, nil
	default:
		c, _ := DB.Do("INCR", "url:clicks:"+id)
		resp := &Url{}
		resp.id = id
		resp.Short = config.BaseURL + id
		resp.Link, _ = redis.String(k, err)
		resp.Clicks = c.(int64)
		UrlCache.Set(id, resp)
		return resp, nil
	}
}

func GetNewUrl(link string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	i, err := GetNewCounter()
	if err != nil {
		return nil, err
	}
	for _, k := range protected {
		for b62_Encode(uint64(i)) == k {
			i, err = GetNewCounter()
			if err != nil {
				return nil, err
			}
		}
	}
	pos := b62_Encode(uint64(i))
	_, err = DB.Do("SET", "url:link:"+pos, link)
	_, err = DB.Do("SET", "url:clicks:"+pos, 0)
	if err != nil {
		return nil, err
	}
	new := &Url{}
	new.id = pos
	new.Link = link
	new.Clicks = 0
	new.Short = config.BaseURL + new.id
	UrlCache.Set(new.id, new)
	log.Printf("Shortened %s to %s", new.Link, config.BaseURL+new.id)
	return new, nil
}

func GetNewCounter() (int64, error) {
	DB := pool.Get()
	defer DB.Close()
	n, err := DB.Do("INCR", "meta:COUNTER")
	if err != nil {
		return 0, err
	}
	return n.(int64), nil
}

func GetSiteStats() SiteStats {
	cc := StatsCache.Get("Stats")
	if cc != nil {
		log.Print("Cache: Site Stats HIT")
		return cc.(SiteStats)
	} else {
		log.Print("Cache: Site Stats MISS")
	}
	k := SiteStats{}
	a, _ := GetTotalClicks()
	b, _ := GetTotalUrls()
	c, _ := GetClicksPerUrl()
	k.Clicks = a
	k.Links = b
	k.ClicksPerUrl = c
	StatsCache.Set("Stats", k)
	return k
}

func newPool() *redis.Pool {
	return redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", config.DBAddress)
		if err != nil {
			return nil, err
		}
		_, err = conn.Do("AUTH", config.DBPassword)
		if err != nil {
			log.Print(err.Error())
			return nil, err
		}
		return conn, nil

	}, 3)
}

func GetTotalUrls() (int, error) {
	db := pool.Get()
	defer db.Close()
	k, err := db.Do("GET", "meta:total:links")
	if err != nil {
		return 0, err
	}
	l, err := redis.Int(k, err)
	if err != nil {
		return 0, err
	}
	return l, err
}

func GetTotalClicks() (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("GET", "meta:total:clicks")
	if err != nil {
		return 0, nil
	}
	j, err := redis.Int(k, err)
	if err != nil {
		return 0, nil
	}
	return j, nil
}
<commit_msg>not as #webscale as I'd hope, looks like I can't do that in a goroutine, at least not without some extra work.<commit_after>package main

import (
	"github.com/garyburd/redigo/redis"
	"log"
	"strings"
)

var (
	protected = []string{"list", "add", "api", "counter", "css", "img", "js"}
)

type Url struct {
	id     string
	Link   string
	Short  string
	Clicks int64
}

type SiteStats struct {
	Clicks       int
	Links        int
	ClicksPerUrl float64
}

func GetUrlById(id string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	cr := UrlCache.Get(id)
	if cr != nil {
		log.Print("UrlCache: Cache HIT!")
		log.Print("Updating click count in goroutine")
		_, err := DB.Do("INCR", "url:clicks:"+id)
		if err != nil {
			log.Print("Error updating click count: ", err.Error())
		}
		return cr.(*Url), nil
	}
	log.Print("UrlCache: Cache Miss, retrieving from DB")
	id = strings.Split(id, ":")[0]
	k, err := DB.Do("GET", "url:link:"+id)
	if err != nil {
		return nil, err
	}
	switch k.(type) {
	case nil:
		return nil, nil
	default:
		c, _ := DB.Do("INCR", "url:clicks:"+id)
		resp := &Url{}
		resp.id = id
		resp.Short = config.BaseURL + id
		resp.Link, _ = redis.String(k, err)
		resp.Clicks = c.(int64)
		UrlCache.Set(id, resp)
		return resp, nil
	}
}

func GetNewUrl(link string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	i, err := GetNewCounter()
	if err != nil {
		return nil, err
	}
	for _, k := range protected {
		for b62_Encode(uint64(i)) == k {
			i, err = GetNewCounter()
			if err != nil {
				return nil, err
			}
		}
	}
	pos := b62_Encode(uint64(i))
	_, err = DB.Do("SET", "url:link:"+pos, link)
	_, err = DB.Do("SET", "url:clicks:"+pos, 0)
	if err != nil {
		return nil, err
	}
	new := &Url{}
	new.id = pos
	new.Link = link
	new.Clicks = 0
	new.Short = config.BaseURL + new.id
	UrlCache.Set(new.id, new)
	log.Printf("Shortened %s to %s", new.Link, config.BaseURL+new.id)
	return new, nil
}

func GetNewCounter() (int64, error) {
	DB := pool.Get()
	defer DB.Close()
	n, err := DB.Do("INCR", "meta:COUNTER")
	if err != nil {
		return 0, err
	}
	return n.(int64), nil
}

func GetSiteStats() SiteStats {
	cc := StatsCache.Get("Stats")
	if cc != nil {
		log.Print("Cache: Site Stats HIT")
		return cc.(SiteStats)
	} else {
		log.Print("Cache: Site Stats MISS")
	}
	k := SiteStats{}
	a, _ := GetTotalClicks()
	b, _ := GetTotalUrls()
	c, _ := GetClicksPerUrl()
	k.Clicks = a
	k.Links = b
	k.ClicksPerUrl = c
	StatsCache.Set("Stats", k)
	return k
}

func newPool() *redis.Pool {
	return redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", config.DBAddress)
		if err != nil {
			return nil, err
		}
		_, err = conn.Do("AUTH", config.DBPassword)
		if err != nil {
			log.Print(err.Error())
			return nil, err
		}
		return conn, nil

	}, 3)
}

func GetTotalUrls() (int, error) {
	db := pool.Get()
	defer db.Close()
	k, err := db.Do("GET", "meta:total:links")
	if err != nil {
		return 0, err
	}
	l, err := redis.Int(k, err)
	if err != nil {
		return 0, err
	}
	return l, err
}

func GetTotalClicks() (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("GET", "meta:total:clicks")
	if err != nil {
		return 0, nil
	}
	j, err := redis.Int(k, err)
	if err != nil {
		return 0, nil
	}
	return j, nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/garyburd/redigo/redis"
	"log"
	"strings"
)

var (
	protected = []string{"list", "add", "api", "counter", "css", "img", "js"}
)

type Url struct {
	id     string
	Link   string
	Short  string
	Clicks int64
}

type SiteStats struct {
	Clicks       int
	Links        int
	ClicksPerUrl float64
}

func GetUrlById(id string, host string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	cr := UrlCache.Get(id)
	if cr != nil {
		log.Print("UrlCache: Cache HIT!")
		log.Print("Updating click count in goroutine")
		go UpdateClickCount(id)
		return cr.(*Url), nil
	}
	log.Print("UrlCache: Cache Miss, retrieving from DB")
	id = strings.Split(id, ":")[0]
	k, err := DB.Do("GET", "url:link:"+id)
	if err != nil {
		return nil, err
	}
	switch k.(type) {
	case nil:
		return nil, nil
	default:
		c, _ := UpdateClickCount(id)
		resp := &Url{}
		resp.id = id
		resp.Short = config.GetBaseUrl(host) + id
		resp.Link, _ = redis.String(k, err)
		resp.Clicks = int64(c)
		UrlCache.Set(id, resp)
		return resp, nil
	}
}

func GetNewUrl(link string, host string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	i, err := GetNewCounter()
	if err != nil {
		return nil, err
	}
	for _, k := range protected {
		for b62_Encode(uint64(i)) == k {
			i, err = GetNewCounter()
			if err != nil {
				return nil, err
			}
		}
	}
	pos := b62_Encode(uint64(i))
	_, err = DB.Do("SET", "url:link:"+pos, link)
	go func(pos string) {
		d := pool.Get()
		defer d.Close()
		_, err := d.Do("SET", "url:clicks:"+pos, 0)
		if err != nil {
			log.Printf("Error settings %s clicks to 0", pos)
		} else {
			log.Printf("%s clicks set to 0", pos)
		}
	}(pos)
	if err != nil {
		return nil, err
	}
	new := &Url{}
	new.id = pos
	new.Link = link
	new.Clicks = 0
	new.Short = config.GetBaseUrl(host) + new.id
	UrlCache.Set(new.id, new)
	log.Printf("Shortened %s to %s", new.Link, config.GetBaseUrl(host)+new.id)
	return new, nil
}

func GetNewCounter() (int64, error) {
	DB := pool.Get()
	defer DB.Close()
	n, err := DB.Do("INCR", "meta:COUNTER")
	if err != nil {
		return 0, err
	}
	return n.(int64), nil
}

func GetSiteStats() SiteStats {
	cc := StatsCache.Get("Stats")
	if cc != nil {
		log.Print("Cache: Site Stats HIT")
		return cc.(SiteStats)
	} else {
		log.Print("Cache: Site Stats MISS")
	}
	k := SiteStats{}
	a, _ := GetTotalClicks()
	b, _ := GetTotalUrls()
	c, _ := GetClicksPerUrl()
	k.Clicks = a
	k.Links = b
	k.ClicksPerUrl = c
	StatsCache.Set("Stats", k)
	return k
}

func newPool() *redis.Pool {
	return redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", config.DBAddress)
		if err != nil {
			return nil, err
		}
		_, err = conn.Do("AUTH", config.DBPassword)
		if err != nil {
			log.Print(err.Error())
			return nil, err
		}
		return conn, nil

	}, 3)
}

func GetTotalUrls() (int, error) {
	db := pool.Get()
	defer db.Close()
	k, err := db.Do("GET", "meta:total:links")
	if err != nil {
		return 0, err
	}
	l, err := redis.Int(k, err)
	if err != nil {
		return 0, err
	}
	return l, err
}

func GetTotalClicks() (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("GET", "meta:total:clicks")
	if err != nil {
		return 0, nil
	}
	j, err := redis.Int(k, err)
	if err != nil {
		return 0, nil
	}
	return j, nil
}

func UpdateClickCount(id string) (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("INCR", "url:clicks:"+id)
	i, err := redis.Int(k, err)
	return i, err
}
<commit_msg>return error before setting click count and fix typo<commit_after>package main

import (
	"github.com/garyburd/redigo/redis"
	"log"
	"strings"
)

var (
	protected = []string{"list", "add", "api", "counter", "css", "img", "js"}
)

type Url struct {
	id     string
	Link   string
	Short  string
	Clicks int64
}

type SiteStats struct {
	Clicks       int
	Links        int
	ClicksPerUrl float64
}

func GetUrlById(id string, host string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	cr := UrlCache.Get(id)
	if cr != nil {
		log.Print("UrlCache: Cache HIT!")
		log.Print("Updating click count in goroutine")
		go UpdateClickCount(id)
		return cr.(*Url), nil
	}
	log.Print("UrlCache: Cache Miss, retrieving from DB")
	id = strings.Split(id, ":")[0]
	k, err := DB.Do("GET", "url:link:"+id)
	if err != nil {
		return nil, err
	}
	switch k.(type) {
	case nil:
		return nil, nil
	default:
		c, _ := UpdateClickCount(id)
		resp := &Url{}
		resp.id = id
		resp.Short = config.GetBaseUrl(host) + id
		resp.Link, _ = redis.String(k, err)
		resp.Clicks = int64(c)
		UrlCache.Set(id, resp)
		return resp, nil
	}
}

func GetNewUrl(link string, host string) (*Url, error) {
	DB := pool.Get()
	defer DB.Close()
	i, err := GetNewCounter()
	if err != nil {
		return nil, err
	}
	for _, k := range protected {
		for b62_Encode(uint64(i)) == k {
			i, err = GetNewCounter()
			if err != nil {
				return nil, err
			}
		}
	}
	pos := b62_Encode(uint64(i))
	_, err = DB.Do("SET", "url:link:"+pos, link)
	if err != nil {
		return nil, err
	}
	go func(pos string) {
		d := pool.Get()
		defer d.Close()
		_, err := d.Do("SET", "url:clicks:"+pos, 0)
		if err != nil {
			log.Printf("Error setting %s clicks to 0", pos)
		} else {
			log.Printf("%s clicks set to 0", pos)
		}
	}(pos)
	new := &Url{}
	new.id = pos
	new.Link = link
	new.Clicks = 0
	new.Short = config.GetBaseUrl(host) + new.id
	UrlCache.Set(new.id, new)
	log.Printf("Shortened %s to %s", new.Link, config.GetBaseUrl(host)+new.id)
	return new, nil
}

func GetNewCounter() (int64, error) {
	DB := pool.Get()
	defer DB.Close()
	n, err := DB.Do("INCR", "meta:COUNTER")
	if err != nil {
		return 0, err
	}
	return n.(int64), nil
}

func GetSiteStats() SiteStats {
	cc := StatsCache.Get("Stats")
	if cc != nil {
		log.Print("Cache: Site Stats HIT")
		return cc.(SiteStats)
	} else {
		log.Print("Cache: Site Stats MISS")
	}
	k := SiteStats{}
	a, _ := GetTotalClicks()
	b, _ := GetTotalUrls()
	c, _ := GetClicksPerUrl()
	k.Clicks = a
	k.Links = b
	k.ClicksPerUrl = c
	StatsCache.Set("Stats", k)
	return k
}

func newPool() *redis.Pool {
	return redis.NewPool(func() (redis.Conn, error) {
		conn, err := redis.Dial("tcp", config.DBAddress)
		if err != nil {
			return nil, err
		}
		_, err = conn.Do("AUTH", config.DBPassword)
		if err != nil {
			log.Print(err.Error())
			return nil, err
		}
		return conn, nil

	}, 3)
}

func GetTotalUrls() (int, error) {
	db := pool.Get()
	defer db.Close()
	k, err := db.Do("GET", "meta:total:links")
	if err != nil {
		return 0, err
	}
	l, err := redis.Int(k, err)
	if err != nil {
		return 0, err
	}
	return l, err
}

func GetTotalClicks() (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("GET", "meta:total:clicks")
	if err != nil {
		return 0, nil
	}
	j, err := redis.Int(k, err)
	if err != nil {
		return 0, nil
	}
	return j, nil
}

func UpdateClickCount(id string) (int, error) {
	DB := pool.Get()
	defer DB.Close()
	k, err := DB.Do("INCR", "url:clicks:"+id)
	i, err := redis.Int(k, err)
	return i, err
}
<|endoftext|>


<commit_before>package policy

import (
	"strings"

	"github.com/weaveworks/flux"
)

const (
	None      = Policy("")
	Ignore    = Policy("ignore")
	Locked    = Policy("locked")
	Automated = Policy("automated")
)

// Policy is an string, denoting the current deployment policy of a service,
// e.g. automated, or locked.
type Policy string

func Boolean(policy Policy) bool {
	for _, p := range []Policy{
		Locked,
		Automated,
		Ignore,
	} {
		if policy == p {
			return true
		}
	}
	return false
}

type Updates map[flux.ServiceID]Update

type Update struct {
	Add    Set `json:"add"`
	Remove Set `json:"remove"`
}

type Set map[Policy]string

func (s Set) String() string {
	var ps []string
	for p, v := range s {
		ps = append(ps, string(p)+":"+v)
	}
	return "{" + strings.Join(ps, ", ") + "}"
}

func (s Set) Add(ps ...Policy) Set {
	s = clone(s)
	for _, p := range ps {
		s[p] = "true"
	}
	return s
}

func (s Set) Set(p Policy, v string) Set {
	s = clone(s)
	s[p] = v
	return s
}

func clone(s Set) Set {
	newMap := Set{}
	for p, v := range s {
		newMap[p] = v
	}
	return newMap
}

func (s Set) Contains(needle Policy) bool {
	for p, _ := range s {
		if p == needle {
			return true
		}
	}
	return false
}

func (s Set) Get(p Policy) (string, bool) {
	v, ok := s[p]
	return v, ok
}

type ServiceMap map[flux.ServiceID]Set

func (s ServiceMap) ToSlice() []flux.ServiceID {
	slice := []flux.ServiceID{}
	for service, _ := range s {
		slice = append(slice, service)
	}
	return slice
}

func (s ServiceMap) Contains(id flux.ServiceID) bool {
	_, ok := s[id]
	return ok
}

func (s ServiceMap) Without(other ServiceMap) ServiceMap {
	newMap := ServiceMap{}
	for k, v := range s {
		if !other.Contains(k) {
			newMap[k] = v
		}
	}
	return newMap
}
<commit_msg>Use a switch statement in policy.Boolean<commit_after>package policy

import (
	"strings"

	"github.com/weaveworks/flux"
)

const (
	None      = Policy("")
	Ignore    = Policy("ignore")
	Locked    = Policy("locked")
	Automated = Policy("automated")
)

// Policy is an string, denoting the current deployment policy of a service,
// e.g. automated, or locked.
type Policy string

func Boolean(policy Policy) bool {
	switch policy {
	case Locked, Automated, Ignore:
		return true
	}
	return false
}

type Updates map[flux.ServiceID]Update

type Update struct {
	Add    Set `json:"add"`
	Remove Set `json:"remove"`
}

type Set map[Policy]string

func (s Set) String() string {
	var ps []string
	for p, v := range s {
		ps = append(ps, string(p)+":"+v)
	}
	return "{" + strings.Join(ps, ", ") + "}"
}

func (s Set) Add(ps ...Policy) Set {
	s = clone(s)
	for _, p := range ps {
		s[p] = "true"
	}
	return s
}

func (s Set) Set(p Policy, v string) Set {
	s = clone(s)
	s[p] = v
	return s
}

func clone(s Set) Set {
	newMap := Set{}
	for p, v := range s {
		newMap[p] = v
	}
	return newMap
}

func (s Set) Contains(needle Policy) bool {
	for p, _ := range s {
		if p == needle {
			return true
		}
	}
	return false
}

func (s Set) Get(p Policy) (string, bool) {
	v, ok := s[p]
	return v, ok
}

type ServiceMap map[flux.ServiceID]Set

func (s ServiceMap) ToSlice() []flux.ServiceID {
	slice := []flux.ServiceID{}
	for service, _ := range s {
		slice = append(slice, service)
	}
	return slice
}

func (s ServiceMap) Contains(id flux.ServiceID) bool {
	_, ok := s[id]
	return ok
}

func (s ServiceMap) Without(other ServiceMap) ServiceMap {
	newMap := ServiceMap{}
	for k, v := range s {
		if !other.Contains(k) {
			newMap[k] = v
		}
	}
	return newMap
}
<|endoftext|>


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

package pretty

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	"github.com/kylelemons/godebug/diff"
)

// A Config represents optional configuration parameters for formatting.
//
// Some options, notably ShortList, dramatically increase the overhead
// of pretty-printing a value.
type Config struct {
	// Verbosity options
	Compact  bool // One-line output. Overrides Diffable.
	Diffable bool // Adds extra newlines for more easily diffable output.

	// Field and value options
	IncludeUnexported   bool // Include unexported fields in output
	PrintStringers      bool // Call String on a fmt.Stringer
	PrintTextMarshalers bool // Call MarshalText on an encoding.TextMarshaler
	SkipZeroFields      bool // Skip struct fields that have a zero value.

	// Output transforms
	ShortList int // Maximum character length for short lists if nonzero.

	// Type-specific overrides
	//
	// Formatter maps a type to a function that will provide a one-line string
	// representation of the input value.  Conceptually:
	//   Formatter[reflect.TypeOf(v)](v) = "v as a string"
	//
	// Note that the first argument need not explicitly match the type, it must
	// merely be callable with it.
	//
	// When processing an input value, if its type exists as a key in Formatter:
	//   1) If the value is nil, no stringification is performed.
	//      This allows overriding of PrintStringers and PrintTextMarshalers.
	//   2) The value will be called with the input as its only argument.
	//      The function must return a string as its first return value.
	//
	// In addition to func literals, two common values for this will be:
	//   fmt.Sprint        (function) func Sprint(...interface{}) string
	//   Type.String         (method) func (Type) String() string
	//
	// Note that neither of these work if the String method is a pointer
	// method and the input will be provided as a value.  In that case,
	// use a function that calls .String on the formal value parameter.
	Formatter map[reflect.Type]interface{}
}

// Default Config objects
var (
	// DefaultFormatter is the default set of overrides for stringification.
	DefaultFormatter = map[reflect.Type]interface{}{
		reflect.TypeOf(time.Time{}):          fmt.Sprint,
		reflect.TypeOf(net.IP{}):             fmt.Sprint,
		reflect.TypeOf((*error)(nil)).Elem(): fmt.Sprint,
	}

	// CompareConfig is the default configuration used for Compare.
	CompareConfig = &Config{
		Diffable:          true,
		IncludeUnexported: true,
		Formatter:         DefaultFormatter,
	}

	// DefaultConfig is the default configuration used for all other top-level functions.
	DefaultConfig = &Config{
		Formatter: DefaultFormatter,
	}
)

func (cfg *Config) fprint(buf *bytes.Buffer, vals ...interface{}) {
	for i, val := range vals {
		if i > 0 {
			buf.WriteByte('\n')
		}
		cfg.val2node(reflect.ValueOf(val)).WriteTo(buf, "", cfg)
	}
}

// Print writes the DefaultConfig representation of the given values to standard output.
func Print(vals ...interface{}) {
	DefaultConfig.Print(vals...)
}

// Print writes the configured presentation of the given values to standard output.
func (cfg *Config) Print(vals ...interface{}) {
	fmt.Println(cfg.Sprint(vals...))
}

// Sprint returns a string representation of the given value according to the DefaultConfig.
func Sprint(vals ...interface{}) string {
	return DefaultConfig.Sprint(vals...)
}

// Sprint returns a string representation of the given value according to cfg.
func (cfg *Config) Sprint(vals ...interface{}) string {
	buf := new(bytes.Buffer)
	cfg.fprint(buf, vals...)
	return buf.String()
}

// Fprint writes the representation of the given value to the writer according to the DefaultConfig.
func Fprint(w io.Writer, vals ...interface{}) (n int64, err error) {
	return DefaultConfig.Fprint(w, vals...)
}

// Fprint writes the representation of the given value to the writer according to the cfg.
func (cfg *Config) Fprint(w io.Writer, vals ...interface{}) (n int64, err error) {
	buf := new(bytes.Buffer)
	cfg.fprint(buf, vals...)
	return buf.WriteTo(w)
}

// Compare returns a string containing a line-by-line unified diff of the
// values in got and want, using the CompareConfig.
//
// Each line in the output is prefixed with '+', '-', or ' ' to indicate if it
// should be added to, removed from, or is correct for the "got" value with
// respect to the "want" value.
func Compare(got, want interface{}) string {
	return CompareConfig.Compare(got, want)
}

// Compare returns a string containing a line-by-line unified diff of the
// values in got and want according to the cfg.
func (cfg *Config) Compare(got, want interface{}) string {
	diffCfg := *cfg
	diffCfg.Diffable = true
	return diff.Diff(cfg.Sprint(got), cfg.Sprint(want))
}
<commit_msg>Change the argument names of pretty.Compare to a and b<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pretty

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	"github.com/kylelemons/godebug/diff"
)

// A Config represents optional configuration parameters for formatting.
//
// Some options, notably ShortList, dramatically increase the overhead
// of pretty-printing a value.
type Config struct {
	// Verbosity options
	Compact  bool // One-line output. Overrides Diffable.
	Diffable bool // Adds extra newlines for more easily diffable output.

	// Field and value options
	IncludeUnexported   bool // Include unexported fields in output
	PrintStringers      bool // Call String on a fmt.Stringer
	PrintTextMarshalers bool // Call MarshalText on an encoding.TextMarshaler
	SkipZeroFields      bool // Skip struct fields that have a zero value.

	// Output transforms
	ShortList int // Maximum character length for short lists if nonzero.

	// Type-specific overrides
	//
	// Formatter maps a type to a function that will provide a one-line string
	// representation of the input value.  Conceptually:
	//   Formatter[reflect.TypeOf(v)](v) = "v as a string"
	//
	// Note that the first argument need not explicitly match the type, it must
	// merely be callable with it.
	//
	// When processing an input value, if its type exists as a key in Formatter:
	//   1) If the value is nil, no stringification is performed.
	//      This allows overriding of PrintStringers and PrintTextMarshalers.
	//   2) The value will be called with the input as its only argument.
	//      The function must return a string as its first return value.
	//
	// In addition to func literals, two common values for this will be:
	//   fmt.Sprint        (function) func Sprint(...interface{}) string
	//   Type.String         (method) func (Type) String() string
	//
	// Note that neither of these work if the String method is a pointer
	// method and the input will be provided as a value.  In that case,
	// use a function that calls .String on the formal value parameter.
	Formatter map[reflect.Type]interface{}
}

// Default Config objects
var (
	// DefaultFormatter is the default set of overrides for stringification.
	DefaultFormatter = map[reflect.Type]interface{}{
		reflect.TypeOf(time.Time{}):          fmt.Sprint,
		reflect.TypeOf(net.IP{}):             fmt.Sprint,
		reflect.TypeOf((*error)(nil)).Elem(): fmt.Sprint,
	}

	// CompareConfig is the default configuration used for Compare.
	CompareConfig = &Config{
		Diffable:          true,
		IncludeUnexported: true,
		Formatter:         DefaultFormatter,
	}

	// DefaultConfig is the default configuration used for all other top-level functions.
	DefaultConfig = &Config{
		Formatter: DefaultFormatter,
	}
)

func (cfg *Config) fprint(buf *bytes.Buffer, vals ...interface{}) {
	for i, val := range vals {
		if i > 0 {
			buf.WriteByte('\n')
		}
		cfg.val2node(reflect.ValueOf(val)).WriteTo(buf, "", cfg)
	}
}

// Print writes the DefaultConfig representation of the given values to standard output.
func Print(vals ...interface{}) {
	DefaultConfig.Print(vals...)
}

// Print writes the configured presentation of the given values to standard output.
func (cfg *Config) Print(vals ...interface{}) {
	fmt.Println(cfg.Sprint(vals...))
}

// Sprint returns a string representation of the given value according to the DefaultConfig.
func Sprint(vals ...interface{}) string {
	return DefaultConfig.Sprint(vals...)
}

// Sprint returns a string representation of the given value according to cfg.
func (cfg *Config) Sprint(vals ...interface{}) string {
	buf := new(bytes.Buffer)
	cfg.fprint(buf, vals...)
	return buf.String()
}

// Fprint writes the representation of the given value to the writer according to the DefaultConfig.
func Fprint(w io.Writer, vals ...interface{}) (n int64, err error) {
	return DefaultConfig.Fprint(w, vals...)
}

// Fprint writes the representation of the given value to the writer according to the cfg.
func (cfg *Config) Fprint(w io.Writer, vals ...interface{}) (n int64, err error) {
	buf := new(bytes.Buffer)
	cfg.fprint(buf, vals...)
	return buf.WriteTo(w)
}

// Compare returns a string containing a line-by-line unified diff of the
// values in a and b, using the CompareConfig.
//
// Each line in the output is prefixed with '+', '-', or ' ' to indicate which
// side it's from. Lines from the a side are marked with a '-', lines from the
// b side are marked with a '+' and lines that are the same on both sides are
// marked with a ' '.
func Compare(a, b interface{}) string {
	return CompareConfig.Compare(a, b)
}

// Compare returns a string containing a line-by-line unified diff of the
// values in got and want according to the cfg.
func (cfg *Config) Compare(a, b interface{}) string {
	diffCfg := *cfg
	diffCfg.Diffable = true
	return diff.Diff(cfg.Sprint(a), cfg.Sprint(b))
}
<|endoftext|>