<commit_before>/*
Copyright 2015 The Kubernetes Authors 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 e2e

import (
	"strconv"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/util"

	. "github.com/onsi/ginkgo"
)

const (
	dynamicConsumptionTimeInSeconds = 30
	staticConsumptionTimeInSeconds  = 3600
	dynamicRequestSizeInMillicores  = 100
	dynamicRequestSizeInMegabytes   = 100
	port                            = 80
	targetPort                      = 8080
	timeoutRC                       = 120 * time.Second
	startServiceTimeout             = time.Minute
	startServiceInterval            = 5 * time.Second
	image                           = "gcr.io/google_containers/resource_consumer:beta"
	rcIsNil                         = "ERROR: replicationController = nil"
)

/*
ResourceConsumer is a tool for testing. It helps create specified usage of CPU or memory (Warnig: memory not supported)
typical use case:
rc.ConsumeCPU(600)
// ... check your assumption here
rc.ConsumeCPU(300)
// ... check your assumption here
*/
type ResourceConsumer struct {
	name                     string
	framework                *Framework
	cpu                      chan int
	mem                      chan int
	stopCPU                  chan int
	stopMem                  chan int
	consumptionTimeInSeconds int
	sleepTime                time.Duration
	requestSizeInMillicores  int
	requestSizeInMegabytes   int
}

func NewDynamicResourceConsumer(name string, replicas, initCPU, initMemory int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	return newResourceConsumer(name, replicas, initCPU, initMemory, dynamicConsumptionTimeInSeconds, dynamicRequestSizeInMillicores, dynamicRequestSizeInMegabytes, cpuLimit, memLimit, framework)
}

func NewStaticResourceConsumer(name string, replicas, initCPU, initMemory int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	return newResourceConsumer(name, replicas, initCPU, initMemory, staticConsumptionTimeInSeconds, initCPU/replicas, initMemory/replicas, cpuLimit, memLimit, framework)
}

/*
NewResourceConsumer creates new ResourceConsumer
initCPU argument is in millicores
initMemory argument is in megabytes
memLimit argument is in megabytes, memLimit is a maximum amount of memory that can be consumed by a single pod
cpuLimit argument is in millicores, cpuLimit is a maximum amount of cpu that can be consumed by a single pod
*/
func newResourceConsumer(name string, replicas, initCPU, initMemory, consumptionTimeInSeconds, requestSizeInMillicores, requestSizeInMegabytes int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	runServiceAndRCForResourceConsumer(framework.Client, framework.Namespace.Name, name, replicas, cpuLimit, memLimit)
	rc := &ResourceConsumer{
		name:                     name,
		framework:                framework,
		cpu:                      make(chan int),
		mem:                      make(chan int),
		stopCPU:                  make(chan int),
		stopMem:                  make(chan int),
		consumptionTimeInSeconds: consumptionTimeInSeconds,
		sleepTime:                time.Duration(consumptionTimeInSeconds) * time.Second,
		requestSizeInMillicores:  requestSizeInMillicores,
		requestSizeInMegabytes:   requestSizeInMegabytes,
	}
	go rc.makeConsumeCPURequests()
	rc.ConsumeCPU(initCPU)
	go rc.makeConsumeMemRequests()
	rc.ConsumeMem(initMemory)
	return rc
}

// ConsumeCPU consumes given number of CPU
func (rc *ResourceConsumer) ConsumeCPU(millicores int) {
	rc.cpu <- millicores
}

// ConsumeMem consumes given number of Mem
func (rc *ResourceConsumer) ConsumeMem(megabytes int) {
	rc.mem <- megabytes
}

func (rc *ResourceConsumer) makeConsumeCPURequests() {
	var count int
	var rest int
	sleepTime := time.Duration(0)
	for {
		select {
		case millicores := <-rc.cpu:
			if rc.requestSizeInMillicores != 0 {
				count = millicores / rc.requestSizeInMillicores
			}
			rest = millicores - count*rc.requestSizeInMillicores
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeCPURequests(count, rc.requestSizeInMillicores, rc.consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeCPURequest(rest, rc.consumptionTimeInSeconds)
			}
			sleepTime = rc.sleepTime
		case <-rc.stopCPU:
			return
		}
	}
}

func (rc *ResourceConsumer) makeConsumeMemRequests() {
	var count int
	var rest int
	sleepTime := time.Duration(0)
	for {
		select {
		case megabytes := <-rc.mem:
			if rc.requestSizeInMegabytes != 0 {
				count = megabytes / rc.requestSizeInMegabytes
			}
			rest = megabytes - count*rc.requestSizeInMegabytes
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeMemRequests(count, rc.requestSizeInMegabytes, rc.consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeMemRequest(rest, rc.consumptionTimeInSeconds)
			}
			sleepTime = rc.sleepTime
		case <-rc.stopMem:
			return
		}
	}
}

func (rc *ResourceConsumer) sendConsumeCPURequests(requests, millicores, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeCPURequest(millicores, durationSec)
	}
}

func (rc *ResourceConsumer) sendConsumeMemRequests(requests, megabytes, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeMemRequest(megabytes, durationSec)
	}
}

// sendOneConsumeCPURequest sends POST request for cpu consumption
func (rc *ResourceConsumer) sendOneConsumeCPURequest(millicores int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeCPU").
		Param("millicores", strconv.Itoa(millicores)).
		Param("durationSec", strconv.Itoa(durationSec)).
		Do().
		Raw()
	expectNoError(err)
}

// sendOneConsumeMemRequest sends POST request for memory consumption
func (rc *ResourceConsumer) sendOneConsumeMemRequest(megabytes int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeMem").
		Param("megabytes", strconv.Itoa(megabytes)).
		Param("durationSec", strconv.Itoa(durationSec)).
		Do().
		Raw()
	expectNoError(err)
}

func (rc *ResourceConsumer) GetReplicas() int {
	replicationController, err := rc.framework.Client.ReplicationControllers(rc.framework.Namespace.Name).Get(rc.name)
	expectNoError(err)
	if replicationController == nil {
		Failf(rcIsNil)
	}
	return replicationController.Status.Replicas
}

func (rc *ResourceConsumer) WaitForReplicas(desiredReplicas int) {
	timeout := 10 * time.Minute
	for start := time.Now(); time.Since(start) < timeout; time.Sleep(20 * time.Second) {
		if desiredReplicas == rc.GetReplicas() {
			Logf("Replication Controller current replicas number is equal to desired replicas number: %d", desiredReplicas)
			return
		} else {
			Logf("Replication Controller current replicas number %d waiting to be %d", rc.GetReplicas(), desiredReplicas)
		}
	}
	Failf("timeout waiting %v for pods size to be %d", timeout, desiredReplicas)
}

func (rc *ResourceConsumer) CleanUp() {
	rc.stopCPU <- 0
	rc.stopMem <- 0
	// Wait some time to ensure all child goroutines are finished.
	time.Sleep(10 * time.Second)
	expectNoError(DeleteRC(rc.framework.Client, rc.framework.Namespace.Name, rc.name))
	expectNoError(rc.framework.Client.Services(rc.framework.Namespace.Name).Delete(rc.name))
}

func runServiceAndRCForResourceConsumer(c *client.Client, ns, name string, replicas int, cpuLimitMillis, memLimitMb int64) {
	_, err := c.Services(ns).Create(&api.Service{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ServiceSpec{
			Ports: []api.ServicePort{{
				Port:       port,
				TargetPort: util.NewIntOrStringFromInt(targetPort),
			}},

			Selector: map[string]string{
				"name": name,
			},
		},
	})
	expectNoError(err)
	config := RCConfig{
		Client:    c,
		Image:     image,
		Name:      name,
		Namespace: ns,
		Timeout:   timeoutRC,
		Replicas:  replicas,
		CpuLimit:  cpuLimitMillis,
		MemLimit:  memLimitMb * 1024 * 1024, // MemLimit is in bytes
	}
	expectNoError(RunRC(config))
	// Make sure endpoints are propagated.
	// TODO(piosz): replace sleep with endpoints watch.
	time.Sleep(10 * time.Second)
}
<commit_msg>Small fixes in autoscaling e2e utils<commit_after>/*
Copyright 2015 The Kubernetes Authors 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 e2e

import (
	"strconv"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/util"

	. "github.com/onsi/ginkgo"
)

const (
	dynamicConsumptionTimeInSeconds = 30
	staticConsumptionTimeInSeconds  = 3600
	dynamicRequestSizeInMillicores  = 100
	dynamicRequestSizeInMegabytes   = 100
	port                            = 80
	targetPort                      = 8080
	timeoutRC                       = 120 * time.Second
	startServiceTimeout             = time.Minute
	startServiceInterval            = 5 * time.Second
	image                           = "gcr.io/google_containers/resource_consumer:beta"
	rcIsNil                         = "ERROR: replicationController = nil"
)

/*
ResourceConsumer is a tool for testing. It helps create specified usage of CPU or memory (Warnig: memory not supported)
typical use case:
rc.ConsumeCPU(600)
// ... check your assumption here
rc.ConsumeCPU(300)
// ... check your assumption here
*/
type ResourceConsumer struct {
	name                     string
	framework                *Framework
	cpu                      chan int
	mem                      chan int
	stopCPU                  chan int
	stopMem                  chan int
	consumptionTimeInSeconds int
	sleepTime                time.Duration
	requestSizeInMillicores  int
	requestSizeInMegabytes   int
}

func NewDynamicResourceConsumer(name string, replicas, initCPU, initMemory int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	return newResourceConsumer(name, replicas, initCPU, initMemory, dynamicConsumptionTimeInSeconds, dynamicRequestSizeInMillicores, dynamicRequestSizeInMegabytes, cpuLimit, memLimit, framework)
}

func NewStaticResourceConsumer(name string, replicas, initCPU, initMemory int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	return newResourceConsumer(name, replicas, initCPU, initMemory, staticConsumptionTimeInSeconds, initCPU/replicas, initMemory/replicas, cpuLimit, memLimit, framework)
}

/*
NewResourceConsumer creates new ResourceConsumer
initCPU argument is in millicores
initMemory argument is in megabytes
memLimit argument is in megabytes, memLimit is a maximum amount of memory that can be consumed by a single pod
cpuLimit argument is in millicores, cpuLimit is a maximum amount of cpu that can be consumed by a single pod
*/
func newResourceConsumer(name string, replicas, initCPU, initMemory, consumptionTimeInSeconds, requestSizeInMillicores, requestSizeInMegabytes int, cpuLimit, memLimit int64, framework *Framework) *ResourceConsumer {
	runServiceAndRCForResourceConsumer(framework.Client, framework.Namespace.Name, name, replicas, cpuLimit, memLimit)
	rc := &ResourceConsumer{
		name:                     name,
		framework:                framework,
		cpu:                      make(chan int),
		mem:                      make(chan int),
		stopCPU:                  make(chan int),
		stopMem:                  make(chan int),
		consumptionTimeInSeconds: consumptionTimeInSeconds,
		sleepTime:                time.Duration(consumptionTimeInSeconds) * time.Second,
		requestSizeInMillicores:  requestSizeInMillicores,
		requestSizeInMegabytes:   requestSizeInMegabytes,
	}
	go rc.makeConsumeCPURequests()
	rc.ConsumeCPU(initCPU)
	go rc.makeConsumeMemRequests()
	rc.ConsumeMem(initMemory)
	return rc
}

// ConsumeCPU consumes given number of CPU
func (rc *ResourceConsumer) ConsumeCPU(millicores int) {
	rc.cpu <- millicores
}

