<commit_before>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	iptablesRules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	return iptablesRules, nil
}

// k8sPodDownHandler cleans up after pod deleted.
func (a *Agent) k8sPodDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infoln("Agent: Entering k8sPodDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)

	return "OK", nil
}

// k8sPodUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) k8sPodUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering k8sPodUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.k8sPodUpHandle(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// index handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
// TODO index should be reserved for an actual index, while this function
// need to be renamed as interfaceHandler and need to respond on it's own url.
func (a *Agent) index(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.interfaceHandle(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// k8sPodUpHandle does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) k8sPodUpHandle(netReq NetworkRequest) error {
	glog.Info("Agent: Entering k8sPodUpHandle()")

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}

	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)

	// Allow ICMP, DHCP and SSH between host and instances.
	inboundChain := chainNames[firewall.InputChainIndex]
	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, inboundRule))

	hostAddr := a.networkConfig.RomanaGW()
	outboundChain := chainNames[firewall.OutputChainIndex]
	outboundRule := firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s -s %s/32 -p udp -m udp --sport 67 --dport 68", outboundChain, hostAddr))

	// forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody("-m comment --comment Outgoing")

	// forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody("-m state --state RELATED,ESTABLISHED")

	fw.SetDefaultRules([]firewall.FirewallRule{inboundRule, outboundRule, forwardInRule, forwardOutRule})

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

// interfaceHandle does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) interfaceHandle(netif NetIf) error {
	glog.Info("Agent: processing request to provision new interface")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}
<commit_msg>Firewall interface for defining default rules (13)<commit_after>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	iptablesRules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	return iptablesRules, nil
}

// k8sPodDownHandler cleans up after pod deleted.
func (a *Agent) k8sPodDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infoln("Agent: Entering k8sPodDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)

	return "OK", nil
}

// k8sPodUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) k8sPodUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering k8sPodUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.k8sPodUpHandle(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// index handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
// TODO index should be reserved for an actual index, while this function
// need to be renamed as interfaceHandler and need to respond on it's own url.
func (a *Agent) index(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.interfaceHandle(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// k8sPodUpHandle does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) k8sPodUpHandle(netReq NetworkRequest) error {
	glog.Info("Agent: Entering k8sPodUpHandle()")

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}

	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)

	// Allow ICMP, DHCP and SSH between host and instances.
	inboundChain := chainNames[firewall.InputChainIndex]
	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, inboundRule))

	hostAddr := a.networkConfig.RomanaGW()
	outboundChain := chainNames[firewall.OutputChainIndex]
	outboundRule := firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s -s %s/32 -p udp -m udp --sport 67 --dport 68", outboundChain, hostAddr))

	// forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody("-m comment --comment Outgoing")

	// forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody("-m state --state RELATED,ESTABLISHED")

	fw.SetDefaultRules([]firewall.FirewallRule{*inboundRule, *outboundRule, *forwardInRule, *forwardOutRule})

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

// interfaceHandle does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) interfaceHandle(netif NetIf) error {
	glog.Info("Agent: processing request to provision new interface")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}
<|endoftext|>


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

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"syscall"

	"github.com/romana/core/common/api"
	"github.com/romana/core/common/client"

	"github.com/docker/libkv/store"
	log "github.com/romana/rlog"
	"github.com/vishvananda/netlink"
)

type ExternalIP struct {
	IP string `json:"ip" form:"ip"`
}

func GetDefaultLink() (netlink.Link, error) {
	defaultR := netlink.Route{}

	routes, err := netlink.RouteList(nil, syscall.AF_INET)
	if err != nil {
		return nil, fmt.Errorf("error finding default route: %s", err)
	}

	for _, r := range routes {
		// If dst/src is not specified for a route, then it
		// means a default route is found which handles packets
		// for everything which is not handled by specific routes.
		if r.Src == nil && r.Dst == nil {
			defaultR = r
			break
		}
	}

	// if default route is not found above, then we have
	// LinkIndex as zero, but LinkIndex start from one, so
	// this will error out, which should be the case, since
	// anyways we couldn't find default link, and other
	// links maybe lo, etc which can't be used and thus
	// error out and return below.
	link, err := netlink.LinkByIndex(defaultR.LinkIndex)
	if err != nil {
		return nil, err
	}
	if link == nil {
		return nil, errors.New("error, could not locate default link for host")
	}

	return link, nil
}

// GetIPs returns all the IPv4 Address attached to link.
func GetIPs(link netlink.Link) ([]string, error) {
	var addresses []string

	addrList, err := netlink.AddrList(link, syscall.AF_INET)
	if err != nil {
		return nil, fmt.Errorf("error finding IP adrress for link (%s): %s",
			link.Attrs().Name, err)
	}

	for _, addr := range addrList {
		addresses = append(addresses, addr.IP.String())
	}

	if len(addresses) < 1 {
		return nil, fmt.Errorf("error finding IP adrress for link (%s)",
			link.Attrs().Name)
	}

	return addresses, nil
}

func linkAddDeleteIP(kvpair *store.KVPairExt, toAdd bool,
	defaultLink netlink.Link, defaultLinkAddressList []string) error {
	var value string
	var IPAddressOnThisNode bool

	if kvpair == nil || (kvpair.Value == "" && kvpair.PrevValue == "") {
		return fmt.Errorf("error retrieving value from the event notification")
	}

	if kvpair.Value != "" {
		value = kvpair.Value
	} else if kvpair.PrevValue != "" {
		value = kvpair.PrevValue
	}

	exposedIP := api.ExposedIPSpec{}
	if err := json.Unmarshal([]byte(value), &exposedIP); err != nil {
		return fmt.Errorf("error retrieving value from the event notification: %s", err)
	}

	if exposedIP.NodeIPAddress == "" {
		return fmt.Errorf("error finding node IP Address for romanaIP")
	}

	if exposedIP.RomanaIP.IP == "" {
		return fmt.Errorf("romana IP error or romana IP not found")
	}

	for i := range defaultLinkAddressList {
		if defaultLinkAddressList[i] == exposedIP.NodeIPAddress {
			IPAddressOnThisNode = true
			break
		}
	}

	if !IPAddressOnThisNode {
		fmt.Printf("romanaIP not for this node, skipping processing it")
		return nil
	}

	ipAddress, err := netlink.ParseAddr(exposedIP.RomanaIP.IP + "/32")
	if err != nil {
		return fmt.Errorf("error parsing romana IP: %s", err)
	}

	if toAdd {
		return netlink.AddrAdd(defaultLink, ipAddress)
	}
	return netlink.AddrDel(defaultLink, ipAddress)
}

func StartRomanaIPSync(ctx context.Context, store *client.Store,
	defaultLink netlink.Link) error {
	var err error

	if store == nil || ctx == nil || defaultLink == nil {
		return fmt.Errorf("error store/context or link empty")
	}

	defaultLinkAddressList, err := GetIPs(defaultLink)
	if err != nil {
		return fmt.Errorf("failed to get default link's IP address: %s\n", err)
	}
	if len(defaultLinkAddressList) < 1 {
		return fmt.Errorf("failed to get default link's IP address")
	}

	go romanaIPWatcher(ctx, store, defaultLink, defaultLinkAddressList)

	return nil
}

func romanaIPWatcher(ctx context.Context, store *client.Store,
	defaultLink netlink.Link, defaultLinkAddressList []string) {

	// TODO: event stream could be broken it store connection is
	// broken so add support for re-watching the events stream here.
	events, err := store.WatchTreeExt(client.DefaultEtcdPrefix+client.RomanaIPPrefix,
		ctx.Done())
	if err != nil {
		log.Println("Error watching kvstore romanaIP keys.")
		return
	}

	for {
		select {
		case pair := <-events:
			switch pair.Action {
			case "create", "set", "update", "compareAndSwap":
				fmt.Printf("creating/updating romanaIP: %#v\n", pair)
				err := linkAddDeleteIP(pair, true, defaultLink, defaultLinkAddressList)
				if err != nil {
					log.Errorf("error adding romanaIP to the link: %s", err)
					continue
				}
			case "delete":
				if pair.Dir {
					// TODO: handle deleting all romanaIPs here.
					log.Printf("deleting ALL romanaIPs: %#v\n", pair)
				} else {
					log.Printf("deleting romanaIP: %#v\n", pair)
					err := linkAddDeleteIP(pair, false, defaultLink, defaultLinkAddressList)
					if err != nil {
						log.Errorf("error deleting romanaIP from the link: %s", err)
						continue
					}
				}
			default:
			}
		case <-ctx.Done():
			fmt.Print("\nStopping romanaIP watcher module.\n")
			return
		}
	}
}
<commit_msg>agent: update log messages to make them more clear.<commit_after>// Copyright (c) 2017 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"syscall"

	"github.com/romana/core/common/api"
	"github.com/romana/core/common/client"

	"github.com/docker/libkv/store"
	log "github.com/romana/rlog"
	"github.com/vishvananda/netlink"
)

type ExternalIP struct {
	IP string `json:"ip" form:"ip"`
}

func GetDefaultLink() (netlink.Link, error) {
	defaultR := netlink.Route{}

	routes, err := netlink.RouteList(nil, syscall.AF_INET)
	if err != nil {
		return nil, fmt.Errorf("error finding default route: %s", err)
	}

	for _, r := range routes {
		// If dst/src is not specified for a route, then it
		// means a default route is found which handles packets
		// for everything which is not handled by specific routes.
		if r.Src == nil && r.Dst == nil {
			defaultR = r
			break
		}
	}

	// if default route is not found above, then we have
	// LinkIndex as zero, but LinkIndex start from one, so
	// this will error out, which should be the case, since
	// anyways we couldn't find default link, and other
	// links maybe lo, etc which can't be used and thus
	// error out and return below.
	link, err := netlink.LinkByIndex(defaultR.LinkIndex)
	if err != nil {
		return nil, err
	}
	if link == nil {
		return nil, errors.New("error, could not locate default link for host")
	}

	return link, nil
}

// GetIPs returns all the IPv4 Address attached to link.
func GetIPs(link netlink.Link) ([]string, error) {
	var addresses []string

	addrList, err := netlink.AddrList(link, syscall.AF_INET)
	if err != nil {
		return nil, fmt.Errorf("error finding IP adrress for link (%s): %s",
			link.Attrs().Name, err)
	}

	for _, addr := range addrList {
		addresses = append(addresses, addr.IP.String())
	}

	if len(addresses) < 1 {
		return nil, fmt.Errorf("error finding IP adrress for link (%s)",
			link.Attrs().Name)
	}

	return addresses, nil
}

