<commit_before>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"fmt"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

const (
	DbHomeRequiredOnlyResource = DbHomeResourceDependencies + `
resource "oci_database_db_home" "test_db_home" {
	#Required
	database {
		#Required
		admin_password = "${var.db_home_database_admin_password}"
		db_name = "${var.db_home_database_db_name}"
	}
	db_system_id = "${oci_database_db_system.test_db_system.id}"
	db_version = "${var.db_home_db_version}"
}
`

	DbHomeResourceConfig = DbHomeResourceDependencies + `
resource "oci_database_db_home" "test_db_home" {
	#Required
	database {
		#Required
		admin_password = "${var.db_home_database_admin_password}"
		backup_id = "${oci_database_backup.test_backup.id}"
		backup_tde_password = "${var.db_home_database_backup_tde_password}"
		db_name = "${var.db_home_database_db_name}"

		#Optional
		character_set = "${var.db_home_database_character_set}"
		db_backup_config {

			#Optional
			auto_backup_enabled = "${var.db_home_database_db_backup_config_auto_backup_enabled}"
		}
		db_workload = "${var.db_home_database_db_workload}"
		defined_tags = "${map("${oci_identity_tag_namespace.tag-namespace1.name}.${oci_identity_tag.tag1.name}", "${var.db_home_database_defined_tags_value}")}"
		freeform_tags = "${var.db_home_database_freeform_tags}"
		ncharacter_set = "${var.db_home_database_ncharacter_set}"
		pdb_name = "${var.db_home_database_pdb_name}"
	}
	db_system_id = "${oci_database_db_system.test_db_system.id}"
	db_version = "${var.db_home_db_version}"

	#Optional
	display_name = "${var.db_home_display_name}"
	source = "${var.db_home_source}"
}
`
	DbHomePropertyVariables = `
variable "db_home_database_admin_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_backup_tde_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_character_set" { default = "AL32UTF8" }
variable "db_home_database_db_backup_config_auto_backup_enabled" { default = false }
variable "db_home_database_db_name" { default = "myTestDb" }
variable "db_home_database_db_workload" { default = "dbWorkload" }
variable "db_home_database_defined_tags_value" { default = "value" }
variable "db_home_database_freeform_tags" { default = {"Department"= "Finance"} }
variable "db_home_database_ncharacter_set" { default = "AL16UTF16" }
variable "db_home_database_pdb_name" { default = "pdbName" }
variable "db_home_display_name" { default = "createdDbHome" }
variable "db_home_source" { default = "DB_BACKUP" }
variable "db_home_db_version" { default = "12.1.0.2" }

`
	DbHomeResourceDependencies = BackupResourceConfig + BackupPropertyVariables
)

func TestDatabaseDbHomeResource_basic(t *testing.T) {
	provider := testAccProvider
	config := testProviderConfig()

	compartmentId := getRequiredEnvSetting("compartment_ocid")
	compartmentIdVariableStr := fmt.Sprintf("variable \"compartment_id\" { default = \"%s\" }\n", compartmentId)

	resourceName := "oci_database_db_home.test_db_home"
	datasourceName := "data.oci_database_db_homes.test_db_homes"

	resource.Test(t, resource.TestCase{
		Providers: map[string]terraform.ResourceProvider{
			"oci": provider,
		},
		Steps: []resource.TestStep{
			// verify create
			{
				ImportState:       true,
				ImportStateVerify: true,
				Config:            config + DbHomePropertyVariables + compartmentIdVariableStr + DbHomeRequiredOnlyResource,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttr(resourceName, "database.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.admin_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_name", "myTestDb"),
					resource.TestCheckResourceAttrSet(resourceName, "db_system_id"),
					resource.TestCheckResourceAttr(resourceName, "db_version", "12.1.0.2"),
				),
			},

			// delete before next create
			{
				Config: config + compartmentIdVariableStr + DbHomeResourceDependencies,
			},
			// verify create with optionals
			{
				Config: config + DbHomePropertyVariables + compartmentIdVariableStr + DbHomeResourceConfig,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttrSet(resourceName, "compartment_id"),
					resource.TestCheckResourceAttr(resourceName, "database.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.admin_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttrSet(resourceName, "database.0.backup_id"),
					resource.TestCheckResourceAttr(resourceName, "database.0.backup_tde_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttr(resourceName, "database.0.character_set", "AL32UTF8"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_backup_config.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_backup_config.0.auto_backup_enabled", "false"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_name", "myTestDb"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_workload", "dbWorkload"),
					resource.TestCheckResourceAttr(resourceName, "database.0.defined_tags.%", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.freeform_tags.%", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.ncharacter_set", "AL16UTF16"),
					resource.TestCheckResourceAttr(resourceName, "database.0.pdb_name", "pdbName"),
					resource.TestCheckResourceAttrSet(resourceName, "db_system_id"),
					resource.TestCheckResourceAttr(resourceName, "db_version", "12.1.0.2"),
					resource.TestCheckResourceAttr(resourceName, "display_name", "createdDbHome"),
					resource.TestCheckResourceAttrSet(resourceName, "id"),
					resource.TestCheckResourceAttr(resourceName, "source", "DB_BACKUP"),
					resource.TestCheckResourceAttrSet(resourceName, "state"),
				),
			},

			// verify datasource
			{
				Config: config + `
variable "db_home_database_admin_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_backup_tde_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_character_set" { default = "AL32UTF8" }
variable "db_home_database_db_backup_config_auto_backup_enabled" { default = false }
variable "db_home_database_db_name" { default = "myTestDb" }
variable "db_home_database_db_workload" { default = "dbWorkload" }
variable "db_home_database_defined_tags_value" { default = "value" }
variable "db_home_database_freeform_tags" { default = {"Department"= "Finance"} }
variable "db_home_database_ncharacter_set" { default = "AL16UTF16" }
variable "db_home_database_pdb_name" { default = "pdbName" }
variable "db_home_display_name" { default = "createdDbHome" }
variable "db_home_source" { default = "DB_BACKUP" }
variable "db_home_db_version" { default = "12.1.0.2" }

data "oci_database_db_homes" "test_db_homes" {
	#Required
	compartment_id = "${var.compartment_id}"
	db_system_id = "${oci_database_db_system.test_db_system.id}"

    filter {
    	name = "id"
    	values = ["${oci_database_db_home.test_db_home.id}"]
    }
}
                ` + compartmentIdVariableStr + DbHomeResourceConfig,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttr(datasourceName, "compartment_id", compartmentId),
					resource.TestCheckResourceAttrSet(datasourceName, "db_system_id"),

					resource.TestCheckResourceAttr(datasourceName, "db_homes.#", "1"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.compartment_id"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.db_system_id"),
					resource.TestCheckResourceAttr(datasourceName, "db_homes.0.db_version", "12.1.0.2"),
					resource.TestCheckResourceAttr(datasourceName, "db_homes.0.display_name", "createdDbHome"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.id"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.state"),
				),
			},
		},
	})
}
<commit_msg>fixing dbHome Test after tag change<commit_after>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"fmt"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

const (
	DbHomeRequiredOnlyResource = DbHomeResourceDependencies + `
resource "oci_database_db_home" "test_db_home" {
	#Required
	database {
		#Required
		admin_password = "${var.db_home_database_admin_password}"
		db_name = "${var.db_home_database_db_name}"
	}
	db_system_id = "${oci_database_db_system.test_db_system.id}"
	db_version = "${var.db_home_db_version}"
}
`

	DbHomeResourceConfig = DbHomeResourceDependencies + `
resource "oci_database_db_home" "test_db_home" {
	#Required
	database {
		#Required
		admin_password = "${var.db_home_database_admin_password}"
		backup_id = "${oci_database_backup.test_backup.id}"
		backup_tde_password = "${var.db_home_database_backup_tde_password}"
		db_name = "${var.db_home_database_db_name}"

		#Optional
		character_set = "${var.db_home_database_character_set}"
		db_backup_config {

			#Optional
			auto_backup_enabled = "${var.db_home_database_db_backup_config_auto_backup_enabled}"
		}
		db_workload = "${var.db_home_database_db_workload}"
		defined_tags = "${map("${oci_identity_tag_namespace.tag-namespace1.name}.${oci_identity_tag.tag1.name}", "${var.db_home_database_defined_tags_value}")}"
		freeform_tags = "${var.db_home_database_freeform_tags}"
		ncharacter_set = "${var.db_home_database_ncharacter_set}"
		pdb_name = "${var.db_home_database_pdb_name}"
	}
	db_system_id = "${oci_database_db_system.test_db_system.id}"
	db_version = "${var.db_home_db_version}"

	#Optional
	display_name = "${var.db_home_display_name}"
	source = "${var.db_home_source}"
}
`
	DbHomePropertyVariables = `
variable "db_home_database_admin_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_backup_tde_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_character_set" { default = "AL32UTF8" }
variable "db_home_database_db_backup_config_auto_backup_enabled" { default = false }
variable "db_home_database_db_name" { default = "myTestDb" }
variable "db_home_database_db_workload" { default = "dbWorkload" }
variable "db_home_database_defined_tags_value" { default = "value" }
variable "db_home_database_freeform_tags" { default = {"Department"= "Finance"} }
variable "db_home_database_ncharacter_set" { default = "AL16UTF16" }
variable "db_home_database_pdb_name" { default = "pdbName" }
variable "db_home_display_name" { default = "createdDbHome" }
variable "db_home_source" { default = "DB_BACKUP" }
variable "db_home_db_version" { default = "12.1.0.2" }

`
	DbHomeResourceDependencies = BackupResourceConfig + BackupPropertyVariables + DefinedTagsDependencies
)

