<commit_before>package main

import (
	"net/http"

	"github.com/sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"github.com/satori/go.uuid"
	"github.com/tokubai/kinu/logger"
	"github.com/tokubai/kinu/resource"
)

const SANDBOX_IMAGE_TYPE = "__sandbox__"

func UploadImageToSandboxHandler(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	w.Header().Set("Content-Type", "application/json")

	r.ParseMultipartForm(0)

	uuidObject, err := uuid.NewV4()
	if err != nil {
		RespondInternalServerError(w, err)
		return
	}
	imageId := uuidObject.String()

	file, _, err := r.FormFile("image")
	if err != nil {
		RespondBadRequest(w, "invalid file")
		return
	}

	err = resource.New(SANDBOX_IMAGE_TYPE, imageId).Store(file)
	if err != nil {
		if _, ok := err.(*ErrInvalidRequest); ok {
			RespondBadRequest(w, err.Error())
		} else {
			RespondInternalServerError(w, err)
		}
	}

	RespondImageUploadSuccessJson(w, SANDBOX_IMAGE_TYPE, imageId)

	logger.WithFields(logrus.Fields{
		"path":   r.URL.Path,
		"params": r.URL.Query(),
		"method": r.Method,
	}).Info("success")
}

func ApplyFromSandboxHandler(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	w.Header().Set("Content-Type", "application/json")

	r.ParseMultipartForm(0)

	sandboxId := r.FormValue("sandbox_id")
	if len(sandboxId) == 0 {
		RespondBadRequest(w, "required sandbox_id")
		return
	}

	imageType := r.FormValue("name")
	if len(imageType) == 0 {
		RespondBadRequest(w, "required name parameter")
		return
	}

	imageId := r.FormValue("id")
	if len(imageId) == 0 {
		RespondBadRequest(w, "required id parameter")
		return
	}

	err := resource.New(SANDBOX_IMAGE_TYPE, sandboxId).MoveTo(imageType, imageId)

	if err != nil {
		RespondInternalServerError(w, err)
		return
	}

	RespondImageUploadSuccessJson(w, imageType, imageId)

	logger.WithFields(logrus.Fields{
		"path":   r.URL.Path,
		"params": r.URL.Query(),
		"method": r.Method,
	}).Info("success")
}
<commit_msg>Simplify variable name<commit_after>package main

import (
	"net/http"

	"github.com/sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"github.com/satori/go.uuid"
	"github.com/tokubai/kinu/logger"
	"github.com/tokubai/kinu/resource"
)

const SANDBOX_IMAGE_TYPE = "__sandbox__"

func UploadImageToSandboxHandler(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	w.Header().Set("Content-Type", "application/json")

	r.ParseMultipartForm(0)

	uid, err := uuid.NewV4()
	if err != nil {
		RespondInternalServerError(w, err)
		return
	}
	imageId := uid.String()

	file, _, err := r.FormFile("image")
	if err != nil {
		RespondBadRequest(w, "invalid file")
		return
	}

	err = resource.New(SANDBOX_IMAGE_TYPE, imageId).Store(file)
	if err != nil {
		if _, ok := err.(*ErrInvalidRequest); ok {
			RespondBadRequest(w, err.Error())
		} else {
			RespondInternalServerError(w, err)
		}
	}

	RespondImageUploadSuccessJson(w, SANDBOX_IMAGE_TYPE, imageId)

	logger.WithFields(logrus.Fields{
		"path":   r.URL.Path,
		"params": r.URL.Query(),
		"method": r.Method,
	}).Info("success")
}

func ApplyFromSandboxHandler(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	w.Header().Set("Content-Type", "application/json")

	r.ParseMultipartForm(0)

	sandboxId := r.FormValue("sandbox_id")
	if len(sandboxId) == 0 {
		RespondBadRequest(w, "required sandbox_id")
		return
	}

	imageType := r.FormValue("name")
	if len(imageType) == 0 {
		RespondBadRequest(w, "required name parameter")
		return
	}

	imageId := r.FormValue("id")
	if len(imageId) == 0 {
		RespondBadRequest(w, "required id parameter")
		return
	}

	err := resource.New(SANDBOX_IMAGE_TYPE, sandboxId).MoveTo(imageType, imageId)

	if err != nil {
		RespondInternalServerError(w, err)
		return
	}

	RespondImageUploadSuccessJson(w, imageType, imageId)

	logger.WithFields(logrus.Fields{
		"path":   r.URL.Path,
		"params": r.URL.Query(),
		"method": r.Method,
	}).Info("success")
}
<|endoftext|>


<commit_before>package tardeploy

