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

package registrydisk

import (
	"io/ioutil"
	"os"
	"os/user"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
)

var _ = Describe("RegistryDisk", func() {
	tmpDir, _ := ioutil.TempDir("", "registrydisktest")
	owner, err := user.Current()
	if err != nil {
		panic(err)
	}

	VerifyDiskType := func(diskExtension string) {
		vm := v1.NewMinimalVM("fake-vm")
		vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
			Type:   "RegistryDisk:v1alpha",
			Device: "disk",
			Source: v1.DiskSource{
				Name: "someimage:v1.2.3.4",
			},
			Target: v1.DiskTarget{
				Device: "vda",
			},
		})

		// create a fake disk file
		volumeMountDir := generateVMBaseDir(vm)
		err = os.MkdirAll(volumeMountDir+"/disk0", 0750)
		Expect(err).ToNot(HaveOccurred())

		filePath := volumeMountDir + "/disk0/disk-image." + diskExtension
		_, err := os.Create(filePath)

		vm, err = MapRegistryDisks(vm)
		Expect(err).ToNot(HaveOccurred())

		// verify file gets renamed by virt-handler to prevent container from
		// removing it before VM is exited
		exists, err := diskutils.FileExists(filePath)
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(false))

		// verify file rename takes place
		exists, err = diskutils.FileExists(filePath + ".virt")
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(true))

		Expect(vm.Spec.Domain.Devices.Disks[0].Type).To(Equal("file"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Target.Device).To(Equal("vda"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver.Type).To(Equal(diskExtension))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source.File).To(Equal(filePath + ".virt"))

		err = CleanupEphemeralDisks(vm)
		exists, err = diskutils.FileExists(volumeMountDir)
		Expect(err).ToNot(HaveOccurred())

		Expect(exists).To(Equal(false))
	}

	BeforeSuite(func() {
		err := SetLocalDirectory(tmpDir)
		if err != nil {
			panic(err)
		}
		SetLocalDataOwner(owner.Username)
	})

	AfterSuite(func() {
		os.RemoveAll(tmpDir)
	})

	Describe("registry-disk", func() {
		Context("verify helper functions", func() {
			table.DescribeTable("by verifying mapping of ",
				func(diskType string) {
					VerifyDiskType(diskType)
				},
				table.Entry("qcow2 disk", "qcow2"),
				table.Entry("raw disk", "raw"),
			)
			It("by verifying error when no disk is present", func() {

				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				vm, err := MapRegistryDisks(vm)
				Expect(err).To(HaveOccurred())
			})
			It("by verifying container generation", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vdb",
					},
				})

				containers, volumes, err := GenerateContainers(vm)
				Expect(err).ToNot(HaveOccurred())

				Expect(len(containers)).To(Equal(2))
				Expect(len(volumes)).To(Equal(2))
			})

			It("by verifying data cleanup", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				volumeMountDir := generateVMBaseDir(vm)
				err = os.MkdirAll(volumeMountDir, 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk0", 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk1", 0755)
				Expect(err).ToNot(HaveOccurred())

				exists, err := diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(true))

				err = CleanupEphemeralDisks(vm)
				exists, err = diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(false))

			})
		})
	})
})
<commit_msg>Unit tests for orphaned registry disk cleanup<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package registrydisk

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/user"

	"k8s.io/client-go/tools/cache"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
)

var _ = Describe("RegistryDisk", func() {
	tmpDir, _ := ioutil.TempDir("", "registrydisktest")
	owner, err := user.Current()
	if err != nil {
		panic(err)
	}

	VerifyDiskType := func(diskExtension string) {
		vm := v1.NewMinimalVM("fake-vm")
		vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
			Type:   "RegistryDisk:v1alpha",
			Device: "disk",
			Source: v1.DiskSource{
				Name: "someimage:v1.2.3.4",
			},
			Target: v1.DiskTarget{
				Device: "vda",
			},
		})

		// create a fake disk file
		volumeMountDir := generateVMBaseDir(vm)
		err = os.MkdirAll(volumeMountDir+"/disk0", 0750)
		Expect(err).ToNot(HaveOccurred())

		filePath := volumeMountDir + "/disk0/disk-image." + diskExtension
		_, err := os.Create(filePath)

		vm, err = MapRegistryDisks(vm)
		Expect(err).ToNot(HaveOccurred())

		// verify file gets renamed by virt-handler to prevent container from
		// removing it before VM is exited
		exists, err := diskutils.FileExists(filePath)
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(false))

		// verify file rename takes place
		exists, err = diskutils.FileExists(filePath + ".virt")
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(true))

		Expect(vm.Spec.Domain.Devices.Disks[0].Type).To(Equal("file"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Target.Device).To(Equal("vda"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver.Type).To(Equal(diskExtension))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source.File).To(Equal(filePath + ".virt"))

		err = CleanupEphemeralDisks(vm)
		exists, err = diskutils.FileExists(volumeMountDir)
		Expect(err).ToNot(HaveOccurred())

		Expect(exists).To(Equal(false))
	}

	BeforeSuite(func() {
		err := SetLocalDirectory(tmpDir)
		if err != nil {
			panic(err)
		}
		SetLocalDataOwner(owner.Username)
	})

	AfterSuite(func() {
		os.RemoveAll(tmpDir)
	})

	Describe("registry-disk", func() {
		Context("verify helper functions", func() {
			table.DescribeTable("by verifying mapping of ",
				func(diskType string) {
					VerifyDiskType(diskType)
				},
				table.Entry("qcow2 disk", "qcow2"),
				table.Entry("raw disk", "raw"),
			)
			It("by verifying error when no disk is present", func() {

				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				vm, err := MapRegistryDisks(vm)
				Expect(err).To(HaveOccurred())
			})
			It("by verifying container generation", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vdb",
					},
				})

				containers, volumes, err := GenerateContainers(vm)
				Expect(err).ToNot(HaveOccurred())

				Expect(len(containers)).To(Equal(2))
				Expect(len(volumes)).To(Equal(2))
			})

			It("by removing unseen registry disk data", func() {
				var domains []string

				domains = append(domains, "fakens1/fakedomain1")
				domains = append(domains, "fakens1/fakedomain2")
				domains = append(domains, "fakens2/fakedomain1")
				domains = append(domains, "fakens2/fakedomain2")
				domains = append(domains, "fakens3/fakedomain1")
				domains = append(domains, "fakens4/fakedomain1")

				for _, dom := range domains {
					err := os.MkdirAll(fmt.Sprintf("%s/%s/some-other-dir", tmpDir, dom), 0755)
					Expect(err).ToNot(HaveOccurred())
					msg := "fake content"
					bytes := []byte(msg)
					err = ioutil.WriteFile(fmt.Sprintf("%s/%s/some-file", tmpDir, dom), bytes, 0644)
					Expect(err).ToNot(HaveOccurred())
				}

				vmStore := cache.NewIndexer(cache.DeletionHandlingMetaNamespaceKeyFunc, cache.Indexers{})

				err := vmStore.Add(v1.NewVMReferenceFromNameWithNS("fakens1", "fakedomain1"))
				Expect(err).ToNot(HaveOccurred())

				err = CleanupOrphanedEphemeralDisks(vmStore)
				Expect(err).ToNot(HaveOccurred())

				// expect this domain to still exist
				_, err = os.Stat(fmt.Sprintf("%s/fakens1/fakedomain1", tmpDir))
				Expect(err).ToNot(HaveOccurred())

				// expect these domains to not exist
				for idx, dom := range domains {
					exists := true
					if idx == 0 {
						continue
					}
					_, err = os.Stat(fmt.Sprintf("%s/%s", tmpDir, dom))
					if os.IsNotExist(err) {
						exists = false
					}
					Expect(exists).To(Equal(false))
				}
			})

			It("by verifying data cleanup", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				volumeMountDir := generateVMBaseDir(vm)
				err = os.MkdirAll(volumeMountDir, 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk0", 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk1", 0755)
				Expect(err).ToNot(HaveOccurred())

				exists, err := diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(true))

				err = CleanupEphemeralDisks(vm)
				exists, err = diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(false))

			})
		})
	})
})
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package registry