func TestDatabaseDbHomeResource_basic(t *testing.T) {
	provider := testAccProvider
	config := testProviderConfig()

	compartmentId := getRequiredEnvSetting("compartment_ocid")
	compartmentIdVariableStr := fmt.Sprintf("variable \"compartment_id\" { default = \"%s\" }\n", compartmentId)

	resourceName := "oci_database_db_home.test_db_home"
	datasourceName := "data.oci_database_db_homes.test_db_homes"

	resource.Test(t, resource.TestCase{
		Providers: map[string]terraform.ResourceProvider{
			"oci": provider,
		},
		Steps: []resource.TestStep{
			// verify create
			{
				ImportState:       true,
				ImportStateVerify: true,
				Config:            config + DbHomePropertyVariables + compartmentIdVariableStr + DbHomeRequiredOnlyResource,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttr(resourceName, "database.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.admin_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_name", "myTestDb"),
					resource.TestCheckResourceAttrSet(resourceName, "db_system_id"),
					resource.TestCheckResourceAttr(resourceName, "db_version", "12.1.0.2"),
				),
			},

			// delete before next create
			{
				Config: config + compartmentIdVariableStr + DbHomeResourceDependencies,
			},
			// verify create with optionals
			{
				Config: config + DbHomePropertyVariables + compartmentIdVariableStr + DbHomeResourceConfig,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttrSet(resourceName, "compartment_id"),
					resource.TestCheckResourceAttr(resourceName, "database.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.admin_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttrSet(resourceName, "database.0.backup_id"),
					resource.TestCheckResourceAttr(resourceName, "database.0.backup_tde_password", "BEstrO0ng_#11"),
					resource.TestCheckResourceAttr(resourceName, "database.0.character_set", "AL32UTF8"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_backup_config.#", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_backup_config.0.auto_backup_enabled", "false"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_name", "myTestDb"),
					resource.TestCheckResourceAttr(resourceName, "database.0.db_workload", "dbWorkload"),
					resource.TestCheckResourceAttr(resourceName, "database.0.defined_tags.%", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.freeform_tags.%", "1"),
					resource.TestCheckResourceAttr(resourceName, "database.0.ncharacter_set", "AL16UTF16"),
					resource.TestCheckResourceAttr(resourceName, "database.0.pdb_name", "pdbName"),
					resource.TestCheckResourceAttrSet(resourceName, "db_system_id"),
					resource.TestCheckResourceAttr(resourceName, "db_version", "12.1.0.2"),
					resource.TestCheckResourceAttr(resourceName, "display_name", "createdDbHome"),
					resource.TestCheckResourceAttrSet(resourceName, "id"),
					resource.TestCheckResourceAttr(resourceName, "source", "DB_BACKUP"),
					resource.TestCheckResourceAttrSet(resourceName, "state"),
				),
			},

			// verify datasource
			{
				Config: config + `
variable "db_home_database_admin_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_backup_tde_password" { default = "BEstrO0ng_#11" }
variable "db_home_database_character_set" { default = "AL32UTF8" }
variable "db_home_database_db_backup_config_auto_backup_enabled" { default = false }
variable "db_home_database_db_name" { default = "myTestDb" }
variable "db_home_database_db_workload" { default = "dbWorkload" }
variable "db_home_database_defined_tags_value" { default = "value" }
variable "db_home_database_freeform_tags" { default = {"Department"= "Finance"} }
variable "db_home_database_ncharacter_set" { default = "AL16UTF16" }
variable "db_home_database_pdb_name" { default = "pdbName" }
variable "db_home_display_name" { default = "createdDbHome" }
variable "db_home_source" { default = "DB_BACKUP" }
variable "db_home_db_version" { default = "12.1.0.2" }

data "oci_database_db_homes" "test_db_homes" {
	#Required
	compartment_id = "${var.compartment_id}"
	db_system_id = "${oci_database_db_system.test_db_system.id}"

    filter {
    	name = "id"
    	values = ["${oci_database_db_home.test_db_home.id}"]
    }
}
                ` + compartmentIdVariableStr + DbHomeResourceConfig,
				Check: resource.ComposeAggregateTestCheckFunc(
					resource.TestCheckResourceAttr(datasourceName, "compartment_id", compartmentId),
					resource.TestCheckResourceAttrSet(datasourceName, "db_system_id"),

					resource.TestCheckResourceAttr(datasourceName, "db_homes.#", "1"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.compartment_id"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.db_system_id"),
					resource.TestCheckResourceAttr(datasourceName, "db_homes.0.db_version", "12.1.0.2"),
					resource.TestCheckResourceAttr(datasourceName, "db_homes.0.display_name", "createdDbHome"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.id"),
					resource.TestCheckResourceAttrSet(datasourceName, "db_homes.0.state"),
				),
			},
		},
	})
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"fmt"
	"regexp"

	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"

	"github.com/rusenask/keel/types"
	"github.com/rusenask/keel/util/version"

	log "github.com/Sirupsen/logrus"
)

// ProviderName - provider name
const ProviderName = "kubernetes"

var versionreg = regexp.MustCompile(`:[^:]*$`)

// Provider - kubernetes provider for auto update
type Provider struct {
	implementer Implementer

	events chan *types.Event
	stop   chan struct{}
}

// NewProvider - create new kubernetes based provider
func NewProvider(implementer Implementer) (*Provider, error) {
	return &Provider{
		implementer: implementer,
		events:      make(chan *types.Event, 100),
		stop:        make(chan struct{}),
	}, nil
}

// Submit - submit event to provider
func (p *Provider) Submit(event types.Event) error {
	p.events <- &event
	return nil
}

// GetName - get provider name
func (p *Provider) GetName() string {
	return ProviderName
}

// Start - starts kubernetes provider, waits for events
func (p *Provider) Start() error {
	return p.startInternal()
}

// Stop - stops kubernetes provider
func (p *Provider) Stop() {
	close(p.stop)
}

func (p *Provider) startInternal() error {
	for {
		select {
		case event := <-p.events:
			log.WithFields(log.Fields{
				"repository": event.Repository.Name,
			}).Info("provider.kubernetes: processing event")
			_, err := p.processEvent(event)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err,
					"image": event.Repository.Name,
					"tag":   event.Repository.Tag,
				}).Error("provider.kubernetes: failed to process event")
			}
		case <-p.stop:
			log.Info("provider.kubernetes: got shutdown signal, stopping...")
			return nil
		}
	}
}

func (p *Provider) processEvent(event *types.Event) (updated []*v1beta1.Deployment, err error) {
	impacted, err := p.impactedDeployments(&event.Repository)
	if err != nil {
		return nil, err
	}

	if len(impacted) == 0 {
		log.WithFields(log.Fields{
			"image": event.Repository.Name,
			"tag":   event.Repository.Tag,
		}).Info("provider.kubernetes: no impacted deployments found for this event")
		return
	}

	return p.updateDeployments(impacted)
}

func (p *Provider) updateDeployments(deployments []v1beta1.Deployment) (updated []*v1beta1.Deployment, err error) {
	for _, deployment := range deployments {
		err := p.implementer.Update(&deployment)
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"namespace":  deployment.Namespace,
				"deployment": deployment.Name,
			}).Error("provider.kubernetes: got error while update deployment")
			continue
		}
		log.WithFields(log.Fields{
			"name":      deployment.Name,
			"namespace": deployment.Namespace,
		}).Info("provider.kubernetes: deployment updated")
		updated = append(updated, &deployment)
	}

	return
}

// getDeployment - helper function to get specific deployment
func (p *Provider) getDeployment(namespace, name string) (*v1beta1.Deployment, error) {
	return p.implementer.Deployment(namespace, name)
}

// gets impacted deployments by changed repository
func (p *Provider) impactedDeployments(repo *types.Repository) ([]v1beta1.Deployment, error) {
	newVersion, err := version.GetVersion(repo.Tag)
	if err != nil {
		return nil, fmt.Errorf("failed to parse version from repository tag, error: %s", err)
	}

	deploymentLists, err := p.deployments()
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Error("provider.kubernetes: failed to get deployment lists")
		return nil, err
	}

	impacted := []v1beta1.Deployment{}

	for _, deploymentList := range deploymentLists {
		for _, deployment := range deploymentList.Items {
			labels := deployment.GetLabels()
			policyStr, ok := labels[types.KeelPolicyLabel]
			// if no policy is set - skipping this deployment
			if !ok {
				continue
			}
			policy := types.ParsePolicy(policyStr)

			log.WithFields(log.Fields{
				"labels":    labels,
				"name":      deployment.Name,
				"namespace": deployment.Namespace,
				"policy":    policy,
			}).Info("provider.kubernetes: keel policy found, checking deployment...")

			shouldUpdateDeployment := false

			for idx, c := range deployment.Spec.Template.Spec.Containers {
				// Remove version if any
				containerImageName := versionreg.ReplaceAllString(c.Image, "")

				log.WithFields(log.Fields{
					"name":              deployment.Name,
					"namespace":         deployment.Namespace,
					"parsed_image_name": containerImageName,
					"target_image_name": repo.Name,
					"target_tag":        repo.Tag,
					"policy":            policy,
					"image":             c.Image,
				}).Info("provider.kubernetes: checking image")

				if containerImageName != repo.Name {
					continue
				}

				currentVersion, err := version.GetVersionFromImageName(c.Image)
				if err != nil {
					log.WithFields(log.Fields{
						"error":       err,
						"image_name":  c.Image,
						"keel_policy": policy,
					}).Error("provider.kubernetes: failed to get image version, is it tagged as semver?")
					continue
				}

				log.WithFields(log.Fields{
					"labels":          labels,
					"name":            deployment.Name,
					"namespace":       deployment.Namespace,
					"image":           c.Image,
					"current_version": currentVersion.String(),
					"policy":          policy,
				}).Info("provider.kubernetes: current image version")

				shouldUpdateContainer, err := version.ShouldUpdate(currentVersion, newVersion, policy)
				if err != nil {
					log.WithFields(log.Fields{
						"error":           err,
						"new_version":     newVersion.String(),
						"current_version": currentVersion.String(),
						"keel_policy":     policy,
					}).Error("provider.kubernetes: got error while checking whether deployment should be updated")
					continue
				}

				log.WithFields(log.Fields{
					"labels":          labels,
					"name":            deployment.Name,
					"namespace":       deployment.Namespace,
					"image":           c.Image,
					"current_version": currentVersion.String(),
					"new_version":     newVersion.String(),
					"policy":          policy,
					"should_update":   shouldUpdateContainer,
				}).Info("provider.kubernetes: checked version, deciding whether to update")

				if shouldUpdateContainer {
					// updating image
					c.Image = fmt.Sprintf("%s:%s", containerImageName, newVersion.String())
					deployment.Spec.Template.Spec.Containers[idx] = c
					// marking this deployment for update
					shouldUpdateDeployment = true

					log.WithFields(log.Fields{
						"parsed_image":     containerImageName,
						"raw_image_name":   c.Image,
						"target_image":     repo.Name,
						"target_image_tag": repo.Tag,
						"policy":           policy,
					}).Info("provider.kubernetes: impacted deployment container found")
				}
			}

			if shouldUpdateDeployment {
				impacted = append(impacted, deployment)
			}
		}
	}

	return impacted, nil
}

func (p *Provider) namespaces() (*v1.NamespaceList, error) {
	return p.implementer.Namespaces()
}

// deployments - gets all deployments
func (p *Provider) deployments() ([]*v1beta1.DeploymentList, error) {
	deployments := []*v1beta1.DeploymentList{}

	n, err := p.namespaces()
	if err != nil {
		return nil, err
	}

	for _, n := range n.Items {
		l, err := p.implementer.Deployments(n.GetName())
		if err != nil {
			log.WithFields(log.Fields{
				"error":     err,
				"namespace": n.GetName(),
			}).Error("provider.kubernetes: failed to list deployments")
			continue
		}
		deployments = append(deployments, l)
	}

	return deployments, nil
}
<commit_msg>providing more info on event process start<commit_after>package kubernetes

import (
	"fmt"
	"regexp"

	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"

	"github.com/rusenask/keel/types"
	"github.com/rusenask/keel/util/version"

	log "github.com/Sirupsen/logrus"
)

// ProviderName - provider name
const ProviderName = "kubernetes"

var versionreg = regexp.MustCompile(`:[^:]*$`)

// Provider - kubernetes provider for auto update
type Provider struct {
	implementer Implementer

	events chan *types.Event
	stop   chan struct{}
}

// NewProvider - create new kubernetes based provider
func NewProvider(implementer Implementer) (*Provider, error) {
	return &Provider{
		implementer: implementer,
		events:      make(chan *types.Event, 100),
		stop:        make(chan struct{}),
	}, nil
}

// Submit - submit event to provider
func (p *Provider) Submit(event types.Event) error {
	p.events <- &event
	return nil
}

// GetName - get provider name
func (p *Provider) GetName() string {
	return ProviderName
}

// Start - starts kubernetes provider, waits for events
func (p *Provider) Start() error {
	return p.startInternal()
}

// Stop - stops kubernetes provider
func (p *Provider) Stop() {
	close(p.stop)
}

func (p *Provider) startInternal() error {
	for {
		select {
		case event := <-p.events:
			log.WithFields(log.Fields{
				"repository": event.Repository.Name,
				"tag":        event.Repository.Tag,
				"registry":   event.Repository.Host,
			}).Info("provider.kubernetes: processing event")
			_, err := p.processEvent(event)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err,
					"image": event.Repository.Name,
					"tag":   event.Repository.Tag,
				}).Error("provider.kubernetes: failed to process event")
			}
		case <-p.stop:
			log.Info("provider.kubernetes: got shutdown signal, stopping...")
			return nil
		}
	}
}

func (p *Provider) processEvent(event *types.Event) (updated []*v1beta1.Deployment, err error) {
	impacted, err := p.impactedDeployments(&event.Repository)
	if err != nil {
		return nil, err
	}

	if len(impacted) == 0 {
		log.WithFields(log.Fields{
			"image": event.Repository.Name,
			"tag":   event.Repository.Tag,
		}).Info("provider.kubernetes: no impacted deployments found for this event")
		return
	}

	return p.updateDeployments(impacted)
}

func (p *Provider) updateDeployments(deployments []v1beta1.Deployment) (updated []*v1beta1.Deployment, err error) {
	for _, deployment := range deployments {
		err := p.implementer.Update(&deployment)
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"namespace":  deployment.Namespace,
				"deployment": deployment.Name,
			}).Error("provider.kubernetes: got error while update deployment")
			continue
		}
		log.WithFields(log.Fields{
			"name":      deployment.Name,
			"namespace": deployment.Namespace,
		}).Info("provider.kubernetes: deployment updated")
		updated = append(updated, &deployment)
	}

	return
}

// getDeployment - helper function to get specific deployment
func (p *Provider) getDeployment(namespace, name string) (*v1beta1.Deployment, error) {
	return p.implementer.Deployment(namespace, name)
}