import (
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

func (configuration *Configuration) beforeRunTrigger(application string) error {
	return configuration.trigger(application, "before")
}

func (configuration *Configuration) afterRunTrigger(application string) error {
	return configuration.trigger(application, "after")
}

func (configuration *Configuration) trigger(application, status string) error {
	log.Infof("Trigger for %s: %s", application, status)
	switch status {
	case "before":
		if "" == configuration.Trigger.Before {
			return nil
		}
		break
	case "after":
		if "" == configuration.Trigger.After {
			return nil
		}
		break
	}
	return errors.New("Trigger not yet implemented")
}
<commit_msg>Running trigger<commit_after>package tardeploy

import (
	"os"
	"os/exec"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

func (configuration *Configuration) beforeRunTrigger(application string) error {
	return configuration.trigger(application, "before")
}

func (configuration *Configuration) afterRunTrigger(application string) error {
	return configuration.trigger(application, "after")
}

func (configuration *Configuration) trigger(application, status string) error {
	var cmd string
	switch status {
	case "before":
		if "" == configuration.Trigger.Before {
			return nil
		}
		cmd = configuration.Trigger.Before
		break
	case "after":
		if "" == configuration.Trigger.After {
			return nil
		}
		cmd = configuration.Trigger.After
		break
	}

	log.Infof("Trigger for %s: %s", application, status)

	command := exec.Command(cmd, application, status)
	command.Stdout = os.Stdout
	command.Stdin = os.Stdin
	command.Stderr = os.Stderr
	var err error
	if err = command.Start(); err != nil {
		return errors.Wrap(err, "could not start command")
	}
	err = command.Wait()
	if err != nil {
		return errors.Wrap(err, "Could not wait for command")
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package flow

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"

	"github.com/redhat-cip/skydive/logging"
)

func LayerFlow(l gopacket.Layer) gopacket.Flow {
	switch l.(type) {
	case gopacket.LinkLayer:
		return l.(gopacket.LinkLayer).LinkFlow()
	case gopacket.NetworkLayer:
		return l.(gopacket.NetworkLayer).NetworkFlow()
	case gopacket.TransportLayer:
		return l.(gopacket.TransportLayer).TransportFlow()
	}
	logging.GetLogger().Critical(fmt.Sprint("Unknown gopacket.Layer %T", l))
	return gopacket.Flow{}
}

type FlowKey struct {
	net, transport uint64
}

func (key FlowKey) fillFromGoPacket(p *gopacket.Packet) FlowKey {
	key.net = LayerFlow((*p).NetworkLayer()).FastHash()
	key.transport = LayerFlow((*p).TransportLayer()).FastHash()
	return key
}

func (key FlowKey) String() string {
	return fmt.Sprintf("%x-%x", key.net, key.transport)
}

func (flow *Flow) fillFromGoPacket(packet *gopacket.Packet) error {
	/* Continue if no ethernet layer */
	ethernetLayer := (*packet).Layer(layers.LayerTypeEthernet)
	_, ok := ethernetLayer.(*layers.Ethernet)
	if !ok {
		return errors.New("Unable to decode the ethernet layer")
	}

	newFlow := false
	fs := flow.GetStatistics()
	now := time.Now().Unix() //(*packet).Metadata().Timestamp.Unix()
	if fs == nil {
		newFlow = true
		fs = NewFlowStatistics()
		fs.Start = now
		fs.newEthernetEndpointStatistics(packet)
		fs.newIPV4EndpointStatistics(packet)
		fs.newTransportEndpointStatistics(packet)
		flow.Statistics = fs
	}
	fs.Last = now
	fs.updateEthernetFromGoPacket(packet)
	fs.updateIPV4FromGoPacket(packet)
	fs.updateTransportFromGoPacket(packet)

	if newFlow {
		hasher := sha1.New()
		path := ""
		for i, layer := range (*packet).Layers() {
			if i > 0 {
				path += "/"
			}
			path += layer.LayerType().String()
		}
		flow.LayersPath = path
		hasher.Write([]byte(flow.LayersPath))

		/* Generate an flow UUID */
		for _, ep := range fs.GetEndpoints() {
			hasher.Write([]byte(ep.AB.Value))
			hasher.Write([]byte(ep.BA.Value))
		}
		flow.UUID = hex.EncodeToString(hasher.Sum(nil))
	}
	return nil
}

func FromData(data []byte) (*Flow, error) {
	flow := new(Flow)

	err := proto.Unmarshal(data, flow)
	if err != nil {
		return nil, err
	}

	return flow, nil
}

func (flow *Flow) GetData() ([]byte, error) {
	data, err := proto.Marshal(flow)
	if err != nil {
		return []byte{}, err
	}

	return data, nil
}

func FLowsFromSFlowSample(ft *FlowTable, sample *layers.SFlowFlowSample, probePath *string) []*Flow {
	flows := []*Flow{}

	for _, rec := range sample.Records {

		/* FIX(safchain): just keeping the raw packet for now */
		record, ok := rec.(layers.SFlowRawPacketFlowRecord)
		if !ok {
			logging.GetLogger().Critical("1st layer is not SFlowRawPacketFlowRecord type")
			continue
		}

		packet := &record.Header
		key := (FlowKey{}).fillFromGoPacket(packet)
		flow, new := ft.GetFlow(key.String(), packet)
		if new {
			flow.ProbeGraphPath = *probePath
		}
		flow.fillFromGoPacket(packet)
		flows = append(flows, flow)
	}

	return flows
}
<commit_msg>fixup null probePath<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package flow

import (
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"

	"github.com/redhat-cip/skydive/logging"
)

func LayerFlow(l gopacket.Layer) gopacket.Flow {
	switch l.(type) {
	case gopacket.LinkLayer:
		return l.(gopacket.LinkLayer).LinkFlow()
	case gopacket.NetworkLayer:
		return l.(gopacket.NetworkLayer).NetworkFlow()
	case gopacket.TransportLayer:
		return l.(gopacket.TransportLayer).TransportFlow()
	}
	logging.GetLogger().Critical(fmt.Sprint("Unknown gopacket.Layer %T", l))
	return gopacket.Flow{}
}

type FlowKey struct {
	net, transport uint64
}

func (key FlowKey) fillFromGoPacket(p *gopacket.Packet) FlowKey {
	key.net = LayerFlow((*p).NetworkLayer()).FastHash()
	key.transport = LayerFlow((*p).TransportLayer()).FastHash()
	return key
}

func (key FlowKey) String() string {
	return fmt.Sprintf("%x-%x", key.net, key.transport)
}

func (flow *Flow) fillFromGoPacket(packet *gopacket.Packet) error {
	/* Continue if no ethernet layer */
	ethernetLayer := (*packet).Layer(layers.LayerTypeEthernet)
	_, ok := ethernetLayer.(*layers.Ethernet)
	if !ok {
		return errors.New("Unable to decode the ethernet layer")
	}

	newFlow := false
	fs := flow.GetStatistics()
	now := time.Now().Unix() //(*packet).Metadata().Timestamp.Unix()
	if fs == nil {
		newFlow = true
		fs = NewFlowStatistics()
		fs.Start = now
		fs.newEthernetEndpointStatistics(packet)
		fs.newIPV4EndpointStatistics(packet)
		fs.newTransportEndpointStatistics(packet)
		flow.Statistics = fs
	}
	fs.Last = now
	fs.updateEthernetFromGoPacket(packet)
	fs.updateIPV4FromGoPacket(packet)
	fs.updateTransportFromGoPacket(packet)

	if newFlow {
		hasher := sha1.New()
		path := ""
		for i, layer := range (*packet).Layers() {
			if i > 0 {
				path += "/"
			}
			path += layer.LayerType().String()
		}
		flow.LayersPath = path
		hasher.Write([]byte(flow.LayersPath))

		/* Generate an flow UUID */
		for _, ep := range fs.GetEndpoints() {
			hasher.Write([]byte(ep.AB.Value))
			hasher.Write([]byte(ep.BA.Value))
		}
		flow.UUID = hex.EncodeToString(hasher.Sum(nil))
	}
	return nil
}

func FromData(data []byte) (*Flow, error) {
	flow := new(Flow)

	err := proto.Unmarshal(data, flow)
	if err != nil {
		return nil, err
	}

	return flow, nil
}

func (flow *Flow) GetData() ([]byte, error) {
	data, err := proto.Marshal(flow)
	if err != nil {
		return []byte{}, err
	}

	return data, nil
}

func FLowsFromSFlowSample(ft *FlowTable, sample *layers.SFlowFlowSample, probePath *string) []*Flow {
	flows := []*Flow{}

	for _, rec := range sample.Records {

		/* FIX(safchain): just keeping the raw packet for now */
		record, ok := rec.(layers.SFlowRawPacketFlowRecord)
		if !ok {
			logging.GetLogger().Critical("1st layer is not SFlowRawPacketFlowRecord type")
			continue
		}

		packet := &record.Header
		key := (FlowKey{}).fillFromGoPacket(packet)
		flow, new := ft.GetFlow(key.String(), packet)
		if new {
			if probePath == nil {
				flow.ProbeGraphPath = "probe-topology-path-not-found"
			} else {
				flow.ProbeGraphPath = *probePath
			}
		}
		flow.fillFromGoPacket(packet)
		flows = append(flows, flow)
	}

	return flows
}
<|endoftext|>


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

// Package dns should not be used directly by google3; instead, use net/dns/go/dns.
package dns

import (
	netpkg "net"
	"strings"
)

func parseDialNetwork(net string) (afnet string, proto int, err error, handled bool) {
	i := strings.LastIndex(net, ":")
	if i < 0 { // no colon
		switch net {
		case "tcp", "tcp4", "tcp6":
		case "udp", "udp4", "udp6":
		default:
			return
		}
		return net, 0, nil, true
	}
	return
}

func resolveNetAddr(op, net, addr string) (afnet string, a netpkg.Addr, err error, handled bool) {
	afnet, _, err, ok := parseDialNetwork(net)
	if !ok {
		return
	}
	if err != nil {
		handled = true
		return
	}
	if op == "dial" && addr == "" {
		return
	}
	switch afnet {
	default:
		return
	case "tcp", "tcp4", "tcp6":
		if addr != "" {
			a, err = ResolveTCPAddr(afnet, addr)
		}
	case "udp", "udp4", "udp6":
		if addr != "" {
			a, err = ResolveUDPAddr(afnet, addr)
		}
	case "ip", "ip4", "ip6":
		if addr != "" {
			a, err = ResolveIPAddr(afnet, addr)
		}
	}
	handled = true
	return
}

// Dial is like net.Dial but handles DNS resolution itself instead of
// in package net.
func Dial(net, addr string) (netpkg.Conn, error) {
	_, addri, err, handled := resolveNetAddr("dial", net, addr)
	if !handled {
		return netpkg.Dial(net, addr)
	}
	if err != nil {
		return nil, err
	}
	return dialAddr(net, addr, addri)
}

func dialAddr(net, addr string, addri netpkg.Addr) (c netpkg.Conn, err error) {
	switch ra := addri.(type) {
	case *netpkg.TCPAddr:
		c, err = netpkg.DialTCP(net, nil, ra)
	case *netpkg.UDPAddr:
		c, err = netpkg.DialUDP(net, nil, ra)
	case *netpkg.IPAddr:
		c, err = netpkg.DialIP(net, nil, ra)
	default:
		err = &netpkg.OpError{"dial", net + " " + addr, nil, netpkg.UnknownNetworkError(net)}
	}
	if err != nil {
		return nil, err
	}
	return
}
<commit_msg>Get rid of reliance on the net package's OpError and error types.  They change in go1.5.<commit_after>// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package dns should not be used directly by google3; instead, use net/dns/go/dns.
package dns

import (
	"errors"
	netpkg "net"
	"strings"
)

func parseDialNetwork(net string) (afnet string, proto int, err error, handled bool) {
	i := strings.LastIndex(net, ":")
	if i < 0 { // no colon
		switch net {
		case "tcp", "tcp4", "tcp6":
		case "udp", "udp4", "udp6":
		default:
			return
		}
		return net, 0, nil, true
	}
	return
}

func resolveNetAddr(op, net, addr string) (afnet string, a netpkg.Addr, err error, handled bool) {
	afnet, _, err, ok := parseDialNetwork(net)
	if !ok {
		return
	}
	if err != nil {
		handled = true
		return
	}
	if op == "dial" && addr == "" {
		return
	}
	switch afnet {
	default:
		return
	case "tcp", "tcp4", "tcp6":
		if addr != "" {
			a, err = ResolveTCPAddr(afnet, addr)
		}
	case "udp", "udp4", "udp6":
		if addr != "" {
			a, err = ResolveUDPAddr(afnet, addr)
		}
	case "ip", "ip4", "ip6":
		if addr != "" {
			a, err = ResolveIPAddr(afnet, addr)
		}
	}
	handled = true
	return
}

// Dial is like net.Dial but handles DNS resolution itself instead of
// in package net.
func Dial(net, addr string) (netpkg.Conn, error) {
	_, addri, err, handled := resolveNetAddr("dial", net, addr)
	if !handled {
		return netpkg.Dial(net, addr)
	}
	if err != nil {
		return nil, err
	}
	return dialAddr(net, addr, addri)
}

func dialAddr(net, addr string, addri netpkg.Addr) (c netpkg.Conn, err error) {
	switch ra := addri.(type) {
	case *netpkg.TCPAddr:
		c, err = netpkg.DialTCP(net, nil, ra)
	case *netpkg.UDPAddr:
		c, err = netpkg.DialUDP(net, nil, ra)
	case *netpkg.IPAddr:
		c, err = netpkg.DialIP(net, nil, ra)
	default:
		err = errors.New("Unknown network address type")
	}
	if err != nil {
		return nil, err
	}
	return
}
<|endoftext|>


<commit_before>// +build !providerless

/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package azure

import (
	"context"
	"fmt"
	"path"
	"strconv"
	"strings"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-07-01/compute"
	"github.com/Azure/go-autorest/autorest/to"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	kwait "k8s.io/apimachinery/pkg/util/wait"
	cloudvolume "k8s.io/cloud-provider/volume"
	volumehelpers "k8s.io/cloud-provider/volume/helpers"
	"k8s.io/klog"
)

const (
	// default IOPS Caps & Throughput Cap (MBps) per https://docs.microsoft.com/en-us/azure/virtual-machines/linux/disks-ultra-ssd
	defaultDiskIOPSReadWrite = 500
	defaultDiskMBpsReadWrite = 100

	diskEncryptionSetIDFormat = "/subscriptions/{subs-id}/resourceGroups/{rg-name}/providers/Microsoft.Compute/diskEncryptionSets/{diskEncryptionSet-name}"
)

//ManagedDiskController : managed disk controller struct
type ManagedDiskController struct {
	common *controllerCommon
}

// ManagedDiskOptions specifies the options of managed disks.
type ManagedDiskOptions struct {
	// The name of the disk.
	DiskName string
	// The size in GB.
	SizeGB int
	// The name of PVC.
	PVCName string
	// The name of resource group.
	ResourceGroup string
	// The AvailabilityZone to create the disk.
	AvailabilityZone string
	// The tags of the disk.
	Tags map[string]string
	// The SKU of storage account.
	StorageAccountType compute.DiskStorageAccountTypes
	// IOPS Caps for UltraSSD disk
	DiskIOPSReadWrite string
	// Throughput Cap (MBps) for UltraSSD disk
	DiskMBpsReadWrite string
	// if SourceResourceID is not empty, then it's a disk copy operation(for snapshot)
	SourceResourceID string
	// The type of source
	SourceType string
	// ResourceId of the disk encryption set to use for enabling encryption at rest.
	DiskEncryptionSetID string
}

//CreateManagedDisk : create managed disk
func (c *ManagedDiskController) CreateManagedDisk(options *ManagedDiskOptions) (string, error) {
	var err error
	klog.V(4).Infof("azureDisk - creating new managed Name:%s StorageAccountType:%s Size:%v", options.DiskName, options.StorageAccountType, options.SizeGB)

	var createZones *[]string
	if len(options.AvailabilityZone) > 0 {
		zoneList := []string{c.common.cloud.GetZoneID(options.AvailabilityZone)}
		createZones = &zoneList
	}

	// insert original tags to newTags
	newTags := make(map[string]*string)
	azureDDTag := "kubernetes-azure-dd"
	newTags["created-by"] = &azureDDTag
	if options.Tags != nil {
		for k, v := range options.Tags {
			// Azure won't allow / (forward slash) in tags
			newKey := strings.Replace(k, "/", "-", -1)
			newValue := strings.Replace(v, "/", "-", -1)
			newTags[newKey] = &newValue
		}
	}

	diskSizeGB := int32(options.SizeGB)
	diskSku := compute.DiskStorageAccountTypes(options.StorageAccountType)

	creationData, err := getValidCreationData(c.common.subscriptionID, options.ResourceGroup, options.SourceResourceID, options.SourceType)
	if err != nil {
		return "", err
	}
	diskProperties := compute.DiskProperties{
		DiskSizeGB:   &diskSizeGB,
		CreationData: &creationData,
	}

	if diskSku == compute.UltraSSDLRS {
		diskIOPSReadWrite := int64(defaultDiskIOPSReadWrite)
		if options.DiskIOPSReadWrite != "" {
			v, err := strconv.Atoi(options.DiskIOPSReadWrite)
			if err != nil {
				return "", fmt.Errorf("AzureDisk - failed to parse DiskIOPSReadWrite: %v", err)
			}
			diskIOPSReadWrite = int64(v)
		}
		diskProperties.DiskIOPSReadWrite = to.Int64Ptr(diskIOPSReadWrite)

		diskMBpsReadWrite := int32(defaultDiskMBpsReadWrite)
		if options.DiskMBpsReadWrite != "" {
			v, err := strconv.Atoi(options.DiskMBpsReadWrite)
			if err != nil {
				return "", fmt.Errorf("AzureDisk - failed to parse DiskMBpsReadWrite: %v", err)
			}
			diskMBpsReadWrite = int32(v)
		}
		diskProperties.DiskMBpsReadWrite = to.Int32Ptr(diskMBpsReadWrite)
	} else {
		if options.DiskIOPSReadWrite != "" {
			return "", fmt.Errorf("AzureDisk - DiskIOPSReadWrite parameter is only applicable in UltraSSD_LRS disk type")
		}
		if options.DiskMBpsReadWrite != "" {
			return "", fmt.Errorf("AzureDisk - DiskMBpsReadWrite parameter is only applicable in UltraSSD_LRS disk type")
		}
	}

	if options.DiskEncryptionSetID != "" {
		if strings.Index(strings.ToLower(options.DiskEncryptionSetID), "/subscriptions/") != 0 {
			return "", fmt.Errorf("AzureDisk - format of DiskEncryptionSetID(%s) is incorrect, correct format: %s", options.DiskEncryptionSetID, diskEncryptionSetIDFormat)
		}
		diskProperties.Encryption = &compute.Encryption{
			DiskEncryptionSetID: &options.DiskEncryptionSetID,
			Type:                compute.EncryptionAtRestWithCustomerKey,
		}
	}

	model := compute.Disk{
		Location: &c.common.location,
		Tags:     newTags,
		Zones:    createZones,
		Sku: &compute.DiskSku{
			Name: diskSku,
		},
		DiskProperties: &diskProperties,
	}

	if options.ResourceGroup == "" {
		options.ResourceGroup = c.common.resourceGroup
	}

	ctx, cancel := getContextWithCancel()
	defer cancel()
	rerr := c.common.cloud.DisksClient.CreateOrUpdate(ctx, options.ResourceGroup, options.DiskName, model)
	if rerr != nil {
		return "", rerr.Error()
	}

	diskID := ""

	err = kwait.ExponentialBackoff(defaultBackOff, func() (bool, error) {
		provisionState, id, err := c.GetDisk(options.ResourceGroup, options.DiskName)
		diskID = id
		// We are waiting for provisioningState==Succeeded
		// We don't want to hand-off managed disks to k8s while they are
		//still being provisioned, this is to avoid some race conditions
		if err != nil {
			return false, err
		}
		if strings.ToLower(provisionState) == "succeeded" {
			return true, nil
		}
		return false, nil
	})

	if err != nil {
		klog.V(2).Infof("azureDisk - created new MD Name:%s StorageAccountType:%s Size:%v but was unable to confirm provisioningState in poll process", options.DiskName, options.StorageAccountType, options.SizeGB)
	} else {
		klog.V(2).Infof("azureDisk - created new MD Name:%s StorageAccountType:%s Size:%v", options.DiskName, options.StorageAccountType, options.SizeGB)
	}

	return diskID, nil
}

//DeleteManagedDisk : delete managed disk
func (c *ManagedDiskController) DeleteManagedDisk(diskURI string) error {
	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		return err
	}

	ctx, cancel := getContextWithCancel()
	defer cancel()

	if _, ok := c.common.diskAttachDetachMap.Load(strings.ToLower(diskURI)); ok {
		return fmt.Errorf("failed to delete disk(%s) since it's in attaching or detaching state", diskURI)
	}

	rerr := c.common.cloud.DisksClient.Delete(ctx, resourceGroup, diskName)
	if rerr != nil {
		return rerr.Error()
	}
	// We don't need poll here, k8s will immediately stop referencing the disk
	// the disk will be eventually deleted - cleanly - by ARM

	klog.V(2).Infof("azureDisk - deleted a managed disk: %s", diskURI)

	return nil
}

// GetDisk return: disk provisionState, diskID, error
func (c *ManagedDiskController) GetDisk(resourceGroup, diskName string) (string, string, error) {
	ctx, cancel := getContextWithCancel()
	defer cancel()

	result, rerr := c.common.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		return "", "", rerr.Error()
	}

	if result.DiskProperties != nil && (*result.DiskProperties).ProvisioningState != nil {
		return *(*result.DiskProperties).ProvisioningState, *result.ID, nil
	}

	return "", "", nil
}

// ResizeDisk Expand the disk to new size
func (c *ManagedDiskController) ResizeDisk(diskURI string, oldSize resource.Quantity, newSize resource.Quantity) (resource.Quantity, error) {
	ctx, cancel := getContextWithCancel()
	defer cancel()

	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		return oldSize, err
	}

	result, rerr := c.common.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		return oldSize, rerr.Error()
	}

	if result.DiskProperties == nil || result.DiskProperties.DiskSizeGB == nil {
		return oldSize, fmt.Errorf("DiskProperties of disk(%s) is nil", diskName)
	}

	// Azure resizes in chunks of GiB (not GB)
	requestGiB := int32(volumehelpers.RoundUpToGiB(newSize))
	newSizeQuant := resource.MustParse(fmt.Sprintf("%dGi", requestGiB))

	klog.V(2).Infof("azureDisk - begin to resize disk(%s) with new size(%d), old size(%v)", diskName, requestGiB, oldSize)
	// If disk already of greater or equal size than requested we return
	if *result.DiskProperties.DiskSizeGB >= requestGiB {
		return newSizeQuant, nil
	}

	result.DiskProperties.DiskSizeGB = &requestGiB

	ctx, cancel = getContextWithCancel()
	defer cancel()
	if rerr := c.common.cloud.DisksClient.CreateOrUpdate(ctx, resourceGroup, diskName, result); rerr != nil {
		return oldSize, rerr.Error()
	}

	klog.V(2).Infof("azureDisk - resize disk(%s) with new size(%d) completed", diskName, requestGiB)

	return newSizeQuant, nil
}

