<commit_before>package worker

import (
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/cloudfoundry-incubator/garden"
)

//go:generate counterfeiter . WorkerProvider

type WorkerProvider interface {
	Workers() ([]Worker, error)
}

var ErrNoWorkers = errors.New("no workers")

type NoCompatibleWorkersError struct {
	Spec    ContainerSpec
	Workers []Worker
}

func (err NoCompatibleWorkersError) Error() string {
	availableWorkers := ""
	for _, worker := range err.Workers {
		availableWorkers += "\n  - " + worker.Description()
	}

	return fmt.Sprintf(
		"no workers satisfying: %s\n\navailable workers: %s",
		err.Spec.Description(),
		availableWorkers,
	)
}

type Pool struct {
	provider WorkerProvider

	rand *rand.Rand
}

func NewPool(provider WorkerProvider) Client {
	return &Pool{
		provider: provider,
		rand:     rand.New(rand.NewSource(time.Now().UnixNano())),
	}
}

func (pool *Pool) CreateContainer(id Identifier, spec ContainerSpec) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	compatibleWorkers := []Worker{}
	for _, worker := range workers {
		if worker.Satisfies(spec) {
			compatibleWorkers = append(compatibleWorkers, worker)
		}
	}

	if len(compatibleWorkers) == 0 {
		return nil, NoCompatibleWorkersError{
			Spec:    spec,
			Workers: workers,
		}
	}

	randomWorker := compatibleWorkers[pool.rand.Intn(len(compatibleWorkers))]

	return randomWorker.CreateContainer(id, spec)
}

func (pool *Pool) FindContainerForIdentifier(id Identifier) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan Container, len(workers))
	multiErrs := make(chan MultipleContainersError, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			container, err := worker.FindContainerForIdentifier(id)
			if err == nil {
				found <- container
			} else if multi, ok := err.(MultipleContainersError); ok {
				multiErrs <- multi
			}
		}(worker)
	}

	wg.Wait()

	totalFound := len(found)
	totalMulti := len(multiErrs)

	if totalMulti != 0 {
		allHandles := []string{}

		for i := 0; i < totalMulti; i++ {
			multiErr := <-multiErrs

			allHandles = append(allHandles, multiErr.Handles...)
		}

		for i := 0; i < totalFound; i++ {
			c := <-found
			allHandles = append(allHandles, c.Handle())
			c.Release()
		}

		return nil, MultipleContainersError{allHandles}
	}

	switch totalFound {
	case 0:
		return nil, ErrContainerNotFound
	case 1:
		return <-found, nil
	default:
		handles := make([]string, totalFound)

		for i := 0; i < totalFound; i++ {
			c := <-found

			handles[i] = c.Handle()

			c.Release()
		}

		return nil, MultipleContainersError{Handles: handles}
	}
}

type workerErrorInfo struct {
	workerName string
	err        error
}

func (pool *Pool) FindContainersForIdentifier(id Identifier) ([]Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan []Container, len(workers))
	errors := make(chan workerErrorInfo, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			containers, err := worker.FindContainersForIdentifier(id)
			found <- containers
			if err != nil {
				errors <- workerErrorInfo{
					workerName: worker.Name(),
					err:        err,
				}
			}
		}(worker)
	}

	wg.Wait()

	totalFound := len(found)

	containers := []Container{}
	for i := 0; i < totalFound; i++ {
		foundContainers := <-found
		for _, c := range foundContainers {
			containers = append(containers, c)
		}
	}

	totalErr := len(errors)
	if len(errors) > 0 {
		multiWorkerError := MultiWorkerError{}

		for i := 0; i < totalErr; i++ {
			e := <-errors
			multiWorkerError.AddError(e.workerName, e.err)
		}
		return containers, multiWorkerError
	}

	return containers, nil
}

type foundContainer struct {
	workerName string
	container  Container
}

func (pool *Pool) LookupContainer(handle string) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan foundContainer, len(workers))
	errors := make(chan workerErrorInfo, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			container, err := worker.LookupContainer(handle)
			if container != nil {
				found <- foundContainer{
					workerName: worker.Name(),
					container:  container,
				}
			}
			if err != nil {
				_, ok := err.(garden.ContainerNotFoundError)
				if !ok {
					errors <- workerErrorInfo{
						workerName: worker.Name(),
						err:        err,
					}
				}
			}
		}(worker)
	}

	wg.Wait()

	totalErrors := len(errors)

	var multiWorkerError *MultiWorkerError
	if totalErrors != 0 {
		allErrors := make(map[string]error, totalErrors)

		for i := 0; i < totalErrors; i++ {
			e := <-errors

			allErrors[e.workerName] = e.err
		}

		multiWorkerError = &MultiWorkerError{allErrors}
	}

	totalFound := len(found)
	switch totalFound {
	case 0:
		if multiWorkerError != nil {
			return nil, *multiWorkerError
		} else {
			return nil, garden.ContainerNotFoundError{}
		}
	case 1:
		return (<-found).container, nil
	default:
		names := make([]string, totalFound)

		for i := 0; i < totalFound; i++ {
			c := <-found
			names[i] = c.workerName
			if c.container != nil {
				c.container.Release()
			}
		}

		return nil, MultipleWorkersFoundContainerError{Names: names}
	}
}

func (pool *Pool) Name() string {
	return "pool"
}

type byActiveContainers []Worker

func (cs byActiveContainers) Len() int { return len(cs) }

func (cs byActiveContainers) Swap(i, j int) { cs[i], cs[j] = cs[j], cs[i] }

func (cs byActiveContainers) Less(i, j int) bool {
	return cs[i].ActiveContainers() < cs[j].ActiveContainers()
}
<commit_msg>Do not export Pool.<commit_after>package worker

import (
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/cloudfoundry-incubator/garden"
)

//go:generate counterfeiter . WorkerProvider

type WorkerProvider interface {
	Workers() ([]Worker, error)
}

var ErrNoWorkers = errors.New("no workers")

type NoCompatibleWorkersError struct {
	Spec    ContainerSpec
	Workers []Worker
}

func (err NoCompatibleWorkersError) Error() string {
	availableWorkers := ""
	for _, worker := range err.Workers {
		availableWorkers += "\n  - " + worker.Description()
	}

	return fmt.Sprintf(
		"no workers satisfying: %s\n\navailable workers: %s",
		err.Spec.Description(),
		availableWorkers,
	)
}

type pool struct {
	provider WorkerProvider

	rand *rand.Rand
}

func NewPool(provider WorkerProvider) Client {
	return &pool{
		provider: provider,
		rand:     rand.New(rand.NewSource(time.Now().UnixNano())),
	}
}

func (pool *pool) CreateContainer(id Identifier, spec ContainerSpec) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	compatibleWorkers := []Worker{}
	for _, worker := range workers {
		if worker.Satisfies(spec) {
			compatibleWorkers = append(compatibleWorkers, worker)
		}
	}

	if len(compatibleWorkers) == 0 {
		return nil, NoCompatibleWorkersError{
			Spec:    spec,
			Workers: workers,
		}
	}

	randomWorker := compatibleWorkers[pool.rand.Intn(len(compatibleWorkers))]

	return randomWorker.CreateContainer(id, spec)
}

func (pool *pool) FindContainerForIdentifier(id Identifier) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan Container, len(workers))
	multiErrs := make(chan MultipleContainersError, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			container, err := worker.FindContainerForIdentifier(id)
			if err == nil {
				found <- container
			} else if multi, ok := err.(MultipleContainersError); ok {
				multiErrs <- multi
			}
		}(worker)
	}

	wg.Wait()

	totalFound := len(found)
	totalMulti := len(multiErrs)

	if totalMulti != 0 {
		allHandles := []string{}

		for i := 0; i < totalMulti; i++ {
			multiErr := <-multiErrs

			allHandles = append(allHandles, multiErr.Handles...)
		}

		for i := 0; i < totalFound; i++ {
			c := <-found
			allHandles = append(allHandles, c.Handle())
			c.Release()
		}

		return nil, MultipleContainersError{allHandles}
	}

	switch totalFound {
	case 0:
		return nil, ErrContainerNotFound
	case 1:
		return <-found, nil
	default:
		handles := make([]string, totalFound)

		for i := 0; i < totalFound; i++ {
			c := <-found

			handles[i] = c.Handle()

			c.Release()
		}

		return nil, MultipleContainersError{Handles: handles}
	}
}

type workerErrorInfo struct {
	workerName string
	err        error
}

func (pool *pool) FindContainersForIdentifier(id Identifier) ([]Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan []Container, len(workers))
	errors := make(chan workerErrorInfo, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			containers, err := worker.FindContainersForIdentifier(id)
			found <- containers
			if err != nil {
				errors <- workerErrorInfo{
					workerName: worker.Name(),
					err:        err,
				}
			}
		}(worker)
	}

	wg.Wait()

	totalFound := len(found)

	containers := []Container{}
	for i := 0; i < totalFound; i++ {
		foundContainers := <-found
		for _, c := range foundContainers {
			containers = append(containers, c)
		}
	}

	totalErr := len(errors)
	if len(errors) > 0 {
		multiWorkerError := MultiWorkerError{}

		for i := 0; i < totalErr; i++ {
			e := <-errors
			multiWorkerError.AddError(e.workerName, e.err)
		}
		return containers, multiWorkerError
	}

	return containers, nil
}

type foundContainer struct {
	workerName string
	container  Container
}

func (pool *pool) LookupContainer(handle string) (Container, error) {
	workers, err := pool.provider.Workers()
	if err != nil {
		return nil, err
	}

	if len(workers) == 0 {
		return nil, ErrNoWorkers
	}

	wg := new(sync.WaitGroup)
	wg.Add(len(workers))

	found := make(chan foundContainer, len(workers))
	errors := make(chan workerErrorInfo, len(workers))

	for _, worker := range workers {
		go func(worker Worker) {
			defer wg.Done()

			container, err := worker.LookupContainer(handle)
			if container != nil {
				found <- foundContainer{
					workerName: worker.Name(),
					container:  container,
				}
			}
			if err != nil {
				_, ok := err.(garden.ContainerNotFoundError)
				if !ok {
					errors <- workerErrorInfo{
						workerName: worker.Name(),
						err:        err,
					}
				}
			}
		}(worker)
	}

	wg.Wait()

	totalErrors := len(errors)

	var multiWorkerError *MultiWorkerError
	if totalErrors != 0 {
		allErrors := make(map[string]error, totalErrors)

		for i := 0; i < totalErrors; i++ {
			e := <-errors

			allErrors[e.workerName] = e.err
		}

		multiWorkerError = &MultiWorkerError{allErrors}
	}

	totalFound := len(found)
	switch totalFound {
	case 0:
		if multiWorkerError != nil {
			return nil, *multiWorkerError
		} else {
			return nil, garden.ContainerNotFoundError{}
		}
	case 1:
		return (<-found).container, nil
	default:
		names := make([]string, totalFound)

		for i := 0; i < totalFound; i++ {
			c := <-found
			names[i] = c.workerName
			if c.container != nil {
				c.container.Release()
			}
		}

		return nil, MultipleWorkersFoundContainerError{Names: names}
	}
}

func (pool *pool) Name() string {
	return "pool"
}

type byActiveContainers []Worker

func (cs byActiveContainers) Len() int { return len(cs) }

func (cs byActiveContainers) Swap(i, j int) { cs[i], cs[j] = cs[j], cs[i] }

func (cs byActiveContainers) Less(i, j int) bool {
	return cs[i].ActiveContainers() < cs[j].ActiveContainers()
}
<|endoftext|>


<commit_before>package clusterservice

/*
Copyright 2017 - 2020 Crunchy Data Solutions, Inc.
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.
*/