func linkAddDeleteIP(kvpair *store.KVPairExt, toAdd bool,
	defaultLink netlink.Link, defaultLinkAddressList []string) error {
	var value string
	var IPAddressOnThisNode bool

	if kvpair == nil || (kvpair.Value == "" && kvpair.PrevValue == "") {
		return fmt.Errorf("error retrieving value from the event notification")
	}

	if kvpair.Value != "" {
		value = kvpair.Value
	} else if kvpair.PrevValue != "" {
		value = kvpair.PrevValue
	}

	exposedIP := api.ExposedIPSpec{}
	if err := json.Unmarshal([]byte(value), &exposedIP); err != nil {
		return fmt.Errorf("error retrieving value from the event notification: %s", err)
	}

	if exposedIP.NodeIPAddress == "" {
		return fmt.Errorf("error finding node IP Address for romanaIP")
	}

	if exposedIP.RomanaIP.IP == "" {
		return fmt.Errorf("romana IP error or romana IP not found")
	}

	for i := range defaultLinkAddressList {
		if defaultLinkAddressList[i] == exposedIP.NodeIPAddress {
			IPAddressOnThisNode = true
			break
		}
	}

	if !IPAddressOnThisNode {
		log.Info("romanaIP not for this node, skipping processing it")
		return nil
	}

	ipAddress, err := netlink.ParseAddr(exposedIP.RomanaIP.IP + "/32")
	if err != nil {
		return fmt.Errorf("error parsing romana IP: %s", err)
	}

	if toAdd {
		return netlink.AddrAdd(defaultLink, ipAddress)
	}
	return netlink.AddrDel(defaultLink, ipAddress)
}

func StartRomanaIPSync(ctx context.Context, store *client.Store,
	defaultLink netlink.Link) error {
	var err error

	if store == nil || ctx == nil || defaultLink == nil {
		return fmt.Errorf("error store/context or link empty")
	}

	defaultLinkAddressList, err := GetIPs(defaultLink)
	if err != nil {
		return fmt.Errorf("failed to get default link's IP address: %s\n", err)
	}
	if len(defaultLinkAddressList) < 1 {
		return fmt.Errorf("failed to get default link's IP address")
	}

	go romanaIPWatcher(ctx, store, defaultLink, defaultLinkAddressList)

	return nil
}

func romanaIPWatcher(ctx context.Context, store *client.Store,
	defaultLink netlink.Link, defaultLinkAddressList []string) {

	// TODO: event stream could be broken it store connection is
	// broken so add support for re-watching the events stream here.
	events, err := store.WatchTreeExt(client.DefaultEtcdPrefix+client.RomanaIPPrefix,
		ctx.Done())
	if err != nil {
		log.Error("Error watching kvstore romanaIP keys.")
		return
	}

	for {
		select {
		case pair := <-events:
			switch pair.Action {
			case "create", "set", "update", "compareAndSwap":
				log.Debugf("creating/updating romanaIP: %#v\n", pair)
				err := linkAddDeleteIP(pair, true, defaultLink, defaultLinkAddressList)
				if err != nil {
					log.Errorf("error adding romanaIP to the link: %s", err)
					continue
				}
			case "delete":
				if pair.Dir {
					// TODO: currently if the whole "/romana/romanaip" kvstore
					// directory is deleted, then we need to delete all romanaIPs,
					// but currently we do nothing here and handle only single
					// romanaIP deletion event below.
					log.Infof("should be deleting ALL romanaIPs(%#v) here, ignoring currently",
						pair)
				} else {
					log.Debugf("deleting romanaIP: %#v\n", pair)
					err := linkAddDeleteIP(pair, false, defaultLink, defaultLinkAddressList)
					if err != nil {
						log.Errorf("error deleting romanaIP from the link: %s", err)
						continue
					}
				}
			default:
				log.Infof("missed romanaIP event type: %s", pair.Action)
			}
		case <-ctx.Done():
			log.Printf("\nStopping romanaIP watcher module.\n")
			return
		}
	}
}
<|endoftext|>


<commit_before>package awsecs

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecs"
)

// a wrapper around an AWS client that makes all the needed calls and just exposes the final results
type ecsClient struct {
	client  *ecs.ECS
	cluster string
}

type ecsInfo struct {
	tasks          map[string]*ecs.Task
	services       map[string]*ecs.Service
	taskServiceMap map[string]string
}

func newClient(cluster string) (*ecsClient, error) {
	sess := session.New()

	region, err := ec2metadata.New(sess).Region()
	if err != nil {
		return nil, err
	}

	return &ecsClient{
		client:  ecs.New(sess, &aws.Config{Region: aws.String(region)}),
		cluster: cluster,
	}, nil
}

// returns a map from deployment ids to service names
func (c ecsClient) getDeploymentMap(services map[string]*ecs.Service) map[string]string {
	results := map[string]string{}
	for serviceName, service := range services {
		for _, deployment := range service.Deployments {
			results[*deployment.Id] = serviceName
		}
	}
	return results
}

// cannot fail as it will attempt to deliver partial results, though that may end up being no results
func (c ecsClient) getServices() map[string]*ecs.Service {
	results := map[string]*ecs.Service{}
	lock := sync.Mutex{} // lock mediates access to results

	group := sync.WaitGroup{}

	err := c.client.ListServicesPages(
		&ecs.ListServicesInput{Cluster: &c.cluster},
		func(page *ecs.ListServicesOutput, lastPage bool) bool {
			// describe each page of 10 (the max for one describe command) concurrently
			group.Add(1)
			serviceArns := page.ServiceArns
			go func() {
				defer group.Done()

				resp, err := c.client.DescribeServices(&ecs.DescribeServicesInput{
					Cluster:  &c.cluster,
					Services: serviceArns,
				})
				if err != nil {
					log.Warnf("Error describing some ECS services, ECS service report may be incomplete: %v", err)
					return
				}

				for _, failure := range resp.Failures {
					log.Warnf("Failed to describe ECS service %s, ECS service report may be incomplete: %s", failure.Arn, failure.Reason)
				}

				lock.Lock()
				for _, service := range resp.Services {
					results[*service.ServiceName] = service
				}
				lock.Unlock()
			}()
			return true
		},
	)
	group.Wait()

	if err != nil {
		log.Warnf("Error listing ECS services, ECS service report may be incomplete: %v", err)
	}
	return results
}

func (c ecsClient) getTasks(taskArns []string) (map[string]*ecs.Task, error) {
	taskPtrs := make([]*string, len(taskArns))
	for i := range taskArns {
		taskPtrs[i] = &taskArns[i]
	}

	// You'd think there's a limit on how many tasks can be described here,
	// but the docs don't mention anything.
	resp, err := c.client.DescribeTasks(&ecs.DescribeTasksInput{
		Cluster: &c.cluster,
		Tasks:   taskPtrs,
	})
	if err != nil {
		return nil, err
	}

	for _, failure := range resp.Failures {
		log.Warnf("Failed to describe ECS task %s, ECS service report may be incomplete: %s", failure.Arn, failure.Reason)
	}

	results := make(map[string]*ecs.Task, len(resp.Tasks))
	for _, task := range resp.Tasks {
		results[*task.TaskArn] = task
	}
	return results, nil
}

// returns a map from task ARNs to service names
func (c ecsClient) getInfo(taskArns []string) (ecsInfo, error) {
	servicesChan := make(chan map[string]*ecs.Service)
	go func() {
		servicesChan <- c.getServices()
	}()

	// do these two fetches in parallel
	tasks, err := c.getTasks(taskArns)
	services := <-servicesChan

	if err != nil {
		return ecsInfo{}, err
	}

	deploymentMap := c.getDeploymentMap(services)

	taskServiceMap := map[string]string{}
	for taskArn, task := range tasks {
		// Note not all tasks map to a deployment, or we could otherwise mismatch due to races.
		// It's safe to just ignore all these cases and consider them "non-service" tasks.
		if serviceName, ok := deploymentMap[*task.StartedBy]; ok {
			taskServiceMap[taskArn] = serviceName
		}
	}

	return ecsInfo{services: services, tasks: tasks, taskServiceMap: taskServiceMap}, nil
}
<commit_msg>ecs reporter: Fix some log lines that were passing *string instead of string<commit_after>package awsecs

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecs"
)

// a wrapper around an AWS client that makes all the needed calls and just exposes the final results
type ecsClient struct {
	client  *ecs.ECS
	cluster string
}

type ecsInfo struct {
	tasks          map[string]*ecs.Task
	services       map[string]*ecs.Service
	taskServiceMap map[string]string
}

func newClient(cluster string) (*ecsClient, error) {
	sess := session.New()

	region, err := ec2metadata.New(sess).Region()
	if err != nil {
		return nil, err
	}

	return &ecsClient{
		client:  ecs.New(sess, &aws.Config{Region: aws.String(region)}),
		cluster: cluster,
	}, nil
}

// returns a map from deployment ids to service names
func (c ecsClient) getDeploymentMap(services map[string]*ecs.Service) map[string]string {
	results := map[string]string{}
	for serviceName, service := range services {
		for _, deployment := range service.Deployments {
			results[*deployment.Id] = serviceName
		}
	}
	return results
}

// cannot fail as it will attempt to deliver partial results, though that may end up being no results
func (c ecsClient) getServices() map[string]*ecs.Service {
	results := map[string]*ecs.Service{}
	lock := sync.Mutex{} // lock mediates access to results

	group := sync.WaitGroup{}

	err := c.client.ListServicesPages(
		&ecs.ListServicesInput{Cluster: &c.cluster},
		func(page *ecs.ListServicesOutput, lastPage bool) bool {
			// describe each page of 10 (the max for one describe command) concurrently
			group.Add(1)
			serviceArns := page.ServiceArns
			go func() {
				defer group.Done()

				resp, err := c.client.DescribeServices(&ecs.DescribeServicesInput{
					Cluster:  &c.cluster,
					Services: serviceArns,
				})
				if err != nil {
					log.Warnf("Error describing some ECS services, ECS service report may be incomplete: %v", err)
					return
				}

				for _, failure := range resp.Failures {
					log.Warnf("Failed to describe ECS service %s, ECS service report may be incomplete: %s", *failure.Arn, *failure.Reason)
				}

				lock.Lock()
				for _, service := range resp.Services {
					results[*service.ServiceName] = service
				}
				lock.Unlock()
			}()
			return true
		},
	)
	group.Wait()

	if err != nil {
		log.Warnf("Error listing ECS services, ECS service report may be incomplete: %v", err)
	}
	return results
}