// gets impacted deployments by changed repository
func (p *Provider) impactedDeployments(repo *types.Repository) ([]v1beta1.Deployment, error) {
	newVersion, err := version.GetVersion(repo.Tag)
	if err != nil {
		return nil, fmt.Errorf("failed to parse version from repository tag, error: %s", err)
	}

	deploymentLists, err := p.deployments()
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Error("provider.kubernetes: failed to get deployment lists")
		return nil, err
	}

	impacted := []v1beta1.Deployment{}

	for _, deploymentList := range deploymentLists {
		for _, deployment := range deploymentList.Items {
			labels := deployment.GetLabels()
			policyStr, ok := labels[types.KeelPolicyLabel]
			// if no policy is set - skipping this deployment
			if !ok {
				continue
			}
			policy := types.ParsePolicy(policyStr)

			log.WithFields(log.Fields{
				"labels":    labels,
				"name":      deployment.Name,
				"namespace": deployment.Namespace,
				"policy":    policy,
			}).Info("provider.kubernetes: keel policy found, checking deployment...")

			shouldUpdateDeployment := false

			for idx, c := range deployment.Spec.Template.Spec.Containers {
				// Remove version if any
				containerImageName := versionreg.ReplaceAllString(c.Image, "")

				log.WithFields(log.Fields{
					"name":              deployment.Name,
					"namespace":         deployment.Namespace,
					"parsed_image_name": containerImageName,
					"target_image_name": repo.Name,
					"target_tag":        repo.Tag,
					"policy":            policy,
					"image":             c.Image,
				}).Info("provider.kubernetes: checking image")

				if containerImageName != repo.Name {
					continue
				}

				currentVersion, err := version.GetVersionFromImageName(c.Image)
				if err != nil {
					log.WithFields(log.Fields{
						"error":       err,
						"image_name":  c.Image,
						"keel_policy": policy,
					}).Error("provider.kubernetes: failed to get image version, is it tagged as semver?")
					continue
				}

				log.WithFields(log.Fields{
					"labels":          labels,
					"name":            deployment.Name,
					"namespace":       deployment.Namespace,
					"image":           c.Image,
					"current_version": currentVersion.String(),
					"policy":          policy,
				}).Info("provider.kubernetes: current image version")

				shouldUpdateContainer, err := version.ShouldUpdate(currentVersion, newVersion, policy)
				if err != nil {
					log.WithFields(log.Fields{
						"error":           err,
						"new_version":     newVersion.String(),
						"current_version": currentVersion.String(),
						"keel_policy":     policy,
					}).Error("provider.kubernetes: got error while checking whether deployment should be updated")
					continue
				}

				log.WithFields(log.Fields{
					"labels":          labels,
					"name":            deployment.Name,
					"namespace":       deployment.Namespace,
					"image":           c.Image,
					"current_version": currentVersion.String(),
					"new_version":     newVersion.String(),
					"policy":          policy,
					"should_update":   shouldUpdateContainer,
				}).Info("provider.kubernetes: checked version, deciding whether to update")

				if shouldUpdateContainer {
					// updating image
					c.Image = fmt.Sprintf("%s:%s", containerImageName, newVersion.String())
					deployment.Spec.Template.Spec.Containers[idx] = c
					// marking this deployment for update
					shouldUpdateDeployment = true

					log.WithFields(log.Fields{
						"parsed_image":     containerImageName,
						"raw_image_name":   c.Image,
						"target_image":     repo.Name,
						"target_image_tag": repo.Tag,
						"policy":           policy,
					}).Info("provider.kubernetes: impacted deployment container found")
				}
			}

			if shouldUpdateDeployment {
				impacted = append(impacted, deployment)
			}
		}
	}

	return impacted, nil
}

func (p *Provider) namespaces() (*v1.NamespaceList, error) {
	return p.implementer.Namespaces()
}

// deployments - gets all deployments
func (p *Provider) deployments() ([]*v1beta1.DeploymentList, error) {
	deployments := []*v1beta1.DeploymentList{}

	n, err := p.namespaces()
	if err != nil {
		return nil, err
	}

	for _, n := range n.Items {
		l, err := p.implementer.Deployments(n.GetName())
		if err != nil {
			log.WithFields(log.Fields{
				"error":     err,
				"namespace": n.GetName(),
			}).Error("provider.kubernetes: failed to list deployments")
			continue
		}
		deployments = append(deployments, l)
	}

	return deployments, nil
}
<|endoftext|>


<commit_before>package dropbox

import (
	"os"
	"testing"

	"github.com/markbates/goth"
	"github.com/stretchr/testify/assert"
)

func provider() *Provider {
	return New(os.Getenv("DROPBOX_KEY"), os.Getenv("DROPBOX_SECRET"), "/foo", "email")
}

func Test_New(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	p := provider()

	a.Equal(p.ClientKey, os.Getenv("DROPBOX_KEY"))
	a.Equal(p.Secret, os.Getenv("DROPBOX_SECRET"))
	a.Equal(p.CallbackURL, "/foo")
}

func Test_Implements_Provider(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	a.Implements((*goth.Provider)(nil), provider())
}

func Test_ImplementsSession(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := Session{}
	a.Implements((*goth.Session)(nil), s)
}

func Test_BeginAuth(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	p := provider()
	session, err := p.BeginAuth("test_state")
	s := session.(*Session)
	a.NoError(err)
	a.Contains(s.AuthURL, "www.dropbox.com/1/oauth2/authorize")
}

func Test_SessionFromJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	p := provider()
	session, err := p.UnmarshalSession(`{"AuthURL":"https://www.dropbox.com/1/oauth2/authorize","Token":"1234567890"}`)
	a.NoError(err)

	s := session.(*Session)
	a.Equal(s.AuthURL, "https://www.dropbox.com/1/oauth2/authorize")
	a.Equal(s.Token, "1234567890")
}

func Test_SessionToJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := &Session{}

	data := s.Marshal()
	a.Equal(data, `{"AuthURL":"","Token":""}`)
}

func Test_GetAuthURL(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := &Session{}

	_, err := s.GetAuthURL()
	a.Error(err)

	s.AuthURL = "/foo"
	url, _ := s.GetAuthURL()
	a.Equal(url, "/foo")
}
<commit_msg>Fix failing tests (for real this time)<commit_after>package dropbox

import (
	"os"
	"testing"

	"github.com/markbates/goth"
	"github.com/stretchr/testify/assert"
)

func provider() *Provider {
	return New(os.Getenv("DROPBOX_KEY"), os.Getenv("DROPBOX_SECRET"), "/foo", "email")
}

func Test_New(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	p := provider()

	a.Equal(p.ClientKey, os.Getenv("DROPBOX_KEY"))
	a.Equal(p.Secret, os.Getenv("DROPBOX_SECRET"))
	a.Equal(p.CallbackURL, "/foo")
}

func Test_Implements_Provider(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	a.Implements((*goth.Provider)(nil), provider())
}

func Test_ImplementsSession(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := &Session{}
	a.Implements((*goth.Session)(nil), s)
}

func Test_BeginAuth(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	p := provider()
	session, err := p.BeginAuth("test_state")
	s := session.(*Session)
	a.NoError(err)
	a.Contains(s.AuthURL, "www.dropbox.com/1/oauth2/authorize")
}

func Test_SessionFromJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	p := provider()
	session, err := p.UnmarshalSession(`{"AuthURL":"https://www.dropbox.com/1/oauth2/authorize","Token":"1234567890"}`)
	a.NoError(err)

	s := session.(*Session)
	a.Equal(s.AuthURL, "https://www.dropbox.com/1/oauth2/authorize")
	a.Equal(s.Token, "1234567890")
}

func Test_SessionToJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := &Session{}

	data := s.Marshal()
	a.Equal(data, `{"AuthURL":"","Token":""}`)
}

func Test_GetAuthURL(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	s := &Session{}

	_, err := s.GetAuthURL()
	a.Error(err)

	s.AuthURL = "/foo"
	url, _ := s.GetAuthURL()
	a.Equal(url, "/foo")
}
<|endoftext|>


<commit_before>package apps

import (
	"ddtxn"
	"ddtxn/dlog"
	"ddtxn/stats"
	"fmt"
	"sync/atomic"
	"time"
)

type Buy struct {
	sp              uint32
	read_rate       int
	nproducts       int
	ncontended_rate int
	nbidders        int
	nworkers        int
	validate        []int32
	lhr             []*stats.LatencyHist
	lhw             []*stats.LatencyHist
}

func InitBuy(s *ddtxn.Store, np, nb, nw, rr int, ncrr float64, ngo int) *Buy {
	b := &Buy{
		nproducts:       np,
		nbidders:        nb,
		nworkers:        nw,
		read_rate:       rr,
		ncontended_rate: int(ncrr * float64(rr)),
		validate:        make([]int32, np),
		lhr:             make([]*stats.LatencyHist, ngo),
		lhw:             make([]*stats.LatencyHist, ngo),
		sp:              uint32(nb / nw),
	}

	for i := 0; i < np; i++ {
		k := ddtxn.ProductKey(i)
		s.CreateKey(k, int32(0), ddtxn.SUM)
	}
	// Uncontended keys
	for i := np; i < nb/10; i++ {
		k := ddtxn.ProductKey(i)
		s.CreateKey(k, int32(0), ddtxn.SUM)
	}
	for i := 0; i < nb; i++ {
		k := ddtxn.UserKey(i)
		s.CreateKey(k, "x", ddtxn.WRITE)
	}
	return b
}

func (b *Buy) SetupLatency(nincr int64, nbuckets int64, ngo int) {
	for i := 0; i < ngo; i++ {
		b.lhr[i] = stats.MakeLatencyHistogram(nincr, nbuckets)
		b.lhw[i] = stats.MakeLatencyHistogram(nincr, nbuckets)
	}
}

// Calls rand 4 times
func (b *Buy) MakeOne(w int, local_seed *uint32, txn *ddtxn.Query) {
	rnd := ddtxn.RandN(local_seed, b.sp)
	lb := int(rnd)
	bidder := lb + w*int(b.sp)
	amt := int32(ddtxn.RandN(local_seed, 10))
	product := int(ddtxn.RandN(local_seed, uint32(b.nproducts)))
	x := int(ddtxn.RandN(local_seed, 100))
	if x < b.read_rate {
		if x > b.ncontended_rate {
			// Contended read
			txn.K1 = ddtxn.ProductKey(product)
		} else {
			// Uncontended read
			txn.K1 = ddtxn.UserKey(bidder)
		}
		txn.TXN = ddtxn.D_READ_ONE
	} else {
		txn.K1 = ddtxn.UserKey(bidder)
		txn.K2 = ddtxn.ProductKey(product)
		txn.A = amt
		txn.TXN = ddtxn.D_BUY
	}
}

func (b *Buy) Add(t ddtxn.Query) {
	if t.TXN == ddtxn.D_BUY || t.TXN == ddtxn.D_BUY_NC {
		x, _ := ddtxn.UndoCKey(t.K2)
		atomic.AddInt32(&b.validate[x], t.A)
	}
}

func (b *Buy) Validate(s *ddtxn.Store, nitr int) bool {
	good := true
	zero_cnt := 0
	for j := 0; j < b.nproducts; j++ {
		var x int32
		k := ddtxn.ProductKey(j)
		v, err := s.Get(k)
		if err != nil {
			if b.validate[j] != 0 {
				fmt.Printf("Validating key %v failed; store: none should have: %v\n", k, b.validate[j])
				good = false
			}
			continue
		}
		x = v.Value().(int32)
		dlog.Printf("Validate: %v %v\n", k, x)
		if x != b.validate[j] {
			fmt.Printf("Validating key %v failed; store: %v should have: %v\n", k, x, b.validate[j])
			good = false
		}
		if x == 0 {
			dlog.Printf("Saying x is zero %v %v\n", x, zero_cnt)
			zero_cnt++
		}
	}
	if zero_cnt == b.nproducts && nitr > 10 {
		fmt.Printf("Bad: all zeroes!\n")
		dlog.Printf("Bad: all zeroes!\n")
		good = false
	}
	return good
}

func (b *Buy) Time(t *ddtxn.Query, txn_end time.Duration, n int) {
	if t.TXN == ddtxn.D_READ_ONE {
		b.lhr[n].AddOne(txn_end.Nanoseconds())
	} else {
		b.lhw[n].AddOne(txn_end.Nanoseconds())
	}
}

func (b *Buy) LatencyString(ngo int) (string, string) {
	for i := 1; i < ngo; i++ {
		b.lhr[0].Combine(b.lhr[i])
		b.lhw[0].Combine(b.lhw[i])
	}
	return fmt.Sprint("Read 25: %v\nRead 50: %v\nRead 75: %v\nRead 99: %v\n", b.lhr[0].GetPercentile(25), b.lhr[0].GetPercentile(50), b.lhr[0].GetPercentile(75), b.lhr[0].GetPercentile(99)), fmt.Sprint("Write 25: %v\nWrite 50: %v\nWrite 75: %v\nWrite 99: %v\n", b.lhw[0].GetPercentile(25), b.lhw[0].GetPercentile(50), b.lhw[0].GetPercentile(75), b.lhw[0].GetPercentile(99))
}
<commit_msg>use a smaller portion of the bidder key space<commit_after>package apps