// ConsumeMem consumes given number of Mem
func (rc *ResourceConsumer) ConsumeMem(megabytes int) {
	rc.mem <- megabytes
}

func (rc *ResourceConsumer) makeConsumeCPURequests() {
	defer GinkgoRecover()
	var count int
	var rest int
	sleepTime := time.Duration(0)
	for {
		select {
		case millicores := <-rc.cpu:
			if rc.requestSizeInMillicores != 0 {
				count = millicores / rc.requestSizeInMillicores
			}
			rest = millicores - count*rc.requestSizeInMillicores
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeCPURequests(count, rc.requestSizeInMillicores, rc.consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeCPURequest(rest, rc.consumptionTimeInSeconds)
			}
			sleepTime = rc.sleepTime
		case <-rc.stopCPU:
			return
		}
	}
}

func (rc *ResourceConsumer) makeConsumeMemRequests() {
	defer GinkgoRecover()
	var count int
	var rest int
	sleepTime := time.Duration(0)
	for {
		select {
		case megabytes := <-rc.mem:
			if rc.requestSizeInMegabytes != 0 {
				count = megabytes / rc.requestSizeInMegabytes
			}
			rest = megabytes - count*rc.requestSizeInMegabytes
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeMemRequests(count, rc.requestSizeInMegabytes, rc.consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeMemRequest(rest, rc.consumptionTimeInSeconds)
			}
			sleepTime = rc.sleepTime
		case <-rc.stopMem:
			return
		}
	}
}

func (rc *ResourceConsumer) sendConsumeCPURequests(requests, millicores, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeCPURequest(millicores, durationSec)
	}
}

func (rc *ResourceConsumer) sendConsumeMemRequests(requests, megabytes, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeMemRequest(megabytes, durationSec)
	}
}

// sendOneConsumeCPURequest sends POST request for cpu consumption
func (rc *ResourceConsumer) sendOneConsumeCPURequest(millicores int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeCPU").
		Param("millicores", strconv.Itoa(millicores)).
		Param("durationSec", strconv.Itoa(durationSec)).
		DoRaw()
	expectNoError(err)
}

// sendOneConsumeMemRequest sends POST request for memory consumption
func (rc *ResourceConsumer) sendOneConsumeMemRequest(megabytes int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeMem").
		Param("megabytes", strconv.Itoa(megabytes)).
		Param("durationSec", strconv.Itoa(durationSec)).
		DoRaw()
	expectNoError(err)
}

func (rc *ResourceConsumer) GetReplicas() int {
	replicationController, err := rc.framework.Client.ReplicationControllers(rc.framework.Namespace.Name).Get(rc.name)
	expectNoError(err)
	if replicationController == nil {
		Failf(rcIsNil)
	}
	return replicationController.Status.Replicas
}

func (rc *ResourceConsumer) WaitForReplicas(desiredReplicas int) {
	timeout := 10 * time.Minute
	for start := time.Now(); time.Since(start) < timeout; time.Sleep(20 * time.Second) {
		if desiredReplicas == rc.GetReplicas() {
			Logf("Replication Controller current replicas number is equal to desired replicas number: %d", desiredReplicas)
			return
		} else {
			Logf("Replication Controller current replicas number %d waiting to be %d", rc.GetReplicas(), desiredReplicas)
		}
	}
	Failf("timeout waiting %v for pods size to be %d", timeout, desiredReplicas)
}

func (rc *ResourceConsumer) CleanUp() {
	rc.stopCPU <- 0
	rc.stopMem <- 0
	// Wait some time to ensure all child goroutines are finished.
	time.Sleep(10 * time.Second)
	expectNoError(DeleteRC(rc.framework.Client, rc.framework.Namespace.Name, rc.name))
	expectNoError(rc.framework.Client.Services(rc.framework.Namespace.Name).Delete(rc.name))
}