import (
	"fmt"
	"strconv"
	"strings"

	crv1 "github.com/crunchydata/postgres-operator/apis/cr/v1"
	"github.com/crunchydata/postgres-operator/apiserver"
	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	"github.com/crunchydata/postgres-operator/config"
	"github.com/crunchydata/postgres-operator/kubeapi"
	"github.com/crunchydata/postgres-operator/util"

	log "github.com/sirupsen/logrus"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// ScaleCluster ...
func ScaleCluster(name, replicaCount, resourcesConfig, storageConfig, nodeLabel,
	ccpImageTag, serviceType, ns, pgouser string) msgs.ClusterScaleResponse {
	var err error

	response := msgs.ClusterScaleResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}

	if name == "all" {
		response.Status.Code = msgs.Error
		response.Status.Msg = "all is not allowed for the scale command"
		return response
	}

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(name).
		Do().Into(&cluster)

	if kerrors.IsNotFound(err) {
		log.Error("no clusters found")
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	if err != nil {
		log.Error("error getting cluster" + err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	var rc int
	rc, err = strconv.Atoi(replicaCount)

	spec := crv1.PgreplicaSpec{}

	//get the resource-config
	if resourcesConfig != "" {
		spec.ContainerResources, _ = apiserver.Pgo.GetContainerResource(resourcesConfig)
	} else {
		defaultContainerResource := apiserver.Pgo.DefaultContainerResources
		if defaultContainerResource != "" {
			spec.ContainerResources, _ = apiserver.Pgo.GetContainerResource(defaultContainerResource)
		}
	}

	//refer to the cluster's replica storage setting by default
	spec.ReplicaStorage = cluster.Spec.ReplicaStorage

	//allow for user override
	if storageConfig != "" {
		spec.ReplicaStorage, _ = apiserver.Pgo.GetStorageSpec(storageConfig)
	}

	spec.UserLabels = cluster.Spec.UserLabels

	if ccpImageTag != "" {
		spec.UserLabels[config.LABEL_CCP_IMAGE_TAG_KEY] = ccpImageTag
	}
	if serviceType != "" {
		if serviceType != config.DEFAULT_SERVICE_TYPE &&
			serviceType != config.NODEPORT_SERVICE_TYPE &&
			serviceType != config.LOAD_BALANCER_SERVICE_TYPE {
			response.Status.Code = msgs.Error
			response.Status.Msg = "error --service-type should be either ClusterIP, NodePort, or LoadBalancer "
			return response
		}
		spec.UserLabels[config.LABEL_SERVICE_TYPE] = serviceType
	}

	var parts []string

	//set replica node lables to blank to start with, then check for overrides
	spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = ""
	spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = ""

	if apiserver.Pgo.Cluster.ReplicaNodeLabel != "" {
		//should have been validated at apiserver startup
		parts = strings.Split(apiserver.Pgo.Cluster.ReplicaNodeLabel, "=")
		spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = parts[0]
		spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = parts[1]
		log.Debug("using pgo.yaml ReplicaNodeLabel for replica creation")
	}

	// validate & parse nodeLabel if exists
	if nodeLabel != "" {

		if err = apiserver.ValidateNodeLabel(nodeLabel); err != nil {
			response.Status.Code = msgs.Error
			response.Status.Msg = err.Error()
			return response
		}

		parts := strings.Split(nodeLabel, "=")
		spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = parts[0]
		spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = parts[1]

		log.Debug("using user entered node label for replica creation")
	}

	labels := make(map[string]string)
	labels[config.LABEL_PG_CLUSTER] = cluster.Spec.Name

	spec.ClusterName = cluster.Spec.Name

	if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	labels[config.LABEL_PGOUSER] = pgouser
	labels[config.LABEL_PG_CLUSTER_IDENTIFIER] = cluster.ObjectMeta.Labels[config.LABEL_PG_CLUSTER_IDENTIFIER]

	for i := 0; i < rc; i++ {

		uniqueName := util.RandStringBytesRmndr(4)
		labels[config.LABEL_NAME] = cluster.Spec.Name + "-" + uniqueName
		spec.Namespace = ns
		spec.Name = labels[config.LABEL_NAME]

		newInstance := &crv1.Pgreplica{
			ObjectMeta: meta_v1.ObjectMeta{
				Name:   labels[config.LABEL_NAME],
				Labels: labels,
			},
			Spec: spec,
			Status: crv1.PgreplicaStatus{
				State:   crv1.PgreplicaStateCreated,
				Message: "Created, not processed yet",
			},
		}

		result := crv1.Pgreplica{}

		err = apiserver.RESTClient.Post().
			Resource(crv1.PgreplicaResourcePlural).
			Namespace(ns).
			Body(newInstance).
			Do().Into(&result)
		if err != nil {
			log.Error(" in creating Pgreplica instance" + err.Error())
		}

		response.Results = append(response.Results, "created Pgreplica "+labels[config.LABEL_NAME])
	}

	return response
}

// ScaleQuery lists the replicas that are in the PostgreSQL cluster
// with information that is helpful in determining which one to fail over to,
// such as the lag behind the replica as well as the timeline
func ScaleQuery(name, ns string) msgs.ScaleQueryResponse {
	var err error

	response := msgs.ScaleQueryResponse{
		Results: make([]msgs.ScaleQueryTargetSpec, 0),
		Status:  msgs.Status{Code: msgs.Ok, Msg: ""},
	}

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(name).
		Do().Into(&cluster)

	// If no clusters are found, return a specific error message,
	// otherwise, pass forward the generic error message that Kubernetes sends
	if kerrors.IsNotFound(err) {
		errorMsg := fmt.Sprintf(`No cluster found for "%s"`, name)
		log.Error(errorMsg)
		response.Status.Code = msgs.Error
		response.Status.Msg = errorMsg
		return response
	} else if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// Get information about the current status of all of the replicas. This is
	// handled by a helper function, that will return the information in a struct
	// with the key elements to help the user understand the current state of the
	// replicas in a cluster
	replicationStatusRequest := util.ReplicationStatusRequest{
		RESTConfig:  apiserver.RESTConfig,
		Clientset:   apiserver.Clientset,
		Namespace:   ns,
		ClusterName: name,
	}

	replicationStatusResponse, err := util.ReplicationStatus(replicationStatusRequest)

	// if an error is return, log the message, and return the response
	if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// indicate in the response whether or not a standby cluster
	response.Standby = cluster.Spec.Standby

	// if there are no results, return the response as is
	if len(replicationStatusResponse.Instances) == 0 {
		return response
	}

	// iterate through response results to create the API response
	for _, instance := range replicationStatusResponse.Instances {
		// create a result for the response
		result := msgs.ScaleQueryTargetSpec{
			Name:           instance.Name,
			Node:           instance.Node,
			Status:         instance.Status,
			ReplicationLag: instance.ReplicationLag,
			Timeline:       instance.Timeline,
		}

		// append the result to the response list
		response.Results = append(response.Results, result)
	}

	return response
}

// ScaleDown ...
func ScaleDown(deleteData bool, clusterName, replicaName, ns string) msgs.ScaleDownResponse {

	var err error

	response := msgs.ScaleDownResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}
	response.Results = make([]string, 0)

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(clusterName).
		Do().Into(&cluster)

	if kerrors.IsNotFound(err) {
		log.Error("no clusters found")
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	if err != nil {
		log.Error("error getting cluster" + err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// dont proceed any further if the cluster is shutdown
	if cluster.Status.State == crv1.PgclusterStateShutdown {
		response.Status.Code = msgs.Error
		response.Status.Msg = "Nothing to scale, the cluster is currently " +
			"shutdown"
		return response
	}

	// selector in the format "pg-cluster=<cluster-name>,pg-ha-scope=<cluster-name>"
	// which will grab the primary and any/all replicas
	selector := fmt.Sprintf("%s=%s,%s=%s", config.LABEL_PG_CLUSTER, clusterName,
		config.LABEL_PGHA_ROLE, "replica")
	replicaList, err := kubeapi.GetPods(apiserver.Clientset, selector, ns)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// check to see if the replica name provided matches the name of any of the
	// replicas found for the cluster
	var replicaNameFound bool
	for _, pod := range replicaList.Items {
		if pod.Labels[config.LABEL_DEPLOYMENT_NAME] == replicaName {
			replicaNameFound = true
			break
		}
	}
	// return an error if the replica name provided does not match the primary or any replicas
	if !replicaNameFound {
		response.Status.Code = msgs.Error
		response.Status.Msg = fmt.Sprintf("Unable to find replica with name %s",
			replicaName)
		return response
	}

	//create the rmdata task which does the cleanup

	clusterPGHAScope := cluster.ObjectMeta.Labels[config.LABEL_PGHA_SCOPE]
	deleteBackups := false
	isReplica := true
	isBackup := false
	taskName := replicaName + "-rmdata"
	err = apiserver.CreateRMDataTask(clusterName, replicaName, taskName, deleteBackups, deleteData, isReplica, isBackup, ns, clusterPGHAScope)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	response.Results = append(response.Results, "deleted replica "+replicaName)
	return response
}
<commit_msg>Move rc var Back to Original Location<commit_after>package clusterservice

/*
Copyright 2017 - 2020 Crunchy Data Solutions, Inc.
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.
*/

import (
	"fmt"
	"strconv"
	"strings"

	crv1 "github.com/crunchydata/postgres-operator/apis/cr/v1"
	"github.com/crunchydata/postgres-operator/apiserver"
	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	"github.com/crunchydata/postgres-operator/config"
	"github.com/crunchydata/postgres-operator/kubeapi"
	"github.com/crunchydata/postgres-operator/util"

	log "github.com/sirupsen/logrus"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// ScaleCluster ...
func ScaleCluster(name, replicaCount, resourcesConfig, storageConfig, nodeLabel,
	ccpImageTag, serviceType, ns, pgouser string) msgs.ClusterScaleResponse {
	var err error

	response := msgs.ClusterScaleResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}

	if name == "all" {
		response.Status.Code = msgs.Error
		response.Status.Msg = "all is not allowed for the scale command"
		return response
	}

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(name).
		Do().Into(&cluster)

	if kerrors.IsNotFound(err) {
		log.Error("no clusters found")
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	if err != nil {
		log.Error("error getting cluster" + err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	spec := crv1.PgreplicaSpec{}

	//get the resource-config
	if resourcesConfig != "" {
		spec.ContainerResources, _ = apiserver.Pgo.GetContainerResource(resourcesConfig)
	} else {
		defaultContainerResource := apiserver.Pgo.DefaultContainerResources
		if defaultContainerResource != "" {
			spec.ContainerResources, _ = apiserver.Pgo.GetContainerResource(defaultContainerResource)
		}
	}

	//refer to the cluster's replica storage setting by default
	spec.ReplicaStorage = cluster.Spec.ReplicaStorage

	//allow for user override
	if storageConfig != "" {
		spec.ReplicaStorage, _ = apiserver.Pgo.GetStorageSpec(storageConfig)
	}

	spec.UserLabels = cluster.Spec.UserLabels

	if ccpImageTag != "" {
		spec.UserLabels[config.LABEL_CCP_IMAGE_TAG_KEY] = ccpImageTag
	}
	if serviceType != "" {
		if serviceType != config.DEFAULT_SERVICE_TYPE &&
			serviceType != config.NODEPORT_SERVICE_TYPE &&
			serviceType != config.LOAD_BALANCER_SERVICE_TYPE {
			response.Status.Code = msgs.Error
			response.Status.Msg = "error --service-type should be either ClusterIP, NodePort, or LoadBalancer "
			return response
		}
		spec.UserLabels[config.LABEL_SERVICE_TYPE] = serviceType
	}

	var parts []string

	//set replica node lables to blank to start with, then check for overrides
	spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = ""
	spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = ""

	if apiserver.Pgo.Cluster.ReplicaNodeLabel != "" {
		//should have been validated at apiserver startup
		parts = strings.Split(apiserver.Pgo.Cluster.ReplicaNodeLabel, "=")
		spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = parts[0]
		spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = parts[1]
		log.Debug("using pgo.yaml ReplicaNodeLabel for replica creation")
	}

	// validate & parse nodeLabel if exists
	if nodeLabel != "" {

		if err = apiserver.ValidateNodeLabel(nodeLabel); err != nil {
			response.Status.Code = msgs.Error
			response.Status.Msg = err.Error()
			return response
		}

		parts := strings.Split(nodeLabel, "=")
		spec.UserLabels[config.LABEL_NODE_LABEL_KEY] = parts[0]
		spec.UserLabels[config.LABEL_NODE_LABEL_VALUE] = parts[1]

		log.Debug("using user entered node label for replica creation")
	}

	labels := make(map[string]string)
	labels[config.LABEL_PG_CLUSTER] = cluster.Spec.Name

	spec.ClusterName = cluster.Spec.Name

	var rc int
	rc, err = strconv.Atoi(replicaCount)
	if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	labels[config.LABEL_PGOUSER] = pgouser
	labels[config.LABEL_PG_CLUSTER_IDENTIFIER] = cluster.ObjectMeta.Labels[config.LABEL_PG_CLUSTER_IDENTIFIER]

	for i := 0; i < rc; i++ {

		uniqueName := util.RandStringBytesRmndr(4)
		labels[config.LABEL_NAME] = cluster.Spec.Name + "-" + uniqueName
		spec.Namespace = ns
		spec.Name = labels[config.LABEL_NAME]

		newInstance := &crv1.Pgreplica{
			ObjectMeta: meta_v1.ObjectMeta{
				Name:   labels[config.LABEL_NAME],
				Labels: labels,
			},
			Spec: spec,
			Status: crv1.PgreplicaStatus{
				State:   crv1.PgreplicaStateCreated,
				Message: "Created, not processed yet",
			},
		}

		result := crv1.Pgreplica{}

		err = apiserver.RESTClient.Post().
			Resource(crv1.PgreplicaResourcePlural).
			Namespace(ns).
			Body(newInstance).
			Do().Into(&result)
		if err != nil {
			log.Error(" in creating Pgreplica instance" + err.Error())
		}

		response.Results = append(response.Results, "created Pgreplica "+labels[config.LABEL_NAME])
	}

	return response
}

// ScaleQuery lists the replicas that are in the PostgreSQL cluster
// with information that is helpful in determining which one to fail over to,
// such as the lag behind the replica as well as the timeline
func ScaleQuery(name, ns string) msgs.ScaleQueryResponse {
	var err error

	response := msgs.ScaleQueryResponse{
		Results: make([]msgs.ScaleQueryTargetSpec, 0),
		Status:  msgs.Status{Code: msgs.Ok, Msg: ""},
	}

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(name).
		Do().Into(&cluster)

	// If no clusters are found, return a specific error message,
	// otherwise, pass forward the generic error message that Kubernetes sends
	if kerrors.IsNotFound(err) {
		errorMsg := fmt.Sprintf(`No cluster found for "%s"`, name)
		log.Error(errorMsg)
		response.Status.Code = msgs.Error
		response.Status.Msg = errorMsg
		return response
	} else if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// Get information about the current status of all of the replicas. This is
	// handled by a helper function, that will return the information in a struct
	// with the key elements to help the user understand the current state of the
	// replicas in a cluster
	replicationStatusRequest := util.ReplicationStatusRequest{
		RESTConfig:  apiserver.RESTConfig,
		Clientset:   apiserver.Clientset,
		Namespace:   ns,
		ClusterName: name,
	}

	replicationStatusResponse, err := util.ReplicationStatus(replicationStatusRequest)

	// if an error is return, log the message, and return the response
	if err != nil {
		log.Error(err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// indicate in the response whether or not a standby cluster
	response.Standby = cluster.Spec.Standby

	// if there are no results, return the response as is
	if len(replicationStatusResponse.Instances) == 0 {
		return response
	}

	// iterate through response results to create the API response
	for _, instance := range replicationStatusResponse.Instances {
		// create a result for the response
		result := msgs.ScaleQueryTargetSpec{
			Name:           instance.Name,
			Node:           instance.Node,
			Status:         instance.Status,
			ReplicationLag: instance.ReplicationLag,
			Timeline:       instance.Timeline,
		}

		// append the result to the response list
		response.Results = append(response.Results, result)
	}

	return response
}

// ScaleDown ...
func ScaleDown(deleteData bool, clusterName, replicaName, ns string) msgs.ScaleDownResponse {

	var err error

	response := msgs.ScaleDownResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}
	response.Results = make([]string, 0)

	cluster := crv1.Pgcluster{}
	err = apiserver.RESTClient.Get().
		Resource(crv1.PgclusterResourcePlural).
		Namespace(ns).
		Name(clusterName).
		Do().Into(&cluster)

	if kerrors.IsNotFound(err) {
		log.Error("no clusters found")
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	if err != nil {
		log.Error("error getting cluster" + err.Error())
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// dont proceed any further if the cluster is shutdown
	if cluster.Status.State == crv1.PgclusterStateShutdown {
		response.Status.Code = msgs.Error
		response.Status.Msg = "Nothing to scale, the cluster is currently " +
			"shutdown"
		return response
	}

	// selector in the format "pg-cluster=<cluster-name>,pg-ha-scope=<cluster-name>"
	// which will grab the primary and any/all replicas
	selector := fmt.Sprintf("%s=%s,%s=%s", config.LABEL_PG_CLUSTER, clusterName,
		config.LABEL_PGHA_ROLE, "replica")
	replicaList, err := kubeapi.GetPods(apiserver.Clientset, selector, ns)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	// check to see if the replica name provided matches the name of any of the
	// replicas found for the cluster
	var replicaNameFound bool
	for _, pod := range replicaList.Items {
		if pod.Labels[config.LABEL_DEPLOYMENT_NAME] == replicaName {
			replicaNameFound = true
			break
		}
	}
	// return an error if the replica name provided does not match the primary or any replicas
	if !replicaNameFound {
		response.Status.Code = msgs.Error
		response.Status.Msg = fmt.Sprintf("Unable to find replica with name %s",
			replicaName)
		return response
	}

	//create the rmdata task which does the cleanup

	clusterPGHAScope := cluster.ObjectMeta.Labels[config.LABEL_PGHA_SCOPE]
	deleteBackups := false
	isReplica := true
	isBackup := false
	taskName := replicaName + "-rmdata"
	err = apiserver.CreateRMDataTask(clusterName, replicaName, taskName, deleteBackups, deleteData, isReplica, isBackup, ns, clusterPGHAScope)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	response.Results = append(response.Results, "deleted replica "+replicaName)
	return response
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/codegangsta/cli"
)

func doApply(c *cli.Context) error {
	inputFilename := c.String("input-filename")
	filterStatusCodeOk := c.Bool("filter-status-code-ok")

	if inputFilename == "" {
		_ = cli.ShowCommandHelp(c, "apply")
		return cli.NewExitError("`input-filename` is a required field.", 1)
	}

	cacheFilename := CacheFilename

	cache, err := LoadCache(cacheFilename)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}

	examples, err := ReadExamples(inputFilename)
	if err != nil {
		return err
	}

	AttachMetaData(cache, examples)
	if filterStatusCodeOk {
		examples = FilterStatusCodeOkExamples(examples)
	}
	model := TrainedModel(examples)

	for _, e := range FilterUnlabeledExamples(examples) {
		e.Score = model.PredictScore(e.Fv)
		fmt.Println(fmt.Sprintf("%0.03f\t%s", e.Score, e.Url))
	}

	cache.Save(cacheFilename)
	return nil
}

var commandApply = cli.Command{
	Name:  "apply",
	Usage: "apply classifier to unlabeled examples",
	Description: `
Apply classifier to unlabeled examples, and print a pair of score and url.
`,
	Action: doApply,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "input-filename"},
		cli.BoolFlag{Name: "filter-status-code-ok", Usage: "Use only examples with status code = 200"},
	},
}
<commit_msg>jsonでの出力をサポート<commit_after>package main

import (
	"fmt"
	"os"

	"encoding/json"

	"github.com/codegangsta/cli"
)

func doApply(c *cli.Context) error {
	inputFilename := c.String("input-filename")
	filterStatusCodeOk := c.Bool("filter-status-code-ok")
	jsonOutput := c.Bool("json-output")

	if inputFilename == "" {
		_ = cli.ShowCommandHelp(c, "apply")
		return cli.NewExitError("`input-filename` is a required field.", 1)
	}

	cacheFilename := CacheFilename

	cache, err := LoadCache(cacheFilename)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}

	examples, err := ReadExamples(inputFilename)
	if err != nil {
		return err
	}

	AttachMetaData(cache, examples)
	if filterStatusCodeOk {
		examples = FilterStatusCodeOkExamples(examples)
	}
	model := TrainedModel(examples)

	for _, e := range FilterUnlabeledExamples(examples) {
		e.Score = model.PredictScore(e.Fv)
		if jsonOutput {
			b, err := json.Marshal(e)
			if err != nil {
				return err
			}
			fmt.Println(string(b))
		} else {
			fmt.Println(fmt.Sprintf("%0.03f\t%s", e.Score, e.Url))
		}
	}

	cache.Save(cacheFilename)
	return nil
}

var commandApply = cli.Command{
	Name:  "apply",
	Usage: "apply classifier to unlabeled examples",
	Description: `
Apply classifier to unlabeled examples, and print a pair of score and url.
`,
	Action: doApply,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "input-filename"},
		cli.BoolFlag{Name: "filter-status-code-ok", Usage: "Use only examples with status code = 200"},
		cli.BoolFlag{Name: "json-output"},
	},
}
<|endoftext|>


<commit_before>// Automatically downloads and configures Steam grid images for all games in a
// given Steam installation.
package main

import (
	"image"
	"image/draw"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"image/jpeg"
	_ "image/png"
)

// User in the local steam installation.
type User struct {
	Name string
	Dir  string
}

// Given the Steam installation dir (NOT the library!), returns all users in
// this computer.
func GetUsers(installationDir string) ([]User, error) {
	userdataDir := filepath.Join(installationDir, "userdata")
	files, err := ioutil.ReadDir(userdataDir)
	if err != nil {
		return nil, err
	}

	users := make([]User, 0)

	for _, userDir := range files {
		userId := userDir.Name()
		userDir := filepath.Join(userdataDir, userId)

		configFile := filepath.Join(userDir, "config", "localconfig.vdf")
		configBytes, err := ioutil.ReadFile(configFile)
		if err != nil {
			return nil, err
		}

		pattern := regexp.MustCompile(`"PersonaName"\s*"(.+?)"`)
		username := pattern.FindStringSubmatch(string(configBytes))[1]
		users = append(users, User{username, userDir})
	}

	return users, nil
}

// Steam profile URL format.
const urlFormat = `http://steamcommunity.com/id/%v/games?tab=all`

// Returns the public Steam profile for a given user, in HTML.
func GetProfile(username string) (string, error) {
	response, err := http.Get(fmt.Sprintf(urlFormat, username))
	if err != nil {
		return "", err
	}

	contentBytes, err := ioutil.ReadAll(response.Body)
	response.Body.Close()
	if err != nil {
		return "", err
	}

	return string(contentBytes), nil
}

// A Steam game in a library. May or may not be installed.
type Game struct {
	// Official Steam id.
	Id string
	// Warning, may contain Unicode characters.
	Name string
	// User created category. May be blank.
	Category string
	// Path for the grid image.
	ImagePath string
}

// Pattern of game declarations in the public profile. It's actually JSON
// inside Javascript, but this way is easier to extract.
const profileGamePattern = `\{"appid":\s*(\d+),\s*"name":\s*"(.+?)"`

// Returns all games from a given user, using both the public profile and local
// files to gather the data. Returns a map of game by ID.
func GetGames(user User) (map[string]*Game, error) {
	profile, err := GetProfile(user.Name)
	if err != nil {
		return nil, err
	}

	// Fetch game list from public profile.
	pattern := regexp.MustCompile(profileGamePattern)
	games := make(map[string]*Game, 0)
	for _, groups := range pattern.FindAllStringSubmatch(profile, -1) {
		gameId := groups[1]
		gameName := groups[2]
		category := ""
		imagePath := ""
		games[gameId] = &Game{gameId, gameName, category, imagePath}
	}

	// Fetch game categories from local file.
	sharedConfFile := filepath.Join(user.Dir, "7", "remote", "sharedconfig.vdf")
	sharedConfBytes, err := ioutil.ReadFile(sharedConfFile)

	sharedConf := string(sharedConfBytes)
	// VDF patterN: "steamid" { "tags { "0" "category" } }
	pattern = regexp.MustCompile(`"([0-9]+)"\s*{[^}]+?"tags"\s*{\s*"0"\s*"([^"]+)"`)
	for _, groups := range pattern.FindAllStringSubmatch(sharedConf, -1) {
		gameId := groups[1]
		category := groups[2]

		game, ok := games[gameId]
		if ok {
			game.Category = category
		} else {
			// If for some reason it wasn't included in the profile, create a new
			// entry for it now. Unfortunately we don't have a name.
			gameName := ""
			games[gameId] = &Game{gameId, gameName, category, ""}
		}
	}

	return games, nil
}

// When all else fails, Google it. Unfortunately this is a deprecated API and
// may go offline at any time. Because this is last resort the number of
// requests shouldn't trigger any punishment.
const googleSearchFormat = `https://ajax.googleapis.com/ajax/services/search/images?v=1.0&q=`

// Returns the first steam grid image URL found by Google search of a given
// game name.
func getGoogleImage(gameName string) (string, error) {
	url := googleSearchFormat + url.QueryEscape("steam grid OR header"+gameName)
	response, err := http.Get(url)
	if err != nil {
		return "", err
	}

	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	response.Body.Close()
	// Again, we could parse JSON. This may be a little too lazy, the pattern
	// is very loose. The order could be wrong, for example.
	pattern := regexp.MustCompile(`"width":"460","height":"215",[^}]+"unescapedUrl":"(.+?)"`)
	matches := pattern.FindStringSubmatch(string(responseBytes))
	if len(matches) >= 1 {
		return matches[0], nil
	} else {
		return "", nil
	}
}

// Tries to fetch a URL, returning the response only if it was positive.
func tryDownload(url string) (*http.Response, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	if response.StatusCode == 404 {
		// Some apps don't have an image and there's nothing we can do.
		return nil, nil
	} else if response.StatusCode > 400 {
		// Other errors should be reported, though.
		return nil, errors.New("Failed to download image " + url + ": " + response.Status)
	}

	return response, nil
}

// Primary URL for downloading grid images.
const akamaiUrlFormat = `https://steamcdn-a.akamaihd.net/steam/apps/%v/header.jpg`

// The subreddit mentions this as primary, but I've found Akamai to contain
// more images and answer faster.
const steamCdnUrlFormat = `http://cdn.steampowered.com/v/gfx/apps/%v/header.jpg`

// Tries to load the grid image for a game from a number of alternative
// sources. Returns the final response received and a flag indicating if it was
// fro ma Google search (useful because we want to log the lower quality
// images).
func getImageAlternatives(game *Game) (response *http.Response, fromSearch bool, err error) {
	response, err = tryDownload(fmt.Sprintf(akamaiUrlFormat, game.Id))
	if err == nil && response != nil {
		return
	}

	response, err = tryDownload(fmt.Sprintf(steamCdnUrlFormat, game.Id))
	if err == nil && response != nil {
		return
	}

	fromSearch = true
	url, err := getGoogleImage(game.Name)
	if err != nil {
		return
	}
	response, err = tryDownload(url)
	if err == nil && response != nil {
		return
	}

	return nil, false, nil
}

// Downloads the grid image for a game into the user's grid directory. Returns
// flags indicating if the operation succeeded and if the image downloaded was
// from a search.
func DownloadImage(game *Game, user User) (found bool, fromSearch bool, err error) {
	gridDir := filepath.Join(user.Dir, "config", "grid")
	filename := filepath.Join(gridDir, game.Id+".jpg")

	game.ImagePath = filename
	if _, err := os.Stat(filename); err == nil {
		// File already exists, skip it.
		return true, false, nil
	}

	response, fromSearch, err := getImageAlternatives(game)
	if response == nil || err != nil {
		return false, false, err
	}

	imageBytes, err := ioutil.ReadAll(response.Body)
	response.Body.Close()
	return true, fromSearch, ioutil.WriteFile(filename, imageBytes, 0666)
}

// Loads an image from a given path.
func loadImage(path string) (img image.Image, err error) {
	reader, err := os.Open(path)
	if err != nil {
		return
	}
	defer reader.Close()

	img, _, err = image.Decode(reader)
	return
}
// Loads the overlays from the given dir, returning a map of name -> image.
func LoadOverlays(dir string) (overlays map[string]image.Image, err error) {

	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return
	}

	overlays = make(map[string]image.Image, 0)

	for _, file := range files {
		img, err := loadImage(filepath.Join(dir, file.Name()))
		if err != nil {
			return nil, err
		}

		// Normalize overlay name.
		name := strings.TrimSuffix(file.Name(), filepath.Ext(file.Name()))
		overlays[strings.ToLower(name)] = img
	}

	return
}

// Applies an overlay to the game image, depending on the category. The
// resulting image is saved over the original.
func ApplyOverlay(game *Game, overlays map[string]image.Image) (err error) {
	if game.ImagePath == "" {
		return nil
	}

	if _, err := os.Stat(game.ImagePath); err != nil {
		// Game has no image, we have to skip it.
		return nil
	}

	// Normalize overlay name.
	categoryName := strings.ToLower(game.Category)

	overlayImage, ok := overlays[categoryName]
	if !ok {
		return
	}

	gameImage, err := loadImage(game.ImagePath)
	if err != nil {
		return err
	}

    result := image.NewRGBA(gameImage.Bounds().Union(overlayImage.Bounds()))
    draw.Draw(result, result.Bounds(), gameImage, image.ZP, draw.Src)
    draw.Draw(result, result.Bounds(), overlayImage, image.Point{0,0}, draw.Over)

	ext := filepath.Ext(game.ImagePath)
	backupPath := strings.TrimSuffix(game.ImagePath, ext) + " (original)" + ext
	if _, err := os.Stat(backupPath); err != nil {
		// Backup doesn't exist, create it.
		err = os.Rename(game.ImagePath, backupPath)
		if err != nil {
			return err
		}
	}

    resultFile, _ := os.Create(game.ImagePath)
    defer resultFile.Close()
    return jpeg.Encode(resultFile, result, &jpeg.Options{90})
}

// Returns the Steam installation directory in Windows. Should work for
// internationalized systems, 32 and 64 bits and users that moved their
// ProgramFiles folder. If a folder is given by program parameter, uses that.
func GetSteamInstallation() (path string, err error) {
	if len(os.Args) == 2 {
		argDir := os.Args[1]
		_, err := os.Stat(argDir)
		if err == nil {
			return argDir, nil
		} else {
			return "", errors.New("Argument must be a valid Steam directory, or empty for auto detection. Got: " + argDir)
		}
	}

	programFiles86Dir := filepath.Join(os.Getenv("ProgramFiles(x86)"), "Steam")
	if _, err = os.Stat(programFiles86Dir); err == nil {
		return programFiles86Dir, nil
	}

	programFilesDir := filepath.Join(os.Getenv("ProgramFiles"), "Steam")
	if _, err = os.Stat(programFilesDir); err == nil {
		return programFilesDir, nil
	}

	return "", errors.New("Could not find Steam installation folder.")
}

// Prints a progress bar, overriding the previous line. It looks like this:
// [=========>        ] (50/100)
func PrintProgress(current int, total int) {
	// \r moves the cursor back to the start of the line.
	fmt.Print("\r[")

	printedHead := false
	for i := 0; i < 40; i++ {
		part := int(float64(i) * (float64(total) / 40.0))
		if part < current {
			fmt.Print("=")
		} else if !printedHead {
			printedHead = true
			fmt.Print(">")
		} else {
			fmt.Print(" ")
		}
	}

	fmt.Printf("] (%v/%v)", current, total)
}

// Prints an error and quits.
func errorAndExit(err error) {
	fmt.Println("An unexpected error occurred:")
	fmt.Println(err)
	os.Stdin.Read(make([]byte, 1))
	os.Exit(1)
}

func main() {
	overlays, err := LoadOverlays("overlays by category")
	if err != nil {
		errorAndExit(err)
	}

	installationDir, err := GetSteamInstallation()
	if err != nil {
		errorAndExit(err)
	}

	users, err := GetUsers(installationDir)
	if err != nil {
		errorAndExit(err)
	}

	for _, user := range users {
		fmt.Printf("Found user %v. Fetching game list from profile...\n\n\n", user.Name)

		games, err := GetGames(user)
		if err != nil {
			errorAndExit(err)
		}

		notFounds := make([]*Game, 0)
		searchFounds := make([]*Game, 0)
		fmt.Printf("Found %v games. Downloading images...\n\n", len(games))

		i := 0
		for _, game := range games {
			i++
			PrintProgress(i, len(games))
			found, fromSearch, err := DownloadImage(game, user)
			if err != nil {
				errorAndExit(err)
			}
			if !found {
				notFounds = append(notFounds, game)
				continue
			}
			if fromSearch {
				searchFounds = append(searchFounds, game)
			}

			err = ApplyOverlay(game, overlays)
			if err != nil {
				errorAndExit(err)
			}
		}
		fmt.Print("\n\n\n")

		if len(notFounds) == 0 && len(searchFounds) == 0 {
			fmt.Println("All grid images downloaded!")
		} else {
			if len(searchFounds) >= 1 {
				fmt.Printf("%v images were found with a Google search and may not be accurate:.\n", len(searchFounds))
				for _, game := range searchFounds {
					fmt.Printf("* %v (steam id %v)\n", game.Name, game.Id)
				}
			}

			fmt.Print("\n\n")

			if len(notFounds) >= 1 {
				fmt.Printf("%v images could not be found:\n", len(notFounds))
				for _, game := range notFounds {
					fmt.Printf("* %v (steam id %v)\n", game.Name, game.Id)
				}
			}
		}
	}

	fmt.Print("\n\n")
	fmt.Println("You can press enter to close this window.")
	os.Stdin.Read(make([]byte, 1))
}
<commit_msg>Ignore when the overlays folder is not found<commit_after>// Automatically downloads and configures Steam grid images for all games in a
// given Steam installation.
package main

import (
	"image"
	"image/draw"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"image/jpeg"
	_ "image/png"
)

// User in the local steam installation.
type User struct {
	Name string
	Dir  string
}

// Given the Steam installation dir (NOT the library!), returns all users in
// this computer.
func GetUsers(installationDir string) ([]User, error) {
	userdataDir := filepath.Join(installationDir, "userdata")
	files, err := ioutil.ReadDir(userdataDir)
	if err != nil {
		return nil, err
	}

	users := make([]User, 0)

	for _, userDir := range files {
		userId := userDir.Name()
		userDir := filepath.Join(userdataDir, userId)

		configFile := filepath.Join(userDir, "config", "localconfig.vdf")
		configBytes, err := ioutil.ReadFile(configFile)
		if err != nil {
			return nil, err
		}

		pattern := regexp.MustCompile(`"PersonaName"\s*"(.+?)"`)
		username := pattern.FindStringSubmatch(string(configBytes))[1]
		users = append(users, User{username, userDir})
	}

	return users, nil
}

// Steam profile URL format.
const urlFormat = `http://steamcommunity.com/id/%v/games?tab=all`

// Returns the public Steam profile for a given user, in HTML.
func GetProfile(username string) (string, error) {
	response, err := http.Get(fmt.Sprintf(urlFormat, username))
	if err != nil {
		return "", err
	}

	contentBytes, err := ioutil.ReadAll(response.Body)
	response.Body.Close()
	if err != nil {
		return "", err
	}

	return string(contentBytes), nil
}

// A Steam game in a library. May or may not be installed.
type Game struct {
	// Official Steam id.
	Id string
	// Warning, may contain Unicode characters.
	Name string
	// User created category. May be blank.
	Category string
	// Path for the grid image.
	ImagePath string
}

// Pattern of game declarations in the public profile. It's actually JSON
// inside Javascript, but this way is easier to extract.
const profileGamePattern = `\{"appid":\s*(\d+),\s*"name":\s*"(.+?)"`

// Returns all games from a given user, using both the public profile and local
// files to gather the data. Returns a map of game by ID.
func GetGames(user User) (map[string]*Game, error) {
	profile, err := GetProfile(user.Name)
	if err != nil {
		return nil, err
	}

	// Fetch game list from public profile.
	pattern := regexp.MustCompile(profileGamePattern)
	games := make(map[string]*Game, 0)
	for _, groups := range pattern.FindAllStringSubmatch(profile, -1) {
		gameId := groups[1]
		gameName := groups[2]
		category := ""
		imagePath := ""
		games[gameId] = &Game{gameId, gameName, category, imagePath}
	}

	// Fetch game categories from local file.
	sharedConfFile := filepath.Join(user.Dir, "7", "remote", "sharedconfig.vdf")
	sharedConfBytes, err := ioutil.ReadFile(sharedConfFile)

	sharedConf := string(sharedConfBytes)
	// VDF patterN: "steamid" { "tags { "0" "category" } }
	pattern = regexp.MustCompile(`"([0-9]+)"\s*{[^}]+?"tags"\s*{\s*"0"\s*"([^"]+)"`)
	for _, groups := range pattern.FindAllStringSubmatch(sharedConf, -1) {
		gameId := groups[1]
		category := groups[2]

		game, ok := games[gameId]
		if ok {
			game.Category = category
		} else {
			// If for some reason it wasn't included in the profile, create a new
			// entry for it now. Unfortunately we don't have a name.
			gameName := ""
			games[gameId] = &Game{gameId, gameName, category, ""}
		}
	}

	return games, nil
}

// When all else fails, Google it. Unfortunately this is a deprecated API and
// may go offline at any time. Because this is last resort the number of
// requests shouldn't trigger any punishment.
const googleSearchFormat = `https://ajax.googleapis.com/ajax/services/search/images?v=1.0&q=`

// Returns the first steam grid image URL found by Google search of a given
// game name.
func getGoogleImage(gameName string) (string, error) {
	url := googleSearchFormat + url.QueryEscape("steam grid OR header"+gameName)
	response, err := http.Get(url)
	if err != nil {
		return "", err
	}

	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	response.Body.Close()
	// Again, we could parse JSON. This may be a little too lazy, the pattern
	// is very loose. The order could be wrong, for example.
	pattern := regexp.MustCompile(`"width":"460","height":"215",[^}]+"unescapedUrl":"(.+?)"`)
	matches := pattern.FindStringSubmatch(string(responseBytes))
	if len(matches) >= 1 {
		return matches[0], nil
	} else {
		return "", nil
	}
}

// Tries to fetch a URL, returning the response only if it was positive.
func tryDownload(url string) (*http.Response, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	if response.StatusCode == 404 {
		// Some apps don't have an image and there's nothing we can do.
		return nil, nil
	} else if response.StatusCode > 400 {
		// Other errors should be reported, though.
		return nil, errors.New("Failed to download image " + url + ": " + response.Status)
	}

	return response, nil
}

// Primary URL for downloading grid images.
const akamaiUrlFormat = `https://steamcdn-a.akamaihd.net/steam/apps/%v/header.jpg`

// The subreddit mentions this as primary, but I've found Akamai to contain
// more images and answer faster.
const steamCdnUrlFormat = `http://cdn.steampowered.com/v/gfx/apps/%v/header.jpg`

// Tries to load the grid image for a game from a number of alternative
// sources. Returns the final response received and a flag indicating if it was
// fro ma Google search (useful because we want to log the lower quality
// images).
func getImageAlternatives(game *Game) (response *http.Response, fromSearch bool, err error) {
	response, err = tryDownload(fmt.Sprintf(akamaiUrlFormat, game.Id))
	if err == nil && response != nil {
		return
	}

	response, err = tryDownload(fmt.Sprintf(steamCdnUrlFormat, game.Id))
	if err == nil && response != nil {
		return
	}

	fromSearch = true
	url, err := getGoogleImage(game.Name)
	if err != nil {
		return
	}
	response, err = tryDownload(url)
	if err == nil && response != nil {
		return
	}

	return nil, false, nil
}

// Downloads the grid image for a game into the user's grid directory. Returns
// flags indicating if the operation succeeded and if the image downloaded was
// from a search.
func DownloadImage(game *Game, user User) (found bool, fromSearch bool, err error) {
	gridDir := filepath.Join(user.Dir, "config", "grid")
	filename := filepath.Join(gridDir, game.Id+".jpg")

	game.ImagePath = filename
	if _, err := os.Stat(filename); err == nil {
		// File already exists, skip it.
		return true, false, nil
	}

	response, fromSearch, err := getImageAlternatives(game)
	if response == nil || err != nil {
		return false, false, err
	}

	imageBytes, err := ioutil.ReadAll(response.Body)
	response.Body.Close()
	return true, fromSearch, ioutil.WriteFile(filename, imageBytes, 0666)
}

// Loads an image from a given path.
func loadImage(path string) (img image.Image, err error) {
	reader, err := os.Open(path)
	if err != nil {
		return
	}
	defer reader.Close()

	img, _, err = image.Decode(reader)
	return
}
// Loads the overlays from the given dir, returning a map of name -> image.
func LoadOverlays(dir string) (overlays map[string]image.Image, err error) {
	overlays = make(map[string]image.Image, 0)

	if _, err = os.Stat(dir); err != nil {
		return overlays, nil
	}

	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return
	}

	for _, file := range files {
		img, err := loadImage(filepath.Join(dir, file.Name()))
		if err != nil {
			return nil, err
		}

		// Normalize overlay name.
		name := strings.TrimSuffix(file.Name(), filepath.Ext(file.Name()))
		overlays[strings.ToLower(name)] = img
	}

	return
}

// Applies an overlay to the game image, depending on the category. The
// resulting image is saved over the original.
func ApplyOverlay(game *Game, overlays map[string]image.Image) (err error) {
	if game.ImagePath == "" {
		return nil
	}

	if _, err := os.Stat(game.ImagePath); err != nil {
		// Game has no image, we have to skip it.
		return nil
	}

	// Normalize overlay name.
	categoryName := strings.ToLower(game.Category)

	overlayImage, ok := overlays[categoryName]
	if !ok {
		return
	}

	gameImage, err := loadImage(game.ImagePath)
	if err != nil {
		return err
	}

    result := image.NewRGBA(gameImage.Bounds().Union(overlayImage.Bounds()))
    draw.Draw(result, result.Bounds(), gameImage, image.ZP, draw.Src)
    draw.Draw(result, result.Bounds(), overlayImage, image.Point{0,0}, draw.Over)

	ext := filepath.Ext(game.ImagePath)
	backupPath := strings.TrimSuffix(game.ImagePath, ext) + " (original)" + ext
	if _, err := os.Stat(backupPath); err != nil {
		// Backup doesn't exist, create it.
		err = os.Rename(game.ImagePath, backupPath)
		if err != nil {
			return err
		}
	}

    resultFile, _ := os.Create(game.ImagePath)
    defer resultFile.Close()
    return jpeg.Encode(resultFile, result, &jpeg.Options{90})
}

// Returns the Steam installation directory in Windows. Should work for
// internationalized systems, 32 and 64 bits and users that moved their
// ProgramFiles folder. If a folder is given by program parameter, uses that.
func GetSteamInstallation() (path string, err error) {
	if len(os.Args) == 2 {
		argDir := os.Args[1]
		_, err := os.Stat(argDir)
		if err == nil {
			return argDir, nil
		} else {
			return "", errors.New("Argument must be a valid Steam directory, or empty for auto detection. Got: " + argDir)
		}
	}

	programFiles86Dir := filepath.Join(os.Getenv("ProgramFiles(x86)"), "Steam")
	if _, err = os.Stat(programFiles86Dir); err == nil {
		return programFiles86Dir, nil
	}

	programFilesDir := filepath.Join(os.Getenv("ProgramFiles"), "Steam")
	if _, err = os.Stat(programFilesDir); err == nil {
		return programFilesDir, nil
	}

	return "", errors.New("Could not find Steam installation folder.")
}

// Prints a progress bar, overriding the previous line. It looks like this:
// [=========>        ] (50/100)
func PrintProgress(current int, total int) {
	// \r moves the cursor back to the start of the line.
	fmt.Print("\r[")

	printedHead := false
	for i := 0; i < 40; i++ {
		part := int(float64(i) * (float64(total) / 40.0))
		if part < current {
			fmt.Print("=")
		} else if !printedHead {
			printedHead = true
			fmt.Print(">")
		} else {
			fmt.Print(" ")
		}
	}

	fmt.Printf("] (%v/%v)", current, total)
}

// Prints an error and quits.
func errorAndExit(err error) {
	fmt.Println("An unexpected error occurred:")
	fmt.Println(err)
	os.Stdin.Read(make([]byte, 1))
	os.Exit(1)
}

func main() {
	overlays, err := LoadOverlays("overlays by category")
	if err != nil {
		errorAndExit(err)
	}

	installationDir, err := GetSteamInstallation()
	if err != nil {
		errorAndExit(err)
	}

	users, err := GetUsers(installationDir)
	if err != nil {
		errorAndExit(err)
	}

	for _, user := range users {
		fmt.Printf("Found user %v. Fetching game list from profile...\n\n\n", user.Name)

		games, err := GetGames(user)
		if err != nil {
			errorAndExit(err)
		}

		notFounds := make([]*Game, 0)
		searchFounds := make([]*Game, 0)
		fmt.Printf("Found %v games. Downloading images...\n\n", len(games))

		i := 0
		for _, game := range games {
			i++
			PrintProgress(i, len(games))
			found, fromSearch, err := DownloadImage(game, user)
			if err != nil {
				errorAndExit(err)
			}
			if !found {
				notFounds = append(notFounds, game)
				continue
			}
			if fromSearch {
				searchFounds = append(searchFounds, game)
			}

			err = ApplyOverlay(game, overlays)
			if err != nil {
				errorAndExit(err)
			}
		}
		fmt.Print("\n\n\n")

		if len(notFounds) == 0 && len(searchFounds) == 0 {
			fmt.Println("All grid images downloaded!")
		} else {
			if len(searchFounds) >= 1 {
				fmt.Printf("%v images were found with a Google search and may not be accurate:.\n", len(searchFounds))
				for _, game := range searchFounds {
					fmt.Printf("* %v (steam id %v)\n", game.Name, game.Id)
				}
			}

			fmt.Print("\n\n")

			if len(notFounds) >= 1 {
				fmt.Printf("%v images could not be found:\n", len(notFounds))
				for _, game := range notFounds {
					fmt.Printf("* %v (steam id %v)\n", game.Name, game.Id)
				}
			}
		}
	}

	fmt.Print("\n\n")
	fmt.Println("You can press enter to close this window.")
	os.Stdin.Read(make([]byte, 1))
}
<|endoftext|>


<commit_before>//wsserver.go
package ws

import (
	"encoding/json"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"github.com/jayaramsankara/gotell/apns"
	"github.com/satori/go.uuid"
	"gopkg.in/redis.v5"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 30 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512

	// Max NUmber of Redis connections maintained for publishing events
	// This is same as the Go Channel buffer size for sending events for redis
	maxRedisConn = 50
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

var clientConnections = map[string][]*wsconnection{}

var redisSender chan *NotifyData

var receiver *redis.PubSub

var logs = log.New(os.Stdout, "INFO", log.LstdFlags)

var Logs = logs

type wsconnection struct {
	// The websocket connection.
	ws       *websocket.Conn
	clientId string
	id       string
	//The redis pubsubs channel
	// Buffered channel of outbound messages.
	//receive *redis.PubSub
	send   chan ([]byte)
	active bool
}

type redisclients struct {
	receiver *redis.PubSub
	sender   *redis.Client
}

type WsMessage struct {
	Message string `json:"message"`
}

type NotifyData struct {
	ClientId string
	Message  string
}

type NotifyResponse struct {
	Status   bool   `json:"status"`
	ClientId string `json:"clientId"`
}

func (conn *wsconnection) Close() {
	conn.write(websocket.CloseMessage, []byte{})
	conn.ws.Close()
}

func (conn *wsconnection) String() string {
	return conn.clientId + " : " + conn.id
}

func (conn *wsconnection) sendMessages() {
	logs.Println("Initing sendMessages for client  " + conn.String())
	clientId := conn.clientId
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		logs.Println("Exiting sendMessages for client  : " + conn.String())
		conn.active = false
		logs.Println("Closing websocket connection for ", conn.String())
		conn.Close()

		currentConnections := clientConnections[clientId]
		if len(currentConnections) == 1 {
			logs.Println("Removing subscription to redis channel for ", conn.String())
			receiver.Unsubscribe(conn.clientId)
		}
		logs.Println("Removing the connection mapped to " + conn.String())
		newConnections := currentConnections[:0]
		for _, ec := range currentConnections {
			if ec.id != conn.id {
				newConnections = append(newConnections, ec)
			}
		}
		clientConnections[clientId] = newConnections
		logs.Println("Exiting sendMessages. Current # conenctions for client " + clientId + " is " + strconv.Itoa(len(clientConnections[clientId])))
	}()

	for {
		select {
		case message, ok := <-conn.send:
			if !ok {
				log.Println("No more messages to be sent to WS connection from channel for  ", conn.String())
				return
			}
			logs.Println("Sending WS message for client  " + conn.String())
			if err := conn.write(websocket.TextMessage, message); err != nil {
				log.Println("Error while sending WS message  ", conn.String(), message, err)
				return
			}
		case <-ticker.C:
			logs.Println("Sending Ping WS message for client  " + conn.String())
			if err := conn.write(websocket.PingMessage, []byte{}); err != nil {
				log.Println("Error while sending WS ping message  ", conn.String(), err)
				return
			}
		}
	}

}

func (conn *wsconnection) receiveMessages() {
	logs.Println("Initing receiveMessages for client  " + conn.String())
	go func() {
		for {
			if _, _, err := conn.ws.NextReader(); err != nil {
				log.Println("Error while receiving WS pong message  ", conn.String(), err)
				conn.active = false
				conn.Close()
				break
			}
		}
	}()

	conn.ws.SetReadLimit(maxMessageSize)
	conn.ws.SetReadDeadline(time.Now().Add(pongWait))
	conn.ws.SetPongHandler(func(string) error {
		logs.Println("PongHandler for client  " + conn.String())
		conn.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
}

// write writes a message with the given message type and payload.
func (conn *wsconnection) write(messageType int, payload []byte) error {
	conn.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return conn.ws.WriteMessage(messageType, payload)
}

func InitPubSub(redisConf *redis.Options) error {
	redisSender = make(chan *NotifyData, redisConf.PoolSize)
	receiver, _ = redis.NewClient(redisConf).Subscribe()

	var publisher = redis.NewClient(redisConf)

	logs.Println("Initialized redis clients for pub and sub.")

	//Init Redis receiver
	go func() {
		for {
			var msgi, err = receiver.Receive()
			if err != nil {
				log.Println("Error while receive message from redis pubsub receiver", err)
				// handle failure.. reinit redis pub sub
				log.Println("Reinitializing Redis PubSub.. and exiting the receive handler routine")
				close(redisSender)
				receiver.Close()
				publisher.Close()
				InitPubSub(redisConf)
				return

			} else {

				switch msg := msgi.(type) {
				case *redis.Subscription:
					logs.Println("Received messagefrom Redis ", msg.Kind, msg.Channel)
				case *redis.Message:
					logs.Println("Received message from Redis ", msg.Payload, msg.Channel)
					go func(clientId string, message string) {
						logs.Println("Attempting to send the message to WS  connections for client ", message, clientId)
						connections := clientConnections[clientId]
						for cnt := range connections {
							logs.Println("Sending the message to WS  client  ", message, connections[cnt].String())
							connections[cnt].send <- []byte(message)
						}

					}(msg.Channel, msg.Payload)

				case *redis.Pong:
					logs.Println("Pong message from Redis : ", msg)
				default:
					log.Printf("Error unknown message from Redis : %#v", msgi)
				}

			}

		}

	}()

	go func(msgsForRedis chan *NotifyData) {
		for {
			data, ok := <-msgsForRedis
			if !ok {
				log.Println("RedisSender channel seems to be closed. Exiting the routine")
				return
			}
			logs.Println("Received notification data from redisSender. ", data.ClientId)
			go func() {
				logs.Println("Publishing message to the Redis ", data.Message, data.ClientId)
				err := publisher.Publish(data.ClientId, data.Message).Err()
				if err != nil {
					log.Println("Error in publishing event to redis", err)
				}
			}()

		}

	}(redisSender)
	return nil
}

func ServeApns(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	deviceToken := vars["devicetoken"]

	//extract body

	decoder := json.NewDecoder(r.Body)
	var data apns.ApnsMessage
	err := decoder.Decode(&data)

	if err != nil {
		log.Println("Error while extracting the apns message.", err)
		w.WriteHeader(http.StatusInternalServerError)

	} else {
		// send data to conn
		logs.Println("Handling apns:  The extracted message is : ", data.Message, data.Badge, data.Sound)
		apns.Notify(&data, deviceToken)
		w.WriteHeader(http.StatusOK)
	}

}

func ServeNotifyCORS(w http.ResponseWriter, r *http.Request) {
	 var headers = r.Header.Get("Access-Control-Request-Headers")

	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST")

	w.Header().Set("Access-Control-Allow-Headers",headers)
	w.WriteHeader(http.StatusOK)
}

//serveNotify receives the API, parses the body and sends the message to the corresponding
// websocket. Returns error if no websocket conn exists for a client id or send fails
func ServeNotify(w http.ResponseWriter, r *http.Request) {

	vars := mux.Vars(r)
	clientId := vars["clientid"]

	//extract body

	decoder := json.NewDecoder(r.Body)
	var data WsMessage
	err := decoder.Decode(&data)

	if err != nil {
		log.Println("Error while extracting the message.", err)
		w.WriteHeader(http.StatusInternalServerError)

	} else {
		// send data to conn
		logs.Println("Handling notify:  The extracted message is : ", data.Message, clientId)
		notifyData := &NotifyData{
			ClientId: clientId,
			Message:  data.Message,
		}
		logs.Println("Handling notify:  Sending notification data to redisSender ", clientId)
		redisSender <- notifyData
		//Check whether clientId is connected.
		currentConnections := clientConnections[clientId]

		notifyResponse := NotifyResponse{
			Status:   (len(currentConnections) > 0),
			ClientId: clientId,
		}
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		err = json.NewEncoder(w).Encode(notifyResponse)
		if err != nil {
			log.Println("Error writing the response.", err)
			w.WriteHeader(http.StatusInternalServerError)
		}

	}

}

// serverWs handles websocket requests from the peer.
func ServeWs(w http.ResponseWriter, r *http.Request) {

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error in accepting websocket connection", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	//new websocket connection
	vars := mux.Vars(r)
	clientId := vars["clientid"]
	logs.Println("Creating new wsconnection for client : " + clientId)

	currentConnections := clientConnections[clientId]
	//Subscribe for a client only for the first connection.
	if len(currentConnections) == 0 {
		// Initialize redis client as subscriber

		err = receiver.Subscribe(clientId)
		if err != nil {
			errMsg := "Failed to subscribe to message channel for " + clientId
			log.Println(errMsg, err)
			ws.WriteMessage(websocket.CloseInternalServerErr, []byte(errMsg))
			ws.Close()
			return
		}
	}

	conn := &wsconnection{active: true, clientId: clientId, id: newUUID(), ws: ws, send: make(chan []byte, 256)}
	logs.Println("Adding clientId-WsConn mapping for client  " + conn.String())

	modifiedConnections := append(currentConnections, conn)
	logs.Println("# connections for  client after appending the new one " + strconv.Itoa(len(modifiedConnections)))

	clientConnections[clientId] = modifiedConnections
	logs.Println("#  updated connections for  client  " + strconv.Itoa(len(clientConnections[clientId])))
	go conn.sendMessages()
	conn.receiveMessages()
}

func newUUID() string {

	u1 := uuid.NewV4()
	return u1.String()
}
<commit_msg>Removed shared map for mainitaining client connections<commit_after>//wsserver.go
package ws

import (
	"encoding/json"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"github.com/jayaramsankara/gotell/apns"
	"github.com/satori/go.uuid"
	"gopkg.in/redis.v5"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 30 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512

	// Max NUmber of Redis connections maintained for publishing events
	// This is same as the Go Channel buffer size for sending events for redis
	maxRedisConn = 50
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

var redisSender chan *NotifyData

var receiver *redis.PubSub

var logs = log.New(os.Stdout, "INFO", log.LstdFlags)

var Logs = logs

var connMgr = initClientConnectionManager()

type wsconnection struct {
	// The websocket connection.
	ws       *websocket.Conn
	clientId string
	id       string
	//The redis pubsubs channel
	// Buffered channel of outbound messages.
	//receive *redis.PubSub
	send   chan ([]byte)
	active bool
}

type redisclients struct {
	receiver *redis.PubSub
	sender   *redis.Client
}

type WsMessage struct {
	Message string `json:"message"`
}

type NotifyData struct {
	ClientId string
	Message  string
}

type NotifyResponse struct {
	Status   bool   `json:"status"`
	ClientId string `json:"clientId"`
}

type connCmd struct {
	Cmd  int
	Data interface{}
}

type newConnCmd struct {
	clientID string
	ws       *websocket.Conn
}

type connectionManager struct {
	incoming chan<- []*wsconnection
	outgoing <-chan []*wsconnection
	command  chan connCmd
}

func (connm *connectionManager) connectionsFor(clientid string) []*wsconnection {
	go func() {
		connm.command <- connCmd{
			Cmd:  0,
			Data: clientid,
		}
	}()
	return <-connm.outgoing

}

func (connm *connectionManager) cleanUpConnection(conn *wsconnection) {
	go func() {
		connm.command <- connCmd{
			Cmd:  2,
			Data: conn,
		}
	}()

}

func (connm *connectionManager) newConnectionFor(clientID string, ws *websocket.Conn) *wsconnection {
	go func() {
		connm.command <- connCmd{
			Cmd: 1,
			Data: newConnCmd{
				clientID,
				ws,
			},
		}
	}()
	newConns := <-connm.outgoing
	return newConns[0]
}

func initClientConnectionManager() connectionManager {
	var clientConnections = map[string][]*wsconnection{}

	outchan := make(chan []*wsconnection)
	cmdchan := make(chan connCmd)

	go func() {
		for {
			select {
			case chanCmd := <-cmdchan:
				switch chanCmd.Cmd {
				case 0:
					clientID := chanCmd.Data.(string)
					connections := clientConnections[clientID]
					outchan <- connections
				case 1:

					data := chanCmd.Data.(newConnCmd)
					clientID := data.clientID
					ws := data.ws
					currentConnections := clientConnections[clientID]
					//Subscribe for a client only for the first connection.
					if len(currentConnections) == 0 {
						// Initialize redis client as subscriber

						err := receiver.Subscribe(clientID)
						if err != nil {
							errMsg := "Failed to subscribe to message channel for " + clientID
							log.Println(errMsg, err)
							conn := &wsconnection{active: false, clientId: clientID, id: newUUID(), ws: ws, send: make(chan []byte, 256)}
							outchan <- []*wsconnection{conn}
							break
						}
					}

					conn := &wsconnection{active: true, clientId: clientID, id: newUUID(), ws: ws, send: make(chan []byte, 256)}
					logs.Println("Adding clientId-WsConn mapping for client  " + conn.String())

					modifiedConnections := append(currentConnections, conn)
					logs.Println("# connections for  client after appending the new one " + strconv.Itoa(len(modifiedConnections)))

					clientConnections[clientID] = modifiedConnections
					logs.Println("#  updated connections for  client  " + strconv.Itoa(len(modifiedConnections)))
					outchan <- []*wsconnection{conn}

				case 2:
					conn := chanCmd.Data.(*wsconnection)
					currentConnections := clientConnections[conn.clientId]
					if len(currentConnections) == 1 {
						logs.Println("Removing subscription to redis channel for ", conn.String())
						receiver.Unsubscribe(conn.clientId)
					}
					logs.Println("Removing the connection mapped to " + conn.String())
					newConnections := currentConnections[:0]
					for _, ec := range currentConnections {
						if ec.id != conn.id {
							newConnections = append(newConnections, ec)
						}
					}
					clientConnections[conn.clientId] = newConnections
					logs.Println(" Current # conenctions for client " + conn.clientId + " is " + strconv.Itoa(len(newConnections)))

				}
			}
		}
	}()

	// go func connectionMmgr(d int) (err) {

	// }()

	return connectionManager{
		outgoing: outchan,
		command:  cmdchan,
	}
}

func (conn *wsconnection) Close() {
	conn.write(websocket.CloseMessage, []byte{})
	conn.ws.Close()
}

func (conn *wsconnection) String() string {
	return conn.clientId + " : " + conn.id
}

func (conn *wsconnection) sendMessages() {
	logs.Println("Initing sendMessages for client  " + conn.String())
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		logs.Println("Exiting sendMessages for client  : " + conn.String())
		conn.active = false
		logs.Println("Closing websocket connection for ", conn.String())
		conn.Close()
		connMgr.cleanUpConnection(conn)
	}()

	for {
		select {
		case message, ok := <-conn.send:
			if !ok {
				log.Println("No more messages to be sent to WS connection from channel for  ", conn.String())
				return
			}
			logs.Println("Sending WS message for client  " + conn.String())
			if err := conn.write(websocket.TextMessage, message); err != nil {
				log.Println("Error while sending WS message  ", conn.String(), message, err)
				return
			}
		case <-ticker.C:
			logs.Println("Sending Ping WS message for client  " + conn.String())
			if err := conn.write(websocket.PingMessage, []byte{}); err != nil {
				log.Println("Error while sending WS ping message  ", conn.String(), err)
				return
			}
		}
	}

}

func (conn *wsconnection) receiveMessages() {
	logs.Println("Initing receiveMessages for client  " + conn.String())
	go func() {
		for {
			if _, _, err := conn.ws.NextReader(); err != nil {
				log.Println("Error while receiving WS pong message  ", conn.String(), err)
				conn.active = false
				conn.Close()
				break
			}
		}
	}()

	conn.ws.SetReadLimit(maxMessageSize)
	conn.ws.SetReadDeadline(time.Now().Add(pongWait))
	conn.ws.SetPongHandler(func(string) error {
		logs.Println("PongHandler for client  " + conn.String())
		conn.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
}

// write writes a message with the given message type and payload.
func (conn *wsconnection) write(messageType int, payload []byte) error {
	conn.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return conn.ws.WriteMessage(messageType, payload)
}

func InitPubSub(redisConf *redis.Options) error {
	redisSender = make(chan *NotifyData, redisConf.PoolSize)
	receiver, _ = redis.NewClient(redisConf).Subscribe()

	var publisher = redis.NewClient(redisConf)

	logs.Println("Initialized redis clients for pub and sub.")

	//Init Redis receiver
	go func() {
		for {
			var msgi, err = receiver.Receive()
			if err != nil {
				log.Println("Error while receive message from redis pubsub receiver", err)
				// handle failure.. reinit redis pub sub
				log.Println("Reinitializing Redis PubSub.. and exiting the receive handler routine")
				close(redisSender)
				receiver.Close()
				publisher.Close()
				InitPubSub(redisConf)
				return

			} else {

				switch msg := msgi.(type) {
				case *redis.Subscription:
					logs.Println("Received messagefrom Redis ", msg.Kind, msg.Channel)
				case *redis.Message:
					logs.Println("Received message from Redis ", msg.Payload, msg.Channel)
					go func(clientId string, message string) {
						logs.Println("Attempting to send the message to WS  connections for client ", message, clientId)
						connections := connMgr.connectionsFor(clientId)
						for cnt := range connections {
							logs.Println("Sending the message to WS  client  ", message, connections[cnt].String())
							connections[cnt].send <- []byte(message)
						}

					}(msg.Channel, msg.Payload)

				case *redis.Pong:
					logs.Println("Pong message from Redis : ", msg)
				default:
					log.Printf("Error unknown message from Redis : %#v", msgi)
				}

			}

		}

	}()

	go func(msgsForRedis chan *NotifyData) {
		for {
			data, ok := <-msgsForRedis
			if !ok {
				log.Println("RedisSender channel seems to be closed. Exiting the routine")
				return
			}
			logs.Println("Received notification data from redisSender. ", data.ClientId)
			go func() {
				logs.Println("Publishing message to the Redis ", data.Message, data.ClientId)
				err := publisher.Publish(data.ClientId, data.Message).Err()
				if err != nil {
					log.Println("Error in publishing event to redis", err)
				}
			}()

		}

	}(redisSender)
	return nil
}

func ServeApns(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	deviceToken := vars["devicetoken"]

	//extract body

	decoder := json.NewDecoder(r.Body)
	var data apns.ApnsMessage
	err := decoder.Decode(&data)

	if err != nil {
		log.Println("Error while extracting the apns message.", err)
		w.WriteHeader(http.StatusInternalServerError)

	} else {
		// send data to conn
		logs.Println("Handling apns:  The extracted message is : ", data.Message, data.Badge, data.Sound)
		apns.Notify(&data, deviceToken)
		w.WriteHeader(http.StatusOK)
	}

}

func ServeNotifyCORS(w http.ResponseWriter, r *http.Request) {
	var headers = r.Header.Get("Access-Control-Request-Headers")

	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST")

	w.Header().Set("Access-Control-Allow-Headers", headers)
	w.WriteHeader(http.StatusOK)
}

//serveNotify receives the API, parses the body and sends the message to the corresponding
// websocket. Returns error if no websocket conn exists for a client id or send fails
func ServeNotify(w http.ResponseWriter, r *http.Request) {

	vars := mux.Vars(r)
	clientId := vars["clientid"]

	//extract body

	decoder := json.NewDecoder(r.Body)
	var data WsMessage
	err := decoder.Decode(&data)

	if err != nil {
		log.Println("Error while extracting the message.", err)
		w.WriteHeader(http.StatusInternalServerError)

	} else {
		// send data to conn
		logs.Println("Handling notify:  The extracted message is : ", data.Message, clientId)
		notifyData := &NotifyData{
			ClientId: clientId,
			Message:  data.Message,
		}
		logs.Println("Handling notify:  Sending notification data to redisSender ", clientId)
		redisSender <- notifyData
		//Check whether clientId is connected.
		currentConnections := connMgr.connectionsFor(clientId)

		notifyResponse := NotifyResponse{
			Status:   (len(currentConnections) > 0),
			ClientId: clientId,
		}
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		err = json.NewEncoder(w).Encode(notifyResponse)
		if err != nil {
			log.Println("Error writing the response.", err)
			w.WriteHeader(http.StatusInternalServerError)
		}

	}

}

// serverWs handles websocket requests from the peer.
func ServeWs(w http.ResponseWriter, r *http.Request) {

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error in accepting websocket connection", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	//new websocket connection
	vars := mux.Vars(r)
	clientID := vars["clientid"]
	logs.Println("Creating new wsconnection for client : " + clientID)
	conn := connMgr.newConnectionFor(clientID, ws)
	if conn.active {
		go conn.sendMessages()
		conn.receiveMessages()
	} else {
		errMsg := "Failed to subscribe to message channel for " + clientID
		ws.WriteMessage(websocket.CloseInternalServerErr, []byte(errMsg))
		ws.Close()
	}

}

func newUUID() string {

	u1 := uuid.NewV4()
	return u1.String()
}
<|endoftext|>


<commit_before>package argot

import (
	"testing"
)

// Step represents a step in a test.
type Step interface {
	Go() error
}

// If we can have a single Step, then we can have a slice of Steps
// representing the order of Steps in a larger unit.
type Steps []Step

// Test runs the steps in order and returns either all the steps, or
// all the steps that did not error, plus the step that errored. Thus
// the results are always a prefix of the Steps.  t can be nil. If t
// is not nil and an error occurs, then t.Fatal will be called. If an
// error occurs, it will be returned.
func (ss Steps) Test(t *testing.T) (results Steps, err error) {
	if t != nil {
		defer func() {
			if err != nil {
				t.Fatalf("Achieved Steps: %v; Error: %v", results, err)
			}
		}()
	}
	for idx, step := range ss {
		err = step.Go()
		if err != nil {
			return ss[:idx], err
		}
	}
	return ss, nil
}

// Steps are useful, but there are times where you don't know ahead of
// time exactly which Steps you wish to run. Consider a test where the
// steps are dependent on the response you receive from a server.
// Thus the advantage of using a chan is that it allows more laziness:
// steps can be responsible for issuing their own subsequent steps.
type StepsChan <-chan Step

// Test runs the steps in order and returns either all the steps, or
// all the steps that did not error, plus the step that errored. Thus
// the results are always a prefix of the Steps.  t can be nil. If t
// is not nil and an error occurs, then t.Fatal will be called. If an
// error occurs, it will be returned.
func (sc StepsChan) Test(t *testing.T) (results Steps, err error) {
	if t != nil {
		defer func() {
			if err != nil {
				t.Fatalf("Achieved Steps: %v; Error: %v", results, err)
			}
		}()
	}
	results = make([]Step, 0, 16)
	for step := range sc {
		results = append(results, step)
		err = step.Go()
		if err != nil {
			return results, err
		}
	}
	return results, nil
}

// StepFunc is the basic type of a Step: a function that takes no
// arguments and returns an error.
type StepFunc func() error

func (sf StepFunc) Go() error {
	return sf()
}

// NamedStep extends StepFunc by adding a name, which is mainly of use
// when formatting a Step.
type NamedStep struct {
	StepFunc
	name string
}

func (ns NamedStep) String() string {
	return ns.name
}

// NewNamedStep creates a NamedStep with the given name and Step
// function.
func NewNamedStep(name string, step StepFunc) *NamedStep {
	return &NamedStep{
		StepFunc: step,
		name:     name,
	}
}

// AnyError is a utility function that returns the first non-nil error
// in the slice, or nil if either the slice or all elements of the
// slice are nil.
func AnyError(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Print results on separate lines.<commit_after>package argot

import (
	"testing"
)

// Step represents a step in a test.
type Step interface {
	Go() error
}

// If we can have a single Step, then we can have a slice of Steps
// representing the order of Steps in a larger unit.
type Steps []Step

// Test runs the steps in order and returns either all the steps, or
// all the steps that did not error, plus the step that errored. Thus
// the results are always a prefix of the Steps.  t can be nil. If t
// is not nil and an error occurs, then t.Fatal will be called. If an
// error occurs, it will be returned.
func (ss Steps) Test(t *testing.T) (results Steps, err error) {
	if t != nil {
		defer func() {
			if err != nil {
				t.Log("Achieved Steps:")
				for _, result := range results {
					t.Logf("  %v", result)
				}
				t.Fatalf("Error: %v", err)
			}
		}()
	}
	for idx, step := range ss {
		err = step.Go()
		if err != nil {
			return ss[:idx], err
		}
	}
	return ss, nil
}

// Steps are useful, but there are times where you don't know ahead of
// time exactly which Steps you wish to run. Consider a test where the
// steps are dependent on the response you receive from a server.
// Thus the advantage of using a chan is that it allows more laziness:
// steps can be responsible for issuing their own subsequent steps.
type StepsChan <-chan Step

// Test runs the steps in order and returns either all the steps, or
// all the steps that did not error, plus the step that errored. Thus
// the results are always a prefix of the Steps.  t can be nil. If t
// is not nil and an error occurs, then t.Fatal will be called. If an
// error occurs, it will be returned.
func (sc StepsChan) Test(t *testing.T) (results Steps, err error) {
	if t != nil {
		defer func() {
			if err != nil {
				t.Log("Achieved Steps:")
				for _, result := range results {
					t.Logf("  %v", result)
				}
				t.Fatalf("Error: %v", err)
			}
		}()
	}
	results = make([]Step, 0, 16)
	for step := range sc {
		results = append(results, step)
		err = step.Go()
		if err != nil {
			return results, err
		}
	}
	return results, nil
}

// StepFunc is the basic type of a Step: a function that takes no
// arguments and returns an error.
type StepFunc func() error

func (sf StepFunc) Go() error {
	return sf()
}

// NamedStep extends StepFunc by adding a name, which is mainly of use
// when formatting a Step.
type NamedStep struct {
	StepFunc
	name string
}

func (ns NamedStep) String() string {
	return ns.name
}

// NewNamedStep creates a NamedStep with the given name and Step
// function.
func NewNamedStep(name string, step StepFunc) *NamedStep {
	return &NamedStep{
		StepFunc: step,
		name:     name,
	}
}

// AnyError is a utility function that returns the first non-nil error
// in the slice, or nil if either the slice or all elements of the
// slice are nil.
func AnyError(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package garbage4

import (
	"os"
	"path"
	"sync"
	"syscall"
	"unsafe"
)

const (
	POSITION_SIZE  = 16
	MAX_STATE_SIZE = 32768
)

var (
	MAX_CHANNEL_NAME_SIZE = 32
)

type State struct {
	sync.RWMutex
	ref      []byte
	file     *os.File
	data     *[MAX_STATE_SIZE]byte
	channels map[string]int
}

type Position struct {
	offset    uint32
	segmentId uint64
}

func loadState(t *Topic) (*State, error) {
	root := t.path
	if err := os.MkdirAll(root, 0700); err != nil {
		return nil, err
	}
	path := path.Join(root, "state.q")
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}

	info, err := file.Stat()
	if err != nil {
		return nil, err
	}

	if size := info.Size(); size == 0 {
		file.Truncate(MAX_STATE_SIZE)
	} else if size != MAX_STATE_SIZE {
		panic("invalid state file size")
	}

	ref, err := syscall.Mmap(int(file.Fd()), 0, MAX_STATE_SIZE, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_SHARED)
	if err != nil {
		file.Close()
		return nil, err
	}

	state := &State{
		ref:      ref,
		file:     file,
		channels: make(map[string]int),
		data:     (*[MAX_STATE_SIZE]byte)(unsafe.Pointer(&ref[0])),
	}

	offset := POSITION_SIZE
	recordSize := offset + MAX_CHANNEL_NAME_SIZE
	for ; (offset+recordSize) < MAX_STATE_SIZE && state.data[offset] != 0; offset += recordSize {
		//deleted channel
		if state.data[offset] == 255 {
			continue
		}
		end := offset
		for ; (end-offset) < MAX_CHANNEL_NAME_SIZE && state.data[end] != 0; end++ {
		}
		println(offset, end, string(state.data[offset:end]))
		state.channels[string(state.data[offset:end])] = offset + MAX_CHANNEL_NAME_SIZE
	}
	return state, nil
}

func (s *State) loadPosition(offset int) *Position {
	return (*Position)(unsafe.Pointer(&s.data[offset]))
}

func (s *State) loadOrCreatePosition(name string) *Position {
	offset, exists := s.channels[name]
	if exists == false {
		//todo check for overflow
		offset = POSITION_SIZE
		for ; offset < MAX_CHANNEL_NAME_SIZE; offset += MAX_CHANNEL_NAME_SIZE + POSITION_SIZE {
			b := s.data[offset]
			if b == 0 || b == 255 {
				break
			}
		}
		copy(s.data[offset:], name)
		offset += MAX_CHANNEL_NAME_SIZE
		s.channels[name] = offset
	}
	println(offset)
	return s.loadPosition(offset)
}

func (s *State) usable(segmentId uint64) bool {
	for _, offset := range s.channels {
		if s.loadPosition(offset).segmentId <= segmentId {
			return true
		}
	}
	return false
}
<commit_msg>small refactoring + remove debug printlns<commit_after>package garbage4

import (
	"os"
	"path"
	"sync"
	"syscall"
	"unsafe"
)

const (
	POSITION_SIZE  = 16
	MAX_STATE_SIZE = 32768
)

var (
	MAX_CHANNEL_NAME_SIZE = 32
)

type State struct {
	sync.RWMutex
	ref      []byte
	file     *os.File
	data     *[MAX_STATE_SIZE]byte
	channels map[string]int
}

type Position struct {
	offset    uint32
	segmentId uint64
}

func loadState(t *Topic) (*State, error) {
	if err := os.MkdirAll(t.path, 0700); err != nil {
		return nil, err
	}
	path := path.Join(t.path, "state.q")
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}

	info, err := file.Stat()
	if err != nil {
		return nil, err
	}

	if size := info.Size(); size == 0 {
		file.Truncate(MAX_STATE_SIZE)
	} else if size != MAX_STATE_SIZE {
		panic("invalid state file size")
	}

	ref, err := syscall.Mmap(int(file.Fd()), 0, MAX_STATE_SIZE, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_SHARED)
	if err != nil {
		file.Close()
		return nil, err
	}

	state := &State{
		ref:      ref,
		file:     file,
		channels: make(map[string]int),
		data:     (*[MAX_STATE_SIZE]byte)(unsafe.Pointer(&ref[0])),
	}

	offset := POSITION_SIZE
	recordSize := offset + MAX_CHANNEL_NAME_SIZE
	for ; (offset+recordSize) < MAX_STATE_SIZE && state.data[offset] != 0; offset += recordSize {
		//deleted channel
		if state.data[offset] == 255 {
			continue
		}
		end := offset
		for ; (end-offset) < MAX_CHANNEL_NAME_SIZE && state.data[end] != 0; end++ {
		}
		state.channels[string(state.data[offset:end])] = offset + MAX_CHANNEL_NAME_SIZE
	}
	return state, nil
}

func (s *State) loadPosition(offset int) *Position {
	return (*Position)(unsafe.Pointer(&s.data[offset]))
}

func (s *State) loadOrCreatePosition(name string) *Position {
	offset, exists := s.channels[name]
	if exists == false {
		//todo check for overflow
		offset = POSITION_SIZE
		for ; offset < MAX_CHANNEL_NAME_SIZE; offset += MAX_CHANNEL_NAME_SIZE + POSITION_SIZE {
			b := s.data[offset]
			if b == 0 || b == 255 {
				break
			}
		}
		copy(s.data[offset:], name)
		offset += MAX_CHANNEL_NAME_SIZE
		s.channels[name] = offset
	}
	return s.loadPosition(offset)
}

func (s *State) usable(segmentId uint64) bool {
	for _, offset := range s.channels {
		if s.loadPosition(offset).segmentId <= segmentId {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"sort"
	"strings"
	"text/template"

	"github.com/google/go-github/v37/github"
	"github.com/jekyll/dashboard"
)

var tmpl = template.Must(template.New(`configuration.go`).Parse(`//go:generate go run ./cmd/generate-project-list
// THIS FILE IS AUTO-GENERATED WITH 'go generate .'
package dashboard

var defaultProjects = []*Project{{"{"}}{{range .Repositories}}
	{
		Name:          {{.Name | printf "%q"}},
		Nwo:           {{.Nwo | printf "%q"}},
		Branch:        {{.Branch | printf "%q"}},
		GemName:       {{.GemName | printf "%q"}},
		GlobalRelayID: {{.GlobalRelayID | printf "%q" }},
	},{{end}}
}
`))

type templateData struct {
	Repositories []*dashboard.Project
}

var gemNames = map[string]string{
	"github-metadata": "jekyll-github-metadata",
}

var additionalProjectNames = map[string]bool{
	"classifier-reborn": true,
	"directory":         true,
	"github-metadata":   true,
}

func relevantProject(name string) bool {
	return strings.HasPrefix(name, "jekyll") || additionalProjectNames[name]
}

func projectGemName(name string) string {
	if gemName, exist := gemNames[name]; exist {
		return gemName
	} else {
		return name
	}
}

func main() {
	flag.Parse()

	client := github.NewClient(http.DefaultClient)
	repositories, _, err := client.Repositories.List(context.Background(), "jekyll", &github.RepositoryListOptions{})
	if err != nil {
		log.Fatalf("unable to list repositories: %v", err)
	}

	repoInfos := make([]*dashboard.Project, 0, len(repositories))

	for _, repository := range repositories {
		if !relevantProject(repository.GetName()) || repository.GetArchived() {
			continue
		}
		info := &dashboard.Project{
			Name:          repository.GetName(),
			Nwo:           repository.GetFullName(),
			Branch:        repository.GetDefaultBranch(),
			GemName:       projectGemName(repository.GetName()),
			GlobalRelayID: repository.GetNodeID(),
			Stars:         repository.GetStargazersCount(),
		}
		repoInfos = append(repoInfos, info)
		log.Printf("repo: %#v", info)
	}

	sort.Slice(repoInfos, func(i, j int) bool {
		return repoInfos[i].Stars > repoInfos[j].Stars
	})

	s := &strings.Builder{}
	err = tmpl.Execute(s, templateData{Repositories: repoInfos})
	if err != nil {
		log.Fatalf("unable to execute template: %v", err)
	}

	path := filepath.Join("configuration.go")
	err = ioutil.WriteFile(path, []byte(s.String()), 0644)
	if err != nil {
		log.Fatalf("unable to write %s: %v", path, err)
	}
}
<commit_msg>Replace use of a map with a switch expression<commit_after>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"sort"
	"strings"
	"text/template"

	"github.com/google/go-github/v37/github"
	"github.com/jekyll/dashboard"
)

var tmpl = template.Must(template.New(`configuration.go`).Parse(`//go:generate go run ./cmd/generate-project-list
// THIS FILE IS AUTO-GENERATED WITH 'go generate .'
package dashboard

var defaultProjects = []*Project{{"{"}}{{range .Repositories}}
	{
		Name:          {{.Name | printf "%q"}},
		Nwo:           {{.Nwo | printf "%q"}},
		Branch:        {{.Branch | printf "%q"}},
		GemName:       {{.GemName | printf "%q"}},
		GlobalRelayID: {{.GlobalRelayID | printf "%q" }},
	},{{end}}
}
`))

type templateData struct {
	Repositories []*dashboard.Project
}

var additionalProjectNames = map[string]bool{
	"classifier-reborn": true,
	"directory":         true,
	"github-metadata":   true,
}

func relevantProject(name string) bool {
	return strings.HasPrefix(name, "jekyll") || additionalProjectNames[name]
}

func projectGemName(name string) string {
	switch name {
	case "github-metadata":
		return "jekyll-github-metadata"
	default:
		return name
	}
}

func main() {
	flag.Parse()

	client := github.NewClient(http.DefaultClient)
	repositories, _, err := client.Repositories.List(context.Background(), "jekyll", &github.RepositoryListOptions{})
	if err != nil {
		log.Fatalf("unable to list repositories: %v", err)
	}

	repoInfos := make([]*dashboard.Project, 0, len(repositories))

	for _, repository := range repositories {
		if !relevantProject(repository.GetName()) || repository.GetArchived() {
			continue
		}
		info := &dashboard.Project{
			Name:          repository.GetName(),
			Nwo:           repository.GetFullName(),
			Branch:        repository.GetDefaultBranch(),
			GemName:       projectGemName(repository.GetName()),
			GlobalRelayID: repository.GetNodeID(),
			Stars:         repository.GetStargazersCount(),
		}
		repoInfos = append(repoInfos, info)
		log.Printf("repo: %#v", info)
	}

	sort.Slice(repoInfos, func(i, j int) bool {
		return repoInfos[i].Stars > repoInfos[j].Stars
	})

	s := &strings.Builder{}
	err = tmpl.Execute(s, templateData{Repositories: repoInfos})
	if err != nil {
		log.Fatalf("unable to execute template: %v", err)
	}

	path := filepath.Join("configuration.go")
	err = ioutil.WriteFile(path, []byte(s.String()), 0644)
	if err != nil {
		log.Fatalf("unable to write %s: %v", path, err)
	}
}
<|endoftext|>