func (c ecsClient) getTasks(taskArns []string) (map[string]*ecs.Task, error) {
	taskPtrs := make([]*string, len(taskArns))
	for i := range taskArns {
		taskPtrs[i] = &taskArns[i]
	}

	// You'd think there's a limit on how many tasks can be described here,
	// but the docs don't mention anything.
	resp, err := c.client.DescribeTasks(&ecs.DescribeTasksInput{
		Cluster: &c.cluster,
		Tasks:   taskPtrs,
	})
	if err != nil {
		return nil, err
	}

	for _, failure := range resp.Failures {
		log.Warnf("Failed to describe ECS task %s, ECS service report may be incomplete: %s", *failure.Arn, *failure.Reason)
	}

	results := make(map[string]*ecs.Task, len(resp.Tasks))
	for _, task := range resp.Tasks {
		results[*task.TaskArn] = task
	}
	return results, nil
}

// returns a map from task ARNs to service names
func (c ecsClient) getInfo(taskArns []string) (ecsInfo, error) {
	servicesChan := make(chan map[string]*ecs.Service)
	go func() {
		servicesChan <- c.getServices()
	}()

	// do these two fetches in parallel
	tasks, err := c.getTasks(taskArns)
	services := <-servicesChan

	if err != nil {
		return ecsInfo{}, err
	}

	deploymentMap := c.getDeploymentMap(services)

	taskServiceMap := map[string]string{}
	for taskArn, task := range tasks {
		// Note not all tasks map to a deployment, or we could otherwise mismatch due to races.
		// It's safe to just ignore all these cases and consider them "non-service" tasks.
		if serviceName, ok := deploymentMap[*task.StartedBy]; ok {
			taskServiceMap[taskArn] = serviceName
		}
	}

	return ecsInfo{services: services, tasks: tasks, taskServiceMap: taskServiceMap}, nil
}
<|endoftext|>


<commit_before>package proc

// #include "threads_darwin.h"
import "C"
import (
	"fmt"
	"unsafe"
)

type OSSpecificDetails struct {
	thread_act C.thread_act_t
	registers  C.x86_thread_state64_t
}

func (t *Thread) Halt() error {
	var kret C.kern_return_t
	kret = C.thread_suspend(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not suspend thread %d", t.Id)
	}
	t.running = false
	return nil
}

func (t *Thread) singleStep() error {
	kret := C.single_step(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not single step")
	}
	t.dbp.trapWait(0)
	kret = C.clear_trap_flag(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not clear CPU trap flag")
	}
	return nil
}

func (t *Thread) resume() error {
	t.running = true
	// TODO(dp) set flag for ptrace stops
	var err error
	t.dbp.execPtraceFunc(func() { err = PtraceCont(t.dbp.Pid, 0) })
	if err == nil {
		return nil
	}
	kret := C.resume_thread(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not continue thread")
	}
	return nil
}

func (t *Thread) blocked() bool {
	// TODO(dp) cache the func pc to remove this lookup
	pc, _ := t.PC()
	fn := t.dbp.goSymTable.PCToFunc(pc)
	if fn != nil && (fn.Name == "runtime.mach_semaphore_wait" || fn.Name == "runtime.usleep") {
		return true
	}
	return false
}

func writeMemory(thread *Thread, addr uintptr, data []byte) (int, error) {
	if len(data) == 0 {
		return 0, nil
	}
	var (
		vm_data = unsafe.Pointer(&data[0])
		vm_addr = C.mach_vm_address_t(addr)
		length  = C.mach_msg_type_number_t(len(data))
	)
	if ret := C.write_memory(thread.dbp.os.task, vm_addr, vm_data, length); ret < 0 {
		return 0, fmt.Errorf("could not write memory")
	}
	return len(data), nil
}

func readMemory(thread *Thread, addr uintptr, data []byte) (int, error) {
	if len(data) == 0 {
		return 0, nil
	}
	var (
		vm_data = unsafe.Pointer(&data[0])
		vm_addr = C.mach_vm_address_t(addr)
		length  = C.mach_msg_type_number_t(len(data))
	)

	ret := C.read_memory(thread.dbp.os.task, vm_addr, vm_data, length)
	if ret < 0 {
		return 0, fmt.Errorf("could not read memory")
	}
	return len(data), nil
}
<commit_msg>Handle thread blocked on kevent<commit_after>package proc

// #include "threads_darwin.h"
import "C"
import (
	"fmt"
	"unsafe"
)

type OSSpecificDetails struct {
	thread_act C.thread_act_t
	registers  C.x86_thread_state64_t
}

func (t *Thread) Halt() error {
	var kret C.kern_return_t
	kret = C.thread_suspend(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not suspend thread %d", t.Id)
	}
	t.running = false
	return nil
}

func (t *Thread) singleStep() error {
	kret := C.single_step(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not single step")
	}
	t.dbp.trapWait(0)
	kret = C.clear_trap_flag(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not clear CPU trap flag")
	}
	return nil
}

func (t *Thread) resume() error {
	t.running = true
	// TODO(dp) set flag for ptrace stops
	var err error
	t.dbp.execPtraceFunc(func() { err = PtraceCont(t.dbp.Pid, 0) })
	if err == nil {
		return nil
	}
	kret := C.resume_thread(t.os.thread_act)
	if kret != C.KERN_SUCCESS {
		return fmt.Errorf("could not continue thread")
	}
	return nil
}

func (t *Thread) blocked() bool {
	// TODO(dp) cache the func pc to remove this lookup
	pc, _ := t.PC()
	fn := t.dbp.goSymTable.PCToFunc(pc)
	if fn == nil {
		return false
	}
	switch fn.Name {
	case "runtime.kevent", "runtime.mach_semaphore_wait", "runtime.usleep":
		return true
	default:
		return false
	}
}

func writeMemory(thread *Thread, addr uintptr, data []byte) (int, error) {
	if len(data) == 0 {
		return 0, nil
	}
	var (
		vm_data = unsafe.Pointer(&data[0])
		vm_addr = C.mach_vm_address_t(addr)
		length  = C.mach_msg_type_number_t(len(data))
	)
	if ret := C.write_memory(thread.dbp.os.task, vm_addr, vm_data, length); ret < 0 {
		return 0, fmt.Errorf("could not write memory")
	}
	return len(data), nil
}

func readMemory(thread *Thread, addr uintptr, data []byte) (int, error) {
	if len(data) == 0 {
		return 0, nil
	}
	var (
		vm_data = unsafe.Pointer(&data[0])
		vm_addr = C.mach_vm_address_t(addr)
		length  = C.mach_msg_type_number_t(len(data))
	)

	ret := C.read_memory(thread.dbp.os.task, vm_addr, vm_data, length)
	if ret < 0 {
		return 0, fmt.Errorf("could not read memory")
	}
	return len(data), nil
}
<|endoftext|>


<commit_before>package processor

import (
	"errors"
	"fmt"
	"log"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/time/rate"

	"gopkg.in/queue.v1"
	"gopkg.in/queue.v1/internal"
)

const consumerBackoff = time.Second
const maxBackoff = 12 * time.Hour

var ErrNotSupported = errors.New("processor: not supported")

type Delayer interface {
	Delay() time.Duration
}

type Stats struct {
	InFlight    uint32
	Deleting    uint32
	Processed   uint32
	Retries     uint32
	Fails       uint32
	AvgDuration time.Duration
}

type Processor struct {
	q   Queuer
	opt *queue.Options

	handler         queue.Handler
	fallbackHandler queue.Handler

	ch chan *queue.Message
	wg sync.WaitGroup

	delBatch *internal.Batcher

	_started uint32
	stop     chan struct{}

	inFlight    uint32
	deleting    uint32
	processed   uint32
	fails       uint32
	retries     uint32
	avgDuration uint32
}

func New(q Queuer, opt *queue.Options) *Processor {
	initOptions(opt)

	p := &Processor{
		q:   q,
		opt: opt,

		ch: make(chan *queue.Message, opt.BufferSize),
	}
	p.setHandler(opt.Handler)
	if opt.FallbackHandler != nil {
		p.setFallbackHandler(opt.FallbackHandler)
	}
	p.delBatch = internal.NewBatcher(opt.Scavengers, p.deleteBatch)
	return p
}

func initOptions(opt *queue.Options) {
	if opt.Workers == 0 {
		opt.Workers = 10 * runtime.NumCPU()
	}
	if opt.Scavengers == 0 {
		opt.Scavengers = runtime.NumCPU() + 1
	}
	if opt.BufferSize == 0 {
		opt.BufferSize = opt.Workers
		if opt.BufferSize > 10 {
			opt.BufferSize = 10
		}
	}
	if opt.RateLimit == 0 {
		opt.RateLimit = rate.Inf
	}
	if opt.Retries == 0 {
		opt.Retries = 10
	}
	if opt.Backoff == 0 {
		opt.Backoff = 3 * time.Second
	}
}

func Start(q Queuer, opt *queue.Options) *Processor {
	p := New(q, opt)
	p.Start()
	return p
}

func (p *Processor) String() string {
	return fmt.Sprintf(
		"Processor<%s workers=%d scavengers=%d buffer=%d>",
		p.q.Name(), p.opt.Workers, p.opt.Scavengers, p.opt.BufferSize,
	)
}

func (p *Processor) Stats() *Stats {
	if p.stopped() {
		return nil
	}
	return &Stats{
		InFlight:    atomic.LoadUint32(&p.inFlight),
		Deleting:    atomic.LoadUint32(&p.deleting),
		Processed:   atomic.LoadUint32(&p.processed),
		Retries:     atomic.LoadUint32(&p.retries),
		Fails:       atomic.LoadUint32(&p.fails),
		AvgDuration: time.Duration(atomic.LoadUint32(&p.avgDuration)) * time.Millisecond,
	}
}

func (p *Processor) setHandler(handler interface{}) {
	p.handler = queue.NewHandler(handler)
}

func (p *Processor) setFallbackHandler(handler interface{}) {
	p.fallbackHandler = queue.NewHandler(handler)
}

func (p *Processor) Add(msg *queue.Message) error {
	p.ch <- msg
	return nil
}

func (p *Processor) Start() error {
	if !atomic.CompareAndSwapUint32(&p._started, 0, 1) {
		return nil
	}

	p.startWorkers()

	p.wg.Add(1)
	go p.messageFetcher()

	return nil
}