func runServiceAndRCForResourceConsumer(c *client.Client, ns, name string, replicas int, cpuLimitMillis, memLimitMb int64) {
	_, err := c.Services(ns).Create(&api.Service{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ServiceSpec{
			Ports: []api.ServicePort{{
				Port:       port,
				TargetPort: util.NewIntOrStringFromInt(targetPort),
			}},

			Selector: map[string]string{
				"name": name,
			},
		},
	})
	expectNoError(err)
	config := RCConfig{
		Client:    c,
		Image:     image,
		Name:      name,
		Namespace: ns,
		Timeout:   timeoutRC,
		Replicas:  replicas,
		CpuLimit:  cpuLimitMillis,
		MemLimit:  memLimitMb * 1024 * 1024, // MemLimit is in bytes
	}
	expectNoError(RunRC(config))
	// Make sure endpoints are propagated.
	// TODO(piosz): replace sleep with endpoints watch.
	time.Sleep(10 * time.Second)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 framework

import (
	"os"
	"path/filepath"
	"time"

	habv1 "github.com/kinvolk/habitat-operator/pkg/apis/habitat/v1"

	appsv1beta1 "k8s.io/api/apps/v1beta1"
	"k8s.io/api/core/v1"
	apiv1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/util/yaml"
)

// CreateHabitat creates a Habitat.
func (f *Framework) CreateHabitat(habitat *habv1.Habitat) error {
	return f.Client.Post().
		Namespace(TestNs).
		Resource(habv1.HabitatResourcePlural).
		Body(habitat).
		Do().
		Error()
}

// WaitForResources waits until numPods are in the "Running" state.
// We wait for pods, because those take the longest to create.
// Waiting for anything else would be already testing.
func (f *Framework) WaitForResources(labelName, habitatName string, numPods int) error {
	return wait.Poll(2*time.Second, 5*time.Minute, func() (bool, error) {
		fs := fields.SelectorFromSet(fields.Set{
			"status.phase": "Running",
		})

		ls := labels.SelectorFromSet(labels.Set{
			labelName: habitatName,
		})

		pods, err := f.KubeClient.CoreV1().Pods(TestNs).List(metav1.ListOptions{FieldSelector: fs.String(), LabelSelector: ls.String()})
		if err != nil {
			return false, err
		}

		if len(pods.Items) != numPods {
			return false, nil
		}

		return true, nil
	})
}

func (f *Framework) WaitForEndpoints(habitatName string) error {
	return wait.Poll(time.Second, time.Minute*5, func() (bool, error) {
		ep, err := f.KubeClient.CoreV1().Endpoints(TestNs).Get(habitatName, metav1.GetOptions{})
		if err != nil {
			return false, err
		}

		if len(ep.Subsets) == 0 && len(ep.Subsets[0].Addresses) == 0 {
			return false, nil
		}

		return true, nil
	})
}

// DeleteHabitat deletes a Habitat as a user would.
func (f *Framework) DeleteHabitat(habitatName string) error {
	return f.Client.Delete().
		Namespace(TestNs).
		Resource(habv1.HabitatResourcePlural).
		Name(habitatName).
		Do().
		Error()
}

// DeleteService delete a Kubernetes service provided.
func (f *Framework) DeleteService(service string) error {
	return f.KubeClient.CoreV1().Services(TestNs).Delete(service, &metav1.DeleteOptions{})
}

func (f *Framework) createRBAC() error {
	// Create Service account.
	sa, err := convertServiceAccount("resources/operator/service-account.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.CoreV1().ServiceAccounts(TestNs).Create(sa)
	if err != nil {
		return err
	}

	// Create cluster role.
	cr, err := convertClusterRole("resources/operator/cluster-role.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.RbacV1().ClusterRoles().Create(cr)
	if err != nil {
		return err
	}

	// Create cluster role bindings.
	crb, err := convertClusterRoleBinding("resources/operator/cluster-role-binding.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.RbacV1().ClusterRoleBindings().Create(crb)
	if err != nil {
		return err
	}

	return nil
}

// convertServiceAccount takes in a path to the YAML file containing the manifest
// It converts it from that file to the ServiceAccount object.
func convertServiceAccount(pathToYaml string) (*apiv1.ServiceAccount, error) {
	sa := apiv1.ServiceAccount{}

	if err := convertToK8sResource(pathToYaml, &sa); err != nil {
		return nil, err
	}

	return &sa, nil
}

// convertClusterRole takes in a path to the YAML file containing the manifest.
// It converts the file to the ClusterRole object.
func convertClusterRole(pathToYaml string) (*rbacv1.ClusterRole, error) {
	cr := rbacv1.ClusterRole{}

	if err := convertToK8sResource(pathToYaml, &cr); err != nil {
		return nil, err
	}

	return &cr, nil
}

// convertClusterRoleBinding takes in a path to the YAML file containing the manifest.
// It converts the file to the ClusterRoleBinding object.
func convertClusterRoleBinding(pathToYaml string) (*rbacv1.ClusterRoleBinding, error) {
	crb := rbacv1.ClusterRoleBinding{}

	if err := convertToK8sResource(pathToYaml, &crb); err != nil {
		return nil, err
	}

	return &crb, nil
}

// ConvertDeployment takes in a path to the YAML file containing the manifest.
// It converts the file to the Deployment object.
func ConvertDeployment(pathToYaml string) (*appsv1beta1.Deployment, error) {
	d := appsv1beta1.Deployment{}

	if err := convertToK8sResource(pathToYaml, &d); err != nil {
		return nil, err
	}

	return &d, nil
}

// ConvertHabitat takes in a path to the YAML file containing the manifest.
// It converts the file to the Habitat object.
func ConvertHabitat(pathToYaml string) (*habv1.Habitat, error) {
	hab := habv1.Habitat{}

	if err := convertToK8sResource(pathToYaml, &hab); err != nil {
		return nil, err
	}

	return &hab, nil
}

// ConvertService takes in a path to the YAML file containing the manifest.
// It converts the file to the Service object.
func ConvertService(pathToYaml string) (*v1.Service, error) {
	s := v1.Service{}

	if err := convertToK8sResource(pathToYaml, &s); err != nil {
		return nil, err
	}

	return &s, nil
}

// ConvertSecret takes in a path to the YAML file containing the manifest.
// It converts the file to the Secret object.
func ConvertSecret(pathToYaml string) (*v1.Secret, error) {
	s := v1.Secret{}

	if err := convertToK8sResource(pathToYaml, &s); err != nil {
		return nil, err
	}

	return &s, nil
}

func convertToK8sResource(pathToYaml string, into interface{}) error {
	manifest, err := pathToOSFile(pathToYaml)
	if err != nil {
		return err
	}

	if err := yaml.NewYAMLToJSONDecoder(manifest).Decode(into); err != nil {
		return err
	}

	return nil
}

// pathToOSFile takes in a path and converts it to a File.
func pathToOSFile(relativePath string) (*os.File, error) {
	path, err := filepath.Abs(relativePath)
	if err != nil {
		return nil, err
	}

	manifest, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	return manifest, nil
}
<commit_msg>Prevent panic in cases when service was not created<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// 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 framework

import (
	"os"
	"path/filepath"
	"time"

	habv1 "github.com/kinvolk/habitat-operator/pkg/apis/habitat/v1"

	appsv1beta1 "k8s.io/api/apps/v1beta1"
	"k8s.io/api/core/v1"
	apiv1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/util/yaml"
)

// CreateHabitat creates a Habitat.
func (f *Framework) CreateHabitat(habitat *habv1.Habitat) error {
	return f.Client.Post().
		Namespace(TestNs).
		Resource(habv1.HabitatResourcePlural).
		Body(habitat).
		Do().
		Error()
}

// WaitForResources waits until numPods are in the "Running" state.
// We wait for pods, because those take the longest to create.
// Waiting for anything else would be already testing.
func (f *Framework) WaitForResources(labelName, habitatName string, numPods int) error {
	return wait.Poll(2*time.Second, 5*time.Minute, func() (bool, error) {
		fs := fields.SelectorFromSet(fields.Set{
			"status.phase": "Running",
		})

		ls := labels.SelectorFromSet(labels.Set{
			labelName: habitatName,
		})

		pods, err := f.KubeClient.CoreV1().Pods(TestNs).List(metav1.ListOptions{FieldSelector: fs.String(), LabelSelector: ls.String()})
		if err != nil {
			return false, err
		}

		if len(pods.Items) != numPods {
			return false, nil
		}

		return true, nil
	})
}

func (f *Framework) WaitForEndpoints(habitatName string) error {
	return wait.Poll(time.Second, time.Minute*5, func() (bool, error) {
		ep, err := f.KubeClient.CoreV1().Endpoints(TestNs).Get(habitatName, metav1.GetOptions{})
		if err != nil {
			return false, err
		}

		if len(ep.Subsets) == 0 {
			return false, nil
		}

		if len(ep.Subsets[0].Addresses) == 0 {
			return false, nil
		}

		return true, nil
	})
}

// DeleteHabitat deletes a Habitat as a user would.
func (f *Framework) DeleteHabitat(habitatName string) error {
	return f.Client.Delete().
		Namespace(TestNs).
		Resource(habv1.HabitatResourcePlural).
		Name(habitatName).
		Do().
		Error()
}

// DeleteService delete a Kubernetes service provided.
func (f *Framework) DeleteService(service string) error {
	return f.KubeClient.CoreV1().Services(TestNs).Delete(service, &metav1.DeleteOptions{})
}

func (f *Framework) createRBAC() error {
	// Create Service account.
	sa, err := convertServiceAccount("resources/operator/service-account.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.CoreV1().ServiceAccounts(TestNs).Create(sa)
	if err != nil {
		return err
	}

	// Create cluster role.
	cr, err := convertClusterRole("resources/operator/cluster-role.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.RbacV1().ClusterRoles().Create(cr)
	if err != nil {
		return err
	}

	// Create cluster role bindings.
	crb, err := convertClusterRoleBinding("resources/operator/cluster-role-binding.yml")
	if err != nil {
		return err
	}
	_, err = f.KubeClient.RbacV1().ClusterRoleBindings().Create(crb)
	if err != nil {
		return err
	}

	return nil
}

// convertServiceAccount takes in a path to the YAML file containing the manifest
// It converts it from that file to the ServiceAccount object.
func convertServiceAccount(pathToYaml string) (*apiv1.ServiceAccount, error) {
	sa := apiv1.ServiceAccount{}

	if err := convertToK8sResource(pathToYaml, &sa); err != nil {
		return nil, err
	}

	return &sa, nil
}

// convertClusterRole takes in a path to the YAML file containing the manifest.
// It converts the file to the ClusterRole object.
func convertClusterRole(pathToYaml string) (*rbacv1.ClusterRole, error) {
	cr := rbacv1.ClusterRole{}

	if err := convertToK8sResource(pathToYaml, &cr); err != nil {
		return nil, err
	}

	return &cr, nil
}

// convertClusterRoleBinding takes in a path to the YAML file containing the manifest.
// It converts the file to the ClusterRoleBinding object.
func convertClusterRoleBinding(pathToYaml string) (*rbacv1.ClusterRoleBinding, error) {
	crb := rbacv1.ClusterRoleBinding{}

	if err := convertToK8sResource(pathToYaml, &crb); err != nil {
		return nil, err
	}

	return &crb, nil
}

// ConvertDeployment takes in a path to the YAML file containing the manifest.
// It converts the file to the Deployment object.
func ConvertDeployment(pathToYaml string) (*appsv1beta1.Deployment, error) {
	d := appsv1beta1.Deployment{}

	if err := convertToK8sResource(pathToYaml, &d); err != nil {
		return nil, err
	}

	return &d, nil
}

// ConvertHabitat takes in a path to the YAML file containing the manifest.
// It converts the file to the Habitat object.
func ConvertHabitat(pathToYaml string) (*habv1.Habitat, error) {
	hab := habv1.Habitat{}

	if err := convertToK8sResource(pathToYaml, &hab); err != nil {
		return nil, err
	}

	return &hab, nil
}

// ConvertService takes in a path to the YAML file containing the manifest.
// It converts the file to the Service object.
func ConvertService(pathToYaml string) (*v1.Service, error) {
	s := v1.Service{}

	if err := convertToK8sResource(pathToYaml, &s); err != nil {
		return nil, err
	}

	return &s, nil
}

// ConvertSecret takes in a path to the YAML file containing the manifest.
// It converts the file to the Secret object.
func ConvertSecret(pathToYaml string) (*v1.Secret, error) {
	s := v1.Secret{}

	if err := convertToK8sResource(pathToYaml, &s); err != nil {
		return nil, err
	}

	return &s, nil
}

func convertToK8sResource(pathToYaml string, into interface{}) error {
	manifest, err := pathToOSFile(pathToYaml)
	if err != nil {
		return err
	}

	if err := yaml.NewYAMLToJSONDecoder(manifest).Decode(into); err != nil {
		return err
	}

	return nil
}

// pathToOSFile takes in a path and converts it to a File.
func pathToOSFile(relativePath string) (*os.File, error) {
	path, err := filepath.Abs(relativePath)
	if err != nil {
		return nil, err
	}

	manifest, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	return manifest, nil
}
<|endoftext|>


<commit_before>package ParquetFile

//go:generate mockgen -destination=../mocks/mock_s3.go -package=mocks github.com/aws/aws-sdk-go/service/s3/s3iface S3API

import (
	"context"
	"errors"
	"fmt"
	"io"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

// S3File is a TODO
type S3File struct {
	ctx    context.Context
	client s3iface.S3API
	offset int64
	whence int

	// write related fields
	writeOpened bool
	writeDone   chan struct{}
	pipeReader  *io.PipeReader
	pipeWriter  *io.PipeWriter
	downloader  *s3manager.Downloader

	// read related fields
	readOpened bool
	fileSize   int64

	BucketName string
	Key        string
}

var errWhence = errors.New("Seek: invalid whence")
var errOffset = errors.New("Seek: invalid offset")
var activeS3Session *session.Session

// NewS3FileWriter is TODO
func NewS3FileWriter(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	if activeS3Session == nil {
		activeS3Session = session.Must(session.NewSession())
	}

	file := &S3File{
		ctx:        ctx,
		client:     s3.New(activeS3Session, cfgs...),
		writeDone:  make(chan struct{}),
		BucketName: bucket,
		Key:        key,
	}

	return file.Create(key)
}

// NewS3FileReader is TODO
func NewS3FileReader(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	if activeS3Session == nil {
		activeS3Session = session.Must(session.NewSession())
	}

	s3Client := s3.New(activeS3Session, cfgs...)
	s3Downloader := s3manager.NewDownloaderWithClient(s3Client)

	file := &S3File{
		ctx:        ctx,
		client:     s3Client,
		downloader: s3Downloader,
		BucketName: bucket,
		Key:        key,
	}

	return file.Open(key)
}

// Seek tracks the offset for the next Read. Has no effect on Write.
func (s *S3File) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case io.SeekStart:
		if offset < 0 || offset > s.fileSize {
			return 0, errOffset
		}
	case io.SeekCurrent:
		currentOffset := s.offset + offset
		if currentOffset < 0 || currentOffset > s.fileSize {
			return 0, errOffset
		}
	case io.SeekEnd:
		if offset > -1 || -offset > s.fileSize {
			return 0, errOffset
		}
	}

	s.offset = offset
	s.whence = whence
	return 0, nil
}

// Read up to len(p) bytes into p and return the number of bytes read.
// not safe for concurrent reads
func (s *S3File) Read(p []byte) (n int, err error) {
	if s.offset >= s.fileSize {
		return 0, io.EOF
	}

	numBytes := len(p)
	getObjRange := buildBytesRange(s.offset, s.whence, numBytes, s.fileSize)
	fmt.Printf("==== byte range %q\n", getObjRange)
	getObj := &s3.GetObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}
	if len(getObjRange) > 0 {
		getObj.Range = aws.String(getObjRange)
	}

	wab := aws.NewWriteAtBuffer(p)
	bytesDownloaded, err := s.downloader.Download(wab, getObj)
	s.offset += bytesDownloaded
	if buf := wab.Bytes(); len(buf) > numBytes {
		// backing buffer reassigned, copy over some of the data
		copy(p, buf)
	}

	fmt.Printf("downloaded %d, new offset %d, %v\n", bytesDownloaded, s.offset, err)

	return int(bytesDownloaded), err
}

// Write is TODO
func (s *S3File) Write(p []byte) (n int, err error) {
	if !s.writeOpened {
		s.openWrite()
	}

	if s.pipeWriter != nil {
		// exit when writer is erroring
		return s.pipeWriter.Write(p)
	}

	return 0, nil
}

// Close is TODO
func (s *S3File) Close() error {
	s.offset = 0

	if s.pipeWriter != nil {
		s.pipeWriter.Close()
	}

	// if s.pipeReader != nil {
	// 	s.pipeReader.Close()
	// }

	// wait for pending uploads
	if s.writeDone != nil {
		<-s.writeDone
	}

	return nil
}

// Open is TODO reader
func (s *S3File) Open(name string) (ParquetFile, error) {
	if !s.readOpened {
		if err := s.openRead(); err != nil {
			return nil, err
		}
	}

	// create a new instance
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		downloader: s.downloader,
		BucketName: s.BucketName,
		Key:        s.Key,
		readOpened: s.readOpened,
		fileSize:   s.fileSize,
		offset:     0,
	}
	return pf, nil
}

// Create is TODO
func (s *S3File) Create(name string) (ParquetFile, error) {
	s.offset = 0
	s.pipeReader = nil
	s.pipeWriter = nil

	return s, nil
}

func (s *S3File) openWrite() {
	pr, pw := io.Pipe()
	s.pipeReader = pr
	s.pipeWriter = pw
	s.writeOpened = true

	uploadParams := &s3manager.UploadInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
		Body:   s.pipeReader,
	}
	uploader := s3manager.NewUploaderWithClient(s.client)

	go func(uploader *s3manager.Uploader) {
		result, err := uploader.Upload(uploadParams)
		fmt.Printf("uploaded: %v %v\n", result, err)
		close(s.writeDone)
	}(uploader)
}

func (s *S3File) openRead() error {
	hoi := &s3.HeadObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}

	hoo, err := s.client.HeadObject(hoi)
	if err != nil {
		return err
	}

	s.readOpened = true
	if hoo.ContentLength != nil {
		s.fileSize = *hoo.ContentLength
	}

	fmt.Printf("opened file, size %d\n", s.fileSize)
	return nil
}

func buildBytesRange(offset int64, whence int, numBytes int, maxBytes int64) string {
	end := offset + int64(numBytes)
	if end > maxBytes {
		end = maxBytes
	}
	var byteRange string
	switch whence {
	case io.SeekStart:
		byteRange = fmt.Sprintf("bytes=%d-%d", offset, end)
	case io.SeekCurrent:
		byteRange = fmt.Sprintf("bytes=%d-%d", offset, end)
	case io.SeekEnd:
		byteRange = fmt.Sprintf("bytes=%d", offset)
	}

	return byteRange
}
<commit_msg>update comments<commit_after>package ParquetFile

//go:generate mockgen -destination=../mocks/mock_s3.go -package=mocks github.com/aws/aws-sdk-go/service/s3/s3iface S3API

import (
	"context"
	"errors"
	"fmt"
	"io"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

// S3File is ParquetFile for AWS S3
type S3File struct {
	ctx    context.Context
	client s3iface.S3API
	offset int64
	whence int

	// write-related fields
	writeOpened bool
	writeDone   chan struct{}
	pipeReader  *io.PipeReader
	pipeWriter  *io.PipeWriter
	downloader  *s3manager.Downloader

	// read-related fields
	readOpened bool
	fileSize   int64

	BucketName string
	Key        string
}

const (
	rangeHeader       = "bytes=%d-%d"
	rangeHeaderSuffix = "bytes=%d"
)

var (
	errWhence        = errors.New("Seek: invalid whence")
	errInvalidOffset = errors.New("Seek: invalid offset")
	errFailedUpload  = errors.New("Write: failed upload")
	activeS3Session  *session.Session
)

// NewS3FileWriter creates an S3 FileWriter, to be used with NewParquetWriter
func NewS3FileWriter(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	if activeS3Session == nil {
		activeS3Session = session.Must(session.NewSession())
	}

	file := &S3File{
		ctx:        ctx,
		client:     s3.New(activeS3Session, cfgs...),
		writeDone:  make(chan struct{}),
		BucketName: bucket,
		Key:        key,
	}

	return file.Create(key)
}

// NewS3FileReader creates an S3 FileReader, to be used with NewParquetReader
func NewS3FileReader(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	if activeS3Session == nil {
		activeS3Session = session.Must(session.NewSession())
	}

	s3Client := s3.New(activeS3Session, cfgs...)
	s3Downloader := s3manager.NewDownloaderWithClient(s3Client)

	file := &S3File{
		ctx:        ctx,
		client:     s3Client,
		downloader: s3Downloader,
		BucketName: bucket,
		Key:        key,
	}

	return file.Open(key)
}

// Seek tracks the offset for the next Read. Has no effect on Write.
func (s *S3File) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case io.SeekStart:
		if offset < 0 || offset > s.fileSize {
			return 0, errInvalidOffset
		}
	case io.SeekCurrent:
		currentOffset := s.offset + offset
		if currentOffset < 0 || currentOffset > s.fileSize {
			return 0, errInvalidOffset
		}
	case io.SeekEnd:
		if offset > -1 || -offset > s.fileSize {
			return 0, errInvalidOffset
		}
	}

	s.offset = offset
	s.whence = whence
	return 0, nil
}