// get resource group name from a managed disk URI, e.g. return {group-name} according to
// /subscriptions/{sub-id}/resourcegroups/{group-name}/providers/microsoft.compute/disks/{disk-id}
// according to https://docs.microsoft.com/en-us/rest/api/compute/disks/get
func getResourceGroupFromDiskURI(diskURI string) (string, error) {
	fields := strings.Split(diskURI, "/")
	if len(fields) != 9 || strings.ToLower(fields[3]) != "resourcegroups" {
		return "", fmt.Errorf("invalid disk URI: %s", diskURI)
	}
	return fields[4], nil
}

// GetLabelsForVolume implements PVLabeler.GetLabelsForVolume
func (c *Cloud) GetLabelsForVolume(ctx context.Context, pv *v1.PersistentVolume) (map[string]string, error) {
	// Ignore if not AzureDisk.
	if pv.Spec.AzureDisk == nil {
		return nil, nil
	}

	// Ignore any volumes that are being provisioned
	if pv.Spec.AzureDisk.DiskName == cloudvolume.ProvisionedVolumeName {
		return nil, nil
	}

	return c.GetAzureDiskLabels(pv.Spec.AzureDisk.DataDiskURI)
}

// GetAzureDiskLabels gets availability zone labels for Azuredisk.
func (c *Cloud) GetAzureDiskLabels(diskURI string) (map[string]string, error) {
	// Get disk's resource group.
	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		klog.Errorf("Failed to get resource group for AzureDisk %q: %v", diskName, err)
		return nil, err
	}

	// Get information of the disk.
	ctx, cancel := getContextWithCancel()
	defer cancel()
	disk, rerr := c.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		klog.Errorf("Failed to get information for AzureDisk %q: %v", diskName, rerr)
		return nil, rerr.Error()
	}

	// Check whether availability zone is specified.
	if disk.Zones == nil || len(*disk.Zones) == 0 {
		klog.V(4).Infof("Azure disk %q is not zoned", diskName)
		return nil, nil
	}

	zones := *disk.Zones
	zoneID, err := strconv.Atoi(zones[0])
	if err != nil {
		return nil, fmt.Errorf("failed to parse zone %v for AzureDisk %v: %v", zones, diskName, err)
	}

	zone := c.makeZone(c.Location, zoneID)
	klog.V(4).Infof("Got zone %q for Azure disk %q", zone, diskName)
	labels := map[string]string{
		v1.LabelZoneRegion:        c.Location,
		v1.LabelZoneFailureDomain: zone,
	}
	return labels, nil
}
<commit_msg>fix: check disk status before disk azure disk<commit_after>// +build !providerless

/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package azure

import (
	"context"
	"fmt"
	"path"
	"strconv"
	"strings"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-07-01/compute"
	"github.com/Azure/go-autorest/autorest/to"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	kwait "k8s.io/apimachinery/pkg/util/wait"
	cloudvolume "k8s.io/cloud-provider/volume"
	volumehelpers "k8s.io/cloud-provider/volume/helpers"
	"k8s.io/klog"
)

const (
	// default IOPS Caps & Throughput Cap (MBps) per https://docs.microsoft.com/en-us/azure/virtual-machines/linux/disks-ultra-ssd
	defaultDiskIOPSReadWrite = 500
	defaultDiskMBpsReadWrite = 100

	diskEncryptionSetIDFormat = "/subscriptions/{subs-id}/resourceGroups/{rg-name}/providers/Microsoft.Compute/diskEncryptionSets/{diskEncryptionSet-name}"
)

//ManagedDiskController : managed disk controller struct
type ManagedDiskController struct {
	common *controllerCommon
}

// ManagedDiskOptions specifies the options of managed disks.
type ManagedDiskOptions struct {
	// The name of the disk.
	DiskName string
	// The size in GB.
	SizeGB int
	// The name of PVC.
	PVCName string
	// The name of resource group.
	ResourceGroup string
	// The AvailabilityZone to create the disk.
	AvailabilityZone string
	// The tags of the disk.
	Tags map[string]string
	// The SKU of storage account.
	StorageAccountType compute.DiskStorageAccountTypes
	// IOPS Caps for UltraSSD disk
	DiskIOPSReadWrite string
	// Throughput Cap (MBps) for UltraSSD disk
	DiskMBpsReadWrite string
	// if SourceResourceID is not empty, then it's a disk copy operation(for snapshot)
	SourceResourceID string
	// The type of source
	SourceType string
	// ResourceId of the disk encryption set to use for enabling encryption at rest.
	DiskEncryptionSetID string
}

//CreateManagedDisk : create managed disk
func (c *ManagedDiskController) CreateManagedDisk(options *ManagedDiskOptions) (string, error) {
	var err error
	klog.V(4).Infof("azureDisk - creating new managed Name:%s StorageAccountType:%s Size:%v", options.DiskName, options.StorageAccountType, options.SizeGB)

	var createZones *[]string
	if len(options.AvailabilityZone) > 0 {
		zoneList := []string{c.common.cloud.GetZoneID(options.AvailabilityZone)}
		createZones = &zoneList
	}

	// insert original tags to newTags
	newTags := make(map[string]*string)
	azureDDTag := "kubernetes-azure-dd"
	newTags["created-by"] = &azureDDTag
	if options.Tags != nil {
		for k, v := range options.Tags {
			// Azure won't allow / (forward slash) in tags
			newKey := strings.Replace(k, "/", "-", -1)
			newValue := strings.Replace(v, "/", "-", -1)
			newTags[newKey] = &newValue
		}
	}

	diskSizeGB := int32(options.SizeGB)
	diskSku := compute.DiskStorageAccountTypes(options.StorageAccountType)

	creationData, err := getValidCreationData(c.common.subscriptionID, options.ResourceGroup, options.SourceResourceID, options.SourceType)
	if err != nil {
		return "", err
	}
	diskProperties := compute.DiskProperties{
		DiskSizeGB:   &diskSizeGB,
		CreationData: &creationData,
	}

	if diskSku == compute.UltraSSDLRS {
		diskIOPSReadWrite := int64(defaultDiskIOPSReadWrite)
		if options.DiskIOPSReadWrite != "" {
			v, err := strconv.Atoi(options.DiskIOPSReadWrite)
			if err != nil {
				return "", fmt.Errorf("AzureDisk - failed to parse DiskIOPSReadWrite: %v", err)
			}
			diskIOPSReadWrite = int64(v)
		}
		diskProperties.DiskIOPSReadWrite = to.Int64Ptr(diskIOPSReadWrite)

		diskMBpsReadWrite := int32(defaultDiskMBpsReadWrite)
		if options.DiskMBpsReadWrite != "" {
			v, err := strconv.Atoi(options.DiskMBpsReadWrite)
			if err != nil {
				return "", fmt.Errorf("AzureDisk - failed to parse DiskMBpsReadWrite: %v", err)
			}
			diskMBpsReadWrite = int32(v)
		}
		diskProperties.DiskMBpsReadWrite = to.Int32Ptr(diskMBpsReadWrite)
	} else {
		if options.DiskIOPSReadWrite != "" {
			return "", fmt.Errorf("AzureDisk - DiskIOPSReadWrite parameter is only applicable in UltraSSD_LRS disk type")
		}
		if options.DiskMBpsReadWrite != "" {
			return "", fmt.Errorf("AzureDisk - DiskMBpsReadWrite parameter is only applicable in UltraSSD_LRS disk type")
		}
	}

	if options.DiskEncryptionSetID != "" {
		if strings.Index(strings.ToLower(options.DiskEncryptionSetID), "/subscriptions/") != 0 {
			return "", fmt.Errorf("AzureDisk - format of DiskEncryptionSetID(%s) is incorrect, correct format: %s", options.DiskEncryptionSetID, diskEncryptionSetIDFormat)
		}
		diskProperties.Encryption = &compute.Encryption{
			DiskEncryptionSetID: &options.DiskEncryptionSetID,
			Type:                compute.EncryptionAtRestWithCustomerKey,
		}
	}

	model := compute.Disk{
		Location: &c.common.location,
		Tags:     newTags,
		Zones:    createZones,
		Sku: &compute.DiskSku{
			Name: diskSku,
		},
		DiskProperties: &diskProperties,
	}

	if options.ResourceGroup == "" {
		options.ResourceGroup = c.common.resourceGroup
	}

	ctx, cancel := getContextWithCancel()
	defer cancel()
	rerr := c.common.cloud.DisksClient.CreateOrUpdate(ctx, options.ResourceGroup, options.DiskName, model)
	if rerr != nil {
		return "", rerr.Error()
	}

	diskID := ""

	err = kwait.ExponentialBackoff(defaultBackOff, func() (bool, error) {
		provisionState, id, err := c.GetDisk(options.ResourceGroup, options.DiskName)
		diskID = id
		// We are waiting for provisioningState==Succeeded
		// We don't want to hand-off managed disks to k8s while they are
		//still being provisioned, this is to avoid some race conditions
		if err != nil {
			return false, err
		}
		if strings.ToLower(provisionState) == "succeeded" {
			return true, nil
		}
		return false, nil
	})

	if err != nil {
		klog.V(2).Infof("azureDisk - created new MD Name:%s StorageAccountType:%s Size:%v but was unable to confirm provisioningState in poll process", options.DiskName, options.StorageAccountType, options.SizeGB)
	} else {
		klog.V(2).Infof("azureDisk - created new MD Name:%s StorageAccountType:%s Size:%v", options.DiskName, options.StorageAccountType, options.SizeGB)
	}

	return diskID, nil
}