import (
	"fmt"
	"net/http"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/health"
)

type HealthyMinionRegistry struct {
	delegate MinionRegistry
	client   health.HTTPGetInterface
	port     int
}

func NewHealthyMinionRegistry(delegate MinionRegistry, client *http.Client) MinionRegistry {
	return &HealthyMinionRegistry{
		delegate: delegate,
		client:   client,
		port:     10250,
	}
}

func (h *HealthyMinionRegistry) makeMinionURL(minion string) string {
	return fmt.Sprintf("http://%s:%d/healthz", minion, h.port)
}

func (h *HealthyMinionRegistry) List() (currentMinions []string, err error) {
	var result []string
	list, err := h.delegate.List()
	if err != nil {
		return result, err
	}
	for _, minion := range list {
		status, err := health.Check(h.makeMinionURL(minion), h.client)
		if err != nil {
			return result, err
		}
		if status == health.Healthy {
			result = append(result, minion)
		}
	}
	return result, nil
}

func (h *HealthyMinionRegistry) Insert(minion string) error {
	return h.delegate.Insert(minion)
}

func (h *HealthyMinionRegistry) Delete(minion string) error {
	return h.delegate.Delete(minion)
}

func (h *HealthyMinionRegistry) Contains(minion string) (bool, error) {
	contains, err := h.delegate.Contains(minion)
	if err != nil {
		return false, err
	}
	if !contains {
		return false, nil
	}
	status, err := health.Check(h.makeMinionURL(minion), h.client)
	if err != nil {
		return false, err
	}
	if status == health.Unhealthy {
		return false, nil
	}
	return true, nil
}
<commit_msg>Fix #663<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package registry

import (
	"fmt"
	"net/http"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/health"
	"github.com/golang/glog"
)

type HealthyMinionRegistry struct {
	delegate MinionRegistry
	client   health.HTTPGetInterface
	port     int
}

func NewHealthyMinionRegistry(delegate MinionRegistry, client *http.Client) MinionRegistry {
	return &HealthyMinionRegistry{
		delegate: delegate,
		client:   client,
		port:     10250,
	}
}

func (h *HealthyMinionRegistry) makeMinionURL(minion string) string {
	return fmt.Sprintf("http://%s:%d/healthz", minion, h.port)
}

func (h *HealthyMinionRegistry) List() (currentMinions []string, err error) {
	var result []string
	list, err := h.delegate.List()
	if err != nil {
		return result, err
	}
	for _, minion := range list {
		status, err := health.Check(h.makeMinionURL(minion), h.client)
		if err != nil {
			glog.Errorf("%s failed health check with error: %s", minion, err)
			continue
		}
		if status == health.Healthy {
			result = append(result, minion)
		}
	}
	return result, nil
}

func (h *HealthyMinionRegistry) Insert(minion string) error {
	return h.delegate.Insert(minion)
}

func (h *HealthyMinionRegistry) Delete(minion string) error {
	return h.delegate.Delete(minion)
}

func (h *HealthyMinionRegistry) Contains(minion string) (bool, error) {
	contains, err := h.delegate.Contains(minion)
	if err != nil {
		return false, err
	}
	if !contains {
		return false, nil
	}
	status, err := health.Check(h.makeMinionURL(minion), h.client)
	if err != nil {
		return false, err
	}
	if status == health.Unhealthy {
		return false, nil
	}
	return true, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package v1

import (
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
)

const (
	// MinExtenderPriority defines the min priority value for extender.
	MinExtenderPriority int64 = 0

	// MaxExtenderPriority defines the max priority value for extender.
	MaxExtenderPriority int64 = 10
)

// ExtenderPreemptionResult represents the result returned by preemption phase of extender.
type ExtenderPreemptionResult struct {
	NodeNameToMetaVictims map[string]*MetaVictims
}

// ExtenderPreemptionArgs represents the arguments needed by the extender to preempt pods on nodes.
type ExtenderPreemptionArgs struct {
	// Pod being scheduled
	Pod *v1.Pod
	// Victims map generated by scheduler preemption phase
	// Only set NodeNameToMetaVictims if ExtenderConfig.NodeCacheCapable == true. Otherwise, only set NodeNameToVictims.
	NodeNameToVictims     map[string]*Victims
	NodeNameToMetaVictims map[string]*MetaVictims
}

// Victims represents:
//   pods:  a group of pods expected to be preempted.
//   numPDBViolations: the count of violations of PodDisruptionBudget
type Victims struct {
	Pods             []*v1.Pod
	NumPDBViolations int64
}

// MetaPod represent identifier for a v1.Pod
type MetaPod struct {
	UID string
}

// MetaVictims represents:
//   pods:  a group of pods expected to be preempted.
//     Only Pod identifiers will be sent and user are expect to get v1.Pod in their own way.
//   numPDBViolations: the count of violations of PodDisruptionBudget
type MetaVictims struct {
	Pods             []*MetaPod
	NumPDBViolations int64
}

// ExtenderArgs represents the arguments needed by the extender to filter/prioritize
// nodes for a pod.
type ExtenderArgs struct {
	// Pod being scheduled
	Pod *v1.Pod
	// List of candidate nodes where the pod can be scheduled; to be populated
	// only if ExtenderConfig.NodeCacheCapable == false
	Nodes *v1.NodeList
	// List of candidate node names where the pod can be scheduled; to be
	// populated only if ExtenderConfig.NodeCacheCapable == true
	NodeNames *[]string
}

// FailedNodesMap represents the filtered out nodes, with node names and failure messages
type FailedNodesMap map[string]string

// ExtenderFilterResult represents the results of a filter call to an extender
type ExtenderFilterResult struct {
	// Filtered set of nodes where the pod can be scheduled; to be populated
	// only if ExtenderConfig.NodeCacheCapable == false
	Nodes *v1.NodeList
	// Filtered set of nodes where the pod can be scheduled; to be populated
	// only if ExtenderConfig.NodeCacheCapable == true
	NodeNames *[]string
	// Filtered out nodes where the pod can't be scheduled and the failure messages
	FailedNodes FailedNodesMap
	// Error message indicating failure
	Error string
}

// ExtenderBindingArgs represents the arguments to an extender for binding a pod to a node.
type ExtenderBindingArgs struct {
	// PodName is the name of the pod being bound
	PodName string
	// PodNamespace is the namespace of the pod being bound
	PodNamespace string
	// PodUID is the UID of the pod being bound
	PodUID types.UID
	// Node selected by the scheduler
	Node string
}

// ExtenderBindingResult represents the result of binding of a pod to a node from an extender.
type ExtenderBindingResult struct {
	// Error message indicating failure
	Error string
}

// HostPriority represents the priority of scheduling to a particular host, higher priority is better.
type HostPriority struct {
	// Name of the host
	Host string
	// Score associated with the host
	Score int64
}

// HostPriorityList declares a []HostPriority type.
type HostPriorityList []HostPriority
<commit_msg>rename ExtenderConfig to Extender<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package v1

import (
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
)

const (
	// MinExtenderPriority defines the min priority value for extender.
	MinExtenderPriority int64 = 0

	// MaxExtenderPriority defines the max priority value for extender.
	MaxExtenderPriority int64 = 10
)

// ExtenderPreemptionResult represents the result returned by preemption phase of extender.
type ExtenderPreemptionResult struct {
	NodeNameToMetaVictims map[string]*MetaVictims
}

// ExtenderPreemptionArgs represents the arguments needed by the extender to preempt pods on nodes.
type ExtenderPreemptionArgs struct {
	// Pod being scheduled
	Pod *v1.Pod
	// Victims map generated by scheduler preemption phase
	// Only set NodeNameToMetaVictims if Extender.NodeCacheCapable == true. Otherwise, only set NodeNameToVictims.
	NodeNameToVictims     map[string]*Victims
	NodeNameToMetaVictims map[string]*MetaVictims
}

// Victims represents:
//   pods:  a group of pods expected to be preempted.
//   numPDBViolations: the count of violations of PodDisruptionBudget
type Victims struct {
	Pods             []*v1.Pod
	NumPDBViolations int64
}

// MetaPod represent identifier for a v1.Pod
type MetaPod struct {
	UID string
}

// MetaVictims represents:
//   pods:  a group of pods expected to be preempted.
//     Only Pod identifiers will be sent and user are expect to get v1.Pod in their own way.
//   numPDBViolations: the count of violations of PodDisruptionBudget
type MetaVictims struct {
	Pods             []*MetaPod
	NumPDBViolations int64
}

// ExtenderArgs represents the arguments needed by the extender to filter/prioritize
// nodes for a pod.
type ExtenderArgs struct {
	// Pod being scheduled
	Pod *v1.Pod
	// List of candidate nodes where the pod can be scheduled; to be populated
	// only if Extender.NodeCacheCapable == false
	Nodes *v1.NodeList
	// List of candidate node names where the pod can be scheduled; to be
	// populated only if Extender.NodeCacheCapable == true
	NodeNames *[]string
}

// FailedNodesMap represents the filtered out nodes, with node names and failure messages
type FailedNodesMap map[string]string

// ExtenderFilterResult represents the results of a filter call to an extender
type ExtenderFilterResult struct {
	// Filtered set of nodes where the pod can be scheduled; to be populated
	// only if Extender.NodeCacheCapable == false
	Nodes *v1.NodeList
	// Filtered set of nodes where the pod can be scheduled; to be populated
	// only if Extender.NodeCacheCapable == true
	NodeNames *[]string
	// Filtered out nodes where the pod can't be scheduled and the failure messages
	FailedNodes FailedNodesMap
	// Error message indicating failure
	Error string
}

// ExtenderBindingArgs represents the arguments to an extender for binding a pod to a node.
type ExtenderBindingArgs struct {
	// PodName is the name of the pod being bound
	PodName string
	// PodNamespace is the namespace of the pod being bound
	PodNamespace string
	// PodUID is the UID of the pod being bound
	PodUID types.UID
	// Node selected by the scheduler
	Node string
}

// ExtenderBindingResult represents the result of binding of a pod to a node from an extender.
type ExtenderBindingResult struct {
	// Error message indicating failure
	Error string
}

// HostPriority represents the priority of scheduling to a particular host, higher priority is better.
type HostPriority struct {
	// Name of the host
	Host string
	// Score associated with the host
	Score int64
}

// HostPriorityList declares a []HostPriority type.
type HostPriorityList []HostPriority
<|endoftext|>


<commit_before>package heartbeat

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	utilwait "k8s.io/apimachinery/pkg/util/wait"
	k8scli "k8s.io/client-go/kubernetes/typed/core/v1"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/log"
	virtutil "kubevirt.io/kubevirt/pkg/util"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	device_manager "kubevirt.io/kubevirt/pkg/virt-handler/device-manager"
)