func (p *Processor) Stop() error {
	return p.StopTimeout(30 * time.Second)
}

func (p *Processor) StopTimeout(timeout time.Duration) error {
	if !atomic.CompareAndSwapUint32(&p._started, 1, 0) {
		return nil
	}
	p.stopWorkers()
	return p.waitWorkers(timeout)
}

func (p *Processor) startWorkers() {
	p.stop = make(chan struct{})
	p.wg.Add(p.opt.Workers)
	for i := 0; i < p.opt.Workers; i++ {
		go p.worker()
	}
}

func (p *Processor) stopWorkers() {
	close(p.stop)
}

func (p *Processor) waitWorkers(timeout time.Duration) error {
	stopped := make(chan struct{})
	go func() {
		p.wg.Wait()
		close(stopped)
	}()

	select {
	case <-time.After(timeout):
		return fmt.Errorf("workers did not stop after %s seconds", timeout)
	case <-stopped:
		return nil
	}
}

func (p *Processor) stopped() bool {
	return atomic.LoadUint32(&p._started) == 0
}

func (p *Processor) Close() error {
	retErr := p.Stop()
	if err := p.delBatch.Close(); err != nil && retErr == nil {
		retErr = err
	}
	return retErr
}

func (p *Processor) ProcessAll() error {
	p.startWorkers()
	var noWork int
	for {
		isIdle := atomic.LoadUint32(&p.inFlight) == 0
		n, err := p.fetchMessages()
		if err != nil && err != ErrNotSupported {
			return err
		}
		if n == 0 && isIdle {
			noWork++
		} else {
			noWork = 0
		}
		if noWork == 2 {
			break
		}
	}
	p.stopWorkers()
	return p.waitWorkers(time.Minute)
}

func (p *Processor) ProcessOne() error {
	msg, err := p.reserveOne()
	if err != nil {
		return err
	}
	atomic.AddUint32(&p.inFlight, 1)
	return p.Process(msg)
}

func (p *Processor) reserveOne() (*queue.Message, error) {
	select {
	case msg := <-p.ch:
		return msg, nil
	default:
	}

	msgs, err := p.q.ReserveN(1)
	if err != nil {
		return nil, err
	}
	if len(msgs) == 0 {
		return nil, errors.New("no messages in queue")
	}
	return &msgs[0], nil
}

func (p *Processor) messageFetcher() {
	defer p.wg.Done()
	for {
		if p.stopped() {
			break
		}

		_, err := p.fetchMessages()
		if err != nil {
			if err == ErrNotSupported {
				break
			}
			log.Printf("%s ReserveN failed: %s (sleeping for %s)", p.q, err, consumerBackoff)
			time.Sleep(consumerBackoff)
			continue
		}
	}
}

func (p *Processor) fetchMessages() (int, error) {
	msgs, err := p.q.ReserveN(p.opt.BufferSize)
	if err != nil {
		return 0, err
	}
	for i := range msgs {
		p.queueMessage(&msgs[i])
	}
	return len(msgs), nil
}

func (p *Processor) worker() {
	defer p.wg.Done()
	for {
		if p.opt.Limiter != nil {
			delay, allow := p.opt.Limiter.AllowRate(p.q.Name(), p.opt.RateLimit)
			if !allow {
				time.Sleep(delay)
				continue
			}
		}

		msg, ok := p.dequeueMessage()
		if !ok {
			break
		}
		p.Process(msg)
	}
}

func (p *Processor) Process(msg *queue.Message) error {
	if msg.Delay > 0 {
		p.release(msg, nil)
		return nil
	}

	start := time.Now()
	err := p.handler.HandleMessage(msg)
	msg.SetValue("err", err)
	p.updateAvgDuration(time.Since(start))

	if err == nil {
		atomic.AddUint32(&p.processed, 1)
		p.delete(msg, nil)
		return nil
	}

	if msg.ReservedCount < p.opt.Retries {
		atomic.AddUint32(&p.retries, 1)
		p.release(msg, err)
	} else {
		atomic.AddUint32(&p.fails, 1)
		p.delete(msg, err)
	}

	if v := msg.Value("err"); v != nil {
		return v.(error)
	}
	return nil
}

func (p *Processor) queueMessage(msg *queue.Message) {
	atomic.AddUint32(&p.inFlight, 1)
	p.ch <- msg
}

func (p *Processor) dequeueMessage() (*queue.Message, bool) {
	select {
	case msg := <-p.ch:
		return msg, true
	case <-p.stop:
		select {
		case msg := <-p.ch:
			return msg, true
		default:
			return nil, false
		}
	}
}

func (p *Processor) release(msg *queue.Message, reason error) {
	delay := p.backoff(msg, reason)

	if reason != nil {
		log.Printf("%s handler failed (retry in %s): %s", p.q, delay, reason)
	}
	if err := p.q.Release(msg, delay); err != nil {
		log.Printf("%s Release failed: %s", p.q, err)
	}

	atomic.AddUint32(&p.inFlight, ^uint32(0))
}

func (p *Processor) backoff(msg *queue.Message, reason error) time.Duration {
	if reason != nil {
		if delayer, ok := reason.(Delayer); ok {
			return delayer.Delay()
		}
	}
	if msg.Delay > 0 {
		return msg.Delay
	}
	return exponentialBackoff(p.opt.Backoff, msg.ReservedCount)
}

func (p *Processor) delete(msg *queue.Message, reason error) {
	if reason != nil {
		log.Printf("%s handler failed: %s", p.q, reason)

		if p.fallbackHandler != nil {
			if err := p.fallbackHandler.HandleMessage(msg); err != nil {
				log.Printf("%s fallback handler failed: %s", p.q, err)
			}
		}
	}

	atomic.AddUint32(&p.inFlight, ^uint32(0))
	atomic.AddUint32(&p.deleting, 1)
	p.delBatch.Add(msg)
}

func (p *Processor) deleteBatch(msgs []*queue.Message) {
	if err := p.q.DeleteBatch(msgs); err != nil {
		log.Printf("%s DeleteBatch failed: %s", p.q, err)
	}
	atomic.AddUint32(&p.deleting, ^uint32(len(msgs)-1))
}

func (p *Processor) updateAvgDuration(dur time.Duration) {
	const decay = float64(1) / 100

	ms := float64(dur / time.Millisecond)
	for {
		avg := atomic.LoadUint32(&p.avgDuration)
		newAvg := uint32((1-decay)*float64(avg) + decay*ms)
		if atomic.CompareAndSwapUint32(&p.avgDuration, avg, newAvg) {
			break
		}
	}
}

func exponentialBackoff(dur time.Duration, retry int) time.Duration {
	dur <<= uint(retry - 1)
	if dur > maxBackoff {
		dur = maxBackoff
	}
	return dur
}
<commit_msg>Small fixes.<commit_after>package processor

import (
	"errors"
	"fmt"
	"log"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/time/rate"

	"gopkg.in/queue.v1"
	"gopkg.in/queue.v1/internal"
)

const consumerBackoff = time.Second
const maxBackoff = 12 * time.Hour
const stopTimeout = time.Minute

var ErrNotSupported = errors.New("processor: not supported")

type Delayer interface {
	Delay() time.Duration
}

type Stats struct {
	InFlight    uint32
	Deleting    uint32
	Processed   uint32
	Retries     uint32
	Fails       uint32
	AvgDuration time.Duration
}

type Processor struct {
	q   Queuer
	opt *queue.Options

	handler         queue.Handler
	fallbackHandler queue.Handler

	ch chan *queue.Message
	wg sync.WaitGroup

	delBatch *internal.Batcher

	_started uint32
	stop     chan struct{}

	inFlight    uint32
	deleting    uint32
	processed   uint32
	fails       uint32
	retries     uint32
	avgDuration uint32
}

func New(q Queuer, opt *queue.Options) *Processor {
	initOptions(opt)

	p := &Processor{
		q:   q,
		opt: opt,

		ch: make(chan *queue.Message, opt.BufferSize),
	}
	p.setHandler(opt.Handler)
	if opt.FallbackHandler != nil {
		p.setFallbackHandler(opt.FallbackHandler)
	}
	p.delBatch = internal.NewBatcher(p.opt.Scavengers, p.deleteBatch)
	return p
}

func initOptions(opt *queue.Options) {
	if opt.Workers == 0 {
		opt.Workers = 10 * runtime.NumCPU()
	}
	if opt.Scavengers == 0 {
		opt.Scavengers = runtime.NumCPU() + 1
	}
	if opt.BufferSize == 0 {
		opt.BufferSize = opt.Workers
		if opt.BufferSize > 10 {
			opt.BufferSize = 10
		}
	}
	if opt.RateLimit == 0 {
		opt.RateLimit = rate.Inf
	}
	if opt.Retries == 0 {
		opt.Retries = 10
	}
	if opt.Backoff == 0 {
		opt.Backoff = 3 * time.Second
	}
}

func Start(q Queuer, opt *queue.Options) *Processor {
	p := New(q, opt)
	p.Start()
	return p
}

func (p *Processor) String() string {
	return fmt.Sprintf(
		"Processor<%s workers=%d scavengers=%d buffer=%d>",
		p.q.Name(), p.opt.Workers, p.opt.Scavengers, p.opt.BufferSize,
	)
}

func (p *Processor) Stats() *Stats {
	if p.stopped() {
		return nil
	}
	return &Stats{
		InFlight:    atomic.LoadUint32(&p.inFlight),
		Deleting:    atomic.LoadUint32(&p.deleting),
		Processed:   atomic.LoadUint32(&p.processed),
		Retries:     atomic.LoadUint32(&p.retries),
		Fails:       atomic.LoadUint32(&p.fails),
		AvgDuration: time.Duration(atomic.LoadUint32(&p.avgDuration)) * time.Millisecond,
	}
}

func (p *Processor) setHandler(handler interface{}) {
	p.handler = queue.NewHandler(handler)
}

func (p *Processor) setFallbackHandler(handler interface{}) {
	p.fallbackHandler = queue.NewHandler(handler)
}

func (p *Processor) Add(msg *queue.Message) error {
	p.queueMessage(msg)
	return nil
}

func (p *Processor) Start() error {
	if !atomic.CompareAndSwapUint32(&p._started, 0, 1) {
		return nil
	}

	p.startWorkers()

	p.wg.Add(1)
	go p.messageFetcher()

	return nil
}

func (p *Processor) Stop() error {
	return p.StopTimeout(30 * time.Second)
}