import (
	"ddtxn"
	"ddtxn/dlog"
	"ddtxn/stats"
	"fmt"
	"sync/atomic"
	"time"
)

type Buy struct {
	sp              uint32
	read_rate       int
	nproducts       int
	ncontended_rate int
	nbidders        int
	nworkers        int
	validate        []int32
	lhr             []*stats.LatencyHist
	lhw             []*stats.LatencyHist
}

func InitBuy(s *ddtxn.Store, np, nb, nw, rr int, ncrr float64, ngo int) *Buy {
	b := &Buy{
		nproducts:       np,
		nbidders:        nb,
		nworkers:        nw,
		read_rate:       rr,
		ncontended_rate: int(ncrr * float64(rr)),
		validate:        make([]int32, np),
		lhr:             make([]*stats.LatencyHist, ngo),
		lhw:             make([]*stats.LatencyHist, ngo),
		sp:              uint32(nb / nw),
	}

	for i := 0; i < np; i++ {
		k := ddtxn.ProductKey(i)
		s.CreateKey(k, int32(0), ddtxn.SUM)
	}
	// Uncontended keys
	for i := np; i < nb/10; i++ {
		k := ddtxn.ProductKey(i)
		s.CreateKey(k, int32(0), ddtxn.SUM)
	}
	for i := 0; i < nb; i++ {
		k := ddtxn.UserKey(i)
		s.CreateKey(k, "x", ddtxn.WRITE)
	}
	return b
}

func (b *Buy) SetupLatency(nincr int64, nbuckets int64, ngo int) {
	for i := 0; i < ngo; i++ {
		b.lhr[i] = stats.MakeLatencyHistogram(nincr, nbuckets)
		b.lhw[i] = stats.MakeLatencyHistogram(nincr, nbuckets)
	}
}

// Calls rand 4 times
func (b *Buy) MakeOne(w int, local_seed *uint32, txn *ddtxn.Query) {
	rnd := ddtxn.RandN(local_seed, b.sp/8)
	lb := int(rnd)
	bidder := lb + w*int(b.sp)
	amt := int32(ddtxn.RandN(local_seed, 10))
	product := int(ddtxn.RandN(local_seed, uint32(b.nproducts)))
	x := int(ddtxn.RandN(local_seed, 100))
	if x < b.read_rate {
		if x > b.ncontended_rate {
			// Contended read
			txn.K1 = ddtxn.ProductKey(product)
		} else {
			// Uncontended read
			txn.K1 = ddtxn.UserKey(bidder)
		}
		txn.TXN = ddtxn.D_READ_ONE
	} else {
		txn.K1 = ddtxn.UserKey(bidder)
		txn.K2 = ddtxn.ProductKey(product)
		txn.A = amt
		txn.TXN = ddtxn.D_BUY
	}
}

func (b *Buy) Add(t ddtxn.Query) {
	if t.TXN == ddtxn.D_BUY || t.TXN == ddtxn.D_BUY_NC {
		x, _ := ddtxn.UndoCKey(t.K2)
		atomic.AddInt32(&b.validate[x], t.A)
	}
}

func (b *Buy) Validate(s *ddtxn.Store, nitr int) bool {
	good := true
	zero_cnt := 0
	for j := 0; j < b.nproducts; j++ {
		var x int32
		k := ddtxn.ProductKey(j)
		v, err := s.Get(k)
		if err != nil {
			if b.validate[j] != 0 {
				fmt.Printf("Validating key %v failed; store: none should have: %v\n", k, b.validate[j])
				good = false
			}
			continue
		}
		x = v.Value().(int32)
		dlog.Printf("Validate: %v %v\n", k, x)
		if x != b.validate[j] {
			fmt.Printf("Validating key %v failed; store: %v should have: %v\n", k, x, b.validate[j])
			good = false
		}
		if x == 0 {
			dlog.Printf("Saying x is zero %v %v\n", x, zero_cnt)
			zero_cnt++
		}
	}
	if zero_cnt == b.nproducts && nitr > 10 {
		fmt.Printf("Bad: all zeroes!\n")
		dlog.Printf("Bad: all zeroes!\n")
		good = false
	}
	return good
}

func (b *Buy) Time(t *ddtxn.Query, txn_end time.Duration, n int) {
	if t.TXN == ddtxn.D_READ_ONE {
		b.lhr[n].AddOne(txn_end.Nanoseconds())
	} else {
		b.lhw[n].AddOne(txn_end.Nanoseconds())
	}
}

func (b *Buy) LatencyString(ngo int) (string, string) {
	for i := 1; i < ngo; i++ {
		b.lhr[0].Combine(b.lhr[i])
		b.lhw[0].Combine(b.lhw[i])
	}
	return fmt.Sprint("Read 25: %v\nRead 50: %v\nRead 75: %v\nRead 99: %v\n", b.lhr[0].GetPercentile(25), b.lhr[0].GetPercentile(50), b.lhr[0].GetPercentile(75), b.lhr[0].GetPercentile(99)), fmt.Sprint("Write 25: %v\nWrite 50: %v\nWrite 75: %v\nWrite 99: %v\n", b.lhw[0].GetPercentile(25), b.lhw[0].GetPercentile(50), b.lhw[0].GetPercentile(75), b.lhw[0].GetPercentile(99))
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"context"
	"github.com/ipfs-search/ipfs-search/index"
	"github.com/olivere/elastic/v6"
	"log"
)

func getElastic(url string) (*elastic.Client, error) {
	el, err := elastic.NewClient(elastic.SetSniff(false), elastic.SetURL(url))
	if err != nil {
		return nil, err
	}

	log.Printf("Connected to ElasticSearch")

	return el, nil
}

func getIndex(ctx context.Context, es *elastic.Client, config *index.Config) (index.Index, error) {
	if config == nil {
		panic("configuration for index nil")
	}

	i := &Index{
		Client: es,
		Config: config,
	}

	// Create index if it doesn't already exists, update if it is different (last parameter, true).
	if err := index.EnsureUpdated(ctx, i, config); err != nil {
		return nil, err
	}

	return i, nil
}

func EnsureIndexes(ctx context.Context, esURL string, configs map[string]*index.Config) (indexes map[string]index.Index, err error) {
	es, err := getElastic(esURL)

	if err != nil {
		return
	}

	indexes = make(map[string]index.Index, len(configs))

	for n, c := range configs {
		indexes[n], err = getIndex(ctx, es, c)
		if err != nil {
			return
		}
	}

	return
}
<commit_msg>Cleanup EnsureIndexes.<commit_after>package elasticsearch

import (
	"context"
	"github.com/ipfs-search/ipfs-search/index"
	"github.com/olivere/elastic/v6"
	"log"
)

func EnsureIndexes(ctx context.Context, esURL string, configs map[string]*index.Config) (indexes map[string]index.Index, err error) {
	es, err := elastic.NewClient(elastic.SetSniff(false), elastic.SetURL(esURL))
	if err != nil {
		return
	}

	log.Printf("Connected to ElasticSearch")

	indexes = make(map[string]index.Index, len(configs))

	for n, c := range configs {
		i := &Index{
			Client: es,
			Config: c,
		}

		if err = index.EnsureUpdated(ctx, i, c); err != nil {
			return
		}

		indexes[n] = i
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/omarqazi/broadcast/datastore"
	"net/http"
	"strings"
	"time"
)

var allChannels map[string]*datastore.Channel = make(map[string]*datastore.Channel)

type PlaylistGenerator struct {
}

func (pl PlaylistGenerator) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	checkCORSHeader(w, r)
	var err error
	channelId := strings.TrimSuffix(r.URL.Path, ".m3u8")
	channel, ok := allChannels[channelId]
	if !ok { // If this is the first time the channel is requested
		channel, err = datastore.GetChannel(channelId)
		if err != nil {
			http.Error(w, "internal server error", 500)
			return
		}

		go channel.AdvanceEvery(5*time.Second, nil)
		allChannels[channelId] = channel
	}

	fmt.Fprintln(w, channel.PlaylistData())
}

func checkCORSHeader(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get("Origin")
	url, err := url.Parse(origin)
	if err != nil {
		errorOut(w, r, "Invalid Origin header")
	}
	host := url.Host
	w.Header().Add("Access-Control-Allow-Origin", origin)
}
<commit_msg>fix errors<commit_after>package main

import (
	"fmt"
	"github.com/omarqazi/broadcast/datastore"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var allChannels map[string]*datastore.Channel = make(map[string]*datastore.Channel)

type PlaylistGenerator struct {
}

func (pl PlaylistGenerator) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	checkCORSHeader(w, r)
	var err error
	channelId := strings.TrimSuffix(r.URL.Path, ".m3u8")
	channel, ok := allChannels[channelId]
	if !ok { // If this is the first time the channel is requested
		channel, err = datastore.GetChannel(channelId)
		if err != nil {
			http.Error(w, "internal server error", 500)
			return
		}

		go channel.AdvanceEvery(5*time.Second, nil)
		allChannels[channelId] = channel
	}

	fmt.Fprintln(w, channel.PlaylistData())
}

func checkCORSHeader(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get("Origin")
	url, err := url.Parse(origin)
	if err != nil {
		http.Error(w, "Invalid origin header", 400)
	}
	host := url.Host
	w.Header().Add("Access-Control-Allow-Origin", host)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/vlifesystems/rulehuntersrv/internal/testhelpers"
	"github.com/vlifesystems/rulehuntersrv/logger"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"syscall"
	"testing"
	"time"
)

func TestSubMain_errors(t *testing.T) {
	tmpDir, err := ioutil.TempDir("", "rulehuntersrv")
	if err != nil {
		t.Fatal("TempDir() couldn't create dir")
	}
	defer os.RemoveAll(tmpDir)

	cases := []struct {
		flags        *cmdFlags
		wantErr      error
		wantExitCode int
	}{
		{
			flags: &cmdFlags{
				user:      "fred",
				configDir: "",
				install:   true,
			},
			wantErr:      errNoConfigDirArg,
			wantExitCode: 1,
		},
		{
			flags: &cmdFlags{
				user:      "fred",
				configDir: tmpDir,
				install:   true,
			},
			wantErr: errConfigLoad{
				filename: filepath.Join(tmpDir, "config.json"),
				err: &os.PathError{
					"open",
					filepath.Join(tmpDir, "config.json"),
					syscall.ENOENT,
				},
			},
			wantExitCode: 1,
		},
	}

	for _, c := range cases {
		l := logger.NewTestLogger()
		exitCode, err := subMain(c.flags, l)
		if exitCode != c.wantExitCode {
			t.Errorf("subMain(%q) exitCode: %d, want: %d",
				c.flags, exitCode, c.wantExitCode)
		}
		if err := checkErrorMatch(err, c.wantErr); err != nil {
			t.Errorf("subMain(%q) %s", c.flags, err)
		}
		if len(l.GetEntries()) != 0 {
			t.Errorf("GetEntries() got: %s, want: {}", l.GetEntries())
		}
	}
}

func TestSubMain(t *testing.T) {
	cases := []struct {
		flags        *cmdFlags
		wantErr      error
		wantExitCode int
		wantEntries  []logger.Entry
	}{
		{
			flags: &cmdFlags{
				user:    "fred",
				install: false,
			},
			wantErr:      nil,
			wantExitCode: 0,
			wantEntries: []logger.Entry{
				{logger.Info, "Waiting for experiments to process"},
			},
		},
	}

	wd, err := os.Getwd()
	if err != nil {
		t.Fatalf("Getwd() err: ", err)
	}
	defer os.Chdir(wd)

	for _, c := range cases {
		configDir, err := testhelpers.BuildConfigDirs()
		if err != nil {
			t.Fatalf("buildConfigDirs() err: %s", err)
		}
		defer os.RemoveAll(configDir)
		c.flags.configDir = configDir

		l := logger.NewTestLogger()
		go func() {
			tryInSeconds := 4
			for i := 0; i < tryInSeconds*5; i++ {
				if reflect.DeepEqual(l.GetEntries(), c.wantEntries) {
					interruptProcess(t)
					return
				}
				time.Sleep(200 * time.Millisecond)
			}
			interruptProcess(t)
		}()

		go func() {
			<-time.After(6 * time.Second)
			t.Fatal("Run() hasn't been stopped")
			if runtime.GOOS == "windows" {
				fmt.Fprintln(os.Stderr, "Run() hasn't been stopped")
				os.Exit(1)
			}
		}()
		if err := os.Chdir(configDir); err != nil {
			t.Fatalf("Chdir() err: %s", err)
		}
		exitCode, err := subMain(c.flags, l)
		if exitCode != c.wantExitCode {
			t.Errorf("subMain(%q) exitCode: %d, want: %d",
				c.flags, exitCode, c.wantExitCode)
		}
		if err := checkErrorMatch(err, c.wantErr); err != nil {
			t.Errorf("subMain(%q) %s", c.flags, err)
		}
		if !reflect.DeepEqual(l.GetEntries(), c.wantEntries) {
			t.Errorf("GetEntries() got: %s, want: %s", l.GetEntries(), c.wantEntries)
		}
	}
}