//DeleteManagedDisk : delete managed disk
func (c *ManagedDiskController) DeleteManagedDisk(diskURI string) error {
	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		return err
	}

	ctx, cancel := getContextWithCancel()
	defer cancel()

	if _, ok := c.common.diskAttachDetachMap.Load(strings.ToLower(diskURI)); ok {
		return fmt.Errorf("failed to delete disk(%s) since it's in attaching or detaching state", diskURI)
	}

	disk, rerr := c.common.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		return rerr.Error()
	}

	if disk.ManagedBy != nil {
		return fmt.Errorf("disk(%s) already attached to node(%s), could not be deleted", diskURI, *disk.ManagedBy)
	}

	rerr = c.common.cloud.DisksClient.Delete(ctx, resourceGroup, diskName)
	if rerr != nil {
		return rerr.Error()
	}
	// We don't need poll here, k8s will immediately stop referencing the disk
	// the disk will be eventually deleted - cleanly - by ARM

	klog.V(2).Infof("azureDisk - deleted a managed disk: %s", diskURI)

	return nil
}

// GetDisk return: disk provisionState, diskID, error
func (c *ManagedDiskController) GetDisk(resourceGroup, diskName string) (string, string, error) {
	ctx, cancel := getContextWithCancel()
	defer cancel()

	result, rerr := c.common.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		return "", "", rerr.Error()
	}

	if result.DiskProperties != nil && (*result.DiskProperties).ProvisioningState != nil {
		return *(*result.DiskProperties).ProvisioningState, *result.ID, nil
	}

	return "", "", nil
}

// ResizeDisk Expand the disk to new size
func (c *ManagedDiskController) ResizeDisk(diskURI string, oldSize resource.Quantity, newSize resource.Quantity) (resource.Quantity, error) {
	ctx, cancel := getContextWithCancel()
	defer cancel()

	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		return oldSize, err
	}

	result, rerr := c.common.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		return oldSize, rerr.Error()
	}

	if result.DiskProperties == nil || result.DiskProperties.DiskSizeGB == nil {
		return oldSize, fmt.Errorf("DiskProperties of disk(%s) is nil", diskName)
	}

	// Azure resizes in chunks of GiB (not GB)
	requestGiB := int32(volumehelpers.RoundUpToGiB(newSize))
	newSizeQuant := resource.MustParse(fmt.Sprintf("%dGi", requestGiB))

	klog.V(2).Infof("azureDisk - begin to resize disk(%s) with new size(%d), old size(%v)", diskName, requestGiB, oldSize)
	// If disk already of greater or equal size than requested we return
	if *result.DiskProperties.DiskSizeGB >= requestGiB {
		return newSizeQuant, nil
	}

	result.DiskProperties.DiskSizeGB = &requestGiB

	ctx, cancel = getContextWithCancel()
	defer cancel()
	if rerr := c.common.cloud.DisksClient.CreateOrUpdate(ctx, resourceGroup, diskName, result); rerr != nil {
		return oldSize, rerr.Error()
	}

	klog.V(2).Infof("azureDisk - resize disk(%s) with new size(%d) completed", diskName, requestGiB)

	return newSizeQuant, nil
}

// get resource group name from a managed disk URI, e.g. return {group-name} according to
// /subscriptions/{sub-id}/resourcegroups/{group-name}/providers/microsoft.compute/disks/{disk-id}
// according to https://docs.microsoft.com/en-us/rest/api/compute/disks/get
func getResourceGroupFromDiskURI(diskURI string) (string, error) {
	fields := strings.Split(diskURI, "/")
	if len(fields) != 9 || strings.ToLower(fields[3]) != "resourcegroups" {
		return "", fmt.Errorf("invalid disk URI: %s", diskURI)
	}
	return fields[4], nil
}

// GetLabelsForVolume implements PVLabeler.GetLabelsForVolume
func (c *Cloud) GetLabelsForVolume(ctx context.Context, pv *v1.PersistentVolume) (map[string]string, error) {
	// Ignore if not AzureDisk.
	if pv.Spec.AzureDisk == nil {
		return nil, nil
	}

	// Ignore any volumes that are being provisioned
	if pv.Spec.AzureDisk.DiskName == cloudvolume.ProvisionedVolumeName {
		return nil, nil
	}

	return c.GetAzureDiskLabels(pv.Spec.AzureDisk.DataDiskURI)
}

// GetAzureDiskLabels gets availability zone labels for Azuredisk.
func (c *Cloud) GetAzureDiskLabels(diskURI string) (map[string]string, error) {
	// Get disk's resource group.
	diskName := path.Base(diskURI)
	resourceGroup, err := getResourceGroupFromDiskURI(diskURI)
	if err != nil {
		klog.Errorf("Failed to get resource group for AzureDisk %q: %v", diskName, err)
		return nil, err
	}

	// Get information of the disk.
	ctx, cancel := getContextWithCancel()
	defer cancel()
	disk, rerr := c.DisksClient.Get(ctx, resourceGroup, diskName)
	if rerr != nil {
		klog.Errorf("Failed to get information for AzureDisk %q: %v", diskName, rerr)
		return nil, rerr.Error()
	}

	// Check whether availability zone is specified.
	if disk.Zones == nil || len(*disk.Zones) == 0 {
		klog.V(4).Infof("Azure disk %q is not zoned", diskName)
		return nil, nil
	}

	zones := *disk.Zones
	zoneID, err := strconv.Atoi(zones[0])
	if err != nil {
		return nil, fmt.Errorf("failed to parse zone %v for AzureDisk %v: %v", zones, diskName, err)
	}

	zone := c.makeZone(c.Location, zoneID)
	klog.V(4).Infof("Got zone %q for Azure disk %q", zone, diskName)
	labels := map[string]string{
		v1.LabelZoneRegion:        c.Location,
		v1.LabelZoneFailureDomain: zone,
	}
	return labels, nil
}
<|endoftext|>


<commit_before>//Go beanstalkd client library
//Copyright(2012) Iwan Budi Kusnanto. See LICENSE for detail
package gobeanstalk

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
)

const (
	MIN_LEN_TO_BUF = 1500 //minimum data len to send using bufio
)

//beanstalkd error
var (
	errOutOfMemory    = errors.New("Out of Memory")
	errInternalError  = errors.New("Internal Error")
	errBadFormat      = errors.New("Bad Format")
	errUnknownCommand = errors.New("Unknown Command")
	errBuried         = errors.New("Buried")
	errExpectedCrlf   = errors.New("Expected CRLF")
	errJobTooBig      = errors.New("Job Too Big")
	errDraining       = errors.New("Draining")
	errDeadlineSoon   = errors.New("Deadline Soon")
	errTimedOut       = errors.New("Timed Out")
	errNotFound       = errors.New("Not Found")
)

//gobeanstalk error
var (
	errInvalidLen = errors.New("Invalid Length")
	errUnknown    = errors.New("Unknown Error")
)

//Connection to beanstalkd
type Conn struct {
	conn      net.Conn
	addr      string
	bufReader *bufio.Reader
	bufWriter *bufio.Writer
}

//create new connection
func NewConn(conn net.Conn, addr string) (*Conn, error) {
	c := new(Conn)
	c.conn = conn
	c.addr = addr
	c.bufReader = bufio.NewReader(conn)
	c.bufWriter = bufio.NewWriter(conn)

	return c, nil
}

//A beanstalkd job
type Job struct {
	Id   uint64
	Body []byte
}

//Create new job
func NewJob(id uint64, body []byte) *Job {
	j := &Job{id, body}
	return j
}

//Connect to beanstalkd server
func Dial(addr string) (*Conn, error) {
	kon, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	c, err := NewConn(kon, addr)

	if err != nil {
		return nil, err
	}

	return c, nil
}