func (p *Processor) StopTimeout(timeout time.Duration) error {
	if !atomic.CompareAndSwapUint32(&p._started, 1, 0) {
		return nil
	}
	return p.stopWorkersTimeout(stopTimeout)
}

func (p *Processor) startWorkers() {
	p.stop = make(chan struct{})
	p.wg.Add(p.opt.Workers)
	for i := 0; i < p.opt.Workers; i++ {
		go p.worker()
	}
}

func (p *Processor) stopWorkersTimeout(timeout time.Duration) error {
	close(p.stop)

	stopped := make(chan struct{})
	go func() {
		p.wg.Wait()
		close(stopped)
	}()

	select {
	case <-time.After(timeout):
		return fmt.Errorf("workers did not stop after %s seconds", timeout)
	case <-stopped:
		return nil
	}
}

func (p *Processor) stopped() bool {
	return atomic.LoadUint32(&p._started) == 0
}

func (p *Processor) ProcessAll() error {
	p.startWorkers()
	defer p.stopWorkersTimeout(stopTimeout)

	var noWork int
	for {
		isIdle := atomic.LoadUint32(&p.inFlight) == 0
		n, err := p.fetchMessages()
		if err != nil && err != ErrNotSupported {
			return err
		}
		if n == 0 && isIdle {
			noWork++
		} else {
			noWork = 0
		}
		if noWork == 2 {
			break
		}
	}
	return nil
}

func (p *Processor) ProcessOne() error {
	msg, err := p.reserveOne()
	if err != nil && err != ErrNotSupported {
		return err
	}
	atomic.AddUint32(&p.inFlight, 1)
	return p.Process(msg)
}

func (p *Processor) reserveOne() (*queue.Message, error) {
	select {
	case msg := <-p.ch:
		return msg, nil
	default:
	}

	msgs, err := p.q.ReserveN(1)
	if err != nil {
		return nil, err
	}
	if len(msgs) == 0 {
		return nil, errors.New("no messages in queue")
	}
	return &msgs[0], nil
}

func (p *Processor) messageFetcher() {
	defer p.wg.Done()
	for {
		if p.stopped() {
			break
		}

		_, err := p.fetchMessages()
		if err != nil {
			if err == ErrNotSupported {
				break
			}
			log.Printf("%s ReserveN failed: %s (sleeping for %s)", p.q, err, consumerBackoff)
			time.Sleep(consumerBackoff)
			continue
		}
	}
}

func (p *Processor) fetchMessages() (int, error) {
	msgs, err := p.q.ReserveN(p.opt.BufferSize)
	if err != nil {
		return 0, err
	}
	for i := range msgs {
		p.queueMessage(&msgs[i])
	}
	return len(msgs), nil
}

func (p *Processor) worker() {
	defer p.wg.Done()
	for {
		if p.opt.Limiter != nil {
			delay, allow := p.opt.Limiter.AllowRate(p.q.Name(), p.opt.RateLimit)
			if !allow {
				time.Sleep(delay)
				continue
			}
		}

		msg, ok := p.dequeueMessage()
		if !ok {
			break
		}
		p.Process(msg)
	}
}

func (p *Processor) Process(msg *queue.Message) error {
	if msg.Delay > 0 {
		p.release(msg, nil)
		return nil
	}

	start := time.Now()
	err := p.handler.HandleMessage(msg)
	msg.SetValue("err", err)
	p.updateAvgDuration(time.Since(start))

	if err == nil {
		atomic.AddUint32(&p.processed, 1)
		p.delete(msg, nil)
		return nil
	}

	if msg.ReservedCount < p.opt.Retries {
		atomic.AddUint32(&p.retries, 1)
		p.release(msg, err)
	} else {
		atomic.AddUint32(&p.fails, 1)
		p.delete(msg, err)
	}

	if v := msg.Value("err"); v != nil {
		return v.(error)
	}
	return nil
}

func (p *Processor) queueMessage(msg *queue.Message) {
	atomic.AddUint32(&p.inFlight, 1)
	p.ch <- msg
}

func (p *Processor) dequeueMessage() (*queue.Message, bool) {
	select {
	case msg := <-p.ch:
		return msg, true
	case <-p.stop:
		select {
		case msg := <-p.ch:
			return msg, true
		default:
			return nil, false
		}
	}
}

func (p *Processor) release(msg *queue.Message, reason error) {
	delay := p.backoff(msg, reason)

	if reason != nil {
		log.Printf("%s handler failed (retry in %s): %s", p.q, delay, reason)
	}
	if err := p.q.Release(msg, delay); err != nil {
		log.Printf("%s Release failed: %s", p.q, err)
	}

	atomic.AddUint32(&p.inFlight, ^uint32(0))
}

func (p *Processor) backoff(msg *queue.Message, reason error) time.Duration {
	if reason != nil {
		if delayer, ok := reason.(Delayer); ok {
			return delayer.Delay()
		}
	}
	if msg.Delay > 0 {
		return msg.Delay
	}
	return exponentialBackoff(p.opt.Backoff, msg.ReservedCount)
}

func (p *Processor) delete(msg *queue.Message, reason error) {
	if reason != nil {
		log.Printf("%s handler failed: %s", p.q, reason)

		if p.fallbackHandler != nil {
			if err := p.fallbackHandler.HandleMessage(msg); err != nil {
				log.Printf("%s fallback handler failed: %s", p.q, err)
			}
		}
	}

	atomic.AddUint32(&p.inFlight, ^uint32(0))
	atomic.AddUint32(&p.deleting, 1)
	p.delBatch.Add(msg)
}

func (p *Processor) deleteBatch(msgs []*queue.Message) {
	if err := p.q.DeleteBatch(msgs); err != nil {
		log.Printf("%s DeleteBatch failed: %s", p.q, err)
	}
	atomic.AddUint32(&p.deleting, ^uint32(len(msgs)-1))
}

func (p *Processor) updateAvgDuration(dur time.Duration) {
	const decay = float64(1) / 100

	ms := float64(dur / time.Millisecond)
	for {
		avg := atomic.LoadUint32(&p.avgDuration)
		newAvg := uint32((1-decay)*float64(avg) + decay*ms)
		if atomic.CompareAndSwapUint32(&p.avgDuration, avg, newAvg) {
			break
		}
	}
}

func exponentialBackoff(dur time.Duration, retry int) time.Duration {
	dur <<= uint(retry - 1)
	if dur > maxBackoff {
		dur = maxBackoff
	}
	return dur
}
<|endoftext|>


<commit_before>package processor

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/krallistic/kafka-operator/controller"
	"github.com/krallistic/kafka-operator/kafka"
	spec "github.com/krallistic/kafka-operator/spec"
	"github.com/krallistic/kafka-operator/util"
	k8sclient "k8s.io/client-go/kubernetes"
)

type Processor struct {
	client          k8sclient.Clientset
	baseBrokerImage string
	util            util.ClientUtil
	tprController   controller.CustomResourceController
	kafkaClusters   map[string]*spec.Kafkacluster
	watchEvents     chan spec.KafkaclusterWatchEvent
	clusterEvents   chan spec.KafkaclusterEvent
	kafkaClient     map[string]*kafka.KafkaUtil
	control         chan int
	errors          chan error
}

func New(client k8sclient.Clientset, image string, util util.ClientUtil, tprClient controller.CustomResourceController, control chan int) (*Processor, error) {
	p := &Processor{
		client:          client,
		baseBrokerImage: image,
		util:            util,
		kafkaClusters:   make(map[string]*spec.Kafkacluster),
		watchEvents:     make(chan spec.KafkaclusterWatchEvent, 100),
		clusterEvents:   make(chan spec.KafkaclusterEvent, 100),
		tprController:   tprClient,
		kafkaClient:     make(map[string]*kafka.KafkaUtil),
		control:         control,
		errors:          make(chan error),
	}
	fmt.Println("Created Processor")
	return p, nil
}

func (p *Processor) Run() error {
	//TODO getListOfAlredyRunningCluster/Refresh
	fmt.Println("Running Processor")
	p.watchKafkaEvents()
	fmt.Println("Watching")
	return nil
}

//We detect basic change through the event type, beyond that we use the API server to find differences.
//Functions compares the KafkaClusterSpec with the real Pods/Services which are there.
//We do that because otherwise we would have to use a local state to track changes.
func (p *Processor) DetectChangeType(event spec.KafkaclusterWatchEvent) spec.KafkaclusterEvent {
	fmt.Println("DetectChangeType: ", event)
	methodLogger := log.WithFields(log.Fields{
		"method":      "DetectChangeType",
		"clusterName": event.Object.Metadata.Name,
		"eventType":   event.Type,
	})
	methodLogger.Info("Detecting type of change in Kafka TPR")

	//TODO multiple changes in one Update? right now we only detect one change
	clusterEvent := spec.KafkaclusterEvent{
		Cluster: event.Object,
	}
	if event.Type == "ADDED" {
		clusterEvent.Type = spec.NEW_CLUSTER
		return clusterEvent
	}
	if event.Type == "DELETED" {
		clusterEvent.Type = spec.DELTE_CLUSTER
		return clusterEvent
		//EVENT type must be modfied now
	} else if p.util.BrokerStSImageUpdate(event.OldObject, event.Object) {
		clusterEvent.Type = spec.CHANGE_IMAGE
		return clusterEvent
	} else if p.util.BrokerStSUpsize(event.OldObject, event.Object) {
		clusterEvent.Type = spec.UPSIZE_CLUSTER
		return clusterEvent
	} else if p.util.BrokerStSDownsize(event.OldObject, event.Object) {
		clusterEvent.Type = spec.DOWNSIZE_CLUSTER
		return clusterEvent
	} else if p.util.BrokerStatefulSetExist(event.Object) {
		clusterEvent.Type = spec.UNKNOWN_CHANGE
		//TODO change to reconsilation event?
		return clusterEvent
	}

	clusterEvent.Type = spec.UNKNOWN_CHANGE
	return clusterEvent
}

func (p *Processor) initKafkaClient(cluster spec.Kafkacluster) error {
	methodLogger := log.WithFields(log.Fields{
		"method":            "initKafkaClient",
		"clusterName":       cluster.Metadata.Name,
		"zookeeperConnectL": cluster.Spec.ZookeeperConnect,
	})
	methodLogger.Info("Creating KafkaCLient for cluster")

	client, err := kafka.New(cluster)
	if err != nil {
		internalErrors.Inc()
		return err
	}

	//TODO can metadata.uuid used? check how that changed
	name := p.GetClusterUUID(cluster)
	p.kafkaClient[name] = client

	methodLogger.Info("Create KakfaClient for cluser")
	return nil
}