/*************************************
 *  Helper functions
 *************************************/

func checkErrorMatch(got, want error) error {
	if got == nil && want == nil {
		return nil
	}
	if got == nil || want == nil {
		return fmt.Errorf("got err: %s, want err: %s", got, want)
	}
	switch x := want.(type) {
	case *os.PathError:
		return checkPathErrorMatch(got, x)
	case errConfigLoad:
		return checkErrConfigLoadMatch(got, x)
	}
	if got.Error() != want.Error() {
		return fmt.Errorf("got err: %s, want err: %s", got, want)
	}
	return nil
}

func checkPathErrorMatch(checkErr error, wantErr error) error {
	cerr, ok := checkErr.(*os.PathError)
	if !ok {
		return fmt.Errorf("got err type: %T, want error type: os.PathError",
			checkErr)
	}
	werr, ok := wantErr.(*os.PathError)
	if !ok {
		panic("wantErr isn't type *os.PathError")
	}
	if cerr.Op != werr.Op {
		return fmt.Errorf("got cerr.Op: %s, want: %s", cerr.Op, werr.Op)
	}
	if filepath.Clean(cerr.Path) != filepath.Clean(werr.Path) {
		return fmt.Errorf("got cerr.Path: %s, want: %s", cerr.Path, werr.Path)
	}
	if cerr.Err != werr.Err {
		return fmt.Errorf("got cerr.Err: %s, want: %s", cerr.Err, werr.Err)
	}
	return nil
}

func checkErrConfigLoadMatch(checkErr error, wantErr errConfigLoad) error {
	cerr, ok := checkErr.(errConfigLoad)
	if !ok {
		return fmt.Errorf("got err type: %T, want error type: errConfigLoad",
			checkErr)
	}
	if filepath.Clean(cerr.filename) != filepath.Clean(wantErr.filename) {
		return fmt.Errorf("got cerr.Path: %s, want: %s",
			cerr.filename, wantErr.filename)
	}
	return checkPathErrorMatch(cerr.err, wantErr.err)
}
<commit_msg>Try locking os thread<commit_after>package main

import (
	"fmt"
	"github.com/vlifesystems/rulehuntersrv/internal/testhelpers"
	"github.com/vlifesystems/rulehuntersrv/logger"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"syscall"
	"testing"
	"time"
)

func init() {
	runtime.LockOSThread()
}

func TestMain(m *testing.M) {
	fmt.Printf("TestMain pid: %d\n", os.Getpid())
	os.Exit(m.Run())
}

func TestSubMain_errors(t *testing.T) {
	tmpDir, err := ioutil.TempDir("", "rulehuntersrv")
	if err != nil {
		t.Fatal("TempDir() couldn't create dir")
	}
	defer os.RemoveAll(tmpDir)

	cases := []struct {
		flags        *cmdFlags
		wantErr      error
		wantExitCode int
	}{
		{
			flags: &cmdFlags{
				user:      "fred",
				configDir: "",
				install:   true,
			},
			wantErr:      errNoConfigDirArg,
			wantExitCode: 1,
		},
		{
			flags: &cmdFlags{
				user:      "fred",
				configDir: tmpDir,
				install:   true,
			},
			wantErr: errConfigLoad{
				filename: filepath.Join(tmpDir, "config.json"),
				err: &os.PathError{
					"open",
					filepath.Join(tmpDir, "config.json"),
					syscall.ENOENT,
				},
			},
			wantExitCode: 1,
		},
	}

	for _, c := range cases {
		l := logger.NewTestLogger()
		exitCode, err := subMain(c.flags, l)
		if exitCode != c.wantExitCode {
			t.Errorf("subMain(%q) exitCode: %d, want: %d",
				c.flags, exitCode, c.wantExitCode)
		}
		if err := checkErrorMatch(err, c.wantErr); err != nil {
			t.Errorf("subMain(%q) %s", c.flags, err)
		}
		if len(l.GetEntries()) != 0 {
			t.Errorf("GetEntries() got: %s, want: {}", l.GetEntries())
		}
	}
}

func TestSubMain(t *testing.T) {
	cases := []struct {
		flags        *cmdFlags
		wantErr      error
		wantExitCode int
		wantEntries  []logger.Entry
	}{
		{
			flags: &cmdFlags{
				user:    "fred",
				install: false,
			},
			wantErr:      nil,
			wantExitCode: 0,
			wantEntries: []logger.Entry{
				{logger.Info, "Waiting for experiments to process"},
			},
		},
	}

	wd, err := os.Getwd()
	if err != nil {
		t.Fatalf("Getwd() err: ", err)
	}
	defer os.Chdir(wd)

	for _, c := range cases {
		configDir, err := testhelpers.BuildConfigDirs()
		if err != nil {
			t.Fatalf("buildConfigDirs() err: %s", err)
		}
		defer os.RemoveAll(configDir)
		c.flags.configDir = configDir

		l := logger.NewTestLogger()
		go func() {
			fmt.Printf("TestSubMain go func pid: %d\n", os.Getpid())
			tryInSeconds := 4
			for i := 0; i < tryInSeconds*5; i++ {
				if reflect.DeepEqual(l.GetEntries(), c.wantEntries) {
					interruptProcess(t)
					return
				}
				time.Sleep(200 * time.Millisecond)
			}
			interruptProcess(t)
		}()

		go func() {
			<-time.After(6 * time.Second)
			t.Fatal("Run() hasn't been stopped")
			if runtime.GOOS == "windows" {
				fmt.Fprintln(os.Stderr, "Run() hasn't been stopped")
				os.Exit(1)
			}
		}()
		if err := os.Chdir(configDir); err != nil {
			t.Fatalf("Chdir() err: %s", err)
		}
		exitCode, err := subMain(c.flags, l)
		if exitCode != c.wantExitCode {
			t.Errorf("subMain(%q) exitCode: %d, want: %d",
				c.flags, exitCode, c.wantExitCode)
		}
		if err := checkErrorMatch(err, c.wantErr); err != nil {
			t.Errorf("subMain(%q) %s", c.flags, err)
		}
		if !reflect.DeepEqual(l.GetEntries(), c.wantEntries) {
			t.Errorf("GetEntries() got: %s, want: %s", l.GetEntries(), c.wantEntries)
		}
	}
}

/*************************************
 *  Helper functions
 *************************************/

func checkErrorMatch(got, want error) error {
	if got == nil && want == nil {
		return nil
	}
	if got == nil || want == nil {
		return fmt.Errorf("got err: %s, want err: %s", got, want)
	}
	switch x := want.(type) {
	case *os.PathError:
		return checkPathErrorMatch(got, x)
	case errConfigLoad:
		return checkErrConfigLoadMatch(got, x)
	}
	if got.Error() != want.Error() {
		return fmt.Errorf("got err: %s, want err: %s", got, want)
	}
	return nil
}

func checkPathErrorMatch(checkErr error, wantErr error) error {
	cerr, ok := checkErr.(*os.PathError)
	if !ok {
		return fmt.Errorf("got err type: %T, want error type: os.PathError",
			checkErr)
	}
	werr, ok := wantErr.(*os.PathError)
	if !ok {
		panic("wantErr isn't type *os.PathError")
	}
	if cerr.Op != werr.Op {
		return fmt.Errorf("got cerr.Op: %s, want: %s", cerr.Op, werr.Op)
	}
	if filepath.Clean(cerr.Path) != filepath.Clean(werr.Path) {
		return fmt.Errorf("got cerr.Path: %s, want: %s", cerr.Path, werr.Path)
	}
	if cerr.Err != werr.Err {
		return fmt.Errorf("got cerr.Err: %s, want: %s", cerr.Err, werr.Err)
	}
	return nil
}

func checkErrConfigLoadMatch(checkErr error, wantErr errConfigLoad) error {
	cerr, ok := checkErr.(errConfigLoad)
	if !ok {
		return fmt.Errorf("got err type: %T, want error type: errConfigLoad",
			checkErr)
	}
	if filepath.Clean(cerr.filename) != filepath.Clean(wantErr.filename) {
		return fmt.Errorf("got cerr.Path: %s, want: %s",
			cerr.filename, wantErr.filename)
	}
	return checkPathErrorMatch(cerr.err, wantErr.err)
}
<|endoftext|>


<commit_before>package redis

import (
	"flag"
	"os"
	"testing"
	"time"

	"github.com/docker/distribution/registry/storage/cache/cachecheck"
	"github.com/garyburd/redigo/redis"
)

var redisAddr string

func init() {
	flag.StringVar(&redisAddr, "test.registry.storage.cache.redis.addr", "", "configure the address of a test instance of redis")
}

// TestRedisLayerInfoCache exercises a live redis instance using the cache
// implementation.
func TestRedisBlobDescriptorCacheProvider(t *testing.T) {
	if redisAddr == "" {
		// fallback to an environement variable
		redisAddr = os.Getenv("TEST_REGISTRY_STORAGE_CACHE_REDIS_ADDR")
	}

	if redisAddr == "" {
		// skip if still not set
		t.Skip("please set -registry.storage.cache.redis to test layer info cache against redis")
	}

	pool := &redis.Pool{
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", redisAddr)
		},
		MaxIdle:   1,
		MaxActive: 2,
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
		Wait: false, // if a connection is not avialable, proceed without cache.
	}

	// Clear the database
	if _, err := pool.Get().Do("FLUSHDB"); err != nil {
		t.Fatalf("unexpected error flushing redis db: %v", err)
	}

	cachecheck.CheckBlobDescriptorCache(t, NewRedisBlobDescriptorCacheProvider(pool))
}
<commit_msg>Fix connection pool exhaustion in Redis tests<commit_after>package redis

import (
	"flag"
	"os"
	"testing"
	"time"

	"github.com/docker/distribution/registry/storage/cache/cachecheck"
	"github.com/garyburd/redigo/redis"
)

var redisAddr string

func init() {
	flag.StringVar(&redisAddr, "test.registry.storage.cache.redis.addr", "", "configure the address of a test instance of redis")
}

// TestRedisLayerInfoCache exercises a live redis instance using the cache
// implementation.
func TestRedisBlobDescriptorCacheProvider(t *testing.T) {
	if redisAddr == "" {
		// fallback to an environement variable
		redisAddr = os.Getenv("TEST_REGISTRY_STORAGE_CACHE_REDIS_ADDR")
	}

	if redisAddr == "" {
		// skip if still not set
		t.Skip("please set -test.registry.storage.cache.redis.addr to test layer info cache against redis")
	}

	pool := &redis.Pool{
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", redisAddr)
		},
		MaxIdle:   1,
		MaxActive: 2,
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
		Wait: false, // if a connection is not avialable, proceed without cache.
	}

	// Clear the database
	conn := pool.Get()
	if _, err := conn.Do("FLUSHDB"); err != nil {
		t.Fatalf("unexpected error flushing redis db: %v", err)
	}
	conn.Close()

	cachecheck.CheckBlobDescriptorCache(t, NewRedisBlobDescriptorCacheProvider(pool))
}
<|endoftext|>


<commit_before>package sampler

import (
	"sync"
	"time"
)

const (
	DefaultIntent = iota
	TimeIntent
)

const (
	StatsType = iota
	CounterType
)

type Fields map[string]interface{}

type Entry struct {
	Metric *Metric
	Time   time.Time
	Fields map[string]interface{}
	Value  int64
}

func (e *Entry) WithField(key string, value interface{}) *Entry {
	e.Fields[key] = value
	return e
}

func (e *Entry) WithFields(fields Fields) *Entry {
	for key, value := range fields {
		e.Fields[key] = value
	}
	return e
}

func (e *Entry) Int(v int) {
	e.Value = int64(v)
	e.Metric.Write(e)
}

func (e *Entry) Int64(v int64) {
	e.Value = v
	e.Metric.Write(e)
}

func (e *Entry) Duration(d time.Duration) {
	e.Value = d.Nanoseconds()
	e.Metric.Intent = TimeIntent
	e.Metric.Write(e)
}

type Metric struct {
	Sampler *Sampler
	Entries []*Entry

	Type   int
	Intent int

	entryMutex sync.Mutex
}