//Watching tube
func (c *Conn) Watch(tubename string) (int, error) {
	cmd := fmt.Sprintf("watch %s\r\n", tubename)

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

/*
Ignore tube.

The "ignore" command is for consumers. It removes the named tube from the
watch list for the current connection
*/
func (c *Conn) Ignore(tubename string) (int, error) {
	//send command and read response string
	cmd := fmt.Sprintf("ignore %s\r\n", tubename)
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	//parse response
	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		if resp == "NOT_IGNORED\r\n" {
			return -1, errors.New("Not Ignored")
		}
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

//Reserve Job
func (c *Conn) Reserve() (*Job, error) {
	//send command and read response
	resp, err := sendGetResp(c, "reserve\r\n")
	if err != nil {
		return nil, err
	}

	//parse response
	var id uint64
	var bodyLen int

	switch {
	case strings.Index(resp, "RESERVED") == 0:
		_, err = fmt.Sscanf(resp, "RESERVED %d %d\r\n", &id, &bodyLen)
		if err != nil {
			return nil, err
		}
	case resp == "DEADLINE_SOON\r\n":
		return nil, errDeadlineSoon
	case resp == "TIMED_OUT\r\n":
		return nil, errTimedOut
	default:
		return nil, parseCommonError(resp)
	}

	//read job body
	body := make([]byte, bodyLen+2) //+2 is for trailing \r\n
	n, err := io.ReadFull(c.bufReader, body)
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	body = body[:n-2] //strip \r\n trail

	return &Job{id, body}, nil
}

//Delete a job
func (c *Conn) Delete(id uint64) error {
	cmd := fmt.Sprintf("delete %d\r\n", id)
	expected := "DELETED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Use tube

The "use" command is for producers. Subsequent put commands will put jobs into
the tube specified by this command. If no use command has been issued, jobs
will be put into the tube named "default".
*/
func (c *Conn) Use(tubename string) error {
	//check parameter
	if len(tubename) > 200 {
		return errInvalidLen
	}

	cmd := fmt.Sprintf("use %s\r\n", tubename)
	expected := fmt.Sprintf("USING %s\r\n", tubename)
	return sendExpectExact(c, cmd, expected)
}

//Put job
func (c *Conn) Put(data []byte, pri, delay, ttr int) (uint64, error) {
	cmd := fmt.Sprintf("put %d %d %d %d\r\n", pri, delay, ttr, len(data))
	cmd = cmd + string(data) + "\r\n"

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return 0, err
	}

	//parse Put response
	switch {
	case strings.Index(resp, "INSERTED") == 0:
		var id uint64
		_, parseErr := fmt.Sscanf(resp, "INSERTED %d\r\n", &id)
		return id, parseErr
	case strings.Index(resp, "BURIED") == 0:
		var id uint64
		fmt.Sscanf(resp, "BURIED %d\r\n", &id)
		return id, errBuried
	case resp == "EXPECTED_CRLF\r\n":
		return 0, errExpectedCrlf
	case resp == "JOB_TOO_BIG\r\n":
		return 0, errJobTooBig
	case resp == "DRAINING\r\n":
		return 0, errDraining
	default:
		return 0, parseCommonError(resp)
	}
	return 0, errUnknown
}

/*
Release a job.

The release command puts a reserved job back into the ready queue (and marks
its state as "ready") to be run by any client. It is normally used when the job
fails because of a transitory error.
	id is the job id to release.
	pri is a new priority to assign to the job.
	delay is an integer number of seconds to wait before putting the job in
		the ready queue. The job will be in the "delayed" state during this time.
*/
func (c *Conn) Release(id uint64, pri, delay int) error {
	cmd := fmt.Sprintf("release %d %d %d\r\n", id, pri, delay)
	expected := "RELEASED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Bury a job.

The bury command puts a job into the "buried" state. Buried jobs are put into a
FIFO linked list and will not be touched by the server again until a client
kicks them with the "kick" command.
	id is the job id to release.
	pri is a new priority to assign to the job.
*/
func (c *Conn) Bury(id uint64, pri int) error {
	cmd := fmt.Sprintf("bury %d %d\r\n", id, pri)
	expected := "BURIED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Touch a job

The "touch" command allows a worker to request more time to work on a job.
This is useful for jobs that potentially take a long time, but you still want
the benefits of a TTR pulling a job away from an unresponsive worker. A worker
may periodically tell the server that it's still alive and processing a job
(e.g. it may do this on DEADLINE_SOON)
*/
func (c *Conn) Touch(id uint64) error {
	cmd := fmt.Sprintf("touch %d\r\n", id)
	expected := "TOUCHED\r\n"
	return sendExpectExact(c, cmd, expected)
}

//send command and expect some exact response
func sendExpectExact(c *Conn, cmd, expected string) error {
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return err
	}

	if resp != expected {
		return parseCommonError(resp)
	}
	return nil
}

//Send command and read response
func sendGetResp(c *Conn, cmd string) (string, error) {
	//_, err := c.conn.Write([]byte(cmd))
	_, err := sendFull(c, []byte(cmd))
	if err != nil {
		return "", err
	}

	//wait for response
	resp, err := c.bufReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	return resp, nil
}

//try to send all of data
//if data len < 1500, it use TCPConn.Write
//if data len >= 1500, it use bufio.Write
func sendFull(c *Conn, data []byte) (int, error) {
	toWrite := data
	totWritten := 0
	var n int
	var err error
	for totWritten < len(data) {
		if len(toWrite) >= MIN_LEN_TO_BUF {
			n, err = c.bufWriter.Write(toWrite)
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
			err = c.bufWriter.Flush()
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
		} else {
			n, err = c.conn.Write(toWrite)
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
		}
		totWritten += n
		toWrite = toWrite[n:]
	}
	return totWritten, nil
}

//parse for Common Error
func parseCommonError(str string) error {
	switch str {
	case "BURIED\r\n":
		return errBuried
	case "NOT_FOUND\r\n":
		return errNotFound
	case "OUT_OF_MEMORY\r\n":
		return errOutOfMemory
	case "INTERNAL_ERROR\r\n":
		return errInternalError
	case "BAD_FORMAT\r\n":
		return errBadFormat
	case "UNKNOWN_COMMAND\r\n":
		return errUnknownCommand
	}
	return errUnknown
}

//Check if it is temporary network error
func isNetTempErr(err error) bool {
	if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
		return true
	}
	return false
}
<commit_msg>Added put-unique command<commit_after>//Go beanstalkd client library
//Copyright(2012) Iwan Budi Kusnanto. See LICENSE for detail
package gobeanstalk

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
)

const (
	MIN_LEN_TO_BUF = 1500 //minimum data len to send using bufio
)

//beanstalkd error
var (
	errOutOfMemory    = errors.New("Out of Memory")
	errInternalError  = errors.New("Internal Error")
	errBadFormat      = errors.New("Bad Format")
	errUnknownCommand = errors.New("Unknown Command")
	errBuried         = errors.New("Buried")
	errExpectedCrlf   = errors.New("Expected CRLF")
	errJobTooBig      = errors.New("Job Too Big")
	errDraining       = errors.New("Draining")
	errDeadlineSoon   = errors.New("Deadline Soon")
	errTimedOut       = errors.New("Timed Out")
	errNotFound       = errors.New("Not Found")
)

//gobeanstalk error
var (
	errInvalidLen = errors.New("Invalid Length")
	errUnknown    = errors.New("Unknown Error")
)

//Connection to beanstalkd
type Conn struct {
	conn      net.Conn
	addr      string
	bufReader *bufio.Reader
	bufWriter *bufio.Writer
}

//create new connection
func NewConn(conn net.Conn, addr string) (*Conn, error) {
	c := new(Conn)
	c.conn = conn
	c.addr = addr
	c.bufReader = bufio.NewReader(conn)
	c.bufWriter = bufio.NewWriter(conn)

	return c, nil
}

//A beanstalkd job
type Job struct {
	Id   uint64
	Body []byte
}

//Create new job
func NewJob(id uint64, body []byte) *Job {
	j := &Job{id, body}
	return j
}

//Connect to beanstalkd server
func Dial(addr string) (*Conn, error) {
	kon, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	c, err := NewConn(kon, addr)

	if err != nil {
		return nil, err
	}

	return c, nil
}

//Watching tube
func (c *Conn) Watch(tubename string) (int, error) {
	cmd := fmt.Sprintf("watch %s\r\n", tubename)

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

/*
Ignore tube.

The "ignore" command is for consumers. It removes the named tube from the
watch list for the current connection
*/
func (c *Conn) Ignore(tubename string) (int, error) {
	//send command and read response string
	cmd := fmt.Sprintf("ignore %s\r\n", tubename)
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return -1, err
	}

	//parse response
	var tubeCount int
	_, err = fmt.Sscanf(resp, "WATCHING %d\r\n", &tubeCount)
	if err != nil {
		if resp == "NOT_IGNORED\r\n" {
			return -1, errors.New("Not Ignored")
		}
		return -1, parseCommonError(resp)
	}
	return tubeCount, nil
}

//Reserve Job
func (c *Conn) Reserve() (*Job, error) {
	//send command and read response
	resp, err := sendGetResp(c, "reserve\r\n")
	if err != nil {
		return nil, err
	}

	//parse response
	var id uint64
	var bodyLen int

	switch {
	case strings.Index(resp, "RESERVED") == 0:
		_, err = fmt.Sscanf(resp, "RESERVED %d %d\r\n", &id, &bodyLen)
		if err != nil {
			return nil, err
		}
	case resp == "DEADLINE_SOON\r\n":
		return nil, errDeadlineSoon
	case resp == "TIMED_OUT\r\n":
		return nil, errTimedOut
	default:
		return nil, parseCommonError(resp)
	}

	//read job body
	body := make([]byte, bodyLen+2) //+2 is for trailing \r\n
	n, err := io.ReadFull(c.bufReader, body)
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	body = body[:n-2] //strip \r\n trail

	return &Job{id, body}, nil
}

//Delete a job
func (c *Conn) Delete(id uint64) error {
	cmd := fmt.Sprintf("delete %d\r\n", id)
	expected := "DELETED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Use tube

The "use" command is for producers. Subsequent put commands will put jobs into
the tube specified by this command. If no use command has been issued, jobs
will be put into the tube named "default".
*/
func (c *Conn) Use(tubename string) error {
	//check parameter
	if len(tubename) > 200 {
		return errInvalidLen
	}

	cmd := fmt.Sprintf("use %s\r\n", tubename)
	expected := fmt.Sprintf("USING %s\r\n", tubename)
	return sendExpectExact(c, cmd, expected)
}

//Put job
func (c *Conn) Put(data []byte, pri, delay, ttr int) (uint64, error) {
	cmd := fmt.Sprintf("put %d %d %d %d\r\n", pri, delay, ttr, len(data))
	cmd = cmd + string(data) + "\r\n"

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return 0, err
	}

	//parse Put response
	switch {
	case strings.Index(resp, "INSERTED") == 0:
		var id uint64
		_, parseErr := fmt.Sscanf(resp, "INSERTED %d\r\n", &id)
		return id, parseErr
	case strings.Index(resp, "BURIED") == 0:
		var id uint64
		fmt.Sscanf(resp, "BURIED %d\r\n", &id)
		return id, errBuried
	case resp == "EXPECTED_CRLF\r\n":
		return 0, errExpectedCrlf
	case resp == "JOB_TOO_BIG\r\n":
		return 0, errJobTooBig
	case resp == "DRAINING\r\n":
		return 0, errDraining
	default:
		return 0, parseCommonError(resp)
	}
	return 0, errUnknown
}

//Put unique job
func (c *Conn) PutUnique(data []byte, pri, delay, ttr int) (uint64, error) {
	cmd := fmt.Sprintf("put-unique %d %d %d %d\r\n", pri, delay, ttr, len(data))
	cmd = cmd + string(data) + "\r\n"

	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return 0, err
	}

	//parse Put response
	switch {
	case strings.Index(resp, "INSERTED") == 0:
		var id uint64
		_, parseErr := fmt.Sscanf(resp, "INSERTED %d\r\n", &id)
		return id, parseErr
	case strings.Index(resp, "BURIED") == 0:
		var id uint64
		fmt.Sscanf(resp, "BURIED %d\r\n", &id)
		return id, errBuried
	case resp == "EXPECTED_CRLF\r\n":
		return 0, errExpectedCrlf
	case resp == "JOB_TOO_BIG\r\n":
		return 0, errJobTooBig
	case resp == "DRAINING\r\n":
		return 0, errDraining
	default:
		return 0, parseCommonError(resp)
	}
	return 0, errUnknown
}

/*
Release a job.

The release command puts a reserved job back into the ready queue (and marks
its state as "ready") to be run by any client. It is normally used when the job
fails because of a transitory error.
	id is the job id to release.
	pri is a new priority to assign to the job.
	delay is an integer number of seconds to wait before putting the job in
		the ready queue. The job will be in the "delayed" state during this time.
*/
func (c *Conn) Release(id uint64, pri, delay int) error {
	cmd := fmt.Sprintf("release %d %d %d\r\n", id, pri, delay)
	expected := "RELEASED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Bury a job.

The bury command puts a job into the "buried" state. Buried jobs are put into a
FIFO linked list and will not be touched by the server again until a client
kicks them with the "kick" command.
	id is the job id to release.
	pri is a new priority to assign to the job.
*/
func (c *Conn) Bury(id uint64, pri int) error {
	cmd := fmt.Sprintf("bury %d %d\r\n", id, pri)
	expected := "BURIED\r\n"
	return sendExpectExact(c, cmd, expected)
}

/*
Touch a job

The "touch" command allows a worker to request more time to work on a job.
This is useful for jobs that potentially take a long time, but you still want
the benefits of a TTR pulling a job away from an unresponsive worker. A worker
may periodically tell the server that it's still alive and processing a job
(e.g. it may do this on DEADLINE_SOON)
*/
func (c *Conn) Touch(id uint64) error {
	cmd := fmt.Sprintf("touch %d\r\n", id)
	expected := "TOUCHED\r\n"
	return sendExpectExact(c, cmd, expected)
}

//send command and expect some exact response
func sendExpectExact(c *Conn, cmd, expected string) error {
	resp, err := sendGetResp(c, cmd)
	if err != nil {
		return err
	}

	if resp != expected {
		return parseCommonError(resp)
	}
	return nil
}

//Send command and read response
func sendGetResp(c *Conn, cmd string) (string, error) {
	//_, err := c.conn.Write([]byte(cmd))
	_, err := sendFull(c, []byte(cmd))
	if err != nil {
		return "", err
	}

	//wait for response
	resp, err := c.bufReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	return resp, nil
}

//try to send all of data
//if data len < 1500, it use TCPConn.Write
//if data len >= 1500, it use bufio.Write
func sendFull(c *Conn, data []byte) (int, error) {
	toWrite := data
	totWritten := 0
	var n int
	var err error
	for totWritten < len(data) {
		if len(toWrite) >= MIN_LEN_TO_BUF {
			n, err = c.bufWriter.Write(toWrite)
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
			err = c.bufWriter.Flush()
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
		} else {
			n, err = c.conn.Write(toWrite)
			if err != nil && !isNetTempErr(err) {
				return totWritten, err
			}
		}
		totWritten += n
		toWrite = toWrite[n:]
	}
	return totWritten, nil
}

//parse for Common Error
func parseCommonError(str string) error {
	switch str {
	case "BURIED\r\n":
		return errBuried
	case "NOT_FOUND\r\n":
		return errNotFound
	case "OUT_OF_MEMORY\r\n":
		return errOutOfMemory
	case "INTERNAL_ERROR\r\n":
		return errInternalError
	case "BAD_FORMAT\r\n":
		return errBadFormat
	case "UNKNOWN_COMMAND\r\n":
		return errUnknownCommand
	}
	return errUnknown
}

//Check if it is temporary network error
func isNetTempErr(err error) bool {
	if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	mongo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	elastic "gopkg.in/olivere/elastic.v3"
	// "log"
	"mongoes/libs"
	"os"
	// "sync/atomic"
	"time"
)

func fatal(e error) {
	fmt.Println(e)
	flag.PrintDefaults()
}

var counts int32 = 0

type Message struct {
	Id       bson.ObjectId
	Document map[string]interface{}
}

func printStats(stats elastic.BulkProcessorStats) {
	fmt.Println("Flushed:", stats.Flushed)
	fmt.Println("Committed:", stats.Committed)
	fmt.Println("Indexed:", stats.Indexed)
	fmt.Println("Created:", stats.Created)
	fmt.Println("Updated:", stats.Updated)
	fmt.Println("Deleted:", stats.Deleted)
	fmt.Println("Succedeed:", stats.Succeeded)
	fmt.Println("Failed:", stats.Failed)

}
func main() {
	var dbName = flag.String("db", "", "Mongodb DB Name")
	var collName = flag.String("collection", "", "Mongodb Collection Name")
	var dbUri = flag.String("dbUri", "localhost:27017", "Mongodb URI")
	var indexName = flag.String("index", "", "ES Index Name")
	var typeName = flag.String("type", "", "ES Type Name")
	var mappingFile = flag.String("mapping", "", "Mapping mongodb field to es")
	flag.Parse()
	if len(*dbName) == 0 || len(*collName) == 0 {
		fatal(errors.New("Please provide db and collection name"))
		return
	}

	if len(*indexName) == 0 {
		indexName = dbName
	}

	if len(*typeName) == 0 {
		typeName = collName
	}

	// Set Tracer
	tracer := libs.NewTracer(os.Stdout)

	// Get connected to mongodb
	tracer.Trace("Connecting to Mongodb at", *dbUri)
	session, err := mongo.Dial(*dbUri)
	if err != nil {
		fatal(err)
		return
	}
	defer session.Close()

	tracer.Trace("Connecting to elasticsearch cluster")
	client, err := elastic.NewClient()
	if err != nil {
		fatal(err)
		return
	}
	client.DeleteIndex(*indexName).Do()
	_, err = client.CreateIndex(*indexName).Do()
	if err != nil {
		fatal(err)
		return
	}
	tracer.Trace("Create Mongodb to ES Mapping")
	rawMapping, _ := libs.ReadMappingJson(*mappingFile)
	esMapping, _ := libs.CreateMapping(rawMapping)
	_, err = client.PutMapping().Index(*indexName).Type(*typeName).BodyJson(esMapping).Do()
	if err != nil {
		fatal(err)
		return
	}
	p := make(map[string]interface{})
	iter := session.DB(*dbName).C(*collName).Find(nil).Iter()
	start := time.Now()
	fmt.Println("Start Indexing MongoDb")
	bulkProcessorService := elastic.NewBulkProcessorService(client).Workers(4).Stats(true)
	bulkProcessor, _ := bulkProcessorService.Do()
	bulkProcessor.Start()
	for iter.Next(&p) {
		var esBody = make(map[string]interface{})
		for k, v := range rawMapping {
			mgoVal, ok := p[k]
			if ok {
				var key = (v.(map[string]interface{}))["es_name"]
				if key == nil {
					key = k
				}
				esBody[key.(string)] = mgoVal
			}
		}
		bulkRequest := elastic.NewBulkIndexRequest().
			Index(*indexName).
			Type(*typeName).
			Id(p["_id"].(bson.ObjectId).Hex()).
			Doc(esBody)

		bulkProcessor.Add(bulkRequest)
	}
	iter.Close()
	elapsed := time.Since(start)
	stats := bulkProcessor.Stats()
	fmt.Println("Finished indexing", stats.Indexed, "documents in", elapsed)
	printStats(stats)
}
<commit_msg>refactor to more controllable bulkservice<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	mongo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	elastic "gopkg.in/olivere/elastic.v3"
	// "log"
	"mongoes/libs"
	"os"
	// "sync/atomic"
	"time"
)