type HeartBeat struct {
	clientset                 k8scli.CoreV1Interface
	deviceManagerController   device_manager.DeviceControllerInterface
	clusterConfig             *virtconfig.ClusterConfig
	host                      string
	cpuManagerPaths           []string
	devicePluginPollIntervall time.Duration
	devicePluginWaitTimeout   time.Duration
}

func NewHeartBeat(clientset k8scli.CoreV1Interface, deviceManager device_manager.DeviceControllerInterface, clusterConfig *virtconfig.ClusterConfig, host string) *HeartBeat {
	return &HeartBeat{
		clientset:               clientset,
		deviceManagerController: deviceManager,
		clusterConfig:           clusterConfig,
		host:                    host,
		// This is a temporary workaround until k8s bug #66525 is resolved
		cpuManagerPaths:           []string{virtutil.CPUManagerPath, virtutil.CPUManagerOS3Path},
		devicePluginPollIntervall: 1 * time.Second,
		devicePluginWaitTimeout:   10 * time.Second,
	}
}

func (h *HeartBeat) Run(heartBeatInterval time.Duration, stopCh chan struct{}) (done chan struct{}) {
	done = make(chan struct{})
	go func() {
		h.heartBeat(heartBeatInterval, stopCh)
		close(done)
	}()
	return done
}

func (h *HeartBeat) heartBeat(heartBeatInterval time.Duration, stopCh chan struct{}) {
	// ensure that the node is synchronized with the actual state
	// especially setting the node to unschedulable if device plugins are not yet ready is very important
	// otherwise workloads get scheduled but are immediately terminated by the kubelet
	h.do()
	// Now wait for 10 seconds for the device plugins  to be initialized
	// This is more than fast enough to be not treated as unschedulable by the cluster
	// and ensures that the cluster gets marked as scheduled as soon as the device plugin is ready
	h.waitForDevicePlugins(stopCh)

	// from now on periodically update the node status
	wait.JitterUntil(h.do, heartBeatInterval, 1.2, true, stopCh)
}

// waitForDevicePlugins gives the device plugins additional time to successfully connect to the kubelet.
// If the connection can not be established it just delays the heartbeat start for devicePluginWaitTimeout.
func (h *HeartBeat) waitForDevicePlugins(stopCh chan struct{}) {
	_ = utilwait.PollImmediate(h.devicePluginPollIntervall, h.devicePluginWaitTimeout, func() (done bool, err error) {
		select {
		case <-stopCh:
			return true, nil
		default:
		}
		return h.deviceManagerController.Initialized(), nil
	})
}

func (h *HeartBeat) do() {
	now, err := json.Marshal(metav1.Now())
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("Can't determine date")
		return
	}

	kubevirtSchedulable := "true"
	if !h.deviceManagerController.Initialized() {
		kubevirtSchedulable = "false"
	}

	var data []byte
	// Label the node if cpu manager is running on it
	// This is a temporary workaround until k8s bug #66525 is resolved
	cpuManagerEnabled := false
	if h.clusterConfig.CPUManagerEnabled() {
		cpuManagerEnabled = h.isCPUManagerEnabled(h.cpuManagerPaths)
	}
	data = []byte(fmt.Sprintf(`{"metadata": { "labels": {"%s": "%s", "%s": "%t"}, "annotations": {"%s": %s}}}`,
		v1.NodeSchedulable, kubevirtSchedulable,
		v1.CPUManager, cpuManagerEnabled,
		v1.VirtHandlerHeartbeat, string(now),
	))
	_, err = h.clientset.Nodes().Patch(context.Background(), h.host, types.StrategicMergePatchType, data, metav1.PatchOptions{})
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("Can't patch node %s", h.host)
		return
	}
	log.DefaultLogger().V(4).Infof("Heartbeat sent")
}

func (h *HeartBeat) isCPUManagerEnabled(cpuManagerPaths []string) bool {
	var cpuManagerOptions map[string]interface{}
	cpuManagerPath, err := detectCPUManagerFile(cpuManagerPaths)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}
	// #nosec No risk for path injection. cpuManagerPath is composed of static values from pkg/util
	content, err := ioutil.ReadFile(cpuManagerPath)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}

	err = json.Unmarshal(content, &cpuManagerOptions)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}

	if v, ok := cpuManagerOptions["policyName"]; ok && v == "static" {
		log.DefaultLogger().V(4).Infof("Node has CPU Manager running")
		return true
	} else {
		log.DefaultLogger().V(4).Infof("Node has CPU Manager not runnning")
		return false
	}
}

func detectCPUManagerFile(cpuManagerPaths []string) (string, error) {
	for _, path := range cpuManagerPaths {
		if _, err := os.Stat(path); err == nil {
			return path, nil
		}
	}
	return "", fmt.Errorf("no cpumanager policy file found")
}
<commit_msg>virt-handler: add comment to describe heartbeat intervals<commit_after>package heartbeat

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	utilwait "k8s.io/apimachinery/pkg/util/wait"
	k8scli "k8s.io/client-go/kubernetes/typed/core/v1"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/log"
	virtutil "kubevirt.io/kubevirt/pkg/util"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	device_manager "kubevirt.io/kubevirt/pkg/virt-handler/device-manager"
)

type HeartBeat struct {
	clientset                 k8scli.CoreV1Interface
	deviceManagerController   device_manager.DeviceControllerInterface
	clusterConfig             *virtconfig.ClusterConfig
	host                      string
	cpuManagerPaths           []string
	devicePluginPollIntervall time.Duration
	devicePluginWaitTimeout   time.Duration
}