func (m *Metric) Entry() *Entry {
	return &Entry{Metric: m, Fields: make(map[string]interface{})}
}

func (m *Metric) WithField(key string, value interface{}) *Entry {
	return m.Entry().WithField(key, value)
}

func (m *Metric) WithFields(fields Fields) *Entry {
	return m.Entry().WithFields(fields)
}

func (m *Metric) Int(v int) {
	m.Entry().Int(v)
}

func (m *Metric) Int64(v int64) {
	m.Entry().Int64(v)
}

func (m *Metric) Duration(d time.Duration) {
	m.Entry().Duration(d)
}

func (m *Metric) Write(e *Entry) {
	m.entryMutex.Lock()
	defer m.entryMutex.Unlock()

	m.Entries = append(m.Entries, e)
}

func (m *Metric) Min() int64 {
	var min int64
	for _, e := range m.Entries {
		if min == 0 || e.Value < min {
			min = e.Value
		}
	}
	return min
}

func (m *Metric) Max() int64 {
	var max int64
	for _, e := range m.Entries {
		if e.Value > max {
			max = e.Value
		}
	}
	return max
}

func (m *Metric) Avg() int64 {
	var sum int64
	for _, e := range m.Entries {
		sum += e.Value
	}
	return sum / int64(len(m.Entries))
}

func (m *Metric) Med() int64 {
	return m.Entries[len(m.Entries)/2].Value
}

type Sampler struct {
	Metrics map[string]*Metric

	MetricMutex sync.Mutex
}

func New() *Sampler {
	return &Sampler{Metrics: make(map[string]*Metric)}
}

func (s *Sampler) Get(name string) *Metric {
	s.MetricMutex.Lock()
	defer s.MetricMutex.Unlock()

	metric, ok := s.Metrics[name]
	if !ok {
		metric = &Metric{Sampler: s}
		s.Metrics[name] = metric
	}
	return metric
}

func (s *Sampler) GetAs(name string, t int) *Metric {
	m := s.Get(name)
	m.Type = t
	return m
}

func (s *Sampler) Stats(name string) *Metric {
	return s.GetAs(name, StatsType)
}

func (s *Sampler) Counter(name string) *Metric {
	return s.GetAs(name, CounterType)
}
<commit_msg>[fix] Divide by zero<commit_after>package sampler

import (
	"sync"
	"time"
)

const (
	DefaultIntent = iota
	TimeIntent
)

const (
	StatsType = iota
	CounterType
)

type Fields map[string]interface{}

type Entry struct {
	Metric *Metric
	Time   time.Time
	Fields map[string]interface{}
	Value  int64
}

func (e *Entry) WithField(key string, value interface{}) *Entry {
	e.Fields[key] = value
	return e
}

func (e *Entry) WithFields(fields Fields) *Entry {
	for key, value := range fields {
		e.Fields[key] = value
	}
	return e
}

func (e *Entry) Int(v int) {
	e.Value = int64(v)
	e.Metric.Write(e)
}

func (e *Entry) Int64(v int64) {
	e.Value = v
	e.Metric.Write(e)
}

func (e *Entry) Duration(d time.Duration) {
	e.Value = d.Nanoseconds()
	e.Metric.Intent = TimeIntent
	e.Metric.Write(e)
}

type Metric struct {
	Sampler *Sampler
	Entries []*Entry

	Type   int
	Intent int

	entryMutex sync.Mutex
}

func (m *Metric) Entry() *Entry {
	return &Entry{Metric: m, Fields: make(map[string]interface{})}
}

func (m *Metric) WithField(key string, value interface{}) *Entry {
	return m.Entry().WithField(key, value)
}

func (m *Metric) WithFields(fields Fields) *Entry {
	return m.Entry().WithFields(fields)
}

func (m *Metric) Int(v int) {
	m.Entry().Int(v)
}

func (m *Metric) Int64(v int64) {
	m.Entry().Int64(v)
}

func (m *Metric) Duration(d time.Duration) {
	m.Entry().Duration(d)
}

func (m *Metric) Write(e *Entry) {
	m.entryMutex.Lock()
	defer m.entryMutex.Unlock()

	m.Entries = append(m.Entries, e)
}

func (m *Metric) Min() int64 {
	var min int64
	for _, e := range m.Entries {
		if min == 0 || e.Value < min {
			min = e.Value
		}
	}
	return min
}

func (m *Metric) Max() int64 {
	var max int64
	for _, e := range m.Entries {
		if e.Value > max {
			max = e.Value
		}
	}
	return max
}

func (m *Metric) Avg() int64 {
	if len(m.Entries) == 0 {
		return 0
	}

	var sum int64
	for _, e := range m.Entries {
		sum += e.Value
	}
	return sum / int64(len(m.Entries))
}

func (m *Metric) Med() int64 {
	return m.Entries[len(m.Entries)/2].Value
}

type Sampler struct {
	Metrics map[string]*Metric

	MetricMutex sync.Mutex
}

func New() *Sampler {
	return &Sampler{Metrics: make(map[string]*Metric)}
}

func (s *Sampler) Get(name string) *Metric {
	s.MetricMutex.Lock()
	defer s.MetricMutex.Unlock()

	metric, ok := s.Metrics[name]
	if !ok {
		metric = &Metric{Sampler: s}
		s.Metrics[name] = metric
	}
	return metric
}

func (s *Sampler) GetAs(name string, t int) *Metric {
	m := s.Get(name)
	m.Type = t
	return m
}

func (s *Sampler) Stats(name string) *Metric {
	return s.GetAs(name, StatsType)
}

func (s *Sampler) Counter(name string) *Metric {
	return s.GetAs(name, CounterType)
}
<|endoftext|>


<commit_before>package plt

import (
	"code.google.com/p/plotinum/vecgfx"
	"image/color"
	"fmt"
	"math"
)

const (
	DefaultFont = "Times-Roman"
)

type Axis struct{
	// Min and Max are the minimum and maximum data
	// coordinates on this axis.
	Min, Max float64

	// Label is the axis label
	Label string

	// LabelStyle is the text style of the label on the axis.
	LabelStyle TextStyle

	// AxisStyle is the style of the axis's line.
	AxisStyle LineStyle

	// Padding between the axis line and the data in inches.
	Padding float64

	// Ticks are the tick marks on the axis.
	Ticks TickMarks
}

// MakeAxis returns a default axis.
func MakeAxis() Axis {
	labelFont, err := MakeFont(DefaultFont, 12)
	if err != nil {
		panic(err)
	}
	return Axis{
		Min: math.Inf(1),
		Max: math.Inf(-1),
		Label: "",
		LabelStyle: TextStyle{
			Color: Black,
			Font: labelFont,
		},
		AxisStyle: LineStyle{
			Color: Black,
			Width: 1.0/64.0,
		},
		Padding: 1.0/8.0,
		Ticks: MakeTickMarks(),
	}
}

// X transfroms the data point x to the drawing coordinate
// for the given drawing area.
func (a *Axis) X(da *DrawArea, x float64) float64 {
	p := (x - a.Min) / (a.Max - a.Min)
	return da.Min.X + p*(da.Max().X - da.Min.X)
}

// Y transforms the data point y to the drawing coordinate
// for the given drawing area.
func (a *Axis) Y(da *DrawArea, y float64) float64 {
	p := (y - a.Min) / (a.Max - a.Min)
	return da.Min.Y + p*(da.Max().Y - da.Min.Y)
}

// height returns the height of the axis in inches
//  if it is drawn as a horizontal axis.
func (a *Axis) height() (h float64) {
	if a.Label != "" {
		h += a.LabelStyle.Font.Extents().Height/vecgfx.PtInch
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		h += a.Ticks.Length + a.Ticks.labelHeight(marks)
	}
	h += a.AxisStyle.Width/2
	h += a.Padding
	return
}

// drawHoriz draws the axis onto the given area.
func (a *Axis) drawHoriz(da *DrawArea) {
	y := da.Min.Y
	if a.Label != "" {
		da.SetTextStyle(a.LabelStyle)
		y += -(a.LabelStyle.Font.Extents().Descent/vecgfx.PtInch * da.DPI())
		da.Text(da.Center().X, y, -0.5, 0, a.Label)
		y += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch * da.DPI()
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		da.SetLineStyle(a.Ticks.MarkStyle)
		da.SetTextStyle(a.Ticks.LabelStyle)
		for _, t := range marks {
			if t.Label == "" {
				continue
			}
			da.Text(a.X(da, t.Value), y, -0.5, 0, t.Label)
		}
		y += a.Ticks.labelHeight(marks) * da.DPI()

		len := a.Ticks.Length*da.DPI()
		for _, t := range marks {
			x := a.X(da, t.Value)
			y1 := y
			if t.Label == "" {
				y1 = y +  len/2
			}
			da.Line([]Point{{x, y1}, {x, y + len}})
		}
		y += len
	}
	da.SetLineStyle(a.AxisStyle)
	da.Line([]Point{{da.Min.X, y}, {da.Max().X, y}})
}

// width returns the width of the axis in inches
//  if it is drawn as a vertically axis.
func (a *Axis) width() (w float64) {
	if a.Label != "" {
		w += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		if lwidth := a.Ticks.labelWidth(marks); lwidth > 0 {
			w += lwidth
			// Add a space after tick labels to separate
			// them from the tick marks
			w += a.Ticks.LabelStyle.Font.Width(" ")/vecgfx.PtInch
		}
		w += a.Ticks.Length
	}
	w += a.AxisStyle.Width/2
	w += a.Padding
	return
}

// drawVert draws the axis onto the given area.
func (a *Axis) drawVert(da *DrawArea) {
	x := da.Min.X
	if a.Label != "" {
		x += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch * da.DPI()
		da.SetTextStyle(a.LabelStyle)
		da.Push()
		da.Rotate(math.Pi/2)
		da.Text(da.Center().Y, -x, -0.5, 0, a.Label)
		da.Pop()
		x += -a.LabelStyle.Font.Extents().Descent/vecgfx.PtInch * da.DPI()
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		da.SetLineStyle(a.Ticks.MarkStyle)
		da.SetTextStyle(a.Ticks.LabelStyle)
		if lwidth := a.Ticks.labelWidth(marks); lwidth > 0 {
			x += lwidth * da.DPI()
			x += a.Ticks.LabelStyle.Font.Width(" ")/vecgfx.PtInch * da.DPI()
		}
		for _, t := range marks {
			if t.Label == "" {
				continue
			}
			da.Text(x, a.Y(da, t.Value), -1, -0.5, t.Label + " ")
		}
		len := a.Ticks.Length*da.DPI()
		for _, t := range marks {
			y := a.Y(da, t.Value)
			x1 := x
			if t.Label == "" {
				x1 = x +  len/2
			}
			da.Line([]Point{{x1, y}, {x + len, y}})
		}
		x += len
	}
	da.SetLineStyle(a.AxisStyle)
	da.Line([]Point{{x, da.Min.Y}, {x, da.Max().Y}})
}

// TickMarks is the style and location of a set of tick marks.
type TickMarks struct {
	// LabelStyle is the text style on the tick labels.
	LabelStyle TextStyle

	// MarkStyle is the style of the tick mark lines.
	MarkStyle LineStyle

	// Length is the length of a major tick mark specified
	// in inches.
	Length float64

	TickMarker
}

// A TickMarker returns a slice of ticks between a given
// range of values.
type TickMarker interface{
	// Marks returns a slice of ticks for the given range.
	Marks(min, max float64) []Tick
}

// A Tick is a single tick mark
type Tick struct {
	Value float64
	Label string
}

// MakeTickMarks returns a TickMarks using the default style
// and TickMarker.
func MakeTickMarks() TickMarks {
	labelFont, err := MakeFont(DefaultFont, 10)
	if err != nil {
		panic(err)
	}
	return TickMarks{
		LabelStyle: TextStyle{
			Color: color.RGBA{A: 255},
			Font: labelFont,
		},
		MarkStyle: LineStyle{
			Color: color.RGBA{A:255},
			Width: 1.0/64.0,
		},
		Length: 1.0/10.0,
		TickMarker: DefaultTicks(struct{}{}),
	}
}
// labelHeight returns the label height in inches.
func (tick TickMarks) labelHeight(ticks []Tick) float64 {
	for _, t := range ticks {
		if t.Label == "" {
			continue
		}
		font := tick.LabelStyle.Font
		return font.Extents().Ascent/vecgfx.PtInch
	}
	return 0
}