// Read up to len(p) bytes into p and return the number of bytes read
func (s *S3File) Read(p []byte) (n int, err error) {
	if s.offset >= s.fileSize {
		return 0, io.EOF
	}

	numBytes := len(p)
	getObjRange := s.getBytesRange(numBytes)
	getObj := &s3.GetObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}
	if len(getObjRange) > 0 {
		getObj.Range = aws.String(getObjRange)
	}

	wab := aws.NewWriteAtBuffer(p)
	bytesDownloaded, err := s.downloader.Download(wab, getObj)
	s.offset += bytesDownloaded
	if buf := wab.Bytes(); len(buf) > numBytes {
		// backing buffer reassigned, copy over some of the data
		copy(p, buf)
	}

	return int(bytesDownloaded), err
}

// Write len(p) bytes from p to the S3 data stream
func (s *S3File) Write(p []byte) (n int, err error) {
	if !s.writeOpened {
		s.openWrite()
	}

	if s.pipeWriter == nil {
		return 0, errFailedUpload
	}

	// prevent further writes upon error
	if _, err := s.pipeWriter.Write(p); err != nil {
		s.pipeWriter.CloseWithError(err)
		s.pipeWriter = nil
		return 0, err
	}

	return 0, nil
}

// Close cleans up any open streams. Will block until
// pending uploads are complete.
func (s *S3File) Close() error {
	s.offset = 0

	if s.pipeWriter != nil {
		s.pipeWriter.Close()
	}

	// wait for pending uploads
	if s.writeDone != nil {
		<-s.writeDone
	}

	if s.pipeReader != nil {
		s.pipeReader.Close()
	}

	return nil
}

// Open creates a new S3 File instance to perform concurrent reads
func (s *S3File) Open(name string) (ParquetFile, error) {
	if !s.readOpened {
		if err := s.openRead(); err != nil {
			return nil, err
		}
	}

	// create a new instance
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		downloader: s.downloader,
		BucketName: s.BucketName,
		Key:        s.Key,
		readOpened: s.readOpened,
		fileSize:   s.fileSize,
		offset:     0,
	}
	return pf, nil
}

// Create creates a new S3 File instance to perform writes
func (s *S3File) Create(name string) (ParquetFile, error) {
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		BucketName: s.BucketName,
		Key:        s.Key,
		writeDone:  make(chan struct{}),
	}
	return pf, nil
}

// openWrite creates an S3 uploader that consumes the Reader end of an io.Pipe.
// Calling Close signals write completion.
func (s *S3File) openWrite() {
	pr, pw := io.Pipe()
	s.pipeReader = pr
	s.pipeWriter = pw
	s.writeOpened = true

	uploadParams := &s3manager.UploadInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
		Body:   s.pipeReader,
	}
	uploader := s3manager.NewUploaderWithClient(s.client)

	go func(uploader *s3manager.Uploader) {
		// upload data and signal done when complete
		uploader.Upload(uploadParams)
		close(s.writeDone)
	}(uploader)
}

// openRead verifies the requested file is accessible and
// tracks the file size
func (s *S3File) openRead() error {
	hoi := &s3.HeadObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}

	hoo, err := s.client.HeadObject(hoi)
	if err != nil {
		return err
	}

	s.readOpened = true
	if hoo.ContentLength != nil {
		s.fileSize = *hoo.ContentLength
	}

	return nil
}

// getBytesRange returns the range request header string
func (s *S3File) getBytesRange(numBytes int) string {
	end := s.offset + int64(numBytes)
	if end > s.fileSize {
		end = s.fileSize
	}
	var byteRange string
	switch s.whence {
	case io.SeekStart:
		byteRange = fmt.Sprintf(rangeHeader, s.offset, end)
	case io.SeekCurrent:
		byteRange = fmt.Sprintf(rangeHeader, s.offset, end)
	case io.SeekEnd:
		byteRange = fmt.Sprintf(rangeHeaderSuffix, s.offset)
	}

	return byteRange
}
<|endoftext|>


<commit_before>package actions

import "testing"

var instance = &TitleExtract{}

var testTable = []struct {
	input  string
	expect string
}{
	{"http://google.com", "Google"},
	{"http://linux.org", "Linux.org"},
}

func TestExtract(t *testing.T) {
	for _, tt := range testTable {
		actual, err := extractTitle(tt.input)
		if err != nil {
			t.Error(err)
		}
		if actual != tt.expect {
			t.Errorf("input %s\nexpected %s\nactual %s\n", tt.input, tt.expect, actual)
		}
	}
}
<commit_msg>benchmark extractTitle<commit_after>package actions

import "testing"

var instance = &TitleExtract{}

var testTable = []struct {
	input  string
	expect string
}{
	{"http://google.com", "Google"},
	{"http://linux.org", "Linux.org"},
}

func TestTitleExtract(t *testing.T) {
	for _, tt := range testTable {
		actual, err := extractTitle(tt.input)
		if err != nil {
			t.Error(err)
		}
		if actual != tt.expect {
			t.Errorf("input %s\nexpected %s\nactual %s\n", tt.input, tt.expect, actual)
		}
	}
}

var result string

func BenchmarkTitleExtract(b *testing.B) {
	var t string
	for i := 0; i < b.N; i++ {
		t, _ = extractTitle("http://google.com")
	}
	result = t
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

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

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

package caddy

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/caddyserver/caddy"
	_ "github.com/caddyserver/caddy/caddyhttp"
	"github.com/caddyserver/caddy/caddytls"
	"github.com/mholt/certmagic"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"

	"github.com/admpub/nging/application/library/msgbox"
)

var (
	DefaultConfig = &Config{
		Agreed:                  true,
		CAUrl:                   certmagic.Default.CA,
		CATimeout:               int64(certmagic.HTTPTimeout),
		DisableHTTPChallenge:    certmagic.Default.DisableHTTPChallenge,
		DisableTLSALPNChallenge: certmagic.Default.DisableTLSALPNChallenge,
		ServerType:              `http`,
		CPU:                     `100%`,
		PidFile:                 `./caddy.pid`,
	}
	DefaultVersion = `2.0.0`
	EnableReload   = true
)

func TrapSignals() {
	caddy.TrapSignals()
}

func Fixed(c *Config) {
	if len(c.CAUrl) == 0 {
		c.CAUrl = DefaultConfig.CAUrl
	}
	if c.CATimeout == 0 {
		c.CATimeout = DefaultConfig.CATimeout
	}
	if len(c.ServerType) == 0 {
		c.ServerType = DefaultConfig.ServerType
	}
	if len(c.CPU) == 0 {
		c.CPU = DefaultConfig.CPU
	}
	pidFile := filepath.Join(echo.Wd(), `data/pid`)
	err := os.MkdirAll(pidFile, os.ModePerm)
	if err != nil {
		log.Println(err)
	}
	pidFile = filepath.Join(pidFile, `caddy.pid`)
	c.PidFile = pidFile
	if len(c.LogFile) == 0 {
		logFile := filepath.Join(echo.Wd(), `data/logs`)
		err := os.MkdirAll(logFile, os.ModePerm)
		if err != nil {
			log.Println(err)
		}
		c.LogFile = filepath.Join(logFile, `caddy.log`)
	} else {
		err := os.MkdirAll(filepath.Dir(c.LogFile), os.ModePerm)
		if err != nil {
			log.Println(err)
		}
	}
	c.appName = `nging`
	c.appVersion = echo.String(`VERSION`, DefaultVersion)
	c.Agreed = true
	c.ctx, c.cancel = context.WithCancel(context.Background())
}

type Config struct {
	Agreed                  bool   `json:"agreed"` //Agree to the CA's Subscriber Agreement
	CAUrl                   string `json:"caURL"`  //URL to certificate authority's ACME server directory
	DisableHTTPChallenge    bool   `json:"disableHTTPChallenge"`
	DisableTLSALPNChallenge bool   `json:"disableTLSALPNChallenge"`
	Caddyfile               string `json:"caddyFile"`  //Caddyfile to load (default caddy.DefaultConfigFile)
	CPU                     string `json:"cpu"`        //CPU cap
	CAEmail                 string `json:"caEmail"`    //Default ACME CA account email address
	CATimeout               int64  `json:"caTimeout"`  //Default ACME CA HTTP timeout
	LogFile                 string `json:"logFile"`    //Process log file
	PidFile                 string `json:"-"`          //Path to write pid file
	Quiet                   bool   `json:"quiet"`      //Quiet mode (no initialization output)
	Revoke                  string `json:"revoke"`     //Hostname for which to revoke the certificate
	ServerType              string `json:"serverType"` //Type of server to run

	//---
	EnvFile string `json:"envFile"` //Path to file with environment variables to load in KEY=VALUE format
	Plugins bool   `json:"plugins"` //List installed plugins
	Version bool   `json:"version"` //Show version

	//---
	appVersion string
	appName    string
	instance   *caddy.Instance
	stopped    bool
	ctx        context.Context
	cancel     context.CancelFunc
}