func fatal(e error) {
	fmt.Println(e)
	flag.PrintDefaults()
}

var counts int32 = 0

type Message struct {
	Id       bson.ObjectId
	Document map[string]interface{}
}

func printStats(stats elastic.BulkProcessorStats) {
	fmt.Println("Flushed:", stats.Flushed)
	fmt.Println("Committed:", stats.Committed)
	fmt.Println("Indexed:", stats.Indexed)
	fmt.Println("Created:", stats.Created)
	fmt.Println("Updated:", stats.Updated)
	fmt.Println("Deleted:", stats.Deleted)
	fmt.Println("Succedeed:", stats.Succeeded)
	fmt.Println("Failed:", stats.Failed)

}
func main() {
	var dbName = flag.String("db", "", "Mongodb DB Name")
	var collName = flag.String("collection", "", "Mongodb Collection Name")
	var dbUri = flag.String("dbUri", "localhost:27017", "Mongodb URI")
	var indexName = flag.String("index", "", "ES Index Name")
	var typeName = flag.String("type", "", "ES Type Name")
	var mappingFile = flag.String("mapping", "", "Mapping mongodb field to es")
	flag.Parse()
	if len(*dbName) == 0 || len(*collName) == 0 {
		fatal(errors.New("Please provide db and collection name"))
		return
	}

	if len(*indexName) == 0 {
		indexName = dbName
	}

	if len(*typeName) == 0 {
		typeName = collName
	}

	// Set Tracer
	tracer := libs.NewTracer(os.Stdout)

	// Get connected to mongodb
	tracer.Trace("Connecting to Mongodb at", *dbUri)
	session, err := mongo.Dial(*dbUri)
	if err != nil {
		fatal(err)
		return
	}
	defer session.Close()

	tracer.Trace("Connecting to elasticsearch cluster")
	client, err := elastic.NewClient()
	if err != nil {
		fatal(err)
		return
	}
	client.DeleteIndex(*indexName).Do()
	_, err = client.CreateIndex(*indexName).Do()
	if err != nil {
		fatal(err)
		return
	}
	tracer.Trace("Create Mongodb to ES Mapping")
	rawMapping, _ := libs.ReadMappingJson(*mappingFile)
	esMapping, _ := libs.CreateMapping(rawMapping)
	_, err = client.PutMapping().Index(*indexName).Type(*typeName).BodyJson(esMapping).Do()
	if err != nil {
		fatal(err)
		return
	}
	p := make(map[string]interface{})
	iter := session.DB(*dbName).C(*collName).Find(nil).Iter()
	start := time.Now()
	fmt.Println("Start Indexing MongoDb")
	bulkService := elastic.NewBulkService(client).Index(*indexName).Type(*typeName)
	// bulkProcessor, _ := bulkProcessorService.Do()
	// bulkProcessor.Start()
	for iter.Next(&p) {
		var esBody = make(map[string]interface{})
		for k, v := range rawMapping {
			mgoVal, ok := p[k]
			if ok {
				var key = (v.(map[string]interface{}))["es_name"]
				if key == nil {
					key = k
				}
				esBody[key.(string)] = mgoVal
			}
		}
		bulkRequest := elastic.NewBulkIndexRequest().
			Index(*indexName).
			Type(*typeName).
			Id(p["_id"].(bson.ObjectId).Hex()).
			Doc(esBody)

		// bulkProcessor.Add(bulkRequest)
		bulkService.Add(bulkRequest)
		if bulkService.NumberOfActions() == 1000 {
			bulkResponse, _ := bulkService.Do()
			counts += int32(len(bulkResponse.Indexed()))
			fmt.Println(counts, " documents indexed")
		}
	}
	iter.Close()
	elapsed := time.Since(start)
	// stats := bulkProcessor.Stats()
	fmt.Println("Finished indexing documents in", elapsed)
	// printStats(stats)
}
<|endoftext|>


<commit_before>package cache

import (
	//"bytes"
	//"encoding/gob"
	"time"

	"github.com/garyburd/redigo/redis"
)

//Cache storage implementation using redis as key/value storage
type RedisCacheStorage struct {
	redisPool      redis.Pool
	ttlReadTimeout int
	cacheAreaa     string
	Serializer     Serializer
}

var enableTTL = true // setup a external config

//recover all cacheregistries of keys
func (s RedisCacheStorage) GetValuesMap(cacheKeys ...string) (map[string]CacheRegistry, error) {

	ttlMapChan := make(chan map[string]int, 1)
	if (enableTTL) {
		go func() {
			defer func() {
				if r := recover(); r != nil {
					log.Critical("Error trying to get ttl for registries %v!", cacheKeys)

					//in case of error, retur an empty map
					ttlMapChan <- make(map[string]int, 0)
				}
			}()

			//put result on channel
			ttlMapChan <- s.GetTTLMap(cacheKeys)
		}()
	}

	mapCacheRegistry := make(map[string]CacheRegistry)

	if len(cacheKeys) <= 0 {
		log.Debug("Nenhuma chave informada para busca. len(arrKeys)=0!")
		return mapCacheRegistry, nil //empty map
	}

	conn := s.redisPool.Get()
	defer conn.Close()
	var err error = nil

	//log.Debug(cacheKeys)

	replyMget, err := conn.Do("MGET", (s.getKeys(cacheKeys))...)
	if err != nil || replyMget == nil {
		log.Error("Error trying to get values from cache %v", err)
		log.Error("Returning an empty registry!")

		return mapCacheRegistry, err // error trying to search cache keys
	}

	arrResults, isArray := replyMget.([]interface{}) //try to convert the returned value to array

	if !isArray {
		log.Error("Value returned by a MGET query is not array for keys %v! No error will be returned!", cacheKeys) //formal check
		return make(map[string]CacheRegistry), nil
	}

	for _, cacheRegistryNotBytes := range arrResults {
		if cacheRegistryNotBytes != nil {

/*
			cacheRegistryBytes, isByteArr := cacheRegistryNotBytes.(string)
			if(isByteArr){
				log.Error("error trying to deserialize! not a byte array")
				return mapCacheRegistry, errors.New("not byte array!")
			}
*/

			cacheRegistryBytes, errBytes := redis.Bytes(cacheRegistryNotBytes, err)
			if errBytes != nil || replyMget == nil {
				return mapCacheRegistry, errBytes
			}

			cacheRegistry := CacheRegistry{}

			interfaceResp, _, errUnm := s.Serializer.UnmarshalMsg(cacheRegistry,cacheRegistryBytes)
			if errUnm!=nil {
				log.Error("error trying to deserialize!",errUnm)
				return mapCacheRegistry, errUnm
			}

			cacheRegistry, isCR := interfaceResp.(CacheRegistry)
			if(!isCR){
				log.Error("error trying to deserialize! object is not a CacheRegistry object type!")
				return mapCacheRegistry, nil
			}

			if err != nil {
				log.Error("Warning!! Error trying to recover data from redis!", err)
			} else {
				if cacheRegistry.Payload == nil {
					log.Error("ATENCAO! NENHUM PAYLOAD FOI RETORNADO DO REDIS!")
				}
				//Everything is alright
				mapCacheRegistry[cacheRegistry.CacheKey] = cacheRegistry
			}
		}
	}

	if (enableTTL) {
		select {
		//wait for ttl channel
		case ttlMap := <-ttlMapChan:
			mapCacheRegistry = s.zipTTL(mapCacheRegistry, ttlMap)
		//in case of timeout, returt an empty map
		case <-time.After(time.Duration(s.ttlReadTimeout) * time.Millisecond):
			log.Warning("Retrieve TTL for cachekeys %v from redis timeout after %dms, continuing without it.", cacheKeys, s.ttlReadTimeout)
			mapCacheRegistry = s.zipTTL(mapCacheRegistry, make(map[string]int, 0))
		}
	}

	return mapCacheRegistry, nil // err=nil by default, if everything is alright
}

//Recover current ttl information about registry
func (s RedisCacheStorage) GetTTL(key string) (int, error) {
	oneItemMap := make(map[string]CacheRegistry, 1)

	oneItemMap[key] = CacheRegistry{key, "", -2 /*not found*/, true, ""}

	respMap, errTTL := s.GetActualTTL(oneItemMap)
	return respMap[key].Ttl, errTTL

}

//Recover current ttl information about registries
func (s RedisCacheStorage) zipTTL(mapCacheRegistry map[string]CacheRegistry, ttlMap map[string]int) map[string]CacheRegistry {
	//prepare a keyval pair array
	for key, cacheRegistry := range mapCacheRegistry {
		if ttl, hasTtl := ttlMap[key]; hasTtl {
			cacheRegistry.Ttl = ttl
		} else {
			cacheRegistry.Ttl = -1
		}
		mapCacheRegistry[key] = cacheRegistry
	}

	return mapCacheRegistry
}

//Recover current ttl information about registries
func (s RedisCacheStorage) GetActualTTL(mapCacheRegistry map[string]CacheRegistry) (map[string]CacheRegistry, error) {

	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for keyMap, cacheRegistry := range mapCacheRegistry {

		respTtl, err := conn.Do("ttl", s.getKey(keyMap))
		log.Debug("TTL %v that came from redis %v", keyMap, respTtl)

		if err != nil {
			log.Error("Error trying to retrieve ttl of key " + keyMap, err)
			cacheRegistry.Ttl = -2
			return mapCacheRegistry, err

		} else {
			intResp, _ := respTtl.(int64)
			cacheRegistry.Ttl = int(intResp)
		}

		mapCacheRegistry[keyMap] = setTTLToPayload(&cacheRegistry)
	}

	return mapCacheRegistry, nil
}

//Recover current ttl information about registries
func (s RedisCacheStorage) GetTTLMap(keys []string) map[string]int {

	ttlMap := make(map[string]int, len(keys))

	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for _, key := range keys {

		respTtl, err := conn.Do("ttl", s.getKey(key))
		log.Debug("TTL %v that came from redis %v", key, respTtl)

		if err != nil {
			log.Error("Error trying to retrieve ttl of key " + key, err)
			ttlMap[key] = -2

		} else {
			intResp, _ := respTtl.(int64)
			ttlMap[key] = int(intResp)
		}

	}

	return ttlMap
}