// labelWidth returns the label width in inches.
func (tick TickMarks) labelWidth(ticks []Tick) float64 {
	maxWidth := 0.0
	for _, t := range ticks {
		if t.Label == "" {
			continue
		}
		w := tick.LabelStyle.Font.Width(t.Label)
		if w > maxWidth {
			maxWidth = w
		}
	}
	return maxWidth/vecgfx.PtInch
}

// A DefalutTicks returns a default set of tick marks within
// the given range.
type DefaultTicks struct{}

// Marks implements the TickMarker Marks method.
func (_ DefaultTicks) Marks(min, max float64) []Tick {
	return []Tick{
			{ Value: min, Label: fmt.Sprintf("%g", min) },
			{ Value: min + (max-min)/4, },
			{ Value: min + (max-min)/2, Label: fmt.Sprintf("%g", min + (max-min)/2) },
			{ Value: min + 3*(max-min)/4, },
			{ Value: max, Label: fmt.Sprintf("%g", max) },
	}
}

// A ConstTicks always returns the same set of tick marks.
type ConstantTicks []Tick

// Marks implements the TickMarker Marks method.
func (tks ConstantTicks) Marks(min, max float64) []Tick {
	return tks
}<commit_msg>Simplify Axis drawing by moving some of the smarts into separate methods.<commit_after>package plt

import (
	"code.google.com/p/plotinum/vecgfx"
	"image/color"
	"fmt"
	"math"
)

const (
	DefaultFont = "Times-Roman"
)

type Axis struct{
	// Min and Max are the minimum and maximum data
	// coordinates on this axis.
	Min, Max float64

	// Label is the axis label
	Label string

	// LabelStyle is the text style of the label on the axis.
	LabelStyle TextStyle

	// AxisStyle is the style of the axis's line.
	AxisStyle LineStyle

	// Padding between the axis line and the data in inches.
	Padding float64

	// Ticks are the tick marks on the axis.
	Ticks TickMarks
}

// MakeAxis returns a default axis.
func MakeAxis() Axis {
	labelFont, err := MakeFont(DefaultFont, 12)
	if err != nil {
		panic(err)
	}
	return Axis{
		Min: math.Inf(1),
		Max: math.Inf(-1),
		Label: "",
		LabelStyle: TextStyle{
			Color: Black,
			Font: labelFont,
		},
		AxisStyle: LineStyle{
			Color: Black,
			Width: 1.0/64.0,
		},
		Padding: 1.0/8.0,
		Ticks: MakeTickMarks(),
	}
}

// X transfroms the data point x to the drawing coordinate
// for the given drawing area.
func (a *Axis) X(da *DrawArea, x float64) float64 {
	p := (x - a.Min) / (a.Max - a.Min)
	return da.Min.X + p*(da.Max().X - da.Min.X)
}

// Y transforms the data point y to the drawing coordinate
// for the given drawing area.
func (a *Axis) Y(da *DrawArea, y float64) float64 {
	p := (y - a.Min) / (a.Max - a.Min)
	return da.Min.Y + p*(da.Max().Y - da.Min.Y)
}

// height returns the height of the axis in inches
//  if it is drawn as a horizontal axis.
func (a *Axis) height() (h float64) {
	if a.Label != "" {
		h += a.LabelStyle.Font.Extents().Height/vecgfx.PtInch
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		h += a.Ticks.Length + a.Ticks.labelHeight(marks)
	}
	h += a.AxisStyle.Width/2
	h += a.Padding
	return
}

// drawHoriz draws the axis onto the given area.
func (a *Axis) drawHoriz(da *DrawArea) {
	y := da.Min.Y
	if a.Label != "" {
		da.SetTextStyle(a.LabelStyle)
		y += -(a.LabelStyle.Font.Extents().Descent/vecgfx.PtInch * da.DPI())
		da.Text(da.Center().X, y, -0.5, 0, a.Label)
		y += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch * da.DPI()
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		da.SetLineStyle(a.Ticks.MarkStyle)
		da.SetTextStyle(a.Ticks.LabelStyle)
		for _, t := range marks {
			if t.minor() {
				continue
			}
			da.Text(a.X(da, t.Value), y, -0.5, 0, t.Label)
		}
		y += a.Ticks.labelHeight(marks) * da.DPI()

		len := a.Ticks.Length*da.DPI()
		for _, t := range marks {
			x := a.X(da, t.Value)
			da.Line([]Point{{x, y + t.lengthOffset(len)}, {x, y + len}})
		}
		y += len
	}
	da.SetLineStyle(a.AxisStyle)
	da.Line([]Point{{da.Min.X, y}, {da.Max().X, y}})
}

// width returns the width of the axis in inches
//  if it is drawn as a vertically axis.
func (a *Axis) width() (w float64) {
	if a.Label != "" {
		w += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		if lwidth := a.Ticks.labelWidth(marks); lwidth > 0 {
			w += lwidth
			// Add a space after tick labels to separate
			// them from the tick marks
			w += a.Ticks.LabelStyle.Font.Width(" ")/vecgfx.PtInch
		}
		w += a.Ticks.Length
	}
	w += a.AxisStyle.Width/2
	w += a.Padding
	return
}

// drawVert draws the axis onto the given area.
func (a *Axis) drawVert(da *DrawArea) {
	x := da.Min.X
	if a.Label != "" {
		x += a.LabelStyle.Font.Extents().Ascent/vecgfx.PtInch * da.DPI()
		da.SetTextStyle(a.LabelStyle)
		da.Push()
		da.Rotate(math.Pi/2)
		da.Text(da.Center().Y, -x, -0.5, 0, a.Label)
		da.Pop()
		x += -a.LabelStyle.Font.Extents().Descent/vecgfx.PtInch * da.DPI()
	}
	marks := a.Ticks.Marks(a.Min, a.Max)
	if len(marks) > 0 {
		da.SetLineStyle(a.Ticks.MarkStyle)
		da.SetTextStyle(a.Ticks.LabelStyle)
		if lwidth := a.Ticks.labelWidth(marks); lwidth > 0 {
			x += lwidth * da.DPI()
			x += a.Ticks.LabelStyle.Font.Width(" ")/vecgfx.PtInch * da.DPI()
		}
		for _, t := range marks {
			if t.minor() {
				continue
			}
			da.Text(x, a.Y(da, t.Value), -1, -0.5, t.Label + " ")
		}
		len := a.Ticks.Length*da.DPI()
		for _, t := range marks {
			y := a.Y(da, t.Value)
			da.Line([]Point{{x + t.lengthOffset(len), y}, {x + len, y}})
		}
		x += len
	}
	da.SetLineStyle(a.AxisStyle)
	da.Line([]Point{{x, da.Min.Y}, {x, da.Max().Y}})
}

// TickMarks specifies the style and location of the tick marks
// on an axis.
type TickMarks struct {
	// LabelStyle is the TextStyle on the tick labels.
	LabelStyle TextStyle

	// MarkStyle is the LineStyle of the tick mark lines.
	MarkStyle LineStyle

	// Length is the length of a major tick mark in inches.
	// Minor tick marks are half of the length of major
	// tick marks.
	Length float64

	// TickMarker locates the tick marks given the
	// minimum and maximum values.
	TickMarker
}

// A TickMarker returns a slice of ticks between a given
// range of values.
type TickMarker interface{
	// Marks returns a slice of ticks for the given range.
	Marks(min, max float64) []Tick
}

// A Tick is a single tick mark
type Tick struct {
	Value float64
	Label string
}

// minor returns true if this is a minor tick mark.
func (t Tick) minor() bool {
	return t.Label == ""
}

// lengthOffset returns an offset that should be added to the
// tick mark's line to accout for its length.  I.e., the start of
// the line for a minor tick mark must be shifted by half of
// the length.
func (t Tick) lengthOffset(len float64) float64 {
	if t.minor() {
		return len/2
	}
	return 0
}

// MakeTickMarks returns a TickMarks using the default style
// and TickMarker.
func MakeTickMarks() TickMarks {
	labelFont, err := MakeFont(DefaultFont, 10)
	if err != nil {
		panic(err)
	}
	return TickMarks{
		LabelStyle: TextStyle{
			Color: color.RGBA{A: 255},
			Font: labelFont,
		},
		MarkStyle: LineStyle{
			Color: color.RGBA{A:255},
			Width: 1.0/64.0,
		},
		Length: 1.0/10.0,
		TickMarker: DefaultTicks(struct{}{}),
	}
}
// labelHeight returns the label height in inches.
func (tick TickMarks) labelHeight(ticks []Tick) float64 {
	for _, t := range ticks {
		if t.minor() {
			continue
		}
		font := tick.LabelStyle.Font
		return font.Extents().Ascent/vecgfx.PtInch
	}
	return 0
}

// labelWidth returns the label width in inches.
func (tick TickMarks) labelWidth(ticks []Tick) float64 {
	maxWidth := 0.0
	for _, t := range ticks {
		if t.minor() {
			continue
		}
		w := tick.LabelStyle.Font.Width(t.Label)
		if w > maxWidth {
			maxWidth = w
		}
	}
	return maxWidth/vecgfx.PtInch
}

// A DefalutTicks returns a default set of tick marks within
// the given range.
type DefaultTicks struct{}

// Marks implements the TickMarker Marks method.
func (_ DefaultTicks) Marks(min, max float64) []Tick {
	return []Tick{
			{ Value: min, Label: fmt.Sprintf("%g", min) },
			{ Value: min + (max-min)/4, },
			{ Value: min + (max-min)/2, Label: fmt.Sprintf("%g", min + (max-min)/2) },
			{ Value: min + 3*(max-min)/4, },
			{ Value: max, Label: fmt.Sprintf("%g", max) },
	}
}

// A ConstTicks always returns the same set of tick marks.
type ConstantTicks []Tick

// Marks implements the TickMarker Marks method.
func (tks ConstantTicks) Marks(min, max float64) []Tick {
	return tks
}<|endoftext|>


<commit_before>package video

import "image"

var color_channel = [12]int{0, 0, 0, 0, 1, 2, 1, 2, 1, 2, 1, 2}

type Macroblock struct {
	macroblock_address_increment int
	macroblock_type              *MacroblockType
	spatial_temporal_weight_code uint32
	frame_motion_type            uint32
	field_motion_type            uint32
	dct_type                     bool

	cpb int
}