func now() string {
	return time.Now().Format(`2006-01-02 15:04:05`)
}

func (c *Config) Start() error {
	caddy.AppName = c.appName
	caddy.AppVersion = c.appVersion
	certmagic.UserAgent = c.appName + "/" + c.appVersion
	c.stopped = false

	// Executes Startup events
	caddy.EmitEvent(caddy.StartupEvent, nil)

	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(c.ServerType)
	if err != nil {
		return err
	}

	if EnableReload {
		c.watchingSignal()
	}

	// Start your engines
	c.instance, err = caddy.Start(caddyfile)
	if err != nil {
		return err
	}

	msgbox.Success(`Caddy`, `Server has been successfully started at `+now())

	// Twiddle your thumbs
	c.instance.Wait()
	return nil
}

// Listen to keypress of "return" and restart the app automatically
func (c *Config) watchingSignal() {
	debug := false
	go func() {
		if debug {
			msgbox.Info(`Caddy`, `listen return ==> `+now())
		}
		in := bufio.NewReader(os.Stdin)
		for {
			select {
			case <-c.ctx.Done():
				return
			default:
				if debug {
					msgbox.Info(`Caddy`, `reading ==> `+now())
				}
				input, _ := in.ReadString(com.LF)
				if input == com.StrLF || input == com.StrCRLF {
					if debug {
						msgbox.Info(`Caddy`, `restart ==> `+now())
					}
					c.Restart()
				} else {
					if debug {
						msgbox.Info(`Caddy`, `waiting ==> `+now())
					}
				}
			}
		}
	}()
}

func (c *Config) Restart() error {
	defer msgbox.Success(`Caddy`, `Server has been successfully reloaded at `+now())
	if c.instance == nil {
		return nil
	}
	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(c.ServerType)
	if err != nil {
		return err
	}
	c.instance, err = c.instance.Restart(caddyfile)
	if err != nil {
		return err
	}
	return nil
}

func (c *Config) Stop() error {
	c.stopped = true
	defer func() {
		c.cancel()
		msgbox.Success(`Caddy`, `Server has been successfully stopped at `+now())
	}()
	if c.instance == nil {
		return nil
	}
	return c.instance.Stop()
}

func (c *Config) Init() *Config {
	certmagic.Default.Agreed = c.Agreed
	certmagic.Default.CA = c.CAUrl
	certmagic.Default.DisableHTTPChallenge = c.DisableHTTPChallenge
	certmagic.Default.DisableTLSALPNChallenge = c.DisableTLSALPNChallenge
	certmagic.Default.Email = c.CAEmail
	certmagic.HTTPTimeout = time.Duration(c.CATimeout)

	caddy.PidFile = c.PidFile
	caddy.Quiet = c.Quiet
	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(c.confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(c.defaultLoader))

	// Set up process log before anything bad happens
	switch c.LogFile {
	case "stdout":
		log.SetOutput(os.Stdout)
	case "stderr":
		log.SetOutput(os.Stderr)
	case "":
		log.SetOutput(ioutil.Discard)
	default:
		log.SetOutput(&lumberjack.Logger{
			Filename:   c.LogFile,
			MaxSize:    100,
			MaxAge:     14,
			MaxBackups: 10,
		})
	}

	//Load all additional envs as soon as possible
	if err := LoadEnvFromFile(c.EnvFile); err != nil {
		mustLogFatalf("%v", err)
	}

	// Check for one-time actions
	if len(c.Revoke) > 0 {
		err := caddytls.Revoke(c.Revoke)
		if err != nil {
			mustLogFatalf(err.Error())
		}
		fmt.Printf("Revoked certificate for %s\n", c.Revoke)
		os.Exit(0)
	}
	if c.Version {
		fmt.Printf("%s %s\n", c.appName, c.appVersion)
		os.Exit(0)
	}
	if c.Plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	err := setCPU(c.CPU)
	if err != nil {
		mustLogFatalf(err.Error())
	}
	return c
}

// confLoader loads the Caddyfile using the -conf flag.
func (c *Config) confLoader(serverType string) (caddy.Input, error) {
	if c.Caddyfile == "" {
		return nil, nil
	}
	if c.Caddyfile == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin, serverType)
	}
	var contents []byte
	if strings.Contains(c.Caddyfile, "*") {
		// Let caddyfile.doImport logic handle the globbed path
		contents = []byte("import " + c.Caddyfile)
	} else {
		var err error
		contents, err = ioutil.ReadFile(c.Caddyfile)
		if err != nil {
			return nil, err
		}
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       c.Caddyfile,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Caddyfile from the current working directory.
func (c *Config) defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			log.Println(err)
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// mustLogFatalf wraps log.Fatalf() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatalf(format string, args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatalf(format, args...)
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
		if numCPU < 1 {
			numCPU = 1
		}
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

// LoadEnvFromFile loads additional envs if file provided and exists
// Envs in file should be in KEY=VALUE format
func LoadEnvFromFile(envFile string) error {
	if envFile == "" {
		return nil
	}

	file, err := os.Open(envFile)
	if err != nil {
		return err
	}
	defer file.Close()

	envMap, err := ParseEnvFile(file)
	if err != nil {
		return err
	}

	for k, v := range envMap {
		if err := os.Setenv(k, v); err != nil {
			return err
		}
	}

	return nil
}

// ParseEnvFile implements parse logic for environment files
func ParseEnvFile(envInput io.Reader) (map[string]string, error) {
	envMap := make(map[string]string)

	scanner := bufio.NewScanner(envInput)
	var line string
	lineNumber := 0

	for scanner.Scan() {
		line = strings.TrimSpace(scanner.Text())
		lineNumber++

		// skip lines starting with comment
		if strings.HasPrefix(line, "#") {
			continue
		}

		// skip empty line
		if len(line) == 0 {
			continue
		}

		fields := strings.SplitN(line, "=", 2)
		if len(fields) != 2 {
			return nil, fmt.Errorf("Can't parse line %d; line should be in KEY=VALUE format", lineNumber)
		}

		if strings.Contains(fields[0], " ") {
			return nil, fmt.Errorf("Can't parse line %d; KEY contains whitespace", lineNumber)
		}

		key := fields[0]
		val := fields[1]

		if key == "" {
			return nil, fmt.Errorf("Can't parse line %d; KEY can't be empty string", lineNumber)
		}
		envMap[key] = val
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return envMap, nil
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

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

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

package caddy

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/caddyserver/caddy"
	_ "github.com/caddyserver/caddy/caddyhttp"
	"github.com/caddyserver/caddy/caddytls"
	"github.com/mholt/certmagic"
	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"

	"github.com/admpub/nging/application/library/msgbox"
)

var (
	DefaultConfig = &Config{
		Agreed:                  true,
		CAUrl:                   certmagic.Default.CA,
		CATimeout:               int64(certmagic.HTTPTimeout),
		DisableHTTPChallenge:    certmagic.Default.DisableHTTPChallenge,
		DisableTLSALPNChallenge: certmagic.Default.DisableTLSALPNChallenge,
		ServerType:              `http`,
		CPU:                     `100%`,
		PidFile:                 `./caddy.pid`,
	}
	DefaultVersion = `2.0.0`
	EnableReload   = true
)

func TrapSignals() {
	caddy.TrapSignals()
}

func Fixed(c *Config) {
	if len(c.CAUrl) == 0 {
		c.CAUrl = DefaultConfig.CAUrl
	}
	if c.CATimeout == 0 {
		c.CATimeout = DefaultConfig.CATimeout
	}
	if len(c.ServerType) == 0 {
		c.ServerType = DefaultConfig.ServerType
	}
	if len(c.CPU) == 0 {
		c.CPU = DefaultConfig.CPU
	}
	pidFile := filepath.Join(echo.Wd(), `data/pid`)
	err := os.MkdirAll(pidFile, os.ModePerm)
	if err != nil {
		log.Println(err)
	}
	pidFile = filepath.Join(pidFile, `caddy.pid`)
	c.PidFile = pidFile
	if len(c.LogFile) == 0 {
		logFile := filepath.Join(echo.Wd(), `data/logs`)
		err := os.MkdirAll(logFile, os.ModePerm)
		if err != nil {
			log.Println(err)
		}
		c.LogFile = filepath.Join(logFile, `caddy.log`)
	} else {
		err := os.MkdirAll(filepath.Dir(c.LogFile), os.ModePerm)
		if err != nil {
			log.Println(err)
		}
	}
	c.appName = `nging`
	c.appVersion = echo.String(`VERSION`, DefaultVersion)
	c.Agreed = true
	c.ctx, c.cancel = context.WithCancel(context.Background())
}

type Config struct {
	Agreed                  bool   `json:"agreed"` //Agree to the CA's Subscriber Agreement
	CAUrl                   string `json:"caURL"`  //URL to certificate authority's ACME server directory
	DisableHTTPChallenge    bool   `json:"disableHTTPChallenge"`
	DisableTLSALPNChallenge bool   `json:"disableTLSALPNChallenge"`
	Caddyfile               string `json:"caddyFile"`  //Caddyfile to load (default caddy.DefaultConfigFile)
	CPU                     string `json:"cpu"`        //CPU cap
	CAEmail                 string `json:"caEmail"`    //Default ACME CA account email address
	CATimeout               int64  `json:"caTimeout"`  //Default ACME CA HTTP timeout
	LogFile                 string `json:"logFile"`    //Process log file
	PidFile                 string `json:"-"`          //Path to write pid file
	Quiet                   bool   `json:"quiet"`      //Quiet mode (no initialization output)
	Revoke                  string `json:"revoke"`     //Hostname for which to revoke the certificate
	ServerType              string `json:"serverType"` //Type of server to run

	//---
	EnvFile string `json:"envFile"` //Path to file with environment variables to load in KEY=VALUE format
	Plugins bool   `json:"plugins"` //List installed plugins
	Version bool   `json:"version"` //Show version

	//---
	appVersion string
	appName    string
	instance   *caddy.Instance
	stopped    bool
	ctx        context.Context
	cancel     context.CancelFunc
}

func now() string {
	return time.Now().Format(`2006-01-02 15:04:05`)
}

func (c *Config) Start() error {
	if !com.FileExists(c.Caddyfile) {
		return errors.New(`not found caddyfile: ` + c.Caddyfile)
	}
	caddy.AppName = c.appName
	caddy.AppVersion = c.appVersion
	certmagic.UserAgent = c.appName + "/" + c.appVersion
	c.stopped = false

	// Executes Startup events
	caddy.EmitEvent(caddy.StartupEvent, nil)

	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(c.ServerType)
	if err != nil {
		return err
	}

	if EnableReload {
		c.watchingSignal()
	}

	// Start your engines
	c.instance, err = caddy.Start(caddyfile)
	if err != nil {
		return err
	}

	msgbox.Success(`Caddy`, `Server has been successfully started at `+now())

	// Twiddle your thumbs
	c.instance.Wait()
	return nil
}

// Listen to keypress of "return" and restart the app automatically
func (c *Config) watchingSignal() {
	debug := false
	go func() {
		if debug {
			msgbox.Info(`Caddy`, `listen return ==> `+now())
		}
		in := bufio.NewReader(os.Stdin)
		for {
			select {
			case <-c.ctx.Done():
				return
			default:
				if debug {
					msgbox.Info(`Caddy`, `reading ==> `+now())
				}
				input, _ := in.ReadString(com.LF)
				if input == com.StrLF || input == com.StrCRLF {
					if debug {
						msgbox.Info(`Caddy`, `restart ==> `+now())
					}
					c.Restart()
				} else {
					if debug {
						msgbox.Info(`Caddy`, `waiting ==> `+now())
					}
				}
			}
		}
	}()
}

func (c *Config) Restart() error {
	defer msgbox.Success(`Caddy`, `Server has been successfully reloaded at `+now())
	if c.instance == nil {
		return nil
	}
	// Get Caddyfile input
	caddyfile, err := caddy.LoadCaddyfile(c.ServerType)
	if err != nil {
		return err
	}
	c.instance, err = c.instance.Restart(caddyfile)
	if err != nil {
		return err
	}
	return nil
}

func (c *Config) Stop() error {
	c.stopped = true
	defer func() {
		c.cancel()
		msgbox.Success(`Caddy`, `Server has been successfully stopped at `+now())
	}()
	if c.instance == nil {
		return nil
	}
	return c.instance.Stop()
}

func (c *Config) Init() *Config {
	certmagic.Default.Agreed = c.Agreed
	certmagic.Default.CA = c.CAUrl
	certmagic.Default.DisableHTTPChallenge = c.DisableHTTPChallenge
	certmagic.Default.DisableTLSALPNChallenge = c.DisableTLSALPNChallenge
	certmagic.Default.Email = c.CAEmail
	certmagic.HTTPTimeout = time.Duration(c.CATimeout)

	caddy.PidFile = c.PidFile
	caddy.Quiet = c.Quiet
	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(c.confLoader))
	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(c.defaultLoader))

	// Set up process log before anything bad happens
	switch c.LogFile {
	case "stdout":
		log.SetOutput(os.Stdout)
	case "stderr":
		log.SetOutput(os.Stderr)
	case "":
		log.SetOutput(ioutil.Discard)
	default:
		log.SetOutput(&lumberjack.Logger{
			Filename:   c.LogFile,
			MaxSize:    100,
			MaxAge:     14,
			MaxBackups: 10,
		})
	}

	//Load all additional envs as soon as possible
	if err := LoadEnvFromFile(c.EnvFile); err != nil {
		mustLogFatalf("%v", err)
	}

	// Check for one-time actions
	if len(c.Revoke) > 0 {
		err := caddytls.Revoke(c.Revoke)
		if err != nil {
			mustLogFatalf(err.Error())
		}
		fmt.Printf("Revoked certificate for %s\n", c.Revoke)
		os.Exit(0)
	}
	if c.Version {
		fmt.Printf("%s %s\n", c.appName, c.appVersion)
		os.Exit(0)
	}
	if c.Plugins {
		fmt.Println(caddy.DescribePlugins())
		os.Exit(0)
	}

	// Set CPU cap
	err := setCPU(c.CPU)
	if err != nil {
		mustLogFatalf(err.Error())
	}
	return c
}