func NewHeartBeat(clientset k8scli.CoreV1Interface, deviceManager device_manager.DeviceControllerInterface, clusterConfig *virtconfig.ClusterConfig, host string) *HeartBeat {
	return &HeartBeat{
		clientset:               clientset,
		deviceManagerController: deviceManager,
		clusterConfig:           clusterConfig,
		host:                    host,
		// This is a temporary workaround until k8s bug #66525 is resolved
		cpuManagerPaths:           []string{virtutil.CPUManagerPath, virtutil.CPUManagerOS3Path},
		devicePluginPollIntervall: 1 * time.Second,
		devicePluginWaitTimeout:   10 * time.Second,
	}
}

func (h *HeartBeat) Run(heartBeatInterval time.Duration, stopCh chan struct{}) (done chan struct{}) {
	done = make(chan struct{})
	go func() {
		h.heartBeat(heartBeatInterval, stopCh)
		close(done)
	}()
	return done
}

func (h *HeartBeat) heartBeat(heartBeatInterval time.Duration, stopCh chan struct{}) {
	// ensure that the node is synchronized with the actual state
	// especially setting the node to unschedulable if device plugins are not yet ready is very important
	// otherwise workloads get scheduled but are immediately terminated by the kubelet
	h.do()
	// Now wait for 10 seconds for the device plugins  to be initialized
	// This is more than fast enough to be not treated as unschedulable by the cluster
	// and ensures that the cluster gets marked as scheduled as soon as the device plugin is ready
	h.waitForDevicePlugins(stopCh)

	// from now on periodically update the node status
	// This sets the heartbeat to:
	// 1 minute with a 1.2 jitter + the time it takes for the heartbeat function to run (sliding == true).
	// So the amount of time between heartbeats randomly varies between 1min and 2min12sec + the heartbeat function execution time.
	wait.JitterUntil(h.do, heartBeatInterval, 1.2, true, stopCh)
}

// waitForDevicePlugins gives the device plugins additional time to successfully connect to the kubelet.
// If the connection can not be established it just delays the heartbeat start for devicePluginWaitTimeout.
func (h *HeartBeat) waitForDevicePlugins(stopCh chan struct{}) {
	_ = utilwait.PollImmediate(h.devicePluginPollIntervall, h.devicePluginWaitTimeout, func() (done bool, err error) {
		select {
		case <-stopCh:
			return true, nil
		default:
		}
		return h.deviceManagerController.Initialized(), nil
	})
}

func (h *HeartBeat) do() {
	now, err := json.Marshal(metav1.Now())
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("Can't determine date")
		return
	}

	kubevirtSchedulable := "true"
	if !h.deviceManagerController.Initialized() {
		kubevirtSchedulable = "false"
	}

	var data []byte
	// Label the node if cpu manager is running on it
	// This is a temporary workaround until k8s bug #66525 is resolved
	cpuManagerEnabled := false
	if h.clusterConfig.CPUManagerEnabled() {
		cpuManagerEnabled = h.isCPUManagerEnabled(h.cpuManagerPaths)
	}
	data = []byte(fmt.Sprintf(`{"metadata": { "labels": {"%s": "%s", "%s": "%t"}, "annotations": {"%s": %s}}}`,
		v1.NodeSchedulable, kubevirtSchedulable,
		v1.CPUManager, cpuManagerEnabled,
		v1.VirtHandlerHeartbeat, string(now),
	))
	_, err = h.clientset.Nodes().Patch(context.Background(), h.host, types.StrategicMergePatchType, data, metav1.PatchOptions{})
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("Can't patch node %s", h.host)
		return
	}
	log.DefaultLogger().V(4).Infof("Heartbeat sent")
}

func (h *HeartBeat) isCPUManagerEnabled(cpuManagerPaths []string) bool {
	var cpuManagerOptions map[string]interface{}
	cpuManagerPath, err := detectCPUManagerFile(cpuManagerPaths)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}
	// #nosec No risk for path injection. cpuManagerPath is composed of static values from pkg/util
	content, err := ioutil.ReadFile(cpuManagerPath)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}

	err = json.Unmarshal(content, &cpuManagerOptions)
	if err != nil {
		log.DefaultLogger().Reason(err).Errorf("failed to set a cpu manager label on host %s", h.host)
		return false
	}

	if v, ok := cpuManagerOptions["policyName"]; ok && v == "static" {
		log.DefaultLogger().V(4).Infof("Node has CPU Manager running")
		return true
	} else {
		log.DefaultLogger().V(4).Infof("Node has CPU Manager not runnning")
		return false
	}
}

func detectCPUManagerFile(cpuManagerPaths []string) (string, error) {
	for _, path := range cpuManagerPaths {
		if _, err := os.Stat(path); err == nil {
			return path, nil
		}
	}
	return "", fmt.Errorf("no cpumanager policy file found")
}
<|endoftext|>


<commit_before>package main

import "fmt"

func main() {
	fmt.Println("Demo2 base grammer")
}
<commit_msg>2015-06-03 lidongxu demo2 基本语法 格式化<commit_after>package main

import "fmt"

func main() {
	fmt.Println("Demo2 base grammer")

	fmt.Println("First line")

	fmt.Println("Second line")

	var age int
	fmt.Println("age=", age)

	var apples = "myApples "
	var oranges = "greenOranges "
	var fruit = apples + oranges
	fmt.Println("fruit=", fruit)
	fmt.Println("apples=", apples, "\t oranges=", oranges)

}
<|endoftext|>


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

package main

import (
	"testing"

	"golang.org/x/net/context"

	speech "cloud.google.com/go/speech/apiv1"
	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
)