func (p *Processor) GetClusterUUID(cluster spec.Kafkacluster) string {
	return cluster.Metadata.Namespace + "-" + cluster.Metadata.Name
}

//Takes in raw Kafka events, lets then detected and the proced to initiate action accoriding to the detected event.
func (p *Processor) processKafkaEvent(currentEvent spec.KafkaclusterEvent) {
	fmt.Println("Recieved Event, proceeding: ", currentEvent)
	methodLogger := log.WithFields(log.Fields{
		"method":                "processKafkaEvent",
		"clusterName":           currentEvent.Cluster.Metadata.Name,
		"KafkaClusterEventType": currentEvent.Type,
	})
	switch currentEvent.Type {
	case spec.NEW_CLUSTER:
		fmt.Println("ADDED")
		clustersTotal.Inc()
		clustersCreated.Inc()
		p.CreateKafkaCluster(currentEvent.Cluster)
		clusterEvent := spec.KafkaclusterEvent{
			Cluster: currentEvent.Cluster,
			Type:    spec.KAKFA_EVENT,
		}

		methodLogger.Info("Init heartbeat type checking...")
		p.Sleep30AndSendEvent(clusterEvent)
		break

	case spec.DELTE_CLUSTER:
		methodLogger.Info("Delete Cluster, deleting all Objects ")
		if p.util.DeleteKafkaCluster(currentEvent.Cluster) != nil {
			//Error while deleting, just resubmit event after wait time.
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

		go func() {
			time.Sleep(time.Duration(currentEvent.Cluster.Spec.BrokerCount) * time.Minute)
			//TODO dynamic sleep, depending till sts is completely scaled down.
			clusterEvent := spec.KafkaclusterEvent{
				Cluster: currentEvent.Cluster,
				Type:    spec.CLEANUP_EVENT,
			}
			p.clusterEvents <- clusterEvent
		}()
		p.util.DeleteOffsetMonitor(currentEvent.Cluster)
		clustersTotal.Dec()
		clustersDeleted.Inc()
	case spec.CHANGE_IMAGE:
		fmt.Println("Change Image, updating StatefulSet should be enough to trigger a new Image Rollout")
		if p.util.UpdateBrokerImage(currentEvent.Cluster) != nil {
			//Error updating
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}
		clustersModified.Inc()
	case spec.UPSIZE_CLUSTER:
		methodLogger.Warn("Upsize Cluster, changing StatefulSet with higher Replicas, no Rebalacing")
		p.util.UpsizeBrokerStS(currentEvent.Cluster)
		clustersModified.Inc()
	case spec.UNKNOWN_CHANGE:
		methodLogger.Warn("Unknown (or unsupported) change occured, doing nothing. Maybe manually check the cluster")
		clustersModified.Inc()
	case spec.DOWNSIZE_CLUSTER:
		fmt.Println("Downsize Cluster")
		//TODO remove poor mans casting :P
		//TODO support Downsizing Multiple Brokers
		brokerToDelete := currentEvent.Cluster.Spec.BrokerCount - 0
		methodLogger.Info("Downsizing Broker, deleting Data on Broker: ", brokerToDelete)

		err := p.util.SetBrokerState(currentEvent.Cluster, brokerToDelete, spec.EMPTY_BROKER)
		if err != nil {
			//just re-try delete event
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

		err = p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)].RemoveTopicsFromBrokers(currentEvent.Cluster, brokerToDelete)
		if err != nil {
			//just re-try delete event
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}
		clustersModified.Inc()
	case spec.CHANGE_ZOOKEEPER_CONNECT:
		methodLogger.Warn("Trying to change zookeeper connect, not supported currently")
		clustersModified.Inc()
	case spec.CLEANUP_EVENT:
		fmt.Println("Recieved CleanupEvent, force delete of StatefuleSet.")
		clustersModified.Inc()
	case spec.KAKFA_EVENT:
		fmt.Println("Kafka Event, heartbeat etc..")
		p.Sleep30AndSendEvent(currentEvent)
	case spec.DOWNSIZE_EVENT:
		methodLogger.Info("Got Downsize Event, checking if all Topics are fully replicated and no topic on to delete cluster")
		//GET CLUSTER TO DELETE
		toDelete, err := p.util.GetBrokersWithState(currentEvent.Cluster, spec.EMPTY_BROKER)
		if err != nil {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
		}
		kafkaClient := p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)]
		topics, err := kafkaClient.GetTopicsOnBroker(currentEvent.Cluster, int32(toDelete))
		if len(topics) > 0 {
			//Move topics from Broker
			methodLogger.Warn("New Topics found on Broker which should be deleted, moving Topics Off", topics)
			for _, topic := range topics {
				kafkaClient.RemoveTopicFromBrokers(currentEvent.Cluster, toDelete, topic)
			}

			break
		}
		if err != nil {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
		}
		//CHECK if all Topics has been moved off
		inSync, err := p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)].AllTopicsInSync()
		if err != nil || !inSync {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

	}
}

func (p *Processor) Sleep30AndSendEvent(currentEvent spec.KafkaclusterEvent) {
	p.SleepAndSendEvent(currentEvent, 30)
}

func (p *Processor) SleepAndSendEvent(currentEvent spec.KafkaclusterEvent, seconds int) {
	go func() {
		time.Sleep(time.Second * time.Duration(seconds))
		p.clusterEvents <- currentEvent
	}()
}

//Creates inside a goroutine a watch channel on the KafkaCLuster Endpoint and distibutes the events.
//control chan used for showdown events from outside
func (p *Processor) watchKafkaEvents() {

	p.tprController.MonitorKafkaEvents(p.watchEvents, p.control)
	fmt.Println("Watching Kafka Events")
	go func() {
		for {

			select {
			case currentEvent := <-p.watchEvents:
				classifiedEvent := p.DetectChangeType(currentEvent)
				p.clusterEvents <- classifiedEvent
			case clusterEvent := <-p.clusterEvents:
				p.processKafkaEvent(clusterEvent)
			case err := <-p.errors:
				println("Error Channel", err)
			case <-p.control:
				fmt.Println("Recieved Something on Control Channel, shutting down: ")
				return
			}
		}
	}()
}

//Create the KafkaCluster, with the following components: Service, Volumes, StatefulSet.
//Maybe move this also into util
func (p *Processor) CreateKafkaCluster(clusterSpec spec.Kafkacluster) {
	fmt.Println("CreatingKafkaCluster", clusterSpec)
	fmt.Println("SPEC: ", clusterSpec.Spec)

	suffix := ".cluster.local:9092"
	brokerNames := make([]string, clusterSpec.Spec.BrokerCount)

	headless_SVC_Name := clusterSpec.Metadata.Name
	round_robing_dns := headless_SVC_Name + suffix
	fmt.Println("Headless Service Name: ", headless_SVC_Name, " Should be accessable through LB: ", round_robing_dns)

	var i int32
	for i = 0; i < clusterSpec.Spec.BrokerCount; i++ {
		brokerNames[i] = "kafka-0." + headless_SVC_Name + suffix
		fmt.Println("Broker", i, " ServiceName: ", brokerNames[i])
	}

	//Create Headless Brokersvc
	//TODO better naming
	p.util.CreateBrokerService(clusterSpec, true)

	//TODO createVolumes

	//CREATE Broker sts
	//Currently we extract name out of spec, maybe move to metadata to be more inline with other k8s komponents.
	p.util.CreateBrokerStatefulSet(clusterSpec)

	p.util.CreateDirectBrokerService(clusterSpec)

	p.initKafkaClient(clusterSpec)

	p.util.DeployOffsetMonitor(clusterSpec)

}
<commit_msg>Processor Changes<commit_after>package processor

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/krallistic/kafka-operator/controller"
	"github.com/krallistic/kafka-operator/kafka"
	spec "github.com/krallistic/kafka-operator/spec"
	"github.com/krallistic/kafka-operator/util"
	k8sclient "k8s.io/client-go/kubernetes"
)

type Processor struct {
	client          k8sclient.Clientset
	baseBrokerImage string
	util            util.ClientUtil
	tprController   controller.CustomResourceController
	kafkaClusters   map[string]*spec.Kafkacluster
	watchEvents     chan spec.KafkaclusterWatchEvent
	clusterEvents   chan spec.KafkaclusterEvent
	kafkaClient     map[string]*kafka.KafkaUtil
	control         chan int
	errors          chan error
}

func New(client k8sclient.Clientset, image string, util util.ClientUtil, tprClient controller.CustomResourceController, control chan int) (*Processor, error) {
	p := &Processor{
		client:          client,
		baseBrokerImage: image,
		util:            util,
		kafkaClusters:   make(map[string]*spec.Kafkacluster),
		watchEvents:     make(chan spec.KafkaclusterWatchEvent, 100),
		clusterEvents:   make(chan spec.KafkaclusterEvent, 100),
		tprController:   tprClient,
		kafkaClient:     make(map[string]*kafka.KafkaUtil),
		control:         control,
		errors:          make(chan error),
	}
	fmt.Println("Created Processor")
	return p, nil
}

func (p *Processor) Run() error {
	//TODO getListOfAlredyRunningCluster/Refresh
	fmt.Println("Running Processor")
	p.watchKafkaEvents()
	fmt.Println("Watching")
	return nil
}

//We detect basic change through the event type, beyond that we use the API server to find differences.
//Functions compares the KafkaClusterSpec with the real Pods/Services which are there.
//We do that because otherwise we would have to use a local state to track changes.
func (p *Processor) DetectChangeType(event spec.KafkaclusterWatchEvent) spec.KafkaclusterEvent {
	fmt.Println("DetectChangeType: ", event)
	methodLogger := log.WithFields(log.Fields{
		"method":      "DetectChangeType",
		"clusterName": event.Object.ObjectMeta.Name,
		"eventType":   event.Type,
	})
	methodLogger.Info("Detecting type of change in Kafka TPR")

	//TODO multiple changes in one Update? right now we only detect one change
	clusterEvent := spec.KafkaclusterEvent{
		Cluster: event.Object,
	}
	if event.Type == "ADDED" {
		clusterEvent.Type = spec.NEW_CLUSTER
		return clusterEvent
	}
	if event.Type == "DELETED" {
		clusterEvent.Type = spec.DELTE_CLUSTER
		return clusterEvent
		//EVENT type must be modfied now
	} else if p.util.BrokerStSImageUpdate(event.OldObject, event.Object) {
		clusterEvent.Type = spec.CHANGE_IMAGE
		return clusterEvent
	} else if p.util.BrokerStSUpsize(event.OldObject, event.Object) {
		clusterEvent.Type = spec.UPSIZE_CLUSTER
		return clusterEvent
	} else if p.util.BrokerStSDownsize(event.OldObject, event.Object) {
		clusterEvent.Type = spec.DOWNSIZE_CLUSTER
		return clusterEvent
	} else if p.util.BrokerStatefulSetExist(event.Object) {
		clusterEvent.Type = spec.UNKNOWN_CHANGE
		//TODO change to reconsilation event?
		return clusterEvent
	}

	clusterEvent.Type = spec.UNKNOWN_CHANGE
	return clusterEvent
}