// confLoader loads the Caddyfile using the -conf flag.
func (c *Config) confLoader(serverType string) (caddy.Input, error) {
	if c.Caddyfile == "" {
		return nil, nil
	}
	if c.Caddyfile == "stdin" {
		return caddy.CaddyfileFromPipe(os.Stdin, serverType)
	}
	var contents []byte
	if strings.Contains(c.Caddyfile, "*") {
		// Let caddyfile.doImport logic handle the globbed path
		contents = []byte("import " + c.Caddyfile)
	} else {
		var err error
		contents, err = ioutil.ReadFile(c.Caddyfile)
		if err != nil {
			return nil, err
		}
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       c.Caddyfile,
		ServerTypeName: serverType,
	}, nil
}

// defaultLoader loads the Caddyfile from the current working directory.
func (c *Config) defaultLoader(serverType string) (caddy.Input, error) {
	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return caddy.CaddyfileInput{
		Contents:       contents,
		Filepath:       caddy.DefaultConfigFile,
		ServerTypeName: serverType,
	}, nil
}

// mustLogFatalf wraps log.Fatalf() in a way that ensures the
// output is always printed to stderr so the user can see it
// if the user is still there, even if the process log was not
// enabled. If this process is an upgrade, however, and the user
// might not be there anymore, this just logs to the process
// log and exits.
func mustLogFatalf(format string, args ...interface{}) {
	if !caddy.IsUpgrade() {
		log.SetOutput(os.Stderr)
	}
	log.Fatalf(format, args...)
}

// setCPU parses string cpu and sets GOMAXPROCS
// according to its value. It accepts either
// a number (e.g. 3) or a percent (e.g. 50%).
func setCPU(cpu string) error {
	var numCPU int

	availCPU := runtime.NumCPU()

	if strings.HasSuffix(cpu, "%") {
		// Percent
		var percent float32
		pctStr := cpu[:len(cpu)-1]
		pctInt, err := strconv.Atoi(pctStr)
		if err != nil || pctInt < 1 || pctInt > 100 {
			return errors.New("invalid CPU value: percentage must be between 1-100")
		}
		percent = float32(pctInt) / 100
		numCPU = int(float32(availCPU) * percent)
		if numCPU < 1 {
			numCPU = 1
		}
	} else {
		// Number
		num, err := strconv.Atoi(cpu)
		if err != nil || num < 1 {
			return errors.New("invalid CPU value: provide a number or percent greater than 0")
		}
		numCPU = num
	}

	if numCPU > availCPU {
		numCPU = availCPU
	}

	runtime.GOMAXPROCS(numCPU)
	return nil
}

// LoadEnvFromFile loads additional envs if file provided and exists
// Envs in file should be in KEY=VALUE format
func LoadEnvFromFile(envFile string) error {
	if envFile == "" {
		return nil
	}

	file, err := os.Open(envFile)
	if err != nil {
		return err
	}
	defer file.Close()

	envMap, err := ParseEnvFile(file)
	if err != nil {
		return err
	}

	for k, v := range envMap {
		if err := os.Setenv(k, v); err != nil {
			return err
		}
	}

	return nil
}

// ParseEnvFile implements parse logic for environment files
func ParseEnvFile(envInput io.Reader) (map[string]string, error) {
	envMap := make(map[string]string)

	scanner := bufio.NewScanner(envInput)
	var line string
	lineNumber := 0

	for scanner.Scan() {
		line = strings.TrimSpace(scanner.Text())
		lineNumber++

		// skip lines starting with comment
		if strings.HasPrefix(line, "#") {
			continue
		}

		// skip empty line
		if len(line) == 0 {
			continue
		}

		fields := strings.SplitN(line, "=", 2)
		if len(fields) != 2 {
			return nil, fmt.Errorf("Can't parse line %d; line should be in KEY=VALUE format", lineNumber)
		}

		if strings.Contains(fields[0], " ") {
			return nil, fmt.Errorf("Can't parse line %d; KEY contains whitespace", lineNumber)
		}

		key := fields[0]
		val := fields[1]

		if key == "" {
			return nil, fmt.Errorf("Can't parse line %d; KEY can't be empty string", lineNumber)
		}
		envMap[key] = val
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return envMap, nil
}
<|endoftext|>


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

// Package batchfossilizer implements a fossilizer that fossilize batches of
// data using a Merkle tree. The evidence will contain the Merkle root, the
// Merkle path, and a timestamp.
package batchfossilizer

import (
	"encoding/gob"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/stratumn/sdk/cs"
	"github.com/stratumn/sdk/cs/evidences"
	"github.com/stratumn/sdk/fossilizer"

	"github.com/stratumn/goprivate/merkle"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "batch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Batch Fossilizer"

	// DefaultInterval is the default interval between batches.
	DefaultInterval = 10 * time.Minute

	// DefaultMaxLeaves if the default maximum number of leaves of a Merkle
	// tree.
	DefaultMaxLeaves = 32 * 1024

	// DefaultMaxSimBatches is the default maximum number of simultaneous
	// batches.
	DefaultMaxSimBatches = 1

	// DefaultArchive is whether to archive completed batches by default.
	DefaultArchive = true

	// DefaultStopBatch is whether to do a batch on stop by default.
	DefaultStopBatch = true

	// DefaultFSync is whether to fsync after saving a hash to disk by
	// default.
	DefaultFSync = false

	// PendingExt is the pending hashes filename extension.
	PendingExt = "pending"

	// DirPerm is the directory's permissions.
	DirPerm = 0600

	// FilePerm is the files's permissions.
	FilePerm = 0600
)

// Config contains configuration options for the fossilizer.
type Config struct {
	// A version string that will be set in the store's information.
	Version string

	// A git commit sha that will be set in the store's information.
	Commit string

	// Interval between batches.
	Interval time.Duration

	// Maximum number of leaves of a Merkle tree.
	MaxLeaves int

	// Maximum number of simultaneous batches.
	MaxSimBatches int

	// Where to store pending hashes.
	// If empty, pending hashes are not saved and will be lost if stopped
	// abruptly.
	Path string

	// Whether to archive completed batches.
	Archive bool

	// Whether to do a batch on stop.
	StopBatch bool

	// Whether to fsync after saving a hash to disk.
	FSync bool
}

// GetInterval returns the configuration's interval or the default value.
func (c *Config) GetInterval() time.Duration {
	if c.Interval > 0 {
		return c.Interval
	}
	return DefaultInterval
}

// GetMaxLeaves returns the configuration's maximum number of leaves of a Merkle
// tree or the default value.
func (c *Config) GetMaxLeaves() int {
	if c.MaxLeaves > 0 {
		return c.MaxLeaves
	}
	return DefaultMaxLeaves
}

// GetMaxSimBatches returns the configuration's maximum number of simultaneous
// batches or the default value.
func (c *Config) GetMaxSimBatches() int {
	if c.MaxSimBatches > 0 {
		return c.MaxSimBatches
	}
	return DefaultMaxSimBatches
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
}

// Fossilizer is the type that
// implements github.com/stratumn/sdk/fossilizer.Adapter.
type Fossilizer struct {
	config      *Config
	startedChan chan chan struct{}
	fossilChan  chan *fossil
	resultChan  chan error
	batchChan   chan *batch
	stopChan    chan error
	semChan     chan struct{}
	resultChans []chan *fossilizer.Result
	waitGroup   sync.WaitGroup
	transformer Transformer
	pending     *batch
	stopping    bool
}

// Transformer is the type of a function to transform results.
type Transformer func(evidence *cs.Evidence, data, meta []byte) (*fossilizer.Result, error)

// New creates an instance of a Fossilizer.
func New(config *Config) (*Fossilizer, error) {
	a := &Fossilizer{
		config:      config,
		startedChan: make(chan chan struct{}),
		fossilChan:  make(chan *fossil),
		resultChan:  make(chan error),
		batchChan:   make(chan *batch, 1),
		stopChan:    make(chan error, 1),
		semChan:     make(chan struct{}, config.GetMaxSimBatches()),
		pending:     newBatch(config.GetMaxLeaves()),
		stopping:    false,
	}

	a.SetTransformer(nil)

	if a.config.Path != "" {
		if err := a.ensurePath(); err != nil {
			return nil, err
		}
		if err := a.recover(); err != nil {
			return nil, err
		}
	}

	return a, nil
}

// GetInfo implements github.com/stratumn/sdk/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	return &Info{
		Name:        Name,
		Description: Description,
		Version:     a.config.Version,
		Commit:      a.config.Commit,
	}, nil
}

// AddResultChan implements
// github.com/stratumn/sdk/fossilizer.Adapter.AddResultChan.
func (a *Fossilizer) AddResultChan(resultChan chan *fossilizer.Result) {
	a.resultChans = append(a.resultChans, resultChan)
}