func TestRecognize(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := send(client, "./quit.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Transcript, "quit"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
}

func TestRecognizeGCS(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := sendGCS(client, "gs://python-docs-samples-tests/speech/audio.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Transcript, "how old is the Brooklyn Bridge"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
}

func TestRecognizeGCSWordTimeOffsets(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := sendGCS(client, "gs://python-docs-samples-tests/speech/audio.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Words[5].Word, "Bridge"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
	if got := result.Alternatives[0].Words[5].StartTime.Seconds; got <= 0 {
		t.Errorf("Word start time: got %f, want positive", got)
	}
	if got := result.Alternatives[0].Words[5].EndTime.Seconds; got <= 0 {
		t.Errorf("Word end time: got %f, want positive", got)
	}
}
<commit_msg>speech/captionasync: fix vet warnings (#284)<commit_after>// Copyright 2016 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package main

import (
	"testing"

	"golang.org/x/net/context"

	speech "cloud.google.com/go/speech/apiv1"
	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
)

func TestRecognize(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := send(client, "./quit.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Transcript, "quit"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
}

func TestRecognizeGCS(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := sendGCS(client, "gs://python-docs-samples-tests/speech/audio.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Transcript, "how old is the Brooklyn Bridge"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
}

func TestRecognizeGCSWordTimeOffsets(t *testing.T) {
	testutil.SystemTest(t)

	ctx := context.Background()
	client, err := speech.NewClient(ctx)
	if err != nil {
		t.Fatal(err)
	}

	opName, err := sendGCS(client, "gs://python-docs-samples-tests/speech/audio.raw")
	if err != nil {
		t.Fatal(err)
	}
	if opName == "" {
		t.Fatal("got no op name; want one")
	}
	resp, err := wait(client, opName)
	if err != nil {
		t.Fatal(err)
	}
	if len(resp.Results) == 0 {
		t.Fatal("got no results; want at least one")
	}
	result := resp.Results[0]
	if len(result.Alternatives) < 1 {
		t.Fatal("got no alternatives; want at least one")
	}
	if got, want := result.Alternatives[0].Words[5].Word, "Bridge"; got != want {
		t.Errorf("Transcript: got %q; want %q", got, want)
	}
	if got := result.Alternatives[0].Words[5].StartTime.Seconds; got <= 0 {
		t.Errorf("Word start time: got %d, want positive", got)
	}
	if got := result.Alternatives[0].Words[5].EndTime.Seconds; got <= 0 {
		t.Errorf("Word end time: got %d, want positive", got)
	}
}
<|endoftext|>


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

package elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	for i := 0; i < b.maxConns; i++ {
		go func() {
			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					// Copy for the potential re-send.
					bufCopy := bytes.NewBuffer(buf.Bytes())
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(bufCopy)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	go func() {
		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	go func() {
		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Delete(index, _type, id string, refresh bool) {
	queryLine := fmt.Sprintf("{\"delete\":{\"_index\":%q,\"_type\":%q,\"_id\":%q,\"refresh\":%t}}\n", index, _type, id, refresh)
	b.bulkChannel <- []byte(queryLine)
	return
}

func (b *BulkIndexer) UpdateWithPartialDoc(index string, _type string, id, ttl string, date *time.Time, partialDoc interface{}, upsert bool, refresh bool) error {

	var data map[string]interface{} = make(map[string]interface{})

	data["doc"] = partialDoc
	if upsert {
		data["doc_as_upsert"] = true
	}
	return b.Update(index, _type, id, ttl, date, data, refresh)
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	response := responseStruct{}

	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			b.numErrors += uint64(len(response.Items))
			return fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<commit_msg>reapply core: add wait group to wait on all goros to terminate properly<commit_after>// Copyright 2013 Matthew Baird
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup

	// wait group to wait on shutdown for all goros
	gorosWg *sync.WaitGroup
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	b.gorosWg.Add(b.maxConns)
	for i := 0; i < b.maxConns; i++ {
		go func() {
			defer b.gorosWg.Done()

			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					// Copy for the potential re-send.
					bufCopy := bytes.NewBuffer(buf.Bytes())
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(bufCopy)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	b.gorosWg.Add(1)
	go func() {
		defer b.gorosWg.Done()

		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	b.gorosWg.Add(1)
	go func() {
		defer b.gorosWg.Done()

		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
	b.gorosWg.Wait()
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Delete(index, _type, id string, refresh bool) {
	queryLine := fmt.Sprintf("{\"delete\":{\"_index\":%q,\"_type\":%q,\"_id\":%q,\"refresh\":%t}}\n", index, _type, id, refresh)
	b.bulkChannel <- []byte(queryLine)
	return
}

func (b *BulkIndexer) UpdateWithPartialDoc(index string, _type string, id, ttl string, date *time.Time, partialDoc interface{}, upsert bool, refresh bool) error {

	var data map[string]interface{} = make(map[string]interface{})

	data["doc"] = partialDoc
	if upsert {
		data["doc_as_upsert"] = true
	}
	return b.Update(index, _type, id, ttl, date, data, refresh)
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	response := responseStruct{}

	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			b.numErrors += uint64(len(response.Items))
			return fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<|endoftext|>


<commit_before>// +build windows

package svc

import (
	"errors"
	"os"
	"syscall"
	"testing"
	"time"

	wsvc "golang.org/x/sys/windows/svc"
)

func setupWinServiceTest(wsf *mockWinServiceFuncs) {
	// wsfWrapper allows signalNotify, svcIsInteractive, and svcRun to be set once.
	// Individual test functions set "wsf" to add behavior.
	wsfWrapper := &mockWinServiceFuncs{
		signalNotify: func(c chan<- os.Signal, sig ...os.Signal) {
			if c == nil {
				panic("os/signal: Notify using nil channel")
			}

			if wsf.signalNotify != nil {
				wsf.signalNotify(c, sig...)
			} else {
				wsf1 := *wsf
				go func() {
					for val := range wsf1.sigChan {
						for _, registeredSig := range sig {
							if val == registeredSig {
								c <- val
							}
						}
					}
				}()
			}
		},
		svcIsInteractive: func() (bool, error) {
			return wsf.svcIsInteractive()
		},
		svcRun: func(name string, handler wsvc.Handler) error {
			return wsf.svcRun(name, handler)
		},
	}

	signalNotify = wsfWrapper.signalNotify
	svcIsAnInteractiveSession = wsfWrapper.svcIsInteractive
	svcRun = wsfWrapper.svcRun
}

type mockWinServiceFuncs struct {
	signalNotify          func(chan<- os.Signal, ...os.Signal)
	svcIsInteractive      func() (bool, error)
	sigChan               chan os.Signal
	svcRun                func(string, wsvc.Handler) error
	ws                    *windowsService
	executeReturnedBool   bool
	executeReturnedUInt32 uint32
	changes               []wsvc.Status
}

func setWindowsServiceFuncs(isInteractive bool, onRunningSendCmd *wsvc.Cmd) (*mockWinServiceFuncs, chan<- wsvc.ChangeRequest) {
	changeRequestChan := make(chan wsvc.ChangeRequest, 4)
	changesChan := make(chan wsvc.Status)
	done := make(chan struct{})

	var wsf *mockWinServiceFuncs
	wsf = &mockWinServiceFuncs{
		sigChan: make(chan os.Signal),
		svcIsInteractive: func() (bool, error) {
			return isInteractive, nil
		},
		svcRun: func(name string, handler wsvc.Handler) error {
			ws, ok := handler.(*windowsService)
			if !ok {
				return errors.New("handler is not a windowsService")
			}
			wsf.ws = ws
			wsf.executeReturnedBool, wsf.executeReturnedUInt32 = handler.Execute(nil, changeRequestChan, changesChan)
			done <- struct{}{}
			return nil
		},
	}

	var currentState wsvc.State

	go func() {
	loop:
		for {
			select {
			case change := <-changesChan:
				wsf.changes = append(wsf.changes, change)
				currentState = change.State

				if change.State == wsvc.Running && onRunningSendCmd != nil {
					changeRequestChan <- wsvc.ChangeRequest{
						Cmd:           *onRunningSendCmd,
						CurrentStatus: wsvc.Status{State: currentState},
					}
				}
			case <-done:
				break loop
			}
		}
	}()

	setupWinServiceTest(wsf)

	return wsf, changeRequestChan
}

func TestWinService_RunWindowsService_NonInteractive(t *testing.T) {
	for _, svcCmd := range []wsvc.Cmd{wsvc.Stop, wsvc.Shutdown} {
		testRunWindowsServiceNonInteractive(t, svcCmd)
	}
}

func testRunWindowsServiceNonInteractive(t *testing.T, svcCmd wsvc.Cmd) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(false, &svcCmd)

	// act
	if err := Run(prg); err != nil {
		t.Fatal(err)
	}

	// assert
	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 3, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.StopPending, changes[2].State)

	equal(t, false, wsf.executeReturnedBool)
	equal(t, uint32(0), wsf.executeReturnedUInt32)

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceNonInteractive_StartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.start = func() error {
		startCalled++
		return errors.New("start error")
	}

	svcStop := wsvc.Stop
	wsf, _ := setWindowsServiceFuncs(false, &svcStop)

	// act
	err := Run(prg)

	// assert
	equal(t, "start error", err.Error())

	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 1, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)

	equal(t, true, wsf.executeReturnedBool)
	equal(t, uint32(1), wsf.executeReturnedUInt32)

	equal(t, "start error", wsf.ws.getError().Error())
}

func TestRunWindowsServiceInteractive_StartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.start = func() error {
		startCalled++
		return errors.New("start error")
	}

	wsf, _ := setWindowsServiceFuncs(true, nil)

	// act
	err := Run(prg)

	// assert
	equal(t, "start error", err.Error())

	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsService_BeforeStartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.init = func(Environment) error {
		initCalled++
		return errors.New("before start error")
	}

	wsf, _ := setWindowsServiceFuncs(false, nil)

	// act
	err := Run(prg)

	// assert
	equal(t, "before start error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsService_IsAnInteractiveSessionError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(false, nil)
	wsf.svcIsInteractive = func() (bool, error) {
		return false, errors.New("IsAnInteractiveSession error")
	}

	// act
	err := Run(prg)

	// assert
	equal(t, "IsAnInteractiveSession error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 0, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsServiceNonInteractive_RunError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	svcStop := wsvc.Stop
	wsf, _ := setWindowsServiceFuncs(false, &svcStop)
	wsf.svcRun = func(name string, handler wsvc.Handler) error {
		ws, ok := handler.(*windowsService)
		if !ok {
			return errors.New("handler is not a windowsService")
		}
		wsf.ws = ws
		return errors.New("wsvc.Run error")
	}

	// act
	err := Run(prg)

	// assert
	equal(t, "wsvc.Run error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceNonInteractive_Interrogate(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, changeRequest := setWindowsServiceFuncs(false, nil)

	time.AfterFunc(50*time.Millisecond, func() {
		// ignored, PausePending won't be in changes slice
		// make sure we don't panic/err on unexpected values
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Pause,
			CurrentStatus: wsvc.Status{State: wsvc.PausePending},
		}
	})

	time.AfterFunc(100*time.Millisecond, func() {
		// handled, Paused will be in changes slice
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Interrogate,
			CurrentStatus: wsvc.Status{State: wsvc.Paused},
		}
	})

	time.AfterFunc(200*time.Millisecond, func() {
		// handled, but CurrentStatus overridden with StopPending;
		// ContinuePending won't be in changes slice
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Stop,
			CurrentStatus: wsvc.Status{State: wsvc.ContinuePending},
		}
	})

	// act
	if err := Run(prg); err != nil {
		t.Fatal(err)
	}

	// assert
	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 4, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.Paused, changes[2].State)
	equal(t, wsvc.StopPending, changes[3].State)

	equal(t, false, wsf.executeReturnedBool)
	equal(t, uint32(0), wsf.executeReturnedUInt32)

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceInteractive_StopError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.stop = func() error {
		stopCalled++
		return errors.New("stop error")
	}

	wsf, _ := setWindowsServiceFuncs(true, nil)

	go func() {
		wsf.sigChan <- os.Interrupt
	}()

	// act
	err := Run(prg)

	// assert
	equal(t, "stop error", err.Error())
	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)
	equal(t, 0, len(wsf.changes))
}

func TestRunWindowsServiceNonInteractive_StopError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.stop = func() error {
		stopCalled++
		return errors.New("stop error")
	}

	shutdownCmd := wsvc.Shutdown
	wsf, _ := setWindowsServiceFuncs(false, &shutdownCmd)

	// act
	err := Run(prg)

	// assert
	changes := wsf.changes

	equal(t, "stop error", err.Error())

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 3, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.StopPending, changes[2].State)

	equal(t, true, wsf.executeReturnedBool)
	equal(t, uint32(2), wsf.executeReturnedUInt32)

	equal(t, "stop error", wsf.ws.getError().Error())
}

func TestDefaultSignalHandling(t *testing.T) {
	signals := []os.Signal{syscall.SIGINT} // default signal handled
	for _, signal := range signals {
		testSignalNotify(t, signal)
	}
}

func TestUserDefinedSignalHandling(t *testing.T) {
	signals := []os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP}
	for _, signal := range signals {
		testSignalNotify(t, signal, signals...)
	}
}

func testSignalNotify(t *testing.T, signal os.Signal, sig ...os.Signal) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(true, nil)

	go func() {
		wsf.sigChan <- signal
	}()

	// act
	if err := Run(prg, sig...); err != nil {
		t.Fatal(err)
	}

	// assert
	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)
	equal(t, 0, len(wsf.changes))
}

func equal(t *testing.T, expected, actual interface{}) {
	if !reflect.DeepEqual(expected, actual) {
		_, file, line, _ := runtime.Caller(1)
		t.Logf("\033[31m%s:%d:\n\n\t   %#v (expected)\n\n\t!= %#v (actual)\033[39m\n\n",
			filepath.Base(file), line, expected, actual)
		t.FailNow()
	}
}

func assertNil(t *testing.T, object interface{}) {
	if !isNil(object) {
		_, file, line, _ := runtime.Caller(1)
		t.Logf("\033[31m%s:%d:\n\n\t   <nil> (expected)\n\n\t!= %#v (actual)\033[39m\n\n",
			filepath.Base(file), line, object)
		t.FailNow()
	}
}

func isNil(object interface{}) bool {
	if object == nil {
		return true
	}

	value := reflect.ValueOf(object)
	kind := value.Kind()
	if kind >= reflect.Chan && kind <= reflect.Slice && value.IsNil() {
		return true
	}

	return false
}
<commit_msg>test: fix windows test imports<commit_after>// +build windows

package svc

import (
	"errors"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"syscall"
	"testing"
	"time"

	wsvc "golang.org/x/sys/windows/svc"
)

func setupWinServiceTest(wsf *mockWinServiceFuncs) {
	// wsfWrapper allows signalNotify, svcIsInteractive, and svcRun to be set once.
	// Individual test functions set "wsf" to add behavior.
	wsfWrapper := &mockWinServiceFuncs{
		signalNotify: func(c chan<- os.Signal, sig ...os.Signal) {
			if c == nil {
				panic("os/signal: Notify using nil channel")
			}

			if wsf.signalNotify != nil {
				wsf.signalNotify(c, sig...)
			} else {
				wsf1 := *wsf
				go func() {
					for val := range wsf1.sigChan {
						for _, registeredSig := range sig {
							if val == registeredSig {
								c <- val
							}
						}
					}
				}()
			}
		},
		svcIsInteractive: func() (bool, error) {
			return wsf.svcIsInteractive()
		},
		svcRun: func(name string, handler wsvc.Handler) error {
			return wsf.svcRun(name, handler)
		},
	}

	signalNotify = wsfWrapper.signalNotify
	svcIsAnInteractiveSession = wsfWrapper.svcIsInteractive
	svcRun = wsfWrapper.svcRun
}

type mockWinServiceFuncs struct {
	signalNotify          func(chan<- os.Signal, ...os.Signal)
	svcIsInteractive      func() (bool, error)
	sigChan               chan os.Signal
	svcRun                func(string, wsvc.Handler) error
	ws                    *windowsService
	executeReturnedBool   bool
	executeReturnedUInt32 uint32
	changes               []wsvc.Status
}

func setWindowsServiceFuncs(isInteractive bool, onRunningSendCmd *wsvc.Cmd) (*mockWinServiceFuncs, chan<- wsvc.ChangeRequest) {
	changeRequestChan := make(chan wsvc.ChangeRequest, 4)
	changesChan := make(chan wsvc.Status)
	done := make(chan struct{})

	var wsf *mockWinServiceFuncs
	wsf = &mockWinServiceFuncs{
		sigChan: make(chan os.Signal),
		svcIsInteractive: func() (bool, error) {
			return isInteractive, nil
		},
		svcRun: func(name string, handler wsvc.Handler) error {
			ws, ok := handler.(*windowsService)
			if !ok {
				return errors.New("handler is not a windowsService")
			}
			wsf.ws = ws
			wsf.executeReturnedBool, wsf.executeReturnedUInt32 = handler.Execute(nil, changeRequestChan, changesChan)
			done <- struct{}{}
			return nil
		},
	}

	var currentState wsvc.State

	go func() {
	loop:
		for {
			select {
			case change := <-changesChan:
				wsf.changes = append(wsf.changes, change)
				currentState = change.State

				if change.State == wsvc.Running && onRunningSendCmd != nil {
					changeRequestChan <- wsvc.ChangeRequest{
						Cmd:           *onRunningSendCmd,
						CurrentStatus: wsvc.Status{State: currentState},
					}
				}
			case <-done:
				break loop
			}
		}
	}()

	setupWinServiceTest(wsf)

	return wsf, changeRequestChan
}

func TestWinService_RunWindowsService_NonInteractive(t *testing.T) {
	for _, svcCmd := range []wsvc.Cmd{wsvc.Stop, wsvc.Shutdown} {
		testRunWindowsServiceNonInteractive(t, svcCmd)
	}
}

func testRunWindowsServiceNonInteractive(t *testing.T, svcCmd wsvc.Cmd) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(false, &svcCmd)

	// act
	if err := Run(prg); err != nil {
		t.Fatal(err)
	}

	// assert
	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 3, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.StopPending, changes[2].State)

	equal(t, false, wsf.executeReturnedBool)
	equal(t, uint32(0), wsf.executeReturnedUInt32)

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceNonInteractive_StartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.start = func() error {
		startCalled++
		return errors.New("start error")
	}

	svcStop := wsvc.Stop
	wsf, _ := setWindowsServiceFuncs(false, &svcStop)

	// act
	err := Run(prg)

	// assert
	equal(t, "start error", err.Error())

	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 1, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)

	equal(t, true, wsf.executeReturnedBool)
	equal(t, uint32(1), wsf.executeReturnedUInt32)

	equal(t, "start error", wsf.ws.getError().Error())
}

func TestRunWindowsServiceInteractive_StartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.start = func() error {
		startCalled++
		return errors.New("start error")
	}

	wsf, _ := setWindowsServiceFuncs(true, nil)

	// act
	err := Run(prg)

	// assert
	equal(t, "start error", err.Error())

	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsService_BeforeStartError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.init = func(Environment) error {
		initCalled++
		return errors.New("before start error")
	}

	wsf, _ := setWindowsServiceFuncs(false, nil)

	// act
	err := Run(prg)

	// assert
	equal(t, "before start error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsService_IsAnInteractiveSessionError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(false, nil)
	wsf.svcIsInteractive = func() (bool, error) {
		return false, errors.New("IsAnInteractiveSession error")
	}

	// act
	err := Run(prg)

	// assert
	equal(t, "IsAnInteractiveSession error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 0, initCalled)

	equal(t, 0, len(changes))
}

func TestRunWindowsServiceNonInteractive_RunError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	svcStop := wsvc.Stop
	wsf, _ := setWindowsServiceFuncs(false, &svcStop)
	wsf.svcRun = func(name string, handler wsvc.Handler) error {
		ws, ok := handler.(*windowsService)
		if !ok {
			return errors.New("handler is not a windowsService")
		}
		wsf.ws = ws
		return errors.New("wsvc.Run error")
	}

	// act
	err := Run(prg)

	// assert
	equal(t, "wsvc.Run error", err.Error())

	changes := wsf.changes

	equal(t, 0, startCalled)
	equal(t, 0, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 0, len(changes))

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceNonInteractive_Interrogate(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, changeRequest := setWindowsServiceFuncs(false, nil)

	time.AfterFunc(50*time.Millisecond, func() {
		// ignored, PausePending won't be in changes slice
		// make sure we don't panic/err on unexpected values
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Pause,
			CurrentStatus: wsvc.Status{State: wsvc.PausePending},
		}
	})

	time.AfterFunc(100*time.Millisecond, func() {
		// handled, Paused will be in changes slice
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Interrogate,
			CurrentStatus: wsvc.Status{State: wsvc.Paused},
		}
	})

	time.AfterFunc(200*time.Millisecond, func() {
		// handled, but CurrentStatus overridden with StopPending;
		// ContinuePending won't be in changes slice
		changeRequest <- wsvc.ChangeRequest{
			Cmd:           wsvc.Stop,
			CurrentStatus: wsvc.Status{State: wsvc.ContinuePending},
		}
	})

	// act
	if err := Run(prg); err != nil {
		t.Fatal(err)
	}

	// assert
	changes := wsf.changes

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 4, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.Paused, changes[2].State)
	equal(t, wsvc.StopPending, changes[3].State)

	equal(t, false, wsf.executeReturnedBool)
	equal(t, uint32(0), wsf.executeReturnedUInt32)

	assertNil(t, wsf.ws.getError())
}