//transfer the ttl information from cacheRegistry to paylaod interface, if it is ExposeTTL
func setTTLToPayload(cacheRegistry *CacheRegistry) CacheRegistry {

	payload := cacheRegistry.Payload

	exposeTTL, hasTtl := payload.(ExposeTTL)

	if hasTtl {
		log.Debug("Transfering ttl from redis (%d seconds) registry to ttl attribute of object %s", cacheRegistry.Ttl, cacheRegistry.CacheKey)
		payload = exposeTTL.SetTtl(cacheRegistry.Ttl) // assure the same type, from set ttl
		cacheRegistry.Payload = payload
		log.Debug("Setting ttl to %v, ttl value %v", cacheRegistry.CacheKey, exposeTTL.GetTtl())
	} else {
		log.Debug("Payload doesn't ExposeTTL %v", cacheRegistry.CacheKey)
	}

	return *cacheRegistry
}

//save informed registries on redis
func (s RedisCacheStorage) SetValues(registries ...CacheRegistry) error {

	var cacheRegistry CacheRegistry
	var index int

	defer func(cacheRegistry *CacheRegistry) {
		if r := recover(); r != nil {
			log.Error("Error trying to save cacheRegistry! recover= %v", r)
		}
	}(&cacheRegistry)

	conn := s.redisPool.Get()
	defer conn.Close()

	keyValPairs := make([]interface{}, 2 * len(registries))

	//prepare a keyval pair array
	for index, cacheRegistry = range registries {

		if len(cacheRegistry.CacheKey) == 0 {
			log.Error("cache key vazio !!!")
			//panic(errors.New("cache key vazio"))
		}

		var bytes = []byte{}
		bytes, err := s.Serializer.MarshalMsg(cacheRegistry,bytes)
		if(err!=nil){
			return err
		}

		if len(bytes) == 0 {
			log.Error("Error trying to decode value for key %v", cacheRegistry.CacheKey)
		}

		keyValPairs[(index * 2)] = s.getKey(cacheRegistry.CacheKey)
		keyValPairs[(index * 2) + 1] = bytes

	}

	_, errDo := conn.Do("MSET", keyValPairs...)
	if errDo != nil {
		log.Error("Error trying to save registry! %v %v", s.getKey(cacheRegistry.CacheKey), errDo)
		return errDo
	} else {
		log.Debug("Updating cache reg key %v ", s.getKey(cacheRegistry.CacheKey))
	}

	errF := conn.Flush()
	if errF != nil {
		log.Error("Error trying to flush connection! %v", errF)
		return errF
	}
	s.SetExpireTTL(registries...)
	return nil
}

//set defined ttl to the cache registries
func (s RedisCacheStorage) SetExpireTTL(cacheRegistries ...CacheRegistry) {
	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for _, cacheRegistry := range cacheRegistries {
		if cacheRegistry.GetTTL() > 0 {
			//log.Debug("Setting ttl to key %s ", cacheRegistry.CacheKey)
			_, err := conn.Do("expire", s.getKey(cacheRegistry.CacheKey), cacheRegistry.GetTTL())
			if err != nil {
				log.Error("Error trying to save cache registry w! %v", err)
				return
			}

		} else {
			log.Debug("TTL for %s, ttl=%d will not be setted! ", s.getKey(cacheRegistry.CacheKey), cacheRegistry.GetTTL())
		}
	}

	err := conn.Flush()
	if err != nil {
		log.Error("Error trying to save cache registry z! %v", err)
		return
	}
}

//delete values from redis
func (s RedisCacheStorage) DeleteValues(cacheKeys ...string) error {

	c := s.redisPool.Get()
	defer func() {
		c.Close()
	}()

	//apply a prefix to cache area
	keys := s.getKeys(cacheKeys)

	reply, err := c.Do("DEL", keys...)
	if err != nil {
		log.Error("Erro ao tentar invalidar registro no cache!", err, reply)
		return err
	}

	return nil
}

//apply a prefix to cache area
func (s RedisCacheStorage) getKey(key string) string {
	var newKey string

	if len(s.cacheAreaa) > 0 {
		newKey = s.cacheAreaa + ":" + key
	} else {
		newKey = key
	}

	return newKey
}

//apply a prefix to cachearea
func (s RedisCacheStorage) getKeys(keys []string) []interface{} {

	newKeys := make([]interface{}, len(keys))

	for index, key := range keys {
		newKey := s.getKey(key)
		newKeys[index] = newKey
	}

	return newKeys
}

//instantiate a new cachestorage redis
func NewRedisCacheStorage(hostPort string, password string, maxIdle int, readTimeout int, ttlReadTimeout int, cacheArea string, serializer Serializer) RedisCacheStorage {

	redisCacheStorage := RedisCacheStorage{
		*newPoolRedis(hostPort, password, maxIdle, readTimeout),
		ttlReadTimeout,
		cacheArea,
		serializer,
	}

	return redisCacheStorage
}