func (p *Processor) initKafkaClient(cluster spec.Kafkacluster) error {
	methodLogger := log.WithFields(log.Fields{
		"method":            "initKafkaClient",
		"clusterName":       cluster.ObjectMeta.Name,
		"zookeeperConnectL": cluster.Spec.ZookeeperConnect,
	})
	methodLogger.Info("Creating KafkaCLient for cluster")

	client, err := kafka.New(cluster)
	if err != nil {
		internalErrors.Inc()
		return err
	}

	//TODO can metadata.uuid used? check how that changed
	name := p.GetClusterUUID(cluster)
	p.kafkaClient[name] = client

	methodLogger.Info("Create KakfaClient for cluser")
	return nil
}

func (p *Processor) GetClusterUUID(cluster spec.Kafkacluster) string {
	return cluster.ObjectMeta.Namespace + "-" + cluster.ObjectMeta.Name
}

//Takes in raw Kafka events, lets then detected and the proced to initiate action accoriding to the detected event.
func (p *Processor) processKafkaEvent(currentEvent spec.KafkaclusterEvent) {
	fmt.Println("Recieved Event, proceeding: ", currentEvent)
	methodLogger := log.WithFields(log.Fields{
		"method":                "processKafkaEvent",
		"clusterName":           currentEvent.Cluster.ObjectMeta.Name,
		"KafkaClusterEventType": currentEvent.Type,
	})
	switch currentEvent.Type {
	case spec.NEW_CLUSTER:
		fmt.Println("ADDED")
		clustersTotal.Inc()
		clustersCreated.Inc()
		p.CreateKafkaCluster(currentEvent.Cluster)
		clusterEvent := spec.KafkaclusterEvent{
			Cluster: currentEvent.Cluster,
			Type:    spec.KAKFA_EVENT,
		}

		methodLogger.Info("Init heartbeat type checking...")
		p.Sleep30AndSendEvent(clusterEvent)
		break

	case spec.DELTE_CLUSTER:
		methodLogger.Info("Delete Cluster, deleting all Objects ")
		if p.util.DeleteKafkaCluster(currentEvent.Cluster) != nil {
			//Error while deleting, just resubmit event after wait time.
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

		go func() {
			time.Sleep(time.Duration(currentEvent.Cluster.Spec.BrokerCount) * time.Minute)
			//TODO dynamic sleep, depending till sts is completely scaled down.
			clusterEvent := spec.KafkaclusterEvent{
				Cluster: currentEvent.Cluster,
				Type:    spec.CLEANUP_EVENT,
			}
			p.clusterEvents <- clusterEvent
		}()
		p.util.DeleteOffsetMonitor(currentEvent.Cluster)
		clustersTotal.Dec()
		clustersDeleted.Inc()
	case spec.CHANGE_IMAGE:
		fmt.Println("Change Image, updating StatefulSet should be enough to trigger a new Image Rollout")
		if p.util.UpdateBrokerImage(currentEvent.Cluster) != nil {
			//Error updating
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}
		clustersModified.Inc()
	case spec.UPSIZE_CLUSTER:
		methodLogger.Warn("Upsize Cluster, changing StatefulSet with higher Replicas, no Rebalacing")
		p.util.UpsizeBrokerStS(currentEvent.Cluster)
		clustersModified.Inc()
	case spec.UNKNOWN_CHANGE:
		methodLogger.Warn("Unknown (or unsupported) change occured, doing nothing. Maybe manually check the cluster")
		clustersModified.Inc()
	case spec.DOWNSIZE_CLUSTER:
		fmt.Println("Downsize Cluster")
		//TODO remove poor mans casting :P
		//TODO support Downsizing Multiple Brokers
		brokerToDelete := currentEvent.Cluster.Spec.BrokerCount - 0
		methodLogger.Info("Downsizing Broker, deleting Data on Broker: ", brokerToDelete)

		err := p.util.SetBrokerState(currentEvent.Cluster, brokerToDelete, spec.EMPTY_BROKER)
		if err != nil {
			//just re-try delete event
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

		err = p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)].RemoveTopicsFromBrokers(currentEvent.Cluster, brokerToDelete)
		if err != nil {
			//just re-try delete event
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}
		clustersModified.Inc()
	case spec.CHANGE_ZOOKEEPER_CONNECT:
		methodLogger.Warn("Trying to change zookeeper connect, not supported currently")
		clustersModified.Inc()
	case spec.CLEANUP_EVENT:
		fmt.Println("Recieved CleanupEvent, force delete of StatefuleSet.")
		p.util.CleanupKafkaCluster(currentEvent.Cluster)
		clustersModified.Inc()
	case spec.KAKFA_EVENT:
		fmt.Println("Kafka Event, heartbeat etc..")
		p.Sleep30AndSendEvent(currentEvent)
	case spec.DOWNSIZE_EVENT:
		methodLogger.Info("Got Downsize Event, checking if all Topics are fully replicated and no topic on to delete cluster")
		//GET CLUSTER TO DELETE
		toDelete, err := p.util.GetBrokersWithState(currentEvent.Cluster, spec.EMPTY_BROKER)
		if err != nil {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
		}
		kafkaClient := p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)]
		topics, err := kafkaClient.GetTopicsOnBroker(currentEvent.Cluster, int32(toDelete))
		if len(topics) > 0 {
			//Move topics from Broker
			methodLogger.Warn("New Topics found on Broker which should be deleted, moving Topics Off", topics)
			for _, topic := range topics {
				kafkaClient.RemoveTopicFromBrokers(currentEvent.Cluster, toDelete, topic)
			}

			break
		}
		if err != nil {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
		}
		//CHECK if all Topics has been moved off
		inSync, err := p.kafkaClient[p.GetClusterUUID(currentEvent.Cluster)].AllTopicsInSync()
		if err != nil || !inSync {
			internalErrors.Inc()
			p.Sleep30AndSendEvent(currentEvent)
			break
		}

	}
}

func (p *Processor) Sleep30AndSendEvent(currentEvent spec.KafkaclusterEvent) {
	p.SleepAndSendEvent(currentEvent, 30)
}

func (p *Processor) SleepAndSendEvent(currentEvent spec.KafkaclusterEvent, seconds int) {
	go func() {
		time.Sleep(time.Second * time.Duration(seconds))
		p.clusterEvents <- currentEvent
	}()
}

//Creates inside a goroutine a watch channel on the KafkaCLuster Endpoint and distibutes the events.
//control chan used for showdown events from outside
func (p *Processor) watchKafkaEvents() {

	p.tprController.MonitorKafkaEvents(p.watchEvents, p.control)
	fmt.Println("Watching Kafka Events")
	go func() {
		for {

			select {
			case currentEvent := <-p.watchEvents:
				classifiedEvent := p.DetectChangeType(currentEvent)
				p.clusterEvents <- classifiedEvent
			case clusterEvent := <-p.clusterEvents:
				p.processKafkaEvent(clusterEvent)
			case err := <-p.errors:
				println("Error Channel", err)
			case <-p.control:
				fmt.Println("Recieved Something on Control Channel, shutting down: ")
				return
			}
		}
	}()
}

// CreateKafkaCluster with the following components: Service, Volumes, StatefulSet.
//Maybe move this also into util
func (p *Processor) CreateKafkaCluster(clusterSpec spec.Kafkacluster) {
	methodLogger := log.WithFields(log.Fields{
		"method":      "CreateKafkaCluster",
		"clusterName": clusterSpec.ObjectMeta.Name,
	})

	err := p.util.CreateBrokerStatefulSet(clusterSpec)
	if err != nil {
		methodLogger.WithField("error", err).Fatal("Cant create statefulset")
	}

	err = p.util.CreateBrokerService(clusterSpec)
	if err != nil {
		methodLogger.WithField("error", err).Fatal("Cant create loadbalacend headless service services")
	}

	err = p.util.CreateDirectBrokerService(clusterSpec)
	if err != nil {
		methodLogger.WithField("error", err).Fatal("Cant create direct broker services")
	}
	//TODO createVolumes here?

	p.initKafkaClient(clusterSpec)

	err = p.util.DeployOffsetMonitor(clusterSpec)
	if err != nil {
		methodLogger.WithField("error", err).Fatal("Cant deploy stats exporter")
	}
}
<|endoftext|>


<commit_before>package db

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"
	"sync"

	"github.com/koron/nvgd/config"
	"github.com/koron/nvgd/ltsv"
	"github.com/koron/nvgd/protocol"
)

var NullReplacement = "(null)"

// Param is connection parameter for the database.
type Param struct {
	// Driver represents driver name for database.
	Driver string `yaml:"driver"`

	// Name represents driver-specific data source name.
	Name string `yaml:"name"`
}

// Config represents configuration for Handler.
type Config map[string]Param

// Handler is database protocol handler.
type Handler struct {
	Config *Config

	l         sync.Mutex
	databases map[string]*sql.DB
}

var dbconfig Config

func init() {
	protocol.MustRegister("db", &Handler{
		Config:    &dbconfig,
		databases: make(map[string]*sql.DB),
	})
	config.RegisterProtocol("db", &dbconfig)
}

// Open creates a database handler.
func (h *Handler) Open(u *url.URL) (io.ReadCloser, error) {
	var (
		name  = u.Host
		query = u.Path
	)
	db, err := h.openDB(name)
	if err != nil {
		return nil, err
	}
	// TODO: sanitize query!
	if strings.HasPrefix(query, "/") {
		query = query[1:]
	}
	fmt.Printf("query=%s\n", query)
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return h.rows2ltsv(rows)
}