func TestRunWindowsServiceInteractive_StopError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.stop = func() error {
		stopCalled++
		return errors.New("stop error")
	}

	wsf, _ := setWindowsServiceFuncs(true, nil)

	go func() {
		wsf.sigChan <- os.Interrupt
	}()

	// act
	err := Run(prg)

	// assert
	equal(t, "stop error", err.Error())
	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)
	equal(t, 0, len(wsf.changes))
}

func TestRunWindowsServiceNonInteractive_StopError(t *testing.T) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)
	prg.stop = func() error {
		stopCalled++
		return errors.New("stop error")
	}

	shutdownCmd := wsvc.Shutdown
	wsf, _ := setWindowsServiceFuncs(false, &shutdownCmd)

	// act
	err := Run(prg)

	// assert
	changes := wsf.changes

	equal(t, "stop error", err.Error())

	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)

	equal(t, 3, len(changes))
	equal(t, wsvc.StartPending, changes[0].State)
	equal(t, wsvc.Running, changes[1].State)
	equal(t, wsvc.StopPending, changes[2].State)

	equal(t, true, wsf.executeReturnedBool)
	equal(t, uint32(2), wsf.executeReturnedUInt32)

	equal(t, "stop error", wsf.ws.getError().Error())
}

func TestDefaultSignalHandling(t *testing.T) {
	signals := []os.Signal{syscall.SIGINT} // default signal handled
	for _, signal := range signals {
		testSignalNotify(t, signal)
	}
}