// Fossilize implements github.com/stratumn/sdk/fossilizer.Adapter.Fossilize.
func (a *Fossilizer) Fossilize(data []byte, meta []byte) error {
	f := fossil{Meta: meta}
	copy(f.Data[:], data)
	a.fossilChan <- &f
	return <-a.resultChan
}

// Start starts the fossilizer.
func (a *Fossilizer) Start() error {
	var (
		interval = a.config.GetInterval()
		timer    = time.NewTimer(interval)
	)

	for {
		select {
		case c := <-a.startedChan:
			c <- struct{}{}
		case f := <-a.fossilChan:
			a.resultChan <- a.fossilize(f)
		case b := <-a.batchChan:
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(interval)
			a.batch(b)
		case <-timer.C:
			timer.Stop()
			timer.Reset(interval)
			if len(a.pending.data) > 0 {
				a.sendBatch()
				log.WithField("interval", interval).Info("Requested new batch because the %s interval was reached")
			} else {
				log.WithField("interval", interval).Info("No batch is needed after the %s interval because there are no pending hashes")
			}
		case err := <-a.stopChan:
			e := a.stop(err)
			a.stopChan <- e
			return e
		}
	}
}

// Stop stops the fossilizer.
func (a *Fossilizer) Stop() {
	a.stopChan <- nil
	<-a.stopChan
}

// Started return a channel that will receive once the fossilizer has started.
func (a *Fossilizer) Started() <-chan struct{} {
	c := make(chan struct{}, 1)
	a.startedChan <- c
	return c
}

// SetTransformer sets a transformer.
func (a *Fossilizer) SetTransformer(t Transformer) {
	if t != nil {
		a.transformer = t
	} else {
		a.transformer = func(evidence *cs.Evidence, data, meta []byte) (*fossilizer.Result, error) {
			return &fossilizer.Result{
				Evidence: *evidence,
				Data:     data,
				Meta:     meta,
			}, nil
		}
	}
}

func (a *Fossilizer) fossilize(f *fossil) error {
	if a.config.Path != "" {
		if a.pending.file == nil {
			if err := a.pending.open(a.pendingPath()); err != nil {
				return err
			}
		}
		if err := f.write(a.pending.encoder); err != nil {
			return err
		}
		if a.config.FSync {
			if err := a.pending.file.Sync(); err != nil {
				return err
			}
		}
	}

	a.pending.append(f)

	if numLeaves, maxLeaves := len(a.pending.data), a.config.GetMaxLeaves(); numLeaves >= maxLeaves {
		a.sendBatch()
		log.WithFields(log.Fields{
			"leaves": numLeaves,
			"max":    maxLeaves,
		}).Info("Requested new batch because the maximum number of leaves was reached")
	}

	return nil
}

func (a *Fossilizer) sendBatch() {
	b := a.pending
	a.pending = newBatch(a.config.GetMaxLeaves())
	a.batchChan <- b
}

func (a *Fossilizer) batch(b *batch) {
	log.Info("Starting batch...")

	a.waitGroup.Add(1)

	go func() {
		defer func() {
			a.waitGroup.Done()
			<-a.semChan
		}()

		a.semChan <- struct{}{}

		tree, err := merkle.NewStaticTree(b.data)
		if err != nil {
			if !a.stopping {
				a.stop(err)
			}
			return
		}

		root := tree.Root()
		log.WithField("root", root).Info("Created tree with Merkle root")

		a.sendEvidence(tree, b.meta)
		log.WithField("root", root).Info("Sent evidence for batch with Merkle root")

		if b.file != nil {
			path := b.file.Name()

			if err := b.close(); err != nil {
				log.WithField("error", err).Warn("Failed to close batch file")
			}

			if a.config.Archive {
				archivePath := filepath.Join(a.config.Path, root.String())
				if err := os.Rename(path, archivePath); err == nil {
					log.WithFields(log.Fields{
						"old": filepath.Base(path),
						"new": filepath.Base(archivePath),
					}).Info("Renamed batch file")
				} else {
					log.WithFields(log.Fields{
						"old":   filepath.Base(path),
						"new":   filepath.Base(archivePath),
						"error": err,
					}).Warn("Failed to rename batch file")
				}
			} else {
				if err := os.Remove(path); err == nil {
					log.WithField("file", filepath.Base(path)).Info("Removed pending hashes file")
				} else {
					log.WithFields(log.Fields{
						"file":  filepath.Base(path),
						"error": err,
					}).Warn("Failed to remove batch file")
				}
			}
		}

		log.WithField("root", root).Info("Finished batch")
	}()
}

func (a *Fossilizer) sendEvidence(tree *merkle.StaticTree, meta [][]byte) {
	for i := 0; i < tree.LeavesLen(); i++ {
		var (
			err  error
			ts   = time.Now().UTC().Unix()
			root = tree.Root()
			leaf = tree.Leaf(i)
			d    = leaf[:]
			m    = meta[i]
			r    *fossilizer.Result
		)

		evidence := cs.Evidence{
			State:    cs.CompleteEvidence,
			Backend:  Name,
			Provider: Name,
			Proof: &evidences.BatchProof{
				Timestamp: ts,
				Root:      root,
				Path:      tree.Path(i),
			},
		}

		if r, err = a.transformer(&evidence, d, m); err != nil {
			log.WithField("error", err).Error("Failed to transform evidence")
		} else {
			for _, c := range a.resultChans {
				c <- r
			}
		}
	}
}

func (a *Fossilizer) stop(err error) error {
	a.stopping = true
	if a.config.StopBatch {
		if len(a.pending.data) > 0 {
			a.batch(a.pending)
			log.Info("Requested final batch for pending hashes")
		} else {
			log.Info("No final batch is needed because there are no pending hashes")
		}
	}

	a.waitGroup.Wait()
	a.SetTransformer(nil)

	if a.pending.file != nil {
		if e := a.pending.file.Close(); e != nil {
			if err == nil {
				err = e
			} else {
				log.WithField("error", err).Error("Failed to close pending batch file")
			}
		}
	}

	return err
}

func (a *Fossilizer) ensurePath() error {
	if err := os.MkdirAll(a.config.Path, DirPerm); err != nil && !os.IsExist(err) {
		return err
	}
	return nil
}

func (a *Fossilizer) recover() error {
	matches, err := filepath.Glob(filepath.Join(a.config.Path, "*."+PendingExt))
	if err != nil {
		return err
	}

	for _, path := range matches {
		file, err := os.OpenFile(path, os.O_RDONLY|os.O_EXCL, FilePerm)
		if err != nil {
			return err
		}
		defer file.Close()

		dec := gob.NewDecoder(file)

		for {
			f, err := newFossilFromDecoder(dec)
			if err == io.EOF {
				break
			}
			if err = a.fossilize(f); err != nil {
				return err
			}
		}

		a.waitGroup.Wait()

		if err := os.Remove(path); err != nil {
			return err
		}

		log.WithField("file", filepath.Base(path)).Info("Recovered pending hashes file")
	}

	return nil
}

func (a *Fossilizer) pendingPath() string {
	filename := fmt.Sprintf("%d.%s", time.Now().UTC().UnixNano(), PendingExt)
	return filepath.Join(a.config.Path, filename)
}
<commit_msg>batchfossilizer: update imports<commit_after>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by the license that can be found in the
// LICENSE file.

// Package batchfossilizer implements a fossilizer that fossilize batches of
// data using a Merkle tree. The evidence will contain the Merkle root, the
// Merkle path, and a timestamp.
package batchfossilizer

import (
	"encoding/gob"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/stratumn/sdk/cs"
	"github.com/stratumn/sdk/cs/evidences"
	"github.com/stratumn/sdk/fossilizer"

	"github.com/stratumn/sdk/merkle"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "batch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Batch Fossilizer"

	// DefaultInterval is the default interval between batches.
	DefaultInterval = 10 * time.Minute

	// DefaultMaxLeaves if the default maximum number of leaves of a Merkle
	// tree.
	DefaultMaxLeaves = 32 * 1024

	// DefaultMaxSimBatches is the default maximum number of simultaneous
	// batches.
	DefaultMaxSimBatches = 1

	// DefaultArchive is whether to archive completed batches by default.
	DefaultArchive = true

	// DefaultStopBatch is whether to do a batch on stop by default.
	DefaultStopBatch = true

	// DefaultFSync is whether to fsync after saving a hash to disk by
	// default.
	DefaultFSync = false

	// PendingExt is the pending hashes filename extension.
	PendingExt = "pending"

	// DirPerm is the directory's permissions.
	DirPerm = 0600

	// FilePerm is the files's permissions.
	FilePerm = 0600
)

// Config contains configuration options for the fossilizer.
type Config struct {
	// A version string that will be set in the store's information.
	Version string

	// A git commit sha that will be set in the store's information.
	Commit string

	// Interval between batches.
	Interval time.Duration

	// Maximum number of leaves of a Merkle tree.
	MaxLeaves int

	// Maximum number of simultaneous batches.
	MaxSimBatches int

	// Where to store pending hashes.
	// If empty, pending hashes are not saved and will be lost if stopped
	// abruptly.
	Path string

	// Whether to archive completed batches.
	Archive bool

	// Whether to do a batch on stop.
	StopBatch bool

	// Whether to fsync after saving a hash to disk.
	FSync bool
}

// GetInterval returns the configuration's interval or the default value.
func (c *Config) GetInterval() time.Duration {
	if c.Interval > 0 {
		return c.Interval
	}
	return DefaultInterval
}

// GetMaxLeaves returns the configuration's maximum number of leaves of a Merkle
// tree or the default value.
func (c *Config) GetMaxLeaves() int {
	if c.MaxLeaves > 0 {
		return c.MaxLeaves
	}
	return DefaultMaxLeaves
}

// GetMaxSimBatches returns the configuration's maximum number of simultaneous
// batches or the default value.
func (c *Config) GetMaxSimBatches() int {
	if c.MaxSimBatches > 0 {
		return c.MaxSimBatches
	}
	return DefaultMaxSimBatches
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
}

// Fossilizer is the type that
// implements github.com/stratumn/sdk/fossilizer.Adapter.
type Fossilizer struct {
	config      *Config
	startedChan chan chan struct{}
	fossilChan  chan *fossil
	resultChan  chan error
	batchChan   chan *batch
	stopChan    chan error
	semChan     chan struct{}
	resultChans []chan *fossilizer.Result
	waitGroup   sync.WaitGroup
	transformer Transformer
	pending     *batch
	stopping    bool
}

// Transformer is the type of a function to transform results.
type Transformer func(evidence *cs.Evidence, data, meta []byte) (*fossilizer.Result, error)

// New creates an instance of a Fossilizer.
func New(config *Config) (*Fossilizer, error) {
	a := &Fossilizer{
		config:      config,
		startedChan: make(chan chan struct{}),
		fossilChan:  make(chan *fossil),
		resultChan:  make(chan error),
		batchChan:   make(chan *batch, 1),
		stopChan:    make(chan error, 1),
		semChan:     make(chan struct{}, config.GetMaxSimBatches()),
		pending:     newBatch(config.GetMaxLeaves()),
		stopping:    false,
	}

	a.SetTransformer(nil)

	if a.config.Path != "" {
		if err := a.ensurePath(); err != nil {
			return nil, err
		}
		if err := a.recover(); err != nil {
			return nil, err
		}
	}

	return a, nil
}

// GetInfo implements github.com/stratumn/sdk/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	return &Info{
		Name:        Name,
		Description: Description,
		Version:     a.config.Version,
		Commit:      a.config.Commit,
	}, nil
}