func (h *Handler) openDB(name string) (*sql.DB, error) {
	h.l.Lock()
	defer h.l.Unlock()
	if db, ok := h.databases[name]; ok {
		return db, nil
	}
	p, ok := (*h.Config)[name]
	if !ok {
		return nil, fmt.Errorf("unknown database: %q", name)
	}
	db, err := sql.Open(p.Driver, p.Name)
	if err != nil {
		return nil, err
	}
	h.databases[name] = db
	return db, nil
}

func (h *Handler) rows2ltsv(rows *sql.Rows) (io.ReadCloser, error) {
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	var (
		buf = &bytes.Buffer{}
		w   = ltsv.NewWriter(buf, cols...)
		n   = len(cols)
	)

	vals := make([]interface{}, n)
	for i := range vals {
		vals[i] = new(sql.NullString)
	}
	strs := make([]string, n)

	for rows.Next() {
		if err := rows.Scan(vals...); err != nil {
			return nil, err
		}
		for i, v := range vals {
			ns := v.(*sql.NullString)
			if ns.Valid {
				strs[i] = ns.String
			} else {
				strs[i] = NullReplacement
			}
		}
		w.Write(strs...)
	}
	return ioutil.NopCloser(buf), nil
}
<commit_msg>wrap queries by transaction for safety<commit_after>package db

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"
	"sync"

	"github.com/koron/nvgd/config"
	"github.com/koron/nvgd/ltsv"
	"github.com/koron/nvgd/protocol"
)

var NullReplacement = "(null)"

// Param is connection parameter for the database.
type Param struct {
	// Driver represents driver name for database.
	Driver string `yaml:"driver"`

	// Name represents driver-specific data source name.
	Name string `yaml:"name"`
}

// Config represents configuration for Handler.
type Config map[string]Param

// Handler is database protocol handler.
type Handler struct {
	Config *Config

	l         sync.Mutex
	databases map[string]*sql.DB
}

var dbconfig Config

func init() {
	protocol.MustRegister("db", &Handler{
		Config:    &dbconfig,
		databases: make(map[string]*sql.DB),
	})
	config.RegisterProtocol("db", &dbconfig)
}

// Open creates a database handler.
func (h *Handler) Open(u *url.URL) (io.ReadCloser, error) {
	var (
		name  = u.Host
		query = u.Path
	)
	db, err := h.openDB(name)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(query, "/") {
		query = query[1:]
	}
	// TODO: sanitize query!
	//fmt.Printf("query=%s\n", query)
	return h.execQuery(db, query)
}

// execQuery executes a query in a transaction which will be rollbacked.
func (h *Handler) execQuery(db *sql.DB, q string) (io.ReadCloser, error) {
	tx, err := db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	rows, err := tx.Query(q)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return h.rows2ltsv(rows)
}

func (h *Handler) openDB(name string) (*sql.DB, error) {
	h.l.Lock()
	defer h.l.Unlock()
	if db, ok := h.databases[name]; ok {
		return db, nil
	}
	p, ok := (*h.Config)[name]
	if !ok {
		return nil, fmt.Errorf("unknown database: %q", name)
	}
	db, err := sql.Open(p.Driver, p.Name)
	if err != nil {
		return nil, err
	}
	h.databases[name] = db
	return db, nil
}

func (h *Handler) rows2ltsv(rows *sql.Rows) (io.ReadCloser, error) {
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	var (
		buf = &bytes.Buffer{}
		w   = ltsv.NewWriter(buf, cols...)
		n   = len(cols)
	)

	vals := make([]interface{}, n)
	for i := range vals {
		vals[i] = new(sql.NullString)
	}
	strs := make([]string, n)

	for rows.Next() {
		if err := rows.Scan(vals...); err != nil {
			return nil, err
		}
		for i, v := range vals {
			ns := v.(*sql.NullString)
			if ns.Valid {
				strs[i] = ns.String
			} else {
				strs[i] = NullReplacement
			}
		}
		w.Write(strs...)
	}
	return ioutil.NopCloser(buf), nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The Terraformer 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"
	"log"
	"strings"

	"github.com/Azure/azure-sdk-for-go/services/dns/mgmt/2018-05-01/dns"
	"github.com/Azure/go-autorest/autorest"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/hashicorp/go-azure-helpers/authentication"
)

type DNSGenerator struct {
	AzureService
}

func (g *DNSGenerator) listRecordSets(resourceGroupName string, zoneName string, top *int32) ([]terraformutils.Resource, error) {
	var resources []terraformutils.Resource
	ctx := context.Background()
	subscriptionID := g.Args["config"].(authentication.Config).SubscriptionID
	RecordSetsClient := dns.NewRecordSetsClient(subscriptionID)
	RecordSetsClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)

	recordSetIterator, err := RecordSetsClient.ListAllByDNSZoneComplete(ctx, resourceGroupName, zoneName, top, "")
	if err != nil {
		return nil, err
	}
	for recordSetIterator.NotDone() {
		recordSet := recordSetIterator.Value()
		// NOTE:
		// Format example: "Microsoft.Network/dnszones/AAAA"
		recordType := *recordSet.Type
		if strings.HasSuffix(recordType, "/A") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_a_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/AAAA") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_aaaa_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/CAA") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_caa_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/CNAME") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_cname_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/NS") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_ns_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/MX") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_mx_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/PTR") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_ptr_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/SRV") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_srv_record",
				g.ProviderName,
				[]string{}))
		} else if strings.HasSuffix(recordType, "/TXT") {
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_txt_record",
				g.ProviderName,
				[]string{}))
		}
		if err := recordSetIterator.Next(); err != nil {
			log.Println(err)
			break
		}

	}
	return resources, nil
}

func (g *DNSGenerator) listAndAddForDNSZone() ([]terraformutils.Resource, error) {
	var resources []terraformutils.Resource
	ctx := context.Background()
	subscriptionID := g.Args["config"].(authentication.Config).SubscriptionID
	DNSZonesClient := dns.NewZonesClient(subscriptionID)
	DNSZonesClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)

	var pageSize int32 = 50
	dnsZoneIterator, err := DNSZonesClient.ListComplete(ctx, &pageSize)
	if err != nil {
		return nil, err
	}
	for dnsZoneIterator.NotDone() {
		zone := dnsZoneIterator.Value()
		resources = append(resources, terraformutils.NewSimpleResource(
			*zone.ID,
			*zone.Name,
			"azurerm_dns_zone",
			g.ProviderName,
			[]string{}))

		id, err := ParseAzureResourceID(*zone.ID)
		if err != nil {
			return nil, err
		}

		records, err := g.listRecordSets(id.ResourceGroup, *zone.Name, &pageSize)
		if err != nil {
			return nil, err
		}
		resources = append(resources, records...)

		if err := dnsZoneIterator.Next(); err != nil {
			log.Println(err)
			break
		}
	}

	return resources, nil
}

func (g *DNSGenerator) InitResources() error {
	functions := []func() ([]terraformutils.Resource, error){
		g.listAndAddForDNSZone,
	}

	for _, f := range functions {
		resources, err := f()
		if err != nil {
			return err
		}
		g.Resources = append(g.Resources, resources...)
	}

	return nil
}
<commit_msg>Fix switchcase syntax<commit_after>// Copyright 2020 The Terraformer 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"
	"log"
	"strings"

	"github.com/Azure/azure-sdk-for-go/services/dns/mgmt/2018-05-01/dns"
	"github.com/Azure/go-autorest/autorest"
	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
	"github.com/hashicorp/go-azure-helpers/authentication"
)

type DNSGenerator struct {
	AzureService
}

func (g *DNSGenerator) listRecordSets(resourceGroupName string, zoneName string, top *int32) ([]terraformutils.Resource, error) {
	var resources []terraformutils.Resource
	ctx := context.Background()
	subscriptionID := g.Args["config"].(authentication.Config).SubscriptionID
	RecordSetsClient := dns.NewRecordSetsClient(subscriptionID)
	RecordSetsClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)

	recordSetIterator, err := RecordSetsClient.ListAllByDNSZoneComplete(ctx, resourceGroupName, zoneName, top, "")
	if err != nil {
		return nil, err
	}
	for recordSetIterator.NotDone() {
		recordSet := recordSetIterator.Value()
		// NOTE:
		// Format example: "Microsoft.Network/dnszones/AAAA"
		recordType := *recordSet.Type
		switch {
		case strings.HasSuffix(recordType, "/A"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_a_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/AAAA"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_aaaa_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/CAA"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_caa_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/CNAME"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_cname_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/NS"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_ns_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/MX"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_mx_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/PTR"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_ptr_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/SRV"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_srv_record",
				g.ProviderName,
				[]string{}))
		case strings.HasSuffix(recordType, "/TXT"):
			resources = append(resources, terraformutils.NewSimpleResource(
				*recordSet.ID,
				*recordSet.Name,
				"azurerm_dns_txt_record",
				g.ProviderName,
				[]string{}))
		}
		if err := recordSetIterator.Next(); err != nil {
			log.Println(err)
			break
		}

	}
	return resources, nil
}

func (g *DNSGenerator) listAndAddForDNSZone() ([]terraformutils.Resource, error) {
	var resources []terraformutils.Resource
	ctx := context.Background()
	subscriptionID := g.Args["config"].(authentication.Config).SubscriptionID
	DNSZonesClient := dns.NewZonesClient(subscriptionID)
	DNSZonesClient.Authorizer = g.Args["authorizer"].(autorest.Authorizer)

	var pageSize int32 = 50
	dnsZoneIterator, err := DNSZonesClient.ListComplete(ctx, &pageSize)
	if err != nil {
		return nil, err
	}
	for dnsZoneIterator.NotDone() {
		zone := dnsZoneIterator.Value()
		resources = append(resources, terraformutils.NewSimpleResource(
			*zone.ID,
			*zone.Name,
			"azurerm_dns_zone",
			g.ProviderName,
			[]string{}))

		id, err := ParseAzureResourceID(*zone.ID)
		if err != nil {
			return nil, err
		}

		records, err := g.listRecordSets(id.ResourceGroup, *zone.Name, &pageSize)
		if err != nil {
			return nil, err
		}
		resources = append(resources, records...)

		if err := dnsZoneIterator.Next(); err != nil {
			log.Println(err)
			break
		}
	}

	return resources, nil
}

func (g *DNSGenerator) InitResources() error {
	functions := []func() ([]terraformutils.Resource, error){
		g.listAndAddForDNSZone,
	}

	for _, f := range functions {
		resources, err := f()
		if err != nil {
			return err
		}
		g.Resources = append(g.Resources, resources...)
	}

	return nil
}
<|endoftext|>