func TestUserDefinedSignalHandling(t *testing.T) {
	signals := []os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP}
	for _, signal := range signals {
		testSignalNotify(t, signal, signals...)
	}
}

func testSignalNotify(t *testing.T, signal os.Signal, sig ...os.Signal) {
	// arrange
	var startCalled, stopCalled, initCalled int
	prg := makeProgram(&startCalled, &stopCalled, &initCalled)

	wsf, _ := setWindowsServiceFuncs(true, nil)

	go func() {
		wsf.sigChan <- signal
	}()

	// act
	if err := Run(prg, sig...); err != nil {
		t.Fatal(err)
	}

	// assert
	equal(t, 1, startCalled)
	equal(t, 1, stopCalled)
	equal(t, 1, initCalled)
	equal(t, 0, len(wsf.changes))
}

func equal(t *testing.T, expected, actual interface{}) {
	if !reflect.DeepEqual(expected, actual) {
		_, file, line, _ := runtime.Caller(1)
		t.Logf("\033[31m%s:%d:\n\n\t   %#v (expected)\n\n\t!= %#v (actual)\033[39m\n\n",
			filepath.Base(file), line, expected, actual)
		t.FailNow()
	}
}

func assertNil(t *testing.T, object interface{}) {
	if !isNil(object) {
		_, file, line, _ := runtime.Caller(1)
		t.Logf("\033[31m%s:%d:\n\n\t   <nil> (expected)\n\n\t!= %#v (actual)\033[39m\n\n",
			filepath.Base(file), line, object)
		t.FailNow()
	}
}

func isNil(object interface{}) bool {
	if object == nil {
		return true
	}

	value := reflect.ValueOf(object)
	kind := value.Kind()
	if kind >= reflect.Chan && kind <= reflect.Slice && value.IsNil() {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go-html-transform/h5"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
)

func runDatTests(ps []string) {
	//for _, p := range ps {
	//}
}

func runTestTests(ps []string) {
	//for _, p := range ps {
	//}
}

func runHtmlTests(ps []string) {
	for _, p := range ps {
		fmt.Println("Attempting to parse file: ", p)
		f, err := os.Open(p)
		if err != nil {
			fmt.Println("Error opening file: ", err)
		}
		parse := h5.NewParser(f)
		err = parse.Parse()
		if err != nil {
			fmt.Println("Error parsing file: ", err)
		} else {
			fmt.Println("Success!!!")
		}
	}
}

type grepSpec map[*regexp.Regexp][]string

func grep(path string, spec grepSpec) error {
	return filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			for re, _ := range spec {
				if re.MatchString(p) {
					spec[re] = append(spec[re], p)
				}
			}
		}
		return nil
	})
}

func main() {
	datRe := regexp.MustCompile("dat$")
	testRe := regexp.MustCompile("test$")
	htmlRe := regexp.MustCompile("html?$")
	spec := make(map[*regexp.Regexp][]string)
	spec[datRe] = []string{}
	spec[testRe] = []string{}
	spec[htmlRe] = []string{}
	err := grep("./", spec)
	if err != nil {
		fmt.Println("Error while grepping: ", err)
	}
	runDatTests(spec[datRe])
	runTestTests(spec[testRe])
	runHtmlTests(spec[htmlRe])
}<commit_msg>Add todo about timing the parsing algorithms.<commit_after>package main

import (
	"code.google.com/p/go-html-transform/h5"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
)

func runDatTests(ps []string) {
	//for _, p := range ps {
	//}
}

func runTestTests(ps []string) {
	//for _, p := range ps {
	//}
}