// AddResultChan implements
// github.com/stratumn/sdk/fossilizer.Adapter.AddResultChan.
func (a *Fossilizer) AddResultChan(resultChan chan *fossilizer.Result) {
	a.resultChans = append(a.resultChans, resultChan)
}

// Fossilize implements github.com/stratumn/sdk/fossilizer.Adapter.Fossilize.
func (a *Fossilizer) Fossilize(data []byte, meta []byte) error {
	f := fossil{Meta: meta}
	copy(f.Data[:], data)
	a.fossilChan <- &f
	return <-a.resultChan
}

// Start starts the fossilizer.
func (a *Fossilizer) Start() error {
	var (
		interval = a.config.GetInterval()
		timer    = time.NewTimer(interval)
	)

	for {
		select {
		case c := <-a.startedChan:
			c <- struct{}{}
		case f := <-a.fossilChan:
			a.resultChan <- a.fossilize(f)
		case b := <-a.batchChan:
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(interval)
			a.batch(b)
		case <-timer.C:
			timer.Stop()
			timer.Reset(interval)
			if len(a.pending.data) > 0 {
				a.sendBatch()
				log.WithField("interval", interval).Info("Requested new batch because the %s interval was reached")
			} else {
				log.WithField("interval", interval).Info("No batch is needed after the %s interval because there are no pending hashes")
			}
		case err := <-a.stopChan:
			e := a.stop(err)
			a.stopChan <- e
			return e
		}
	}
}

// Stop stops the fossilizer.
func (a *Fossilizer) Stop() {
	a.stopChan <- nil
	<-a.stopChan
}

// Started return a channel that will receive once the fossilizer has started.
func (a *Fossilizer) Started() <-chan struct{} {
	c := make(chan struct{}, 1)
	a.startedChan <- c
	return c
}

// SetTransformer sets a transformer.
func (a *Fossilizer) SetTransformer(t Transformer) {
	if t != nil {
		a.transformer = t
	} else {
		a.transformer = func(evidence *cs.Evidence, data, meta []byte) (*fossilizer.Result, error) {
			return &fossilizer.Result{
				Evidence: *evidence,
				Data:     data,
				Meta:     meta,
			}, nil
		}
	}
}

func (a *Fossilizer) fossilize(f *fossil) error {
	if a.config.Path != "" {
		if a.pending.file == nil {
			if err := a.pending.open(a.pendingPath()); err != nil {
				return err
			}
		}
		if err := f.write(a.pending.encoder); err != nil {
			return err
		}
		if a.config.FSync {
			if err := a.pending.file.Sync(); err != nil {
				return err
			}
		}
	}

	a.pending.append(f)

	if numLeaves, maxLeaves := len(a.pending.data), a.config.GetMaxLeaves(); numLeaves >= maxLeaves {
		a.sendBatch()
		log.WithFields(log.Fields{
			"leaves": numLeaves,
			"max":    maxLeaves,
		}).Info("Requested new batch because the maximum number of leaves was reached")
	}

	return nil
}

func (a *Fossilizer) sendBatch() {
	b := a.pending
	a.pending = newBatch(a.config.GetMaxLeaves())
	a.batchChan <- b
}

func (a *Fossilizer) batch(b *batch) {
	log.Info("Starting batch...")

	a.waitGroup.Add(1)

	go func() {
		defer func() {
			a.waitGroup.Done()
			<-a.semChan
		}()

		a.semChan <- struct{}{}

		tree, err := merkle.NewStaticTree(b.data)
		if err != nil {
			if !a.stopping {
				a.stop(err)
			}
			return
		}

		root := tree.Root()
		log.WithField("root", root).Info("Created tree with Merkle root")

		a.sendEvidence(tree, b.meta)
		log.WithField("root", root).Info("Sent evidence for batch with Merkle root")

		if b.file != nil {
			path := b.file.Name()

			if err := b.close(); err != nil {
				log.WithField("error", err).Warn("Failed to close batch file")
			}

			if a.config.Archive {
				archivePath := filepath.Join(a.config.Path, root.String())
				if err := os.Rename(path, archivePath); err == nil {
					log.WithFields(log.Fields{
						"old": filepath.Base(path),
						"new": filepath.Base(archivePath),
					}).Info("Renamed batch file")
				} else {
					log.WithFields(log.Fields{
						"old":   filepath.Base(path),
						"new":   filepath.Base(archivePath),
						"error": err,
					}).Warn("Failed to rename batch file")
				}
			} else {
				if err := os.Remove(path); err == nil {
					log.WithField("file", filepath.Base(path)).Info("Removed pending hashes file")
				} else {
					log.WithFields(log.Fields{
						"file":  filepath.Base(path),
						"error": err,
					}).Warn("Failed to remove batch file")
				}
			}
		}

		log.WithField("root", root).Info("Finished batch")
	}()
}

func (a *Fossilizer) sendEvidence(tree *merkle.StaticTree, meta [][]byte) {
	for i := 0; i < tree.LeavesLen(); i++ {
		var (
			err  error
			ts   = time.Now().UTC().Unix()
			root = tree.Root()
			leaf = tree.Leaf(i)
			d    = leaf[:]
			m    = meta[i]
			r    *fossilizer.Result
		)

		evidence := cs.Evidence{
			State:    cs.CompleteEvidence,
			Backend:  Name,
			Provider: Name,
			Proof: &evidences.BatchProof{
				Timestamp: ts,
				Root:      root,
				Path:      tree.Path(i),
			},
		}

		if r, err = a.transformer(&evidence, d, m); err != nil {
			log.WithField("error", err).Error("Failed to transform evidence")
		} else {
			for _, c := range a.resultChans {
				c <- r
			}
		}
	}
}

func (a *Fossilizer) stop(err error) error {
	a.stopping = true
	if a.config.StopBatch {
		if len(a.pending.data) > 0 {
			a.batch(a.pending)
			log.Info("Requested final batch for pending hashes")
		} else {
			log.Info("No final batch is needed because there are no pending hashes")
		}
	}

	a.waitGroup.Wait()
	a.SetTransformer(nil)

	if a.pending.file != nil {
		if e := a.pending.file.Close(); e != nil {
			if err == nil {
				err = e
			} else {
				log.WithField("error", err).Error("Failed to close pending batch file")
			}
		}
	}

	return err
}

func (a *Fossilizer) ensurePath() error {
	if err := os.MkdirAll(a.config.Path, DirPerm); err != nil && !os.IsExist(err) {
		return err
	}
	return nil
}

func (a *Fossilizer) recover() error {
	matches, err := filepath.Glob(filepath.Join(a.config.Path, "*."+PendingExt))
	if err != nil {
		return err
	}

	for _, path := range matches {
		file, err := os.OpenFile(path, os.O_RDONLY|os.O_EXCL, FilePerm)
		if err != nil {
			return err
		}
		defer file.Close()

		dec := gob.NewDecoder(file)

		for {
			f, err := newFossilFromDecoder(dec)
			if err == io.EOF {
				break
			}
			if err = a.fossilize(f); err != nil {
				return err
			}
		}

		a.waitGroup.Wait()

		if err := os.Remove(path); err != nil {
			return err
		}

		log.WithField("file", filepath.Base(path)).Info("Recovered pending hashes file")
	}

	return nil
}

func (a *Fossilizer) pendingPath() string {
	filename := fmt.Sprintf("%d.%s", time.Now().UTC().UnixNano(), PendingExt)
	return filepath.Join(a.config.Path, filename)
}
<|endoftext|>


<commit_before>package instructions

import (
	"jvmgo/jvm/rtda"
	rtc "jvmgo/jvm/rtda/class"
)

// Create new array
type newarray struct {
	atype uint8
}

func (self *newarray) fetchOperands(decoder *InstructionDecoder) {
	self.atype = decoder.readUint8()
}
func (self *newarray) Execute(frame *rtda.Frame) {
	stack := frame.OperandStack()
	count := stack.PopInt()
	if count < 0 {
		frame.Thread().ThrowNegativeArraySizeException()
		return
	}

	arr := rtc.NewPrimitiveArray(self.atype, uint(count), frame.ClassLoader())
	stack.PushRef(arr)
}

// Create new array of reference
type anewarray struct{ Index16Instruction }

func (self *anewarray) Execute(frame *rtda.Frame) {
	cp := frame.ConstantPool()
	kClass := cp.GetConstant(self.index).(*rtc.ConstantClass)
	componentClass := kClass.Class()

	if componentClass.InitializationNotStarted() {
		thread := frame.Thread()
		frame.SetNextPC(thread.PC()) // undo anewarray
		thread.InitClass(componentClass)
		return
	}

	stack := frame.OperandStack()
	count := stack.PopInt()
	if count < 0 {
		frame.Thread().ThrowNegativeArraySizeException()
		return
	}

	arr := rtc.NewRefArray(componentClass, uint(count))
	stack.PushRef(arr)
}

// Create new multidimensional array
type multianewarray struct {
	index      uint16
	dimensions uint8
}

func (self *multianewarray) fetchOperands(decoder *InstructionDecoder) {
	self.index = decoder.readUint16()
	self.dimensions = decoder.readUint8()
}
func (self *multianewarray) Execute(frame *rtda.Frame) {
	// todo
	panic("todo multianewarray")
}
<commit_msg>multianewarray..<commit_after>package instructions

import (
	"jvmgo/jvm/rtda"
	rtc "jvmgo/jvm/rtda/class"
)

// Create new array
type newarray struct {
	atype uint8
}

func (self *newarray) fetchOperands(decoder *InstructionDecoder) {
	self.atype = decoder.readUint8()
}
func (self *newarray) Execute(frame *rtda.Frame) {
	stack := frame.OperandStack()
	count := stack.PopInt()
	if count < 0 {
		frame.Thread().ThrowNegativeArraySizeException()
		return
	}

	arr := rtc.NewPrimitiveArray(self.atype, uint(count), frame.ClassLoader())
	stack.PushRef(arr)
}

// Create new array of reference
type anewarray struct{ Index16Instruction }

func (self *anewarray) Execute(frame *rtda.Frame) {
	cp := frame.ConstantPool()
	kClass := cp.GetConstant(self.index).(*rtc.ConstantClass)
	componentClass := kClass.Class()

	if componentClass.InitializationNotStarted() {
		thread := frame.Thread()
		frame.SetNextPC(thread.PC()) // undo anewarray
		thread.InitClass(componentClass)
		return
	}

	stack := frame.OperandStack()
	count := stack.PopInt()
	if count < 0 {
		frame.Thread().ThrowNegativeArraySizeException()
		return
	}

	arr := rtc.NewRefArray(componentClass, uint(count))
	stack.PushRef(arr)
}

// Create new multidimensional array
type multianewarray struct {
	index      uint16
	dimensions uint8
}

func (self *multianewarray) fetchOperands(decoder *InstructionDecoder) {
	self.index = decoder.readUint16()
	self.dimensions = decoder.readUint8()
}
func (self *multianewarray) Execute(frame *rtda.Frame) {
	cp := frame.ConstantPool()
	kClass := cp.GetConstant(uint(self.index)).(*rtc.ConstantClass)
	componentClass := kClass.Class()

	if componentClass.InitializationNotStarted() {
		thread := frame.Thread()
		frame.SetNextPC(thread.PC()) // undo anewarray
		thread.InitClass(componentClass)
		return
	}

	stack := frame.OperandStack()
	counts := stack.PopTops(uint(self.dimensions))
	// todo
	println(counts)
	panic("todo multianewarray")
}
<|endoftext|>