func (br *VideoSequence) macroblock(
	// location
	mb_address, mb_row int,
	// dct predictors
	dcp *dcDctPredictors, resetDCPredictors dcDctPredictorResetter,
	mvd *motionVectorData,
	qsc *uint32,
	frameSlice *image.YCbCr) (int, error) {

	mb := Macroblock{}

	for {
		if nextbits, err := br.Peek32(11); err != nil {
			return 0, err
		} else if nextbits == 0x08 { // 0000 0001 000
			br.Trash(11)
			mb.macroblock_address_increment += 33
		} else {
			break
		}
	}

	if incr, err := macroblockAddressIncrementDecoder.Decode(br); err != nil {
		return 0, err
	} else {
		mb.macroblock_address_increment += incr
	}

	if br.PictureHeader.picture_coding_type == PFrame && mb.macroblock_address_increment > 1 {
		copy_macroblocks(mb_row, mb_address+1, mb.macroblock_address_increment-1, frameSlice, br.frameStore.past)
	}

	// Reset dcDctPredictors: whenever a macroblock is skipped. (7.2.1)
	if mb.macroblock_address_increment > 1 {
		resetDCPredictors()
	}

	// Reset motion vector predictors: P-picture with a skipped macroblock (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_address_increment > 1 {
		mvd.reset()
	}

	if err := br.macroblock_mode(&mb); err != nil {
		return 0, err
	}

	// Reset dcDctPredictors: whenever a non-intra macroblock is decoded. (7.2.1)
	if mb.macroblock_type.macroblock_intra == false {
		resetDCPredictors()
	}

	// Reset motion vector predictors: intra macroblock without concealment motion vectors (7.6.4.3)
	if mb.macroblock_type.macroblock_intra == true &&
		br.PictureCodingExtension.concealment_motion_vectors == false {
		mvd.reset()
	}

	// Reset motion vector predictors: non-intra P-picture with no forward motion vectors (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_type.macroblock_intra == false &&
		mb.macroblock_type.macroblock_motion_forward == false {
		mvd.reset()
	}

	if mb.macroblock_type.macroblock_quant {
		if mb_qsc, err := br.Read32(5); err != nil {
			return 0, err
		} else {
			*qsc = mb_qsc
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		(mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors) {
		if err := br.motion_vectors(0, &mb, mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_motion_backward {
		if err := br.motion_vectors(1, &mb, mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors {
		if err := marker_bit(br); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		if cpb, err := coded_block_pattern(br, br.SequenceExtension.chroma_format); err != nil {
			return 0, nil
		} else {
			mb.cpb = cpb
		}
	}

	var block_count int
	switch br.SequenceExtension.chroma_format {
	case ChromaFormat_420:
		block_count = 6
	case ChromaFormat_422:
		block_count = 8
	case ChromaFormat_444:
		block_count = 12
	}

	mb_address += mb.macroblock_address_increment
	pattern_code := mb.decodePatternCode(br.SequenceExtension.chroma_format)

	var b block

	for i := 0; i < block_count; i++ {
		cc := color_channel[i]

		if pattern_code[i] {
			if err := b.read(br, dcp, br.PictureCodingExtension.intra_vlc_format, cc, mb.macroblock_type.macroblock_intra); err != nil {
				return 0, err
			}
			b.decode_block(br, cc, *qsc, mb.macroblock_type.macroblock_intra)
			b.idct()
		} else {
			b.empty()
		}

		br.motion_compensation(mvd, i, mb_row, mb_address, &mb, &b)
		updateFrameSlice(i, mb_address, mb.dct_type, frameSlice, &b)
	}

	return mb_address, nil
}

func updateFrameSlice(i, mb_address int, interlaced bool, frameSlice *image.YCbCr, b *block) {

	var cb clampedblock
	b.clamp(&cb)

	var (
		base_i  int
		channel []uint8
		stride  int
	)

	// channel switch
	switch i {
	case 0, 1, 2, 3:
		channel = frameSlice.Y
	case 4:
		channel = frameSlice.Cb
	case 5:
		channel = frameSlice.Cr
	}

	// base address and stride switch
	switch i {
	case 0, 1, 2, 3:
		stride = frameSlice.YStride
		base_i = mb_address * 16
	case 4, 5:
		stride = frameSlice.CStride
		base_i = mb_address * 8
	}

	// position switch
	if interlaced {
		// Field DCT coding alternates lines from each block:
		//
		//  <-8px-> <-8px->
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		switch i {
		case 0, 1, 2, 3:
			base_i += (i & 1) << 3
			base_i += ((i & 2) >> 1) * stride
			stride *= 2
		}
	} else {
		// Frame DCT coding are mapped in the follow order:
		//
		//  <-8px-> <-8px->
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───2───│───3───
		//  ───2───│───3───
		//  ───2───│───3───
		switch i {
		case 0, 1, 2, 3:
			base_i += (i & 1) << 3            // horiztonal positioning
			base_i += ((i & 2) << 2) * stride // vertical positioning
		}
	}

	// perform copy
	for y := 0; y < 8; y++ {
		si := y * 8
		di := base_i + (y * stride)
		copy(channel[di:di+8], cb[si:si+8])
	}

}

type PatternCode [12]bool

func (mb *Macroblock) decodePatternCode(chroma_format chromaFormat) (pattern_code PatternCode) {
	for i := 0; i < 12; i++ {
		if mb.macroblock_type.macroblock_intra {
			pattern_code[i] = true
		} else {
			pattern_code[i] = false
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		for i := 0; i < 6; i++ {
			if mask := 1 << uint(5-i); mb.cpb&mask == mask {
				pattern_code[i] = true
			}
		}

		if chroma_format == ChromaFormat_422 || chroma_format == ChromaFormat_444 {
			panic("unsupported: coded block pattern chroma format")
		}
	}

	return
}

func (br *VideoSequence) macroblock_mode(mb *Macroblock) (err error) {

	var typeDecoder macroblockTypeDecoderFn
	switch br.PictureHeader.picture_coding_type {
	case IntraCoded:
		typeDecoder = macroblockTypeDecoder.IFrame
	case PredictiveCoded:
		typeDecoder = macroblockTypeDecoder.PFrame
	case BidirectionallyPredictiveCoded:
		typeDecoder = macroblockTypeDecoder.BFrame
	default:
		panic("not implemented: macroblock type decoder")
	}

	mb.macroblock_type, err = typeDecoder(br)
	if err != nil {
		return err
	}

	if mb.macroblock_type.spatial_temporal_weight_code_flag &&
		false /* ( spatial_temporal_weight_code_table_index != ‘00’) */ {
		mb.spatial_temporal_weight_code, err = br.Read32(2)
		if err != nil {
			return err
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		mb.macroblock_type.macroblock_motion_backward {
		if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture {
			if br.PictureCodingExtension.frame_pred_frame_dct == 0 {
				mb.frame_motion_type, err = br.Read32(2)
				if err != nil {
					return err
				}
			}
		} else {
			mb.field_motion_type, err = br.Read32(2)
			if err != nil {
				return err
			}
		}
	}

	if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture &&
		br.PictureCodingExtension.frame_pred_frame_dct == 0 &&
		(mb.macroblock_type.macroblock_intra || mb.macroblock_type.macroblock_pattern) {
		mb.dct_type, err = br.ReadBit() //dct_type 1 uimsbf
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>reduce clampedblock allocations<commit_after>package video

import "image"

var color_channel = [12]int{0, 0, 0, 0, 1, 2, 1, 2, 1, 2, 1, 2}

type Macroblock struct {
	macroblock_address_increment int
	macroblock_type              *MacroblockType
	spatial_temporal_weight_code uint32
	frame_motion_type            uint32
	field_motion_type            uint32
	dct_type                     bool

	cpb int
}

func (br *VideoSequence) macroblock(
	// location
	mb_address, mb_row int,
	// dct predictors
	dcp *dcDctPredictors, resetDCPredictors dcDctPredictorResetter,
	mvd *motionVectorData,
	qsc *uint32,
	frameSlice *image.YCbCr) (int, error) {

	mb := Macroblock{}

	for {
		if nextbits, err := br.Peek32(11); err != nil {
			return 0, err
		} else if nextbits == 0x08 { // 0000 0001 000
			br.Trash(11)
			mb.macroblock_address_increment += 33
		} else {
			break
		}
	}

	if incr, err := macroblockAddressIncrementDecoder.Decode(br); err != nil {
		return 0, err
	} else {
		mb.macroblock_address_increment += incr
	}

	if br.PictureHeader.picture_coding_type == PFrame && mb.macroblock_address_increment > 1 {
		copy_macroblocks(mb_row, mb_address+1, mb.macroblock_address_increment-1, frameSlice, br.frameStore.past)
	}

	// Reset dcDctPredictors: whenever a macroblock is skipped. (7.2.1)
	if mb.macroblock_address_increment > 1 {
		resetDCPredictors()
	}

	// Reset motion vector predictors: P-picture with a skipped macroblock (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_address_increment > 1 {
		mvd.reset()
	}

	if err := br.macroblock_mode(&mb); err != nil {
		return 0, err
	}

	// Reset dcDctPredictors: whenever a non-intra macroblock is decoded. (7.2.1)
	if mb.macroblock_type.macroblock_intra == false {
		resetDCPredictors()
	}

	// Reset motion vector predictors: intra macroblock without concealment motion vectors (7.6.4.3)
	if mb.macroblock_type.macroblock_intra == true &&
		br.PictureCodingExtension.concealment_motion_vectors == false {
		mvd.reset()
	}

	// Reset motion vector predictors: non-intra P-picture with no forward motion vectors (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_type.macroblock_intra == false &&
		mb.macroblock_type.macroblock_motion_forward == false {
		mvd.reset()
	}

	if mb.macroblock_type.macroblock_quant {
		if mb_qsc, err := br.Read32(5); err != nil {
			return 0, err
		} else {
			*qsc = mb_qsc
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		(mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors) {
		if err := br.motion_vectors(0, &mb, mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_motion_backward {
		if err := br.motion_vectors(1, &mb, mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors {
		if err := marker_bit(br); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		if cpb, err := coded_block_pattern(br, br.SequenceExtension.chroma_format); err != nil {
			return 0, nil
		} else {
			mb.cpb = cpb
		}
	}

	var block_count int
	switch br.SequenceExtension.chroma_format {
	case ChromaFormat_420:
		block_count = 6
	case ChromaFormat_422:
		block_count = 8
	case ChromaFormat_444:
		block_count = 12
	}

	mb_address += mb.macroblock_address_increment
	pattern_code := mb.decodePatternCode(br.SequenceExtension.chroma_format)

	var b block
	var cb clampedblock

	for i := 0; i < block_count; i++ {
		cc := color_channel[i]

		if pattern_code[i] {
			if err := b.read(br, dcp, br.PictureCodingExtension.intra_vlc_format, cc, mb.macroblock_type.macroblock_intra); err != nil {
				return 0, err
			}
			b.decode_block(br, cc, *qsc, mb.macroblock_type.macroblock_intra)
			b.idct()
		} else {
			b.empty()
		}

		br.motion_compensation(mvd, i, mb_row, mb_address, &mb, &b)
		b.clamp(&cb)
		updateFrameSlice(i, mb_address, mb.dct_type, frameSlice, &cb)
	}

	return mb_address, nil
}

func updateFrameSlice(i, mb_address int, interlaced bool, frameSlice *image.YCbCr, cb *clampedblock) {

	var (
		base_i  int
		channel []uint8
		stride  int
	)

	// channel switch
	switch i {
	case 0, 1, 2, 3:
		channel = frameSlice.Y
	case 4:
		channel = frameSlice.Cb
	case 5:
		channel = frameSlice.Cr
	}

	// base address and stride switch
	switch i {
	case 0, 1, 2, 3:
		stride = frameSlice.YStride
		base_i = mb_address * 16
	case 4, 5:
		stride = frameSlice.CStride
		base_i = mb_address * 8
	}

	// position switch
	if interlaced {
		// Field DCT coding alternates lines from each block:
		//
		//  <-8px-> <-8px->
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───0───│───1───
		//  ───2───│───3───
		switch i {
		case 0, 1, 2, 3:
			base_i += (i & 1) << 3
			base_i += ((i & 2) >> 1) * stride
			stride *= 2
		}
	} else {
		// Frame DCT coding are mapped in the follow order:
		//
		//  <-8px-> <-8px->
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───0───│───1───
		//  ───2───│───3───
		//  ───2───│───3───
		//  ───2───│───3───
		//  ───2───│───3───
		switch i {
		case 0, 1, 2, 3:
			base_i += (i & 1) << 3            // horiztonal positioning
			base_i += ((i & 2) << 2) * stride // vertical positioning
		}
	}

	// perform copy
	for y := 0; y < 8; y++ {
		si := y * 8
		di := base_i + (y * stride)
		copy(channel[di:di+8], cb[si:si+8])
	}

}

type PatternCode [12]bool

func (mb *Macroblock) decodePatternCode(chroma_format chromaFormat) (pattern_code PatternCode) {
	for i := 0; i < 12; i++ {
		if mb.macroblock_type.macroblock_intra {
			pattern_code[i] = true
		} else {
			pattern_code[i] = false
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		for i := 0; i < 6; i++ {
			if mask := 1 << uint(5-i); mb.cpb&mask == mask {
				pattern_code[i] = true
			}
		}

		if chroma_format == ChromaFormat_422 || chroma_format == ChromaFormat_444 {
			panic("unsupported: coded block pattern chroma format")
		}
	}

	return
}

func (br *VideoSequence) macroblock_mode(mb *Macroblock) (err error) {

	var typeDecoder macroblockTypeDecoderFn
	switch br.PictureHeader.picture_coding_type {
	case IntraCoded:
		typeDecoder = macroblockTypeDecoder.IFrame
	case PredictiveCoded:
		typeDecoder = macroblockTypeDecoder.PFrame
	case BidirectionallyPredictiveCoded:
		typeDecoder = macroblockTypeDecoder.BFrame
	default:
		panic("not implemented: macroblock type decoder")
	}

	mb.macroblock_type, err = typeDecoder(br)
	if err != nil {
		return err
	}

	if mb.macroblock_type.spatial_temporal_weight_code_flag &&
		false /* ( spatial_temporal_weight_code_table_index != ‘00’) */ {
		mb.spatial_temporal_weight_code, err = br.Read32(2)
		if err != nil {
			return err
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		mb.macroblock_type.macroblock_motion_backward {
		if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture {
			if br.PictureCodingExtension.frame_pred_frame_dct == 0 {
				mb.frame_motion_type, err = br.Read32(2)
				if err != nil {
					return err
				}
			}
		} else {
			mb.field_motion_type, err = br.Read32(2)
			if err != nil {
				return err
			}
		}
	}

	if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture &&
		br.PictureCodingExtension.frame_pred_frame_dct == 0 &&
		(mb.macroblock_type.macroblock_intra || mb.macroblock_type.macroblock_pattern) {
		mb.dct_type, err = br.ReadBit() //dct_type 1 uimsbf
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>