func runHtmlTests(ps []string) {
	// TODO(jwall): with timings?
	for _, p := range ps {
		fmt.Println("Attempting to parse file: ", p)
		f, err := os.Open(p)
		if err != nil {
			fmt.Println("Error opening file: ", err)
		}
		parse := h5.NewParser(f)
		err = parse.Parse()
		if err != nil {
			fmt.Println("Error parsing file: ", err)
		} else {
			fmt.Println("Success!!!")
		}
	}
}

type grepSpec map[*regexp.Regexp][]string

func grep(path string, spec grepSpec) error {
	return filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			for re, _ := range spec {
				if re.MatchString(p) {
					spec[re] = append(spec[re], p)
				}
			}
		}
		return nil
	})
}

func main() {
	datRe := regexp.MustCompile("dat$")
	testRe := regexp.MustCompile("test$")
	htmlRe := regexp.MustCompile("html?$")
	spec := make(map[*regexp.Regexp][]string)
	spec[datRe] = []string{}
	spec[testRe] = []string{}
	spec[htmlRe] = []string{}
	err := grep("./", spec)
	if err != nil {
		fmt.Println("Error while grepping: ", err)
	}
	runDatTests(spec[datRe])
	runTestTests(spec[testRe])
	runHtmlTests(spec[htmlRe])
}<|endoftext|>


<commit_before>package catalog

import (
	"github.com/coreos-inc/operator-client/pkg/client"
	"github.com/coreos/go-semver/semver"
)

// Service Catalog

// 1. InstallDeclaration loop

// Strawpeople:
//
// ```yaml
// declare:
//  - Vault AppType 1.0.0 (ref by sha?)
// approval: manual/automatic
// status: Unresolved
//```
//
//```yaml
// declare:
//  - Vault AppType (ref by sha?)
// approval: manual/automatic
// resolved:
//   - Vault AppType 1.0.0
//   - Vault OpVer
//   - VaultService CRD
//   - Etcd AppType
//   - Etcd OpVer
//   - EtcdCluster CRD
// status: resolved
//```

// States: Unresolved, Resolved, Approved, Complete

// A. Watches for new InstallDeclarations in a namespace
//     i. If Unresolved, attempt to resolve those resources and write them back to the
//        InstallDeclaration
//    ii. If Resolved, wait for state to be Approved
//   iii. If approval is set to automatic, state is transitioned to Approved
//    iv. If Approved, creates all resolved resources, reports back status
//     v. If Complete, nothing

// 2. Subscription loop

//```yaml
// type: Subscription
// declare:
//  - Vault Apptype
// source: quay
// package: vault
// channel: stable
// approval: manual/automatic
// status:
//   current: v1.0.0
// ---
// type: CatalogSource
// url: quay.io/catalog
// name: quay
//```

// A. Watches for Subscription objects
//      i. If no InstallDeclaration exists for the Subscription, creates it
//     ii. Checks CatalogSource for updates
//    iii. If newer version is available in the channel and is greater than current, creates an
//         InstallDeclaration for it.

// TEMP - to be defined elsewhere
type Subscription interface {
	AppType() string
	CurrentVersion() semver.Version
}

type InstallDeclaration struct{}

type Catalog interface {
	FetchLatestVersion(apptype string) (*semver.Version, error) // TODO location? appid?
	FetchInstallDeclarationForAppVersion(apptype string, ver *semver.Version) (InstallDeclaration, error)
	ResolveDependencies(decl InstallDeclaration) error
}

// TEMP - to be filled in with k8s and logic
type Installer interface {
	CheckForInstallDeclaration(sub Subscription) (bool, error)
	CreateInstallDeclaration(InstallDeclaration) error
}

type SubscriptionController struct {
	catalog   Catalog
	client    client.Interface
	installer Installer
}

func (ctl *SubscriptionController) installDeclarationForSubscription(sub Subscription) (*InstallDeclaration, error) {
	decl, err := ctl.catalog.FetchInstallDeclarationForAppVersion(sub.AppType(), sub.CurrentVersion())
	if err != nil {
		return nil, err
	}
	return &decl, ctl.installer.CreateInstallDeclaration(decl)
}

func (ctl *SubscriptionController) HandleNewSubscription(sub Subscription) error {
	ok, err := ctl.installer.CheckForInstallDeclaration(sub)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}
	_, err = ctl.installDeclarationForSubscription(sub)
	return err
}

func (ctl *SubscriptionController) CheckCatalogForUpdate(sub Subscription) error {
	ver, err := ctl.catalog.FetchLatestVersion(sub.AppType())
	if err != nil {
		return err
	}
	currVer := sub.CurrentVersion()
	if !currVer.LessThan(ver) {
		return nil
	}
	_, err = ctl.installDeclarationForSubscription(sub)
	return err
}
<commit_msg>doc(catalog): add comment on temporary type decl<commit_after>package catalog

import (
	"github.com/coreos-inc/operator-client/pkg/client"
	"github.com/coreos/go-semver/semver"
)

// Service Catalog

// 1. InstallDeclaration loop

// Strawpeople:
//
// ```yaml
// declare:
//  - Vault AppType 1.0.0 (ref by sha?)
// approval: manual/automatic
// status: Unresolved
//```
//
//```yaml
// declare:
//  - Vault AppType (ref by sha?)
// approval: manual/automatic
// resolved:
//   - Vault AppType 1.0.0
//   - Vault OpVer
//   - VaultService CRD
//   - Etcd AppType
//   - Etcd OpVer
//   - EtcdCluster CRD
// status: resolved
//```

// States: Unresolved, Resolved, Approved, Complete

// A. Watches for new InstallDeclarations in a namespace
//     i. If Unresolved, attempt to resolve those resources and write them back to the
//        InstallDeclaration
//    ii. If Resolved, wait for state to be Approved
//   iii. If approval is set to automatic, state is transitioned to Approved
//    iv. If Approved, creates all resolved resources, reports back status
//     v. If Complete, nothing

// 2. Subscription loop

//```yaml
// type: Subscription
// declare:
//  - Vault Apptype
// source: quay
// package: vault
// channel: stable
// approval: manual/automatic
// status:
//   current: v1.0.0
// ---
// type: CatalogSource
// url: quay.io/catalog
// name: quay
//```

// A. Watches for Subscription objects
//      i. If no InstallDeclaration exists for the Subscription, creates it
//     ii. Checks CatalogSource for updates
//    iii. If newer version is available in the channel and is greater than current, creates an
//         InstallDeclaration for it.

// TEMP - to be defined elsewhere
type Subscription interface {
	AppType() string
	CurrentVersion() semver.Version
}

// TEMP - to be defined elsewhere
type InstallDeclaration struct{}

type Catalog interface {
	FetchLatestVersion(apptype string) (*semver.Version, error) // TODO location? appid?
	FetchInstallDeclarationForAppVersion(apptype string, ver *semver.Version) (InstallDeclaration, error)
	ResolveDependencies(decl InstallDeclaration) error
}

// TEMP - to be filled in with k8s and logic
type Installer interface {
	CheckForInstallDeclaration(sub Subscription) (bool, error)
	CreateInstallDeclaration(InstallDeclaration) error
}

type SubscriptionController struct {
	catalog   Catalog
	client    client.Interface
	installer Installer
}

func (ctl *SubscriptionController) installDeclarationForSubscription(sub Subscription) (*InstallDeclaration, error) {
	decl, err := ctl.catalog.FetchInstallDeclarationForAppVersion(sub.AppType(), sub.CurrentVersion())
	if err != nil {
		return nil, err
	}
	return &decl, ctl.installer.CreateInstallDeclaration(decl)
}

func (ctl *SubscriptionController) HandleNewSubscription(sub Subscription) error {
	ok, err := ctl.installer.CheckForInstallDeclaration(sub)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}
	_, err = ctl.installDeclarationForSubscription(sub)
	return err
}

func (ctl *SubscriptionController) CheckCatalogForUpdate(sub Subscription) error {
	ver, err := ctl.catalog.FetchLatestVersion(sub.AppType())
	if err != nil {
		return err
	}
	currVer := sub.CurrentVersion()
	if !currVer.LessThan(ver) {
		return nil
	}
	_, err = ctl.installDeclarationForSubscription(sub)
	return err
}
<|endoftext|>