//create a redis connection pool
func newPoolRedis(server, password string, maxIdle int, readTimeout int) *redis.Pool {

	return &redis.Pool{
		MaxIdle:     maxIdle,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {

			c, err := redis.Dial("tcp", server, redis.DialReadTimeout(time.Duration(readTimeout) * time.Millisecond))
			if err != nil {
				log.Error("Erro ao tentar se conectar ao redis! ", err)
				return nil, err
			}

			return c, err
		},

		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}
<commit_msg>disabling ttl<commit_after>package cache

import (
	//"bytes"
	//"encoding/gob"
	"time"

	"github.com/garyburd/redigo/redis"
)

//Cache storage implementation using redis as key/value storage
type RedisCacheStorage struct {
	redisPool      redis.Pool
	ttlReadTimeout int
	cacheAreaa     string
	Serializer     Serializer // usually SerializerGOB implementation
}

var _=SerializerGOB{} // this is the usual serializer used above!!

var enableTTL = false // setup a external config

//recover all cacheregistries of keys
func (s RedisCacheStorage) GetValuesMap(cacheKeys ...string) (map[string]CacheRegistry, error) {

	ttlMapChan := make(chan map[string]int, 1)
	if (enableTTL) {
		go func() {
			defer func() {
				if r := recover(); r != nil {
					log.Critical("Error trying to get ttl for registries %v!", cacheKeys)

					//in case of error, retur an empty map
					ttlMapChan <- make(map[string]int, 0)
				}
			}()

			//put result on channel
			ttlMapChan <- s.GetTTLMap(cacheKeys)
		}()
	}

	mapCacheRegistry := make(map[string]CacheRegistry)

	if len(cacheKeys) <= 0 {
		log.Debug("Nenhuma chave informada para busca. len(arrKeys)=0!")
		return mapCacheRegistry, nil //empty map
	}

	conn := s.redisPool.Get()
	defer conn.Close()
	var err error = nil

	//log.Debug(cacheKeys)

	replyMget, err := conn.Do("MGET", (s.getKeys(cacheKeys))...)
	if err != nil || replyMget == nil {
		log.Error("Error trying to get values from cache %v", err)
		log.Error("Returning an empty registry!")

		return mapCacheRegistry, err // error trying to search cache keys
	}

	arrResults, isArray := replyMget.([]interface{}) //try to convert the returned value to array

	if !isArray {
		log.Error("Value returned by a MGET query is not array for keys %v! No error will be returned!", cacheKeys) //formal check
		return make(map[string]CacheRegistry), nil
	}

	for _, cacheRegistryNotBytes := range arrResults {
		if cacheRegistryNotBytes != nil {

/*
			cacheRegistryBytes, isByteArr := cacheRegistryNotBytes.(string)
			if(isByteArr){
				log.Error("error trying to deserialize! not a byte array")
				return mapCacheRegistry, errors.New("not byte array!")
			}
*/

			cacheRegistryBytes, errBytes := redis.Bytes(cacheRegistryNotBytes, err)
			if errBytes != nil || replyMget == nil {
				return mapCacheRegistry, errBytes
			}

			cacheRegistry := CacheRegistry{}

			interfaceResp, _, errUnm := s.Serializer.UnmarshalMsg(cacheRegistry,cacheRegistryBytes)
			if errUnm!=nil {
				log.Error("error trying to deserialize!",errUnm)
				return mapCacheRegistry, errUnm
			}

			cacheRegistry, isCR := interfaceResp.(CacheRegistry)
			if(!isCR){
				log.Error("error trying to deserialize! object is not a CacheRegistry object type!")
				return mapCacheRegistry, nil
			}

			if err != nil {
				log.Error("Warning!! Error trying to recover data from redis!", err)
			} else {
				if cacheRegistry.Payload == nil {
					log.Error("ATENCAO! NENHUM PAYLOAD FOI RETORNADO DO REDIS!")
				}
				//Everything is alright
				mapCacheRegistry[cacheRegistry.CacheKey] = cacheRegistry
			}
		}
	}

	if (enableTTL) {
		select {
		//wait for ttl channel
		case ttlMap := <-ttlMapChan:
			mapCacheRegistry = s.zipTTL(mapCacheRegistry, ttlMap)
		//in case of timeout, returt an empty map
		case <-time.After(time.Duration(s.ttlReadTimeout) * time.Millisecond):
			log.Warning("Retrieve TTL for cachekeys %v from redis timeout after %dms, continuing without it.", cacheKeys, s.ttlReadTimeout)
			mapCacheRegistry = s.zipTTL(mapCacheRegistry, make(map[string]int, 0))
		}
	}

	return mapCacheRegistry, nil // err=nil by default, if everything is alright
}

//Recover current ttl information about registry
func (s RedisCacheStorage) GetTTL(key string) (int, error) {
	oneItemMap := make(map[string]CacheRegistry, 1)

	oneItemMap[key] = CacheRegistry{key, "", -2 /*not found*/, true, ""}

	respMap, errTTL := s.GetActualTTL(oneItemMap)
	return respMap[key].Ttl, errTTL

}

//Recover current ttl information about registries
func (s RedisCacheStorage) zipTTL(mapCacheRegistry map[string]CacheRegistry, ttlMap map[string]int) map[string]CacheRegistry {
	//prepare a keyval pair array
	for key, cacheRegistry := range mapCacheRegistry {
		if ttl, hasTtl := ttlMap[key]; hasTtl {
			cacheRegistry.Ttl = ttl
		} else {
			cacheRegistry.Ttl = -1
		}
		mapCacheRegistry[key] = cacheRegistry
	}

	return mapCacheRegistry
}

//Recover current ttl information about registries
func (s RedisCacheStorage) GetActualTTL(mapCacheRegistry map[string]CacheRegistry) (map[string]CacheRegistry, error) {

	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for keyMap, cacheRegistry := range mapCacheRegistry {

		respTtl, err := conn.Do("ttl", s.getKey(keyMap))
		log.Debug("TTL %v that came from redis %v", keyMap, respTtl)

		if err != nil {
			log.Error("Error trying to retrieve ttl of key " + keyMap, err)
			cacheRegistry.Ttl = -2
			return mapCacheRegistry, err

		} else {
			intResp, _ := respTtl.(int64)
			cacheRegistry.Ttl = int(intResp)
		}

		mapCacheRegistry[keyMap] = setTTLToPayload(&cacheRegistry)
	}

	return mapCacheRegistry, nil
}

//Recover current ttl information about registries
func (s RedisCacheStorage) GetTTLMap(keys []string) map[string]int {

	ttlMap := make(map[string]int, len(keys))

	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for _, key := range keys {

		respTtl, err := conn.Do("ttl", s.getKey(key))
		log.Debug("TTL %v that came from redis %v", key, respTtl)

		if err != nil {
			log.Error("Error trying to retrieve ttl of key " + key, err)
			ttlMap[key] = -2

		} else {
			intResp, _ := respTtl.(int64)
			ttlMap[key] = int(intResp)
		}

	}

	return ttlMap
}

//transfer the ttl information from cacheRegistry to paylaod interface, if it is ExposeTTL
func setTTLToPayload(cacheRegistry *CacheRegistry) CacheRegistry {

	payload := cacheRegistry.Payload

	exposeTTL, hasTtl := payload.(ExposeTTL)

	if hasTtl {
		log.Debug("Transfering ttl from redis (%d seconds) registry to ttl attribute of object %s", cacheRegistry.Ttl, cacheRegistry.CacheKey)
		payload = exposeTTL.SetTtl(cacheRegistry.Ttl) // assure the same type, from set ttl
		cacheRegistry.Payload = payload
		log.Debug("Setting ttl to %v, ttl value %v", cacheRegistry.CacheKey, exposeTTL.GetTtl())
	} else {
		log.Debug("Payload doesn't ExposeTTL %v", cacheRegistry.CacheKey)
	}

	return *cacheRegistry
}

//save informed registries on redis
func (s RedisCacheStorage) SetValues(registries ...CacheRegistry) error {

	var cacheRegistry CacheRegistry
	var index int

	defer func(cacheRegistry *CacheRegistry) {
		if r := recover(); r != nil {
			log.Error("Error trying to save cacheRegistry! recover= %v", r)
		}
	}(&cacheRegistry)

	conn := s.redisPool.Get()
	defer conn.Close()

	keyValPairs := make([]interface{}, 2 * len(registries))

	//prepare a keyval pair array
	for index, cacheRegistry = range registries {

		if len(cacheRegistry.CacheKey) == 0 {
			log.Error("cache key vazio !!!")
			//panic(errors.New("cache key vazio"))
		}

		var bytes = []byte{}
		bytes, err := s.Serializer.MarshalMsg(cacheRegistry,bytes)
		if(err!=nil){
			return err
		}

		if len(bytes) == 0 {
			log.Error("Error trying to decode value for key %v", cacheRegistry.CacheKey)
		}

		keyValPairs[(index * 2)] = s.getKey(cacheRegistry.CacheKey)
		keyValPairs[(index * 2) + 1] = bytes

	}

	_, errDo := conn.Do("MSET", keyValPairs...)
	if errDo != nil {
		log.Error("Error trying to save registry! %v %v", s.getKey(cacheRegistry.CacheKey), errDo)
		return errDo
	} else {
		log.Debug("Updating cache reg key %v ", s.getKey(cacheRegistry.CacheKey))
	}

	errF := conn.Flush()
	if errF != nil {
		log.Error("Error trying to flush connection! %v", errF)
		return errF
	}
	s.SetExpireTTL(registries...)
	return nil
}

//set defined ttl to the cache registries
func (s RedisCacheStorage) SetExpireTTL(cacheRegistries ...CacheRegistry) {
	conn := s.redisPool.Get()
	defer conn.Close()

	//prepare a keyval pair array
	for _, cacheRegistry := range cacheRegistries {
		if cacheRegistry.GetTTL() > 0 {
			//log.Debug("Setting ttl to key %s ", cacheRegistry.CacheKey)
			_, err := conn.Do("expire", s.getKey(cacheRegistry.CacheKey), cacheRegistry.GetTTL())
			if err != nil {
				log.Error("Error trying to save cache registry w! %v", err)
				return
			}

		} else {
			log.Debug("TTL for %s, ttl=%d will not be setted! ", s.getKey(cacheRegistry.CacheKey), cacheRegistry.GetTTL())
		}
	}

	err := conn.Flush()
	if err != nil {
		log.Error("Error trying to save cache registry z! %v", err)
		return
	}
}

//delete values from redis
func (s RedisCacheStorage) DeleteValues(cacheKeys ...string) error {

	c := s.redisPool.Get()
	defer func() {
		c.Close()
	}()

	//apply a prefix to cache area
	keys := s.getKeys(cacheKeys)

	reply, err := c.Do("DEL", keys...)
	if err != nil {
		log.Error("Erro ao tentar invalidar registro no cache!", err, reply)
		return err
	}

	return nil
}

//apply a prefix to cache area
func (s RedisCacheStorage) getKey(key string) string {
	var newKey string

	if len(s.cacheAreaa) > 0 {
		newKey = s.cacheAreaa + ":" + key
	} else {
		newKey = key
	}

	return newKey
}

//apply a prefix to cachearea
func (s RedisCacheStorage) getKeys(keys []string) []interface{} {

	newKeys := make([]interface{}, len(keys))

	for index, key := range keys {
		newKey := s.getKey(key)
		newKeys[index] = newKey
	}

	return newKeys
}

//instantiate a new cachestorage redis
func NewRedisCacheStorage(hostPort string, password string, maxIdle int, readTimeout int, ttlReadTimeout int, cacheArea string, serializer Serializer) RedisCacheStorage {

	redisCacheStorage := RedisCacheStorage{
		*newPoolRedis(hostPort, password, maxIdle, readTimeout),
		ttlReadTimeout,
		cacheArea,
		serializer,
	}

	return redisCacheStorage
}

//create a redis connection pool
func newPoolRedis(server, password string, maxIdle int, readTimeout int) *redis.Pool {

	return &redis.Pool{
		MaxIdle:     maxIdle,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {

			c, err := redis.Dial("tcp", server, redis.DialReadTimeout(time.Duration(readTimeout) * time.Millisecond))
			if err != nil {
				log.Error("Erro ao tentar se conectar ao redis! ", err)
				return nil, err
			}

			return c, err
		},

		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}
<|endoftext|>


<commit_before>package gopdf

import (
	"fmt"
	"io"
	"log"
)

type cacheContentImage struct {
	withMask         bool
	maskAngle        float64
	imageAngle       float64
	verticalFlip     bool
	horizontalFlip   bool
	index            int
	x                float64
	y                float64
	pageHeight       float64
	rect             Rect
	crop             *CropOptions
	extGStateIndexes []int
}

func (c *cacheContentImage) write(writer io.Writer, protection *PDFProtection) error {
	width := c.rect.W
	height := c.rect.H

	var angle float64
	if c.withMask {
		angle = 0
	} else {
		angle = c.imageAngle
	}

	if angle != 0 {
		w := c.rect.W
		h := c.rect.H

		if c.crop != nil {
			w = c.crop.Width
			h = c.crop.Height
		}

		x := c.x + w/2
		y := c.y + h/2

		cacheRotate := cacheContentRotate{
			x:          x,
			y:          y,
			pageHeight: c.pageHeight,
			angle:      angle,
		}
		if err := cacheRotate.write(writer, protection); err != nil {
			return err
		}

		defer func() {
			resetCacheRotate := cacheContentRotate{isReset: true}

			if err := resetCacheRotate.write(writer, protection); err != nil {
				log.Println(err)
			}
		}()
	}

	contentStream := "q\n"

	for _, extGStateIndex := range c.extGStateIndexes {
		contentStream += fmt.Sprintf("/GS%d gs\n", extGStateIndex)
	}

	if c.horizontalFlip || c.verticalFlip {
		fh := "1"
		if c.horizontalFlip {
			fh = "-1"
		}

		fv := "1"
		if c.verticalFlip {
			fv = "-1"
		}

		contentStream += fmt.Sprintf("%s 0 0 %s 0 0 cm\n", fh, fv)
	}

	if c.crop != nil {
		clippingX := c.x
		if c.horizontalFlip {
			clippingX = -clippingX - c.crop.Width
		}

		clippingY := c.pageHeight - (c.y + c.crop.Height)
		if c.verticalFlip {
			clippingY = -clippingY - c.crop.Height
		}

		startX := c.x - c.crop.X
		if c.horizontalFlip {
			startX = -startX - width
		}

		startY := c.pageHeight - (c.y - c.crop.Y + c.rect.H)
		if c.verticalFlip {
			startY = -startY - height
		}

		contentStream += fmt.Sprintf("%0.2f %0.2f %0.2f %0.2f re W* n\n", clippingX, clippingY, c.crop.Width, c.crop.Height)

		var rotateMat string
		if c.maskAngle != 0 {
			x := c.x + width/2
			y := c.y + height/2

			rotateMat = computeRotateTransformationMatrix(x, y, c.imageAngle, c.pageHeight)
		}

		contentStream += fmt.Sprintf("q\n %s %0.2f 0 0\n %0.2f %0.2f %0.2f cm /I%d Do \nQ\n", rotateMat, width, height, startX, startY, c.index+1)
	} else {
		x := c.x
		y := c.pageHeight - (c.y + height)

		if c.horizontalFlip {
			x = -x - width
		}

		if c.verticalFlip {
			y = -y - height
		}

		var rotateMat string
		if c.maskAngle != 0 {
			rotatedX := c.x + width/2
			rotatedY := c.y + height/2

			rotateMat = computeRotateTransformationMatrix(rotatedX, rotatedY, c.maskAngle+c.imageAngle, c.pageHeight)
		}

		contentStream += fmt.Sprintf("q\n %s %0.2f 0 0\n %0.2f %0.2f %0.2f cm\n /I%d Do \nQ\n", rotateMat, width, height, x, y, c.index+1)
	}

	contentStream += "Q\n"

	if _, err := io.WriteString(writer, contentStream); err != nil {
		return err
	}

	return nil
}
<commit_msg>refactor<commit_after>package gopdf

import (
	"fmt"
	"io"
)

type cacheContentImage struct {
	withMask         bool
	maskAngle        float64
	imageAngle       float64
	verticalFlip     bool
	horizontalFlip   bool
	index            int
	x                float64
	y                float64
	pageHeight       float64
	rect             Rect
	crop             *CropOptions
	extGStateIndexes []int
}

func (c *cacheContentImage) openImageRotateTrMt(writer io.Writer, protection *PDFProtection) error {
	w := c.rect.W
	h := c.rect.H

	if c.crop != nil {
		w = c.crop.Width
		h = c.crop.Height
	}

	x := c.x + w/2
	y := c.y + h/2

	cacheRotate := cacheContentRotate{
		x:          x,
		y:          y,
		pageHeight: c.pageHeight,
		angle:      c.imageAngle,
	}
	if err := cacheRotate.write(writer, protection); err != nil {
		return err
	}

	return nil
}

func (c *cacheContentImage) closeImageRotateTrMt(writer io.Writer, protection *PDFProtection) error {
	resetCacheRotate := cacheContentRotate{isReset: true}

	return resetCacheRotate.write(writer, protection)
}

func (c *cacheContentImage) computeMaskImageRotateTrMt() string {
	x := c.x + c.rect.W/2
	y := c.y + c.rect.H/2

	rotateMat := computeRotateTransformationMatrix(x, y, c.maskAngle+c.imageAngle, c.pageHeight)

	return rotateMat
}

func (c *cacheContentImage) write(writer io.Writer, protection *PDFProtection) error {
	width := c.rect.W
	height := c.rect.H

	if !c.withMask {
		if err := c.openImageRotateTrMt(writer, protection); err != nil {
			return err
		}

		defer c.closeImageRotateTrMt(writer, protection)
	}

	contentStream := "q\n"

	for _, extGStateIndex := range c.extGStateIndexes {
		contentStream += fmt.Sprintf("/GS%d gs\n", extGStateIndex)
	}

	if c.horizontalFlip || c.verticalFlip {
		fh := "1"
		if c.horizontalFlip {
			fh = "-1"
		}

		fv := "1"
		if c.verticalFlip {
			fv = "-1"
		}

		contentStream += fmt.Sprintf("%s 0 0 %s 0 0 cm\n", fh, fv)
	}

	if c.crop != nil {
		clippingX := c.x
		if c.horizontalFlip {
			clippingX = -clippingX - c.crop.Width
		}

		clippingY := c.pageHeight - (c.y + c.crop.Height)
		if c.verticalFlip {
			clippingY = -clippingY - c.crop.Height
		}

		startX := c.x - c.crop.X
		if c.horizontalFlip {
			startX = -startX - width
		}

		startY := c.pageHeight - (c.y - c.crop.Y + c.rect.H)
		if c.verticalFlip {
			startY = -startY - height
		}

		contentStream += fmt.Sprintf("%0.2f %0.2f %0.2f %0.2f re W* n\n", clippingX, clippingY, c.crop.Width, c.crop.Height)

		var rotateMat string
		if c.maskAngle != 0 {
			rotateMat = c.computeMaskImageRotateTrMt()
		}

		contentStream += fmt.Sprintf("q\n %s %0.2f 0 0\n %0.2f %0.2f %0.2f cm /I%d Do \nQ\n", rotateMat, width, height, startX, startY, c.index+1)
	} else {
		x := c.x
		y := c.pageHeight - (c.y + height)

		if c.horizontalFlip {
			x = -x - width
		}

		if c.verticalFlip {
			y = -y - height
		}

		var rotateMat string
		if c.maskAngle != 0 {
			rotateMat = c.computeMaskImageRotateTrMt()
		}

		contentStream += fmt.Sprintf("q\n %s %0.2f 0 0\n %0.2f %0.2f %0.2f cm\n /I%d Do \nQ\n", rotateMat, width, height, x, y, c.index+1)
	}

	contentStream += "Q\n"

	if _, err := io.WriteString(writer, contentStream); err != nil {
		return err
	}

	return nil
}
<|endoftext|>