<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsVpnConnectionRoute() *schema.Resource {
	return &schema.Resource{
		// You can't update a route. You can just delete one and make
		// a new one.
		Create: resourceAwsVpnConnectionRouteCreate,
		Read:   resourceAwsVpnConnectionRouteRead,
		Delete: resourceAwsVpnConnectionRouteDelete,

		Schema: map[string]*schema.Schema{
			"destination_cidr_block": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"vpn_connection_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
		},
	}
}

func resourceAwsVpnConnectionRouteCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	createOpts := &ec2.CreateVpnConnectionRouteInput{
		DestinationCidrBlock: aws.String(d.Get("destination_cidr_block").(string)),
		VpnConnectionId:      aws.String(d.Get("vpn_connection_id").(string)),
	}

	// Create the route.
	log.Printf("[DEBUG] Creating VPN connection route")
	_, err := conn.CreateVpnConnectionRoute(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating VPN connection route: %s", err)
	}

	// Store the ID by the only two data we have available to us.
	d.SetId(fmt.Sprintf("%s:%s", *createOpts.DestinationCidrBlock, *createOpts.VpnConnectionId))

	return resourceAwsVpnConnectionRouteRead(d, meta)
}

func resourceAwsVpnConnectionRouteRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	cidrBlock, vpnConnectionId := resourceAwsVpnConnectionRouteParseId(d.Id())

	routeFilters := []*ec2.Filter{
		&ec2.Filter{
			Name:   aws.String("route.destination-cidr-block"),
			Values: []*string{aws.String(cidrBlock)},
		},
		&ec2.Filter{
			Name:   aws.String("vpn-connection-id"),
			Values: []*string{aws.String(vpnConnectionId)},
		},
	}

	// Technically, we know everything there is to know about the route
	// from its ID, but we still want to catch cases where it changes
	// outside of terraform and results in a stale state file. Hence,
	// conduct a read.
	resp, err := conn.DescribeVpnConnections(&ec2.DescribeVpnConnectionsInput{
		Filters: routeFilters,
	})
	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidVpnConnectionID.NotFound" {
			d.SetId("")
			return nil
		} else {
			log.Printf("[ERROR] Error finding VPN connection route: %s", err)
			return err
		}
	}
	if resp == nil || len(resp.VpnConnections) == 0 {
		// This is kind of a weird edge case. I'd rather return an error
		// instead of just blindly setting the ID to ""... since I don't know
		// what might cause this.
		return fmt.Errorf("No VPN connections returned")
	}

	vpnConnection := resp.VpnConnections[0]

	var found bool
	for _, r := range vpnConnection.Routes {
		if *r.DestinationCidrBlock == cidrBlock && *r.State != "deleted" {
			d.Set("destination_cidr_block", *r.DestinationCidrBlock)
			d.Set("vpn_connection_id", *vpnConnection.VpnConnectionId)
			found = true
		}
	}
	if !found {
		// Something other than terraform eliminated the route.
		d.SetId("")
	}

	return nil
}

func resourceAwsVpnConnectionRouteDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	_, err := conn.DeleteVpnConnectionRoute(&ec2.DeleteVpnConnectionRouteInput{
		DestinationCidrBlock: aws.String(d.Get("destination_cidr_block").(string)),
		VpnConnectionId:      aws.String(d.Get("vpn_connection_id").(string)),
	})
	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidVpnConnectionID.NotFound" {
			d.SetId("")
			return nil
		} else {
			log.Printf("[ERROR] Error deleting VPN connection route: %s", err)
			return err
		}
	}

	return nil
}

func resourceAwsVpnConnectionRouteParseId(id string) (string, string) {
	parts := strings.SplitN(id, ":", 2)
	return parts[0], parts[1]
}
<commit_msg>r/vpn_connection_route: Wait until route is available/deleted<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsVpnConnectionRoute() *schema.Resource {
	return &schema.Resource{
		// You can't update a route. You can just delete one and make
		// a new one.
		Create: resourceAwsVpnConnectionRouteCreate,
		Read:   resourceAwsVpnConnectionRouteRead,
		Delete: resourceAwsVpnConnectionRouteDelete,

		Schema: map[string]*schema.Schema{
			"destination_cidr_block": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"vpn_connection_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
		},
	}
}

func resourceAwsVpnConnectionRouteCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	cidrBlock := d.Get("destination_cidr_block").(string)
	vpnConnectionId := d.Get("vpn_connection_id").(string)
	createOpts := &ec2.CreateVpnConnectionRouteInput{
		DestinationCidrBlock: aws.String(cidrBlock),
		VpnConnectionId:      aws.String(vpnConnectionId),
	}

	// Create the route.
	log.Printf("[DEBUG] Creating VPN connection route")
	_, err := conn.CreateVpnConnectionRoute(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating VPN connection route: %s", err)
	}

	// Store the ID by the only two data we have available to us.
	d.SetId(fmt.Sprintf("%s:%s", *createOpts.DestinationCidrBlock, *createOpts.VpnConnectionId))

	stateConf := resource.StateChangeConf{
		Pending: []string{"pending"},
		Target:  []string{"available"},
		Timeout: 15 * time.Second,
		Refresh: func() (interface{}, string, error) {
			route, err := findConnectionRoute(conn, cidrBlock, vpnConnectionId)
			if err != nil {
				return 42, "", err
			}
			return route, *route.State, nil
		},
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return err
	}

	return resourceAwsVpnConnectionRouteRead(d, meta)
}

func resourceAwsVpnConnectionRouteRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	cidrBlock, vpnConnectionId := resourceAwsVpnConnectionRouteParseId(d.Id())

	route, err := findConnectionRoute(conn, cidrBlock, vpnConnectionId)
	if err != nil {
		return err
	}
	if route == nil {
		// Something other than terraform eliminated the route.
		d.SetId("")
	}

	return nil
}

func resourceAwsVpnConnectionRouteDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	cidrBlock := d.Get("destination_cidr_block").(string)
	vpnConnectionId := d.Get("vpn_connection_id").(string)
	_, err := conn.DeleteVpnConnectionRoute(&ec2.DeleteVpnConnectionRouteInput{
		DestinationCidrBlock: aws.String(cidrBlock),
		VpnConnectionId:      aws.String(vpnConnectionId),
	})
	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidVpnConnectionID.NotFound" {
			d.SetId("")
			return nil
		}
		log.Printf("[ERROR] Error deleting VPN connection route: %s", err)
		return err
	}

	stateConf := resource.StateChangeConf{
		Pending: []string{"pending", "available", "deleting"},
		Target:  []string{"deleted"},
		Timeout: 15 * time.Second,
		Refresh: func() (interface{}, string, error) {
			route, err := findConnectionRoute(conn, cidrBlock, vpnConnectionId)
			if err != nil {
				return 42, "", err
			}
			if route == nil {
				return 42, "deleted", nil
			}
			return route, *route.State, nil
		},
	}
	_, err = stateConf.WaitForState()
	if err != nil {
		return err
	}

	return nil
}

func findConnectionRoute(conn *ec2.EC2, cidrBlock, vpnConnectionId string) (*ec2.VpnStaticRoute, error) {
	resp, err := conn.DescribeVpnConnections(&ec2.DescribeVpnConnectionsInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name:   aws.String("route.destination-cidr-block"),
				Values: []*string{aws.String(cidrBlock)},
			},
			&ec2.Filter{
				Name:   aws.String("vpn-connection-id"),
				Values: []*string{aws.String(vpnConnectionId)},
			},
		},
	})
	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidVpnConnectionID.NotFound" {
			return nil, nil
		}
		return nil, err
	}
	if resp == nil || len(resp.VpnConnections) == 0 {
		return nil, nil
	}
	vpnConnection := resp.VpnConnections[0]

	for _, r := range vpnConnection.Routes {
		if *r.DestinationCidrBlock == cidrBlock && *r.State != "deleted" {
			return r, nil
		}
	}
	return nil, nil
}

func resourceAwsVpnConnectionRouteParseId(id string) (string, string) {
	parts := strings.SplitN(id, ":", 2)
	return parts[0], parts[1]
}
<|endoftext|>


<commit_before>// Copyright 2016 The etcd 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 command

import (
	"context"
	"fmt"
	"strconv"

	v3 "github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

// NewLeaseCommand returns the cobra command for "lease".
func NewLeaseCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "lease <subcommand>",
		Short: "Lease related commands",
	}

	lc.AddCommand(NewLeaseGrantCommand())
	lc.AddCommand(NewLeaseRevokeCommand())
	lc.AddCommand(NewLeaseTimeToLiveCommand())
	lc.AddCommand(NewLeaseListCommand())
	lc.AddCommand(NewLeaseKeepAliveCommand())

	return lc
}

// NewLeaseGrantCommand returns the cobra command for "lease grant".
func NewLeaseGrantCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "grant <ttl>",
		Short: "Creates leases",

		Run: leaseGrantCommandFunc,
	}

	return lc
}

// leaseGrantCommandFunc executes the "lease grant" command.
func leaseGrantCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease grant command needs TTL argument."))
	}

	ttl, err := strconv.ParseInt(args[0], 10, 64)
	if err != nil {
		ExitWithError(ExitBadArgs, fmt.Errorf("bad TTL (%v)", err))
	}

	ctx, cancel := commandCtx(cmd)
	resp, err := mustClientFromCmd(cmd).Grant(ctx, ttl)
	cancel()
	if err != nil {
		ExitWithError(ExitError, fmt.Errorf("failed to grant lease (%v)\n", err))
	}
	display.Grant(*resp)
}

// NewLeaseRevokeCommand returns the cobra command for "lease revoke".
func NewLeaseRevokeCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "revoke <leaseID>",
		Short: "Revokes leases",

		Run: leaseRevokeCommandFunc,
	}

	return lc
}

// leaseRevokeCommandFunc executes the "lease grant" command.
func leaseRevokeCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease revoke command needs 1 argument"))
	}

	id := leaseFromArgs(args[0])
	ctx, cancel := commandCtx(cmd)
	resp, err := mustClientFromCmd(cmd).Revoke(ctx, id)
	cancel()
	if err != nil {
		ExitWithError(ExitError, fmt.Errorf("failed to revoke lease (%v)\n", err))
	}
	display.Revoke(id, *resp)
}

var timeToLiveKeys bool

// NewLeaseTimeToLiveCommand returns the cobra command for "lease timetolive".
func NewLeaseTimeToLiveCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "timetolive <leaseID> [options]",
		Short: "Get lease information",

		Run: leaseTimeToLiveCommandFunc,
	}
	lc.Flags().BoolVar(&timeToLiveKeys, "keys", false, "Get keys attached to this lease")

	return lc
}

// leaseTimeToLiveCommandFunc executes the "lease timetolive" command.
func leaseTimeToLiveCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease timetolive command needs lease ID as argument"))
	}
	var opts []v3.LeaseOption
	if timeToLiveKeys {
		opts = append(opts, v3.WithAttachedKeys())
	}
	resp, rerr := mustClientFromCmd(cmd).TimeToLive(context.TODO(), leaseFromArgs(args[0]), opts...)
	if rerr != nil {
		ExitWithError(ExitBadConnection, rerr)
	}
	display.TimeToLive(*resp, timeToLiveKeys)
}

// NewLeaseListCommand returns the cobra command for "lease list".
func NewLeaseListCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "list",
		Short: "List all active leases",
		Run:   leaseListCommandFunc,
	}
	return lc
}

// leaseListCommandFunc executes the "lease list" command.
func leaseListCommandFunc(cmd *cobra.Command, args []string) {
	resp, rerr := mustClientFromCmd(cmd).Leases(context.TODO())
	if rerr != nil {
		ExitWithError(ExitBadConnection, rerr)
	}
	display.Leases(*resp)
}

// NewLeaseKeepAliveCommand returns the cobra command for "lease keep-alive".
func NewLeaseKeepAliveCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "keep-alive <leaseID>",
		Short: "Keeps leases alive (renew)",

		Run: leaseKeepAliveCommandFunc,
	}

	return lc
}

// leaseKeepAliveCommandFunc executes the "lease keep-alive" command.
func leaseKeepAliveCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease keep-alive command needs lease ID as argument"))
	}

	id := leaseFromArgs(args[0])
	respc, kerr := mustClientFromCmd(cmd).KeepAlive(context.TODO(), id)
	if kerr != nil {
		ExitWithError(ExitBadConnection, kerr)
	}

	for resp := range respc {
		display.KeepAlive(*resp)
	}

	if _, ok := (display).(*simplePrinter); ok {
		fmt.Printf("lease %016x expired or revoked.\n", id)
	}
}

func leaseFromArgs(arg string) v3.LeaseID {
	id, err := strconv.ParseInt(arg, 16, 64)
	if err != nil {
		ExitWithError(ExitBadArgs, fmt.Errorf("bad lease ID arg (%v), expecting ID in Hex", err))
	}
	return v3.LeaseID(id)
}
<commit_msg>etcdctl v3: adds the --once option to the lease keep-alive command<commit_after>// Copyright 2016 The etcd 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 command

import (
	"context"
	"fmt"
	"strconv"

	v3 "github.com/coreos/etcd/clientv3"

	"github.com/spf13/cobra"
)

// NewLeaseCommand returns the cobra command for "lease".
func NewLeaseCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "lease <subcommand>",
		Short: "Lease related commands",
	}

	lc.AddCommand(NewLeaseGrantCommand())
	lc.AddCommand(NewLeaseRevokeCommand())
	lc.AddCommand(NewLeaseTimeToLiveCommand())
	lc.AddCommand(NewLeaseListCommand())
	lc.AddCommand(NewLeaseKeepAliveCommand())

	return lc
}

// NewLeaseGrantCommand returns the cobra command for "lease grant".
func NewLeaseGrantCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "grant <ttl>",
		Short: "Creates leases",

		Run: leaseGrantCommandFunc,
	}

	return lc
}

// leaseGrantCommandFunc executes the "lease grant" command.
func leaseGrantCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease grant command needs TTL argument."))
	}

	ttl, err := strconv.ParseInt(args[0], 10, 64)
	if err != nil {
		ExitWithError(ExitBadArgs, fmt.Errorf("bad TTL (%v)", err))
	}

	ctx, cancel := commandCtx(cmd)
	resp, err := mustClientFromCmd(cmd).Grant(ctx, ttl)
	cancel()
	if err != nil {
		ExitWithError(ExitError, fmt.Errorf("failed to grant lease (%v)\n", err))
	}
	display.Grant(*resp)
}

// NewLeaseRevokeCommand returns the cobra command for "lease revoke".
func NewLeaseRevokeCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "revoke <leaseID>",
		Short: "Revokes leases",

		Run: leaseRevokeCommandFunc,
	}

	return lc
}

// leaseRevokeCommandFunc executes the "lease grant" command.
func leaseRevokeCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease revoke command needs 1 argument"))
	}

	id := leaseFromArgs(args[0])
	ctx, cancel := commandCtx(cmd)
	resp, err := mustClientFromCmd(cmd).Revoke(ctx, id)
	cancel()
	if err != nil {
		ExitWithError(ExitError, fmt.Errorf("failed to revoke lease (%v)\n", err))
	}
	display.Revoke(id, *resp)
}

var timeToLiveKeys bool

// NewLeaseTimeToLiveCommand returns the cobra command for "lease timetolive".
func NewLeaseTimeToLiveCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "timetolive <leaseID> [options]",
		Short: "Get lease information",

		Run: leaseTimeToLiveCommandFunc,
	}
	lc.Flags().BoolVar(&timeToLiveKeys, "keys", false, "Get keys attached to this lease")

	return lc
}

// leaseTimeToLiveCommandFunc executes the "lease timetolive" command.
func leaseTimeToLiveCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease timetolive command needs lease ID as argument"))
	}
	var opts []v3.LeaseOption
	if timeToLiveKeys {
		opts = append(opts, v3.WithAttachedKeys())
	}
	resp, rerr := mustClientFromCmd(cmd).TimeToLive(context.TODO(), leaseFromArgs(args[0]), opts...)
	if rerr != nil {
		ExitWithError(ExitBadConnection, rerr)
	}
	display.TimeToLive(*resp, timeToLiveKeys)
}

// NewLeaseListCommand returns the cobra command for "lease list".
func NewLeaseListCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "list",
		Short: "List all active leases",
		Run:   leaseListCommandFunc,
	}
	return lc
}

// leaseListCommandFunc executes the "lease list" command.
func leaseListCommandFunc(cmd *cobra.Command, args []string) {
	resp, rerr := mustClientFromCmd(cmd).Leases(context.TODO())
	if rerr != nil {
		ExitWithError(ExitBadConnection, rerr)
	}
	display.Leases(*resp)
}

var (
	leaseKeepAliveOnce bool
)

// NewLeaseKeepAliveCommand returns the cobra command for "lease keep-alive".
func NewLeaseKeepAliveCommand() *cobra.Command {
	lc := &cobra.Command{
		Use:   "keep-alive [options] <leaseID>",
		Short: "Keeps leases alive (renew)",

		Run: leaseKeepAliveCommandFunc,
	}

	lc.Flags().BoolVar(&leaseKeepAliveOnce, "once", false, "Resets the keep-alive time to its original value and exits immediately")

	return lc
}

// leaseKeepAliveCommandFunc executes the "lease keep-alive" command.
func leaseKeepAliveCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		ExitWithError(ExitBadArgs, fmt.Errorf("lease keep-alive command needs lease ID as argument"))
	}

	id := leaseFromArgs(args[0])

	if leaseKeepAliveOnce {
		respc, kerr := mustClientFromCmd(cmd).KeepAliveOnce(context.TODO(), id)
		if kerr != nil {
			ExitWithError(ExitBadConnection, kerr)
		}
		display.KeepAlive(*respc)
		return
	}

	respc, kerr := mustClientFromCmd(cmd).KeepAlive(context.TODO(), id)
	if kerr != nil {
		ExitWithError(ExitBadConnection, kerr)
	}
	for resp := range respc {
		display.KeepAlive(*resp)
	}

	if _, ok := (display).(*simplePrinter); ok {
		fmt.Printf("lease %016x expired or revoked.\n", id)
	}
}

func leaseFromArgs(arg string) v3.LeaseID {
	id, err := strconv.ParseInt(arg, 16, 64)
	if err != nil {
		ExitWithError(ExitBadArgs, fmt.Errorf("bad lease ID arg (%v), expecting ID in Hex", err))
	}
	return v3.LeaseID(id)
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd 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 v2auth

import (
	"context"
	"encoding/json"
	"path"

	"go.etcd.io/etcd/etcdserver"
	"go.etcd.io/etcd/etcdserver/api/v2error"
	"go.etcd.io/etcd/etcdserver/etcdserverpb"

	"go.uber.org/zap"
)

func (s *store) ensureAuthDirectories() error {
	if s.ensuredOnce {
		return nil
	}
	for _, res := range []string{StorePermsPrefix, StorePermsPrefix + "/users/", StorePermsPrefix + "/roles/"} {
		ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
		defer cancel()
		pe := false
		rr := etcdserverpb.Request{
			Method:    "PUT",
			Path:      res,
			Dir:       true,
			PrevExist: &pe,
		}
		_, err := s.server.Do(ctx, rr)
		if err != nil {
			if e, ok := err.(*v2error.Error); ok {
				if e.ErrorCode == v2error.EcodeNodeExist {
					continue
				}
			}
			if s.lg != nil {
				s.lg.Warn(
					"failed to create auth directories",
					zap.Error(err),
				)
			} else {
				plog.Errorf("failed to create auth directories in the store (%v)", err)
			}
			return err
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	pe := false
	rr := etcdserverpb.Request{
		Method:    "PUT",
		Path:      StorePermsPrefix + "/enabled",
		Val:       "false",
		PrevExist: &pe,
	}
	_, err := s.server.Do(ctx, rr)
	if err != nil {
		if e, ok := err.(*v2error.Error); ok {
			if e.ErrorCode == v2error.EcodeNodeExist {
				s.ensuredOnce = true
				return nil
			}
		}
		return err
	}
	s.ensuredOnce = true
	return nil
}

func (s *store) enableAuth() error {
	_, err := s.updateResource("/enabled", true)
	return err
}
func (s *store) disableAuth() error {
	_, err := s.updateResource("/enabled", false)
	return err
}

func (s *store) detectAuth() bool {
	if s.server == nil {
		return false
	}
	value, err := s.requestResource("/enabled", false)
	if err != nil {
		if e, ok := err.(*v2error.Error); ok {
			if e.ErrorCode == v2error.EcodeKeyNotFound {
				return false
			}
		}
		if s.lg != nil {
			s.lg.Warn(
				"failed to detect auth settings",
				zap.Error(err),
			)
		} else {
			plog.Errorf("failed to detect auth settings (%s)", err)
		}
		return false
	}

	var u bool
	err = json.Unmarshal([]byte(*value.Event.Node.Value), &u)
	if err != nil {
		if s.lg != nil {
			s.lg.Warn(
				"internal bookkeeping value for enabled isn't valid JSON",
				zap.Error(err),
			)
		} else {
			plog.Errorf("internal bookkeeping value for enabled isn't valid JSON (%v)", err)
		}
		return false
	}
	return u
}

func (s *store) requestResource(res string, quorum bool) (etcdserver.Response, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	p := path.Join(StorePermsPrefix, res)
	method := "GET"
	if quorum {
		method = "QGET"
	}
	rr := etcdserverpb.Request{
		Method: method,
		Path:   p,
		Dir:    false, // TODO: always false?
	}
	return s.server.Do(ctx, rr)
}

func (s *store) updateResource(res string, value interface{}) (etcdserver.Response, error) {
	return s.setResource(res, value, true)
}
func (s *store) createResource(res string, value interface{}) (etcdserver.Response, error) {
	return s.setResource(res, value, false)
}
func (s *store) setResource(res string, value interface{}, prevexist bool) (etcdserver.Response, error) {
	err := s.ensureAuthDirectories()
	if err != nil {
		return etcdserver.Response{}, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	data, err := json.Marshal(value)
	if err != nil {
		return etcdserver.Response{}, err
	}
	p := path.Join(StorePermsPrefix, res)
	rr := etcdserverpb.Request{
		Method:    "PUT",
		Path:      p,
		Val:       string(data),
		PrevExist: &prevexist,
	}
	return s.server.Do(ctx, rr)
}

func (s *store) deleteResource(res string) error {
	err := s.ensureAuthDirectories()
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	pex := true
	p := path.Join(StorePermsPrefix, res)
	_, err = s.server.Do(ctx, etcdserverpb.Request{
		Method:    "DELETE",
		Path:      p,
		PrevExist: &pex,
	})
	return err
}
<commit_msg>api/v2auth: remove defer in loop.<commit_after>// Copyright 2015 The etcd 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 v2auth

import (
	"context"
	"encoding/json"
	"path"

	"go.etcd.io/etcd/etcdserver"
	"go.etcd.io/etcd/etcdserver/api/v2error"
	"go.etcd.io/etcd/etcdserver/etcdserverpb"

	"go.uber.org/zap"
)

func (s *store) ensureAuthDirectories() error {
	if s.ensuredOnce {
		return nil
	}
	for _, res := range []string{StorePermsPrefix, StorePermsPrefix + "/users/", StorePermsPrefix + "/roles/"} {
		ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
		pe := false
		rr := etcdserverpb.Request{
			Method:    "PUT",
			Path:      res,
			Dir:       true,
			PrevExist: &pe,
		}
		_, err := s.server.Do(ctx, rr)
		cancel()
		if err != nil {
			if e, ok := err.(*v2error.Error); ok {
				if e.ErrorCode == v2error.EcodeNodeExist {
					continue
				}
			}
			if s.lg != nil {
				s.lg.Warn(
					"failed to create auth directories",
					zap.Error(err),
				)
			} else {
				plog.Errorf("failed to create auth directories in the store (%v)", err)
			}
			return err
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	pe := false
	rr := etcdserverpb.Request{
		Method:    "PUT",
		Path:      StorePermsPrefix + "/enabled",
		Val:       "false",
		PrevExist: &pe,
	}
	_, err := s.server.Do(ctx, rr)
	if err != nil {
		if e, ok := err.(*v2error.Error); ok {
			if e.ErrorCode == v2error.EcodeNodeExist {
				s.ensuredOnce = true
				return nil
			}
		}
		return err
	}
	s.ensuredOnce = true
	return nil
}

func (s *store) enableAuth() error {
	_, err := s.updateResource("/enabled", true)
	return err
}
func (s *store) disableAuth() error {
	_, err := s.updateResource("/enabled", false)
	return err
}

func (s *store) detectAuth() bool {
	if s.server == nil {
		return false
	}
	value, err := s.requestResource("/enabled", false)
	if err != nil {
		if e, ok := err.(*v2error.Error); ok {
			if e.ErrorCode == v2error.EcodeKeyNotFound {
				return false
			}
		}
		if s.lg != nil {
			s.lg.Warn(
				"failed to detect auth settings",
				zap.Error(err),
			)
		} else {
			plog.Errorf("failed to detect auth settings (%s)", err)
		}
		return false
	}

	var u bool
	err = json.Unmarshal([]byte(*value.Event.Node.Value), &u)
	if err != nil {
		if s.lg != nil {
			s.lg.Warn(
				"internal bookkeeping value for enabled isn't valid JSON",
				zap.Error(err),
			)
		} else {
			plog.Errorf("internal bookkeeping value for enabled isn't valid JSON (%v)", err)
		}
		return false
	}
	return u
}

func (s *store) requestResource(res string, quorum bool) (etcdserver.Response, error) {
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	p := path.Join(StorePermsPrefix, res)
	method := "GET"
	if quorum {
		method = "QGET"
	}
	rr := etcdserverpb.Request{
		Method: method,
		Path:   p,
		Dir:    false, // TODO: always false?
	}
	return s.server.Do(ctx, rr)
}

func (s *store) updateResource(res string, value interface{}) (etcdserver.Response, error) {
	return s.setResource(res, value, true)
}
func (s *store) createResource(res string, value interface{}) (etcdserver.Response, error) {
	return s.setResource(res, value, false)
}
func (s *store) setResource(res string, value interface{}, prevexist bool) (etcdserver.Response, error) {
	err := s.ensureAuthDirectories()
	if err != nil {
		return etcdserver.Response{}, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	data, err := json.Marshal(value)
	if err != nil {
		return etcdserver.Response{}, err
	}
	p := path.Join(StorePermsPrefix, res)
	rr := etcdserverpb.Request{
		Method:    "PUT",
		Path:      p,
		Val:       string(data),
		PrevExist: &prevexist,
	}
	return s.server.Do(ctx, rr)
}

func (s *store) deleteResource(res string) error {
	err := s.ensureAuthDirectories()
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
	defer cancel()
	pex := true
	p := path.Join(StorePermsPrefix, res)
	_, err = s.server.Do(ctx, etcdserverpb.Request{
		Method:    "DELETE",
		Path:      p,
		PrevExist: &pex,
	})
	return err
}
<|endoftext|>


<commit_before>package main

import (
	neural "./neural"
	problems "./problems"
	"fmt"
	sdl "github.com/veandco/go-sdl2/sdl"
	"math"
	"runtime"
	"sort"
	"time"
)

type FBird struct {
	bird  *problems.Bird
	brain *neural.Net
	bestX float64
	dead  bool
}

type Flock []FBird

func (birds Flock) Len() int {
	return len(birds)
}

func (birds Flock) Less(c, r int) bool {
	return birds[c].bestX > birds[r].bestX
}

func (birds Flock) Swap(c, r int) {
	birds[c], birds[r] = birds[r], birds[c]
}

// will check if going from pos to next will collide
func checkFlock(birds Flock, lvl *problems.Level) {

	collide := func(aX, bX, cX float64) bool {
		// c.X == d.X
		return aX-1 <= cX && bX+1 >= cX
	}

	hSize := float64(problems.PylonHole / 2)

	for c := range birds {
		if birds[c].bird.Pos.Y >= lvl.GetSize().Y || birds[c].bird.Pos.Y < 1 {
			// hit ceeling or floor
			birds[c].dead = true
			continue
		}

		pylon := lvl.ClosestPylon(&birds[c].bird.Pos)
		if birds[c].bird.Pos.Y >= pylon.Y-hSize && birds[c].bird.Pos.Y <= pylon.Y+hSize {
			// in the pylon hole
			continue
		}

		if birds[c].bird.Pos.Y > pylon.Y {
			// bottom pylon segment
			birds[c].dead = collide(birds[c].bird.Pos.X, birds[c].bird.NextPos.X, pylon.X)
		} else {
			// top pylon segment
			birds[c].dead = collide(birds[c].bird.Pos.X, birds[c].bird.NextPos.X, pylon.X)
		}
	}

}

func thnikFlock(birds Flock, lvl *problems.Level) {
	wg := make(chan struct{}, len(birds))

	thinkBird := func(c int) {
		birds[c].bestX = math.Max(birds[c].bird.Pos.X, birds[c].bestX)
		next := lvl.FirstPylonAfter(&birds[c].bird.Pos)
		diffY := next.Y - birds[c].bird.Pos.Y
		diffX := next.X - birds[c].bird.Pos.X

		birds[c].brain.Stimulate(0, diffY)
		birds[c].brain.Stimulate(1, diffX)
		birds[c].brain.Stimulate(2, birds[c].bird.Vel.Y)

		birds[c].brain.Step()
		if birds[c].brain.ValueOf(7) > 0.75 {
			birds[c].bird.Vel.Y = -0.4
		}

		birds[c].brain.Clear()
		wg <- struct{}{}
	}

	for c := 0; c < len(birds); c++ {
		go thinkBird(c)
	}

	for c := 0; c < len(birds); c++ {
		<-wg
	}
}

func mutateFlock(birds Flock, lvl *problems.Level) {
	sort.Sort(birds)

	randNet := func() *neural.Net {
		return birds[int(neural.RandMax(float64(len(birds))))].brain
	}

	best := birds[0].brain

	// TODO move dead check out of this loop
	// TODO check if the bird jumps trough the pylon - kill
	for c := range birds {
		if birds[c].dead {
			birds[c].dead = false
			birds[c].bird.Pos = *lvl.NewBirdPos()
			birds[c].bird.Vel = *problems.NewVector(0.1, 0)

			birds[c].brain = neural.Cross(best, randNet())

			if neural.Chance(0.1) {
				birds[c].brain.Mutate(0.33)
			}
		} else {
			birds[c].bird.Pos = birds[c].bird.NextPos
		}
	}

}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	doDraw := true

	W := 1500
	H := 800
	LVL_W := W * 50
	fmt.Println(W, H)
	var FPS float64 = 60.0

	lvl := problems.NewLevel(LVL_W, H)

	sdl.Init(sdl.INIT_EVERYTHING)

	window, err := sdl.CreateWindow("test", sdl.WINDOWPOS_UNDEFINED, sdl.WINDOWPOS_UNDEFINED,
		int(W), int(H), sdl.WINDOW_SHOWN)
	if err != nil {
		panic(err)
	}
	defer window.Destroy()

	surface, err := window.GetSurface()
	if err != nil {
		panic(err)
	}

	rect := sdl.Rect{0, 0, 5, 5}
	clearRect := sdl.Rect{0, 0, int32(W), int32(H)}
	surface.FillRect(&clearRect, 0xffffffff)

	bcount := 10000

	nets := make([]*neural.Net, bcount, bcount)
	for c := range nets {
		nets[c] = neural.NewNet(8)

		// diffY- to hidden
		*nets[c].Synapse(0, 3) = 0.0
		*nets[c].Synapse(0, 4) = 0.0
		*nets[c].Synapse(0, 5) = 0.0
		*nets[c].Synapse(0, 6) = 0.0

		// diffX- to hidden
		*nets[c].Synapse(1, 3) = 0.0
		*nets[c].Synapse(1, 4) = 0.0
		*nets[c].Synapse(1, 5) = 0.0
		*nets[c].Synapse(1, 6) = 0.0

		// velY - to hidden
		*nets[c].Synapse(2, 3) = 0.0
		*nets[c].Synapse(2, 4) = 0.0
		*nets[c].Synapse(2, 5) = 0.0
		*nets[c].Synapse(2, 6) = 0.0

		// hidden to output
		*nets[c].Synapse(3, 7) = 0.0
		*nets[c].Synapse(4, 7) = 0.0
		*nets[c].Synapse(5, 7) = 0.0
		*nets[c].Synapse(6, 7) = 0.0

		nets[c].Randomize()
	}

	lvl.AddBirds(bcount)
	flock := make(Flock, bcount)
	for c := 0; c < bcount; c++ {
		flock[c].bird = (*lvl.GetBirds())[c]
		flock[c].brain = nets[c]
		flock[c].bestX = 0
	}

	offset := 0
	step := 65

	frame := 0
	var frameTime float64 = 1000 / FPS
	start := time.Now()
	for {

		frame++
		thnikFlock(flock, lvl)

		visible := func(x float64) bool {
			return x >= float64(offset) && x < float64(offset+W)
		}

		toScreen := func(x float64) float64 {
			return float64(x - float64(offset))
		}

		brds := *lvl.GetBirds()
		for _, brd := range brds {
			if !visible(brd.Pos.X) {
				continue
			}

			rect.X = int32(toScreen(brd.Pos.X))
			rect.Y = int32(brd.Pos.Y)
			rect.W = 5
			rect.H = 5
			if doDraw {
				surface.FillRect(&rect, 0xffff0000)
			}
		}

		hSize := float64(problems.PylonHole) / 2.0
		for _, pylon := range lvl.GetPylons() {
			if !visible(pylon.X) {
				continue
			}

			rect.X = int32(toScreen(pylon.X))
			rect.Y = int32(0)
			rect.W = 5

			// top part
			rect.H = int32(pylon.Y - hSize)
			if doDraw {
				surface.FillRect(&rect, 0xff00ff00)
			}

			// bottom part
			rect.Y = int32(pylon.Y + hSize)
			rect.H = int32(float64(H) - (pylon.Y + hSize))
			if doDraw {
				surface.FillRect(&rect, 0xff00ff00)
			}

			rect.Y = int32(pylon.Y)
			rect.W = 3
			rect.H = 3
			if doDraw {
				surface.FillRect(&rect, 0xff0000ff)
			}
		}

		elapsed := time.Since(start)

		if doDraw && frameTime < 1000.0/FPS {
			time.Sleep(time.Millisecond * time.Duration(1000.0/FPS))
		}

		start = time.Now()

		frameTime = frameTime*0.9 + float64(elapsed.Nanoseconds())*0.1

		if frame > 60 {
			frame = 0
			fmt.Printf("fps last: %s\tfps average %f\tcompletion %f%%\n", elapsed, frameTime/1000000.0, flock[0].bestX/float64(LVL_W)*100.0)
		}

		window.UpdateSurface()
		lvl.Step(float64(elapsed.Nanoseconds()) / 1000000.0)
		checkFlock(flock, lvl)
		// lvl.DoStep()

		if doDraw {
			surface.FillRect(&clearRect, 0xffffffff)
		}

		stop := false
		for event := sdl.PollEvent(); event != nil; event = sdl.PollEvent() {
			switch t := event.(type) {
			case *sdl.QuitEvent:
				stop = true
			case *sdl.KeyDownEvent:
				switch t.Keysym.Sym {
				case sdl.K_LEFT:
					offset = int(math.Max(0, float64(offset-step)))
				case sdl.K_RIGHT:
					offset = int(math.Min(float64(LVL_W-W), float64(offset+step)))
				case sdl.K_SPACE:
					doDraw = !doDraw
					surface.FillRect(&clearRect, 0xffaaaaaa)
				}
			}
		}

		if stop {
			break
		}

		mutateFlock(flock, lvl)
	}

	sdl.Quit()
}
<commit_msg>Bird best results penlized on mutation<commit_after>package main

import (
	neural "./neural"
	problems "./problems"
	"fmt"
	sdl "github.com/veandco/go-sdl2/sdl"
	"math"
	"runtime"
	"sort"
	"time"
)

type FBird struct {
	bird  *problems.Bird
	brain *neural.Net
	bestX float64
	dead  bool
}

type Flock []FBird

func (birds Flock) Len() int {
	return len(birds)
}

func (birds Flock) Less(c, r int) bool {
	return birds[c].bestX > birds[r].bestX
}

func (birds Flock) Swap(c, r int) {
	birds[c], birds[r] = birds[r], birds[c]
}

// will check if going from pos to next will collide
func checkFlock(birds Flock, lvl *problems.Level) {

	collide := func(aX, bX, cX float64) bool {
		// c.X == d.X
		return aX-1 <= cX && bX+1 >= cX
	}

	hSize := float64(problems.PylonHole / 2)

	for c := range birds {
		if birds[c].bird.Pos.Y >= lvl.GetSize().Y || birds[c].bird.Pos.Y < 1 {
			// hit ceeling or floor
			birds[c].dead = true
			continue
		}

		pylon := lvl.ClosestPylon(&birds[c].bird.Pos)
		if birds[c].bird.Pos.Y >= pylon.Y-hSize && birds[c].bird.Pos.Y <= pylon.Y+hSize {
			// in the pylon hole
			continue
		}

		if birds[c].bird.Pos.Y > pylon.Y {
			// bottom pylon segment
			birds[c].dead = collide(birds[c].bird.Pos.X, birds[c].bird.NextPos.X, pylon.X)
		} else {
			// top pylon segment
			birds[c].dead = collide(birds[c].bird.Pos.X, birds[c].bird.NextPos.X, pylon.X)
		}
	}

}

func thnikFlock(birds Flock, lvl *problems.Level) {
	wg := make(chan struct{}, len(birds))

	thinkBird := func(c int) {
		next := lvl.FirstPylonAfter(&birds[c].bird.Pos)
		diffY := next.Y - birds[c].bird.Pos.Y
		diffX := next.X - birds[c].bird.Pos.X

		birds[c].brain.Stimulate(0, diffY)
		birds[c].brain.Stimulate(1, diffX)
		birds[c].brain.Stimulate(2, birds[c].bird.Vel.Y)

		birds[c].brain.Step()
		if birds[c].brain.ValueOf(7) > 0.75 {
			birds[c].bird.Vel.Y = -0.4
		}

		birds[c].brain.Clear()
		wg <- struct{}{}
	}

	for c := 0; c < len(birds); c++ {
		go thinkBird(c)
	}

	for c := 0; c < len(birds); c++ {
		<-wg
	}
}

func mutateFlock(birds Flock, lvl *problems.Level) {
	sort.Sort(birds)

	randNet := func() *neural.Net {
		return birds[int(neural.RandMax(float64(len(birds))))].brain
	}

	best := birds[0].brain

	for c := range birds {
		if birds[c].dead {
			birds[c].dead = false
			birds[c].bird.Pos = *lvl.NewBirdPos()
			birds[c].bird.Vel = *problems.NewVector(0.1, 0)

			birds[c].brain = neural.Cross(best, randNet())

			if neural.Chance(0.1) {
				// penalize best achievement due to mutation
				birds[c].bestX *= 0.66
				birds[c].brain.Mutate(0.33)
			}
		} else {
			birds[c].bird.Pos = birds[c].bird.NextPos
			birds[c].bestX = math.Max(birds[c].bird.Pos.X, birds[c].bestX)
		}
	}

}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	doDraw := true

	W := 1500
	H := 800
	LVL_W := W * 50
	fmt.Println(W, H)
	var FPS float64 = 60.0

	lvl := problems.NewLevel(LVL_W, H)

	sdl.Init(sdl.INIT_EVERYTHING)

	window, err := sdl.CreateWindow("test", sdl.WINDOWPOS_UNDEFINED, sdl.WINDOWPOS_UNDEFINED,
		int(W), int(H), sdl.WINDOW_SHOWN)
	if err != nil {
		panic(err)
	}
	defer window.Destroy()

	surface, err := window.GetSurface()
	if err != nil {
		panic(err)
	}

	rect := sdl.Rect{0, 0, 5, 5}
	clearRect := sdl.Rect{0, 0, int32(W), int32(H)}
	surface.FillRect(&clearRect, 0xffffffff)

	bcount := 10000

	nets := make([]*neural.Net, bcount, bcount)
	for c := range nets {
		nets[c] = neural.NewNet(8)

		// diffY- to hidden
		*nets[c].Synapse(0, 3) = 0.0
		*nets[c].Synapse(0, 4) = 0.0
		*nets[c].Synapse(0, 5) = 0.0
		*nets[c].Synapse(0, 6) = 0.0

		// diffX- to hidden
		*nets[c].Synapse(1, 3) = 0.0
		*nets[c].Synapse(1, 4) = 0.0
		*nets[c].Synapse(1, 5) = 0.0
		*nets[c].Synapse(1, 6) = 0.0

		// velY - to hidden
		*nets[c].Synapse(2, 3) = 0.0
		*nets[c].Synapse(2, 4) = 0.0
		*nets[c].Synapse(2, 5) = 0.0
		*nets[c].Synapse(2, 6) = 0.0

		// hidden to output
		*nets[c].Synapse(3, 7) = 0.0
		*nets[c].Synapse(4, 7) = 0.0
		*nets[c].Synapse(5, 7) = 0.0
		*nets[c].Synapse(6, 7) = 0.0

		nets[c].Randomize()
	}

	lvl.AddBirds(bcount)
	flock := make(Flock, bcount)
	for c := 0; c < bcount; c++ {
		flock[c].bird = (*lvl.GetBirds())[c]
		flock[c].brain = nets[c]
		flock[c].bestX = 0
	}

	offset := 0
	step := 65

	frame := 0
	var frameTime float64 = 1000 / FPS
	start := time.Now()
	for {

		frame++
		thnikFlock(flock, lvl)

		visible := func(x float64) bool {
			return x >= float64(offset) && x < float64(offset+W)
		}

		toScreen := func(x float64) float64 {
			return float64(x - float64(offset))
		}

		brds := *lvl.GetBirds()
		for _, brd := range brds {
			if !visible(brd.Pos.X) {
				continue
			}

			rect.X = int32(toScreen(brd.Pos.X))
			rect.Y = int32(brd.Pos.Y)
			rect.W = 5
			rect.H = 5
			if doDraw {
				surface.FillRect(&rect, 0xffff0000)
			}
		}

		hSize := float64(problems.PylonHole) / 2.0
		for _, pylon := range lvl.GetPylons() {
			if !visible(pylon.X) {
				continue
			}

			rect.X = int32(toScreen(pylon.X))
			rect.Y = int32(0)
			rect.W = 5

			// top part
			rect.H = int32(pylon.Y - hSize)
			if doDraw {
				surface.FillRect(&rect, 0xff00ff00)
			}

			// bottom part
			rect.Y = int32(pylon.Y + hSize)
			rect.H = int32(float64(H) - (pylon.Y + hSize))
			if doDraw {
				surface.FillRect(&rect, 0xff00ff00)
			}

			rect.Y = int32(pylon.Y)
			rect.W = 3
			rect.H = 3
			if doDraw {
				surface.FillRect(&rect, 0xff0000ff)
			}
		}

		elapsed := time.Since(start)

		if doDraw && frameTime < 1000.0/FPS {
			time.Sleep(time.Millisecond * time.Duration(1000.0/FPS))
		}

		start = time.Now()

		frameTime = frameTime*0.9 + float64(elapsed.Nanoseconds())*0.1

		if frame > 60 {
			frame = 0
			fmt.Printf("fps last: %s\tfps average %f\tcompletion %f%%\n", elapsed, frameTime/1000000.0, flock[0].bestX/float64(LVL_W)*100.0)
		}

		window.UpdateSurface()
		lvl.Step(float64(elapsed.Nanoseconds()) / 1000000.0)
		checkFlock(flock, lvl)
		// lvl.DoStep()

		if doDraw {
			surface.FillRect(&clearRect, 0xffffffff)
		}

		stop := false
		for event := sdl.PollEvent(); event != nil; event = sdl.PollEvent() {
			switch t := event.(type) {
			case *sdl.QuitEvent:
				stop = true
			case *sdl.KeyDownEvent:
				switch t.Keysym.Sym {
				case sdl.K_LEFT:
					offset = int(math.Max(0, float64(offset-step)))
				case sdl.K_RIGHT:
					offset = int(math.Min(float64(LVL_W-W), float64(offset+step)))
				case sdl.K_SPACE:
					doDraw = !doDraw
					surface.FillRect(&clearRect, 0xffaaaaaa)
				}
			}
		}

		if stop {
			break
		}

		mutateFlock(flock, lvl)
	}

	sdl.Quit()
}
<|endoftext|>


<commit_before>package main

import (
	//"bufio"
	//"fmt"
	"io/ioutil"
	"strings"
	//"os"
)

func main() {
	//println("hello world")
	cpu := read_cpu_info()
	//println(strings.Split(cpu, "\n"))
	for _, element := range strings.Split(cpu, "\n") {
		println(element)
		println()
	}
}

func read_cpu_info() string {
	data, err := ioutil.ReadFile("/proc/stat")
	if err != nil {
		panic(err)
	}
	return string(data)
}
<commit_msg>add split_on_newline function<commit_after>package main

import (
	//"bufio"
	//"fmt"
	"io/ioutil"
	"strings"
	//"os"
)

func main() {
	//println("hello world")
	cpu := read_cpu_info()
	for _, element := range split_on_newline(cpu) {
		println(element)
		println()
	}
}

func read_cpu_info() string {
	data, err := ioutil.ReadFile("/proc/stat")
	if err != nil {
		panic(err)
	}
	return string(data)
}

func split_on_newline(str string) []string {
	return strings.Split(str, "\n")
}
<|endoftext|>


<commit_before>// timeutil provides a set of time utilities including comparisons,
// conversion to "DT8" int32 and "DT14" int64 formats and other
// capabilities.
package timeutil

import (
	"fmt"
	"strings"
	"time"
)

// More predefined layouts for use in Time.Format and time.Parse.
const (
	DT14               = "20060102150405"
	DT8                = "20060102"
	DT6                = "200601"
	RFC3339FullDate    = "2006-01-02"
	ISO8601            = "2006-01-02T15:04:05Z0700"
	ISO8601TZHour      = "2006-01-02T15:04:05Z07"
	ISO8601MilliNoTZ   = "2006-01-02T15:04:05.000"
	ISO8601CompactZ    = "20060102T150405Z0700"
	ISO8601CompactNoTZ = "20060102T150405"
	ISO8601YM          = "2006-01"
	InsightlyApiQuery  = "_1/_2/2006 _3:04:05 PM"
	SQLTimestamp       = "2006-01-02 15:04:05" // MySQL, BigQuery, etc.
	DateMDYSlash       = "01/02/2006"
	DateDMYHM2         = "02:01:06 15:04" // GMT time in format dd:mm:yy hh:mm
)

const (
	RFC3339Min         = "0000-01-01T00:00:00Z"
	RFC3339Max         = "9999-12-31T23:59:59Z"
	RFC3339Zero        = "0001-01-01T00:00:00Z"
	RFC3339ZeroUnix    = "1970-01-01T00:00:00Z"
	RFC3339YMDZeroUnix = int64(-62135596800)
)

// Reformat a time string from one format to another
func FromTo(value, fromLayout, toLayout string) (string, error) {
	t, err := time.Parse(fromLayout, strings.TrimSpace(value))
	if err != nil {
		return "", err
	}
	return t.Format(toLayout), nil
}

// ParseOrZero returns a parsed time.Time or the RFC-3339 zero time.
func ParseOrZero(layout, value string) time.Time {
	t, err := time.Parse(layout, value)
	if err != nil {
		return TimeRFC3339Zero()
	}
	return t
}

// ParseFirst attempts to parse a string with a set of layouts.
func ParseFirst(layouts []string, value string) (time.Time, error) {
	value = strings.TrimSpace(value)
	if len(value) == 0 || len(layouts) == 0 {
		return time.Now(), fmt.Errorf(
			"Requires value [%v] and at least one layout [%v]", value, strings.Join(layouts, ","))
	}
	for _, layout := range layouts {
		layout = strings.TrimSpace(layout)
		if len(layout) == 0 {
			continue
		}
		if dt, err := time.Parse(layout, value); err == nil {
			return dt, nil
		}
	}
	return time.Now(), fmt.Errorf("Cannot parse time [%v] with layouts [%v]",
		value, strings.Join(layouts, ","))
}

var FormatMap = map[string]string{
	"RFC3339":    time.RFC3339,
	"RFC3339YMD": RFC3339FullDate,
	"ISO8601YM":  ISO8601YM,
}

func GetFormat(formatName string) (string, error) {
	format, ok := FormatMap[strings.TrimSpace(formatName)]
	if !ok {
		return "", fmt.Errorf("Format Not Found: %v", format)
	}
	return format, nil
}

// FormatQuarter takes quarter time and formats it using "Q# YYYY".
func FormatQuarter(t time.Time) string {
	return fmt.Sprintf("Q%d %d", MonthToQuarter(uint8(t.Month())), t.Year())
}

// FormatQuarter takes quarter time and formats it using "Q# YYYY".
func FormatQuarterYYYYQ(t time.Time) string {
	return fmt.Sprintf("%d Q%d", t.Year(), MonthToQuarter(uint8(t.Month())))
}

func TimeRFC3339Min() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339Min)
	return t0
}

func TimeRFC3339Zero() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339Zero)
	return t0
}

func TimeRFC3339ZeroUnix() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339ZeroUnix)
	return t0
}

func IsZeroAny(u time.Time) bool { return TimeIsZeroAny(u) }

func TimeIsZeroAny(u time.Time) bool {
	if u.Equal(TimeRFC3339Zero()) ||
		u.Equal(TimeRFC3339Min()) ||
		u.Equal(TimeRFC3339ZeroUnix()) {
		return true
	}
	return false
}

type RFC3339YMDTime struct{ time.Time }

type ISO8601NoTzMilliTime struct{ time.Time }

func (t *RFC3339YMDTime) UnmarshalJSON(buf []byte) error {
	tt, isNil, err := timeUnmarshalJSON(buf, RFC3339FullDate)
	if err != nil || isNil {
		return err
	}
	t.Time = tt
	return nil
}

func (t RFC3339YMDTime) MarshalJSON() ([]byte, error) {
	return timeMarshalJSON(t.Time, RFC3339FullDate)
}

func (t *ISO8601NoTzMilliTime) UnmarshalJSON(buf []byte) error {
	tt, isNil, err := timeUnmarshalJSON(buf, ISO8601MilliNoTZ)
	if err != nil || isNil {
		return err
	}
	t.Time = tt
	return nil
}

func (t ISO8601NoTzMilliTime) MarshalJSON() ([]byte, error) {
	return timeMarshalJSON(t.Time, ISO8601MilliNoTZ)
}

func timeUnmarshalJSON(buf []byte, layout string) (time.Time, bool, error) {
	str := string(buf)
	isNil := true
	if str == "null" || str == "\"\"" {
		return time.Time{}, isNil, nil
	}
	tt, err := time.Parse(layout, strings.Trim(str, `"`))
	if err != nil {
		return time.Time{}, false, err
	}
	return tt, false, nil
}

func timeMarshalJSON(t time.Time, layout string) ([]byte, error) {
	return []byte(`"` + t.Format(layout) + `"`), nil
}
<commit_msg>add timeutil.IsZero() convenience function<commit_after>// timeutil provides a set of time utilities including comparisons,
// conversion to "DT8" int32 and "DT14" int64 formats and other
// capabilities.
package timeutil

import (
	"fmt"
	"strings"
	"time"
)

// More predefined layouts for use in Time.Format and time.Parse.
const (
	DT14               = "20060102150405"
	DT8                = "20060102"
	DT6                = "200601"
	RFC3339FullDate    = "2006-01-02"
	ISO8601            = "2006-01-02T15:04:05Z0700"
	ISO8601TZHour      = "2006-01-02T15:04:05Z07"
	ISO8601MilliNoTZ   = "2006-01-02T15:04:05.000"
	ISO8601CompactZ    = "20060102T150405Z0700"
	ISO8601CompactNoTZ = "20060102T150405"
	ISO8601YM          = "2006-01"
	InsightlyApiQuery  = "_1/_2/2006 _3:04:05 PM"
	SQLTimestamp       = "2006-01-02 15:04:05" // MySQL, BigQuery, etc.
	DateMDYSlash       = "01/02/2006"
	DateDMYHM2         = "02:01:06 15:04" // GMT time in format dd:mm:yy hh:mm
)

const (
	RFC3339Min         = "0000-01-01T00:00:00Z"
	RFC3339Max         = "9999-12-31T23:59:59Z"
	RFC3339Zero        = "0001-01-01T00:00:00Z"
	RFC3339ZeroUnix    = "1970-01-01T00:00:00Z"
	RFC3339YMDZeroUnix = int64(-62135596800)
)

// Reformat a time string from one format to another
func FromTo(value, fromLayout, toLayout string) (string, error) {
	t, err := time.Parse(fromLayout, strings.TrimSpace(value))
	if err != nil {
		return "", err
	}
	return t.Format(toLayout), nil
}

// ParseOrZero returns a parsed time.Time or the RFC-3339 zero time.
func ParseOrZero(layout, value string) time.Time {
	t, err := time.Parse(layout, value)
	if err != nil {
		return TimeRFC3339Zero()
	}
	return t
}

// ParseFirst attempts to parse a string with a set of layouts.
func ParseFirst(layouts []string, value string) (time.Time, error) {
	value = strings.TrimSpace(value)
	if len(value) == 0 || len(layouts) == 0 {
		return time.Now(), fmt.Errorf(
			"Requires value [%v] and at least one layout [%v]", value, strings.Join(layouts, ","))
	}
	for _, layout := range layouts {
		layout = strings.TrimSpace(layout)
		if len(layout) == 0 {
			continue
		}
		if dt, err := time.Parse(layout, value); err == nil {
			return dt, nil
		}
	}
	return time.Now(), fmt.Errorf("Cannot parse time [%v] with layouts [%v]",
		value, strings.Join(layouts, ","))
}

var FormatMap = map[string]string{
	"RFC3339":    time.RFC3339,
	"RFC3339YMD": RFC3339FullDate,
	"ISO8601YM":  ISO8601YM,
}

func GetFormat(formatName string) (string, error) {
	format, ok := FormatMap[strings.TrimSpace(formatName)]
	if !ok {
		return "", fmt.Errorf("Format Not Found: %v", format)
	}
	return format, nil
}

// FormatQuarter takes quarter time and formats it using "Q# YYYY".
func FormatQuarter(t time.Time) string {
	return fmt.Sprintf("Q%d %d", MonthToQuarter(uint8(t.Month())), t.Year())
}

// FormatQuarter takes quarter time and formats it using "Q# YYYY".
func FormatQuarterYYYYQ(t time.Time) string {
	return fmt.Sprintf("%d Q%d", t.Year(), MonthToQuarter(uint8(t.Month())))
}

func TimeRFC3339Min() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339Min)
	return t0
}

func TimeRFC3339Zero() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339Zero)
	return t0
}

func TimeRFC3339ZeroUnix() time.Time {
	t0, _ := time.Parse(time.RFC3339, RFC3339ZeroUnix)
	return t0
}

func IsZero(u time.Time) bool { return TimeIsZeroAny(u) }

func IsZeroAny(u time.Time) bool { return TimeIsZeroAny(u) }

func TimeIsZeroAny(u time.Time) bool {
	if u.Equal(TimeRFC3339Zero()) ||
		u.Equal(TimeRFC3339Min()) ||
		u.Equal(TimeRFC3339ZeroUnix()) {
		return true
	}
	return false
}

type RFC3339YMDTime struct{ time.Time }

type ISO8601NoTzMilliTime struct{ time.Time }

func (t *RFC3339YMDTime) UnmarshalJSON(buf []byte) error {
	tt, isNil, err := timeUnmarshalJSON(buf, RFC3339FullDate)
	if err != nil || isNil {
		return err
	}
	t.Time = tt
	return nil
}

func (t RFC3339YMDTime) MarshalJSON() ([]byte, error) {
	return timeMarshalJSON(t.Time, RFC3339FullDate)
}

func (t *ISO8601NoTzMilliTime) UnmarshalJSON(buf []byte) error {
	tt, isNil, err := timeUnmarshalJSON(buf, ISO8601MilliNoTZ)
	if err != nil || isNil {
		return err
	}
	t.Time = tt
	return nil
}

func (t ISO8601NoTzMilliTime) MarshalJSON() ([]byte, error) {
	return timeMarshalJSON(t.Time, ISO8601MilliNoTZ)
}

func timeUnmarshalJSON(buf []byte, layout string) (time.Time, bool, error) {
	str := string(buf)
	isNil := true
	if str == "null" || str == "\"\"" {
		return time.Time{}, isNil, nil
	}
	tt, err := time.Parse(layout, strings.Trim(str, `"`))
	if err != nil {
		return time.Time{}, false, err
	}
	return tt, false, nil
}

func timeMarshalJSON(t time.Time, layout string) ([]byte, error) {
	return []byte(`"` + t.Format(layout) + `"`), nil
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"os/exec"
	"strconv"
	"time"
)

func generator() {
	generatorLock.Lock()
	defer generatorLock.Unlock()
	logger.Info("updating new list...")
	for _, tunnel := range tunnels {
		ipt.ClearChain("mangle", tunnel.ChainName)
	}
	var weightSum, i int
	for _, tunnel := range tunnels {
		if tunnel.status {
			weightSum += tunnel.Weight
		}
	}
	for _, tunnel := range tunnels {
		if tunnel.status {
			for j := 0; j < tunnel.Weight; j++ {
				logger.Debug("iptables -t mangle -A %s -m statistic --mode nth --every %s --packet %s -j MARK --set-xmark %s", tunnel.ChainName, strconv.Itoa(weightSum), strconv.Itoa(i), tunnel.Mark)
				ipt.Append("mangle", tunnel.ChainName, "-m", "statistic", "--mode", "nth", "--every", strconv.Itoa(weightSum), "--packet", strconv.Itoa(i), "-j", "MARK", "--set-xmark", tunnel.Mark)
				i++
			}
		}
	}
	logger.Trace("updating finished")
}

func monitor() {
	pin.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		logger.Trace("ping %s %f ms", addr.String(), rtt.Seconds())
		tunInfo := tunnels[ip2tunnel[addr.String()]]
		tunInfo.lastlive = time.Now()
		if tunInfo.status == false {
			logger.Info("connection recover %s(%s)", ip2tunnel[addr.String()], tunInfo.Ip)
			tunInfo.status = true
			if tunInfo.RecoverCommand != "" {
				exec.Command("bash", "-c", tunInfo.RecoverCommand)
			}
			generator()
		}
	}
	pin.OnIdle = func() {}
	pin.MaxRTT = time.Second * time.Duration(rtt)
	go func() {
		for {
			pin.Run()
			time.Sleep(time.Second * time.Duration(rtt))
			<-pin.Done()
		}
	}()

	for {
		regenerate := false
		for tunName, tunInfo := range tunnels {
			if tunInfo.lastlive.Add(time.Second * time.Duration(determineTime)).Before(time.Now()) {
				if tunInfo.status {
					logger.Info("connection lost %s(%s)", tunName, tunInfo.Ip)
					tunInfo.status = false
					regenerate = true
					if tunInfo.DownCommand != "" {
						exec.Command("bash", "-c", tunInfo.DownCommand)
					}
				}
			}
		}
		if regenerate {
			generator()
		}
		time.Sleep(time.Second * time.Duration(detectingDuration))
	}
}
<commit_msg>fix balance serial in different chain<commit_after>package main

import (
	"net"
	"os/exec"
	"strconv"
	"time"
)

func generator() {
	generatorLock.Lock()
	defer generatorLock.Unlock()
	logger.Info("updating new list...")
	for _, tunnel := range tunnels {
		ipt.ClearChain("mangle", tunnel.ChainName)
	}
	var weightSum, i int
	var chainSerials map[string]int
	var chainSerialsTotal map[string]int
	chainSerials = make(map[string]int)
	chainSerialsTotal = make(map[string]int)
	for _, tunnel := range tunnels {
		if tunnel.status {
			weightSum += tunnel.Weight
		}
	}
	for _, tunnel := range tunnels {
		if tunnel.status {
			for j := 0; j < tunnel.Weight; j++ {
				_, exists := chainSerialsTotal[tunnel.ChainName]
				if !exists {
					chainSerialsTotal[tunnel.ChainName] = 0
				}
				chainSerialsTotal[tunnel.ChainName]++
				i++
			}
		}
	}
	for _, tunnel := range tunnels {
		if tunnel.status {
			for j := 0; j < tunnel.Weight; j++ {
				_, exists := chainSerials[tunnel.ChainName]
				if !exists {
					chainSerials[tunnel.ChainName] = 0
				}
				logger.Debug("iptables -t mangle -A %s -m statistic --mode nth --every %s --packet %s -j MARK --set-xmark %s", tunnel.ChainName, strconv.Itoa(weightSum), strconv.Itoa(i), tunnel.Mark)
				ipt.Append("mangle", tunnel.ChainName, "-m", "statistic", "--mode", "nth", "--every", strconv.Itoa(chainSerialsTotal[tunnel.ChainName]), "--packet", strconv.Itoa(chainSerials[tunnel.ChainName]), "-j", "MARK", "--set-xmark", tunnel.Mark)
				chainSerials[tunnel.ChainName]++
				i++
			}
		}
	}
	logger.Trace("updating finished")
}

func monitor() {
	pin.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		logger.Trace("ping %s %f ms", addr.String(), rtt.Seconds())
		tunInfo := tunnels[ip2tunnel[addr.String()]]
		tunInfo.lastlive = time.Now()
		if tunInfo.status == false {
			logger.Info("connection recover %s(%s)", ip2tunnel[addr.String()], tunInfo.Ip)
			tunInfo.status = true
			if tunInfo.RecoverCommand != "" {
				exec.Command("bash", "-c", tunInfo.RecoverCommand)
			}
			generator()
		}
	}
	pin.OnIdle = func() {}
	pin.MaxRTT = time.Second * time.Duration(rtt)
	go func() {
		for {
			pin.Run()
			time.Sleep(time.Second * time.Duration(rtt))
			<-pin.Done()
		}
	}()

	for {
		regenerate := false
		for tunName, tunInfo := range tunnels {
			if tunInfo.lastlive.Add(time.Second * time.Duration(determineTime)).Before(time.Now()) {
				if tunInfo.status {
					logger.Info("connection lost %s(%s)", tunName, tunInfo.Ip)
					tunInfo.status = false
					regenerate = true
					if tunInfo.DownCommand != "" {
						exec.Command("bash", "-c", tunInfo.DownCommand)
					}
				}
			}
		}
		if regenerate {
			generator()
		}
		time.Sleep(time.Second * time.Duration(detectingDuration))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"sort"
	"strconv"
	"strings"
	"time"

	"flag"

	"fmt"

	"os"

	"github.com/Shopify/sarama"
	"github.com/fatih/color"
	"github.com/olekukonko/tablewriter"
	"github.com/samuel/go-zookeeper/zk"
)

var (
	zkAddr            = flag.String("zkAddr", "", "zookeeper address")
	basePath          = flag.String("basePath", "/kafka", "kafka base path in zookeeper")
	brokers           = flag.String("brokers", "localhost:9092", "brokers' address")
	topic             = flag.String("topic", "trigger", "topic name")
	group             = flag.String("group", "default", "consumer group name")
	version           = flag.String("version", "0.10.0.1", "kafka version. min version is 0.8.2.0")
	lagThreshold      = flag.Int("lagThreshold", 1000, "alarm lag threshold for partition")
	totalLagThreshold = flag.Int("totalLagThreshold", 5000, "alarm total lag threshold for topic")
	interval          = flag.Duration("duration", time.Minute, "check interval time")
	informEmail       = flag.String("email", "xxx@xxxxxx.com", "inform user email")
	smtpHost          = flag.String("smtp.host", "smtp.sina.com", "smtp host for sending alarms")
	smtpPort          = flag.Int("smtp.port", 25, "smtp port for sending alarms")
	smtpUser          = flag.String("smtp.user", "kafka_monitor", "smtp user for sending alarms")
	smtpPassword      = flag.String("smtp.password", "xxxxxx", "smtp user password for sending alarms")
)

var (
	maybeProblem       = false
	lastTriggeredTime  time.Time
	mergeAlertDuration = 10 * time.Minute
	kafkaVersions      = kafkaVersion()
)

func kafkaVersion() map[string]sarama.KafkaVersion {
	m := make(map[string]sarama.KafkaVersion)
	m["0.8.2.0"] = sarama.V0_8_2_0
	m["0.8.2.1"] = sarama.V0_8_2_1
	m["0.8.2.2"] = sarama.V0_8_2_2
	m["0.9.0.0"] = sarama.V0_9_0_0
	m["0.9.0.1"] = sarama.V0_9_0_1
	m["0.10.0.0"] = sarama.V0_10_0_0
	m["0.10.0.1"] = sarama.V0_10_0_1
	m["0.10.1.0"] = sarama.V0_10_1_0
	return m
}

func main() {
	flag.Parse()

	for {
		check()
	}
}

func check() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("check error: %v", r)
			alert(*informEmail, *brokers, *topic, []byte(fmt.Sprintf("%v", r)), *smtpHost, *smtpPort, *smtpUser, *smtpPassword)
		}
	}()

	kafkaBrokers := strings.Split(*brokers, ",")
	v := kafkaVersions[*version]
	client := NewSaramaClient(kafkaBrokers, v)

	var buf bytes.Buffer

	ticker := time.NewTicker(*interval)
	for range ticker.C {
		buf.Reset()
		partitions, err := client.Partitions(*topic)
		if err != nil {
			fmt.Printf("failed to get partitions for topic=%s, err=%v\n", *topic, err)
			panic(err)
		}

		writablePartitions, err := client.WritablePartitions(*topic)
		if err != nil {
			fmt.Printf("failed to get writable partitions for topic=%s, err=%v\n", *topic, err)
			panic(err)
		}

		if len(partitions) != len(writablePartitions) {
			buf.WriteString("some partitions are not writable\n")
			buf.WriteString(fmt.Sprintf("all partitions: %v\n", partitions))
			buf.WriteString(fmt.Sprintf("writable partitions: %v\n", writablePartitions))
			//TODO print unwritable partitions
			maybeProblem = true
		}

		fmt.Printf("Time: %s\n", time.Now().Format("2006-01-02 15:04:05"))
		fmt.Printf("Brokers: %s\nVersion: %s\n", color.GreenString(*brokers), color.GreenString(*version))
		fmt.Printf("Topic: %s, Group: %s, Partitions: %s\n\n",
			color.GreenString(*topic), color.GreenString(*group), color.GreenString(strconv.Itoa(len(partitions))))

		//partition info
		var c *zk.Conn
		if *zkAddr != "" {
			c, _, err = zk.Connect(strings.Split(*zkAddr, ","), 5*time.Second)
			if err != nil {
				panic(err)
			}
		}

		infos := GetPartitionInfo(client, *topic, partitions, c, *basePath)
		if len(infos) > 0 {
			table := tablewriter.NewWriter(&buf)
			table.SetHeader([]string{"partition", "leader address", "leader", "replicas", "isr"})
			for _, info := range infos {
				replicas := fmt.Sprintf("%v", info.Replicas)
				replicas = replicas[1 : len(replicas)-1]
				replicas = compareString(replicas)
				isr := compareString(info.Isr)

				table.Append([]string{strconv.Itoa(int(info.Partition)), info.LeaderAddress, strconv.Itoa(int(info.Leader)), "[" + replicas + "]", "[" + isr + "]"})

				if replicas != isr {
					maybeProblem = true
				}
			}

			table.SetAlignment(tablewriter.ALIGN_LEFT)
			table.Render()

			buf.WriteString("\n\n")
		}
		//lag
		if offsets, err := FetchOffsets(client, *topic, *group); err == nil {

			table := tablewriter.NewWriter(&buf)
			table.SetHeader([]string{"partition", "end of log", "group offset", "lag"})

			var totalLag int64
			for _, info := range offsets {
				lag := info.PartitionOffset - info.GroupOffset
				if int(lag) > *lagThreshold {
					maybeProblem = true
				}

				totalLag += lag
				table.Append([]string{strconv.Itoa(int(info.Partition)), strconv.Itoa(int(info.PartitionOffset)),
					strconv.Itoa(int(info.GroupOffset)), convertLag(lag, *lagThreshold)})
			}

			if int(totalLag) > *totalLagThreshold {
				maybeProblem = true
			}

			table.SetFooter([]string{"", "", "Total", convertLag(totalLag, *totalLagThreshold)})
			table.SetAlignment(tablewriter.ALIGN_LEFT)
			table.SetFooterAlignment(tablewriter.ALIGN_LEFT)
			table.Render()
		}

		//output tables to stdout
		bytes := buf.Bytes()
		os.Stdout.Write(bytes)

		if maybeProblem && time.Since(lastTriggeredTime) > mergeAlertDuration {
			lastTriggeredTime = time.Now()
			alert(*informEmail, *brokers, *topic, bytes, *smtpHost, *smtpPort, *smtpUser, *smtpPassword)
		}

		maybeProblem = false
	}
}

func compareString(s string) string {
	s1 := strings.Split(s, " ")
	sort.Sort(sort.StringSlice(s1))

	return strings.Join(s1, " ")
}
func convertLag(lag int64, threshold int) string {
	lagStr := strconv.Itoa(int(lag))
	if int(lag) > threshold {
		lagStr = color.RedString(lagStr)
	}

	return lagStr
}

func NewSaramaClient(brokers []string, version sarama.KafkaVersion) sarama.Client {
	config := sarama.NewConfig()
	config.Version = version
	config.Consumer.Return.Errors = true
	config.Metadata.RefreshFrequency = 1 * time.Minute
	config.Metadata.Retry.Max = 10
	config.Net.MaxOpenRequests = 10
	config.Net.DialTimeout = 10 * time.Second
	config.Net.ReadTimeout = 10 * time.Second
	config.Net.WriteTimeout = 10 * time.Second

	client, err := sarama.NewClient(brokers, config)

	if err != nil {
		panic("Failed to start client: " + err.Error())
	}

	return client
}
<commit_msg>add broker/version info in email<commit_after>package main

import (
	"bytes"
	"sort"
	"strconv"
	"strings"
	"time"

	"flag"

	"fmt"

	"os"

	"github.com/Shopify/sarama"
	"github.com/fatih/color"
	"github.com/olekukonko/tablewriter"
	"github.com/samuel/go-zookeeper/zk"
)

var (
	zkAddr            = flag.String("zkAddr", "", "zookeeper address")
	basePath          = flag.String("basePath", "/kafka", "kafka base path in zookeeper")
	brokers           = flag.String("brokers", "localhost:9092", "brokers' address")
	topic             = flag.String("topic", "trigger", "topic name")
	group             = flag.String("group", "default", "consumer group name")
	version           = flag.String("version", "0.10.0.1", "kafka version. min version is 0.8.2.0")
	lagThreshold      = flag.Int("lagThreshold", 1000, "alarm lag threshold for partition")
	totalLagThreshold = flag.Int("totalLagThreshold", 5000, "alarm total lag threshold for topic")
	interval          = flag.Duration("duration", time.Minute, "check interval time")
	informEmail       = flag.String("email", "xxx@xxxxxx.com", "inform user email")
	smtpHost          = flag.String("smtp.host", "smtp.sina.com", "smtp host for sending alarms")
	smtpPort          = flag.Int("smtp.port", 25, "smtp port for sending alarms")
	smtpUser          = flag.String("smtp.user", "kafka_monitor", "smtp user for sending alarms")
	smtpPassword      = flag.String("smtp.password", "xxxxxx", "smtp user password for sending alarms")
)

var (
	maybeProblem       = false
	lastTriggeredTime  time.Time
	mergeAlertDuration = 10 * time.Minute
	kafkaVersions      = kafkaVersion()
)

func kafkaVersion() map[string]sarama.KafkaVersion {
	m := make(map[string]sarama.KafkaVersion)
	m["0.8.2.0"] = sarama.V0_8_2_0
	m["0.8.2.1"] = sarama.V0_8_2_1
	m["0.8.2.2"] = sarama.V0_8_2_2
	m["0.9.0.0"] = sarama.V0_9_0_0
	m["0.9.0.1"] = sarama.V0_9_0_1
	m["0.10.0.0"] = sarama.V0_10_0_0
	m["0.10.0.1"] = sarama.V0_10_0_1
	m["0.10.1.0"] = sarama.V0_10_1_0
	return m
}

func main() {
	flag.Parse()

	for {
		check()
	}
}

func check() {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("check error: %v", r)
			alert(*informEmail, *brokers, *topic, []byte(fmt.Sprintf("%v", r)), *smtpHost, *smtpPort, *smtpUser, *smtpPassword)
		}
	}()

	kafkaBrokers := strings.Split(*brokers, ",")
	v := kafkaVersions[*version]
	client := NewSaramaClient(kafkaBrokers, v)

	var buf bytes.Buffer

	ticker := time.NewTicker(*interval)
	for range ticker.C {
		buf.Reset()
		partitions, err := client.Partitions(*topic)
		if err != nil {
			fmt.Printf("failed to get partitions for topic=%s, err=%v\n", *topic, err)
			panic(err)
		}

		writablePartitions, err := client.WritablePartitions(*topic)
		if err != nil {
			fmt.Printf("failed to get writable partitions for topic=%s, err=%v\n", *topic, err)
			panic(err)
		}

		if len(partitions) != len(writablePartitions) {
			buf.WriteString("some partitions are not writable\n")
			buf.WriteString(fmt.Sprintf("all partitions: %v\n", partitions))
			buf.WriteString(fmt.Sprintf("writable partitions: %v\n", writablePartitions))
			//TODO print unwritable partitions
			maybeProblem = true
		}

		buf.WriteString(fmt.Sprintf("Time: %s\n", time.Now().Format("2006-01-02 15:04:05")))
		buf.WriteString(fmt.Sprintf("Brokers: %s\nVersion: %s\n", color.GreenString(*brokers), color.GreenString(*version)))
		buf.WriteString(fmt.Sprintf("Topic: %s, Group: %s, Partitions: %s\n\n",
			color.GreenString(*topic), color.GreenString(*group), color.GreenString(strconv.Itoa(len(partitions)))))

		fmt.Println(string(buf.Bytes()))

		//partition info
		var c *zk.Conn
		if *zkAddr != "" {
			c, _, err = zk.Connect(strings.Split(*zkAddr, ","), 5*time.Second)
			if err != nil {
				panic(err)
			}
		}

		infos := GetPartitionInfo(client, *topic, partitions, c, *basePath)
		if len(infos) > 0 {
			table := tablewriter.NewWriter(&buf)
			table.SetHeader([]string{"partition", "leader address", "leader", "replicas", "isr"})
			for _, info := range infos {
				replicas := fmt.Sprintf("%v", info.Replicas)
				replicas = replicas[1 : len(replicas)-1]
				replicas = compareString(replicas)
				isr := compareString(info.Isr)

				table.Append([]string{strconv.Itoa(int(info.Partition)), info.LeaderAddress, strconv.Itoa(int(info.Leader)), "[" + replicas + "]", "[" + isr + "]"})

				if replicas != isr {
					maybeProblem = true
				}
			}

			table.SetAlignment(tablewriter.ALIGN_LEFT)
			table.Render()

			buf.WriteString("\n\n")
		}
		//lag
		if offsets, err := FetchOffsets(client, *topic, *group); err == nil {

			table := tablewriter.NewWriter(&buf)
			table.SetHeader([]string{"partition", "end of log", "group offset", "lag"})

			var totalLag int64
			for _, info := range offsets {
				lag := info.PartitionOffset - info.GroupOffset
				if int(lag) > *lagThreshold {
					maybeProblem = true
				}

				totalLag += lag
				table.Append([]string{strconv.Itoa(int(info.Partition)), strconv.Itoa(int(info.PartitionOffset)),
					strconv.Itoa(int(info.GroupOffset)), convertLag(lag, *lagThreshold)})
			}

			if int(totalLag) > *totalLagThreshold {
				maybeProblem = true
			}

			table.SetFooter([]string{"", "", "Total", convertLag(totalLag, *totalLagThreshold)})
			table.SetAlignment(tablewriter.ALIGN_LEFT)
			table.SetFooterAlignment(tablewriter.ALIGN_LEFT)
			table.Render()
		}

		//output tables to stdout
		bytes := buf.Bytes()
		os.Stdout.Write(bytes)

		if maybeProblem && time.Since(lastTriggeredTime) > mergeAlertDuration {
			lastTriggeredTime = time.Now()
			alert(*informEmail, *brokers, *topic, bytes, *smtpHost, *smtpPort, *smtpUser, *smtpPassword)
		}

		maybeProblem = false
	}
}

func compareString(s string) string {
	s1 := strings.Split(s, " ")
	sort.Sort(sort.StringSlice(s1))

	return strings.Join(s1, " ")
}
func convertLag(lag int64, threshold int) string {
	lagStr := strconv.Itoa(int(lag))
	if int(lag) > threshold {
		lagStr = color.RedString(lagStr)
	}

	return lagStr
}

func NewSaramaClient(brokers []string, version sarama.KafkaVersion) sarama.Client {
	config := sarama.NewConfig()
	config.Version = version
	config.Consumer.Return.Errors = true
	config.Metadata.RefreshFrequency = 1 * time.Minute
	config.Metadata.Retry.Max = 10
	config.Net.MaxOpenRequests = 10
	config.Net.DialTimeout = 10 * time.Second
	config.Net.ReadTimeout = 10 * time.Second
	config.Net.WriteTimeout = 10 * time.Second

	client, err := sarama.NewClient(brokers, config)

	if err != nil {
		panic("Failed to start client: " + err.Error())
	}

	return client
}
<|endoftext|>


<commit_before>/*
Package msg creates a simple messaging package with flexible formatting that
can be used for stdout and/or logging.

Here is an example use.

    import (
      "jlinoff/utils/msg"
      "io"
      "os"
    )

    // My package logger.
    var log *msg.Object

    // Initialize it at startup.
    func init() {
      // Only write to stdout.
      w := []io.Writer{os.Stdout}

      // The name of my package. It is only used %pkg is specified in the
      // format string.
      n := "MyPackage"

      // Format string. Note that i could use %utc instead of %time to get
      // UTC time.
      f := `%pkg %(-27)time %(-7)type %file %line - %msg`

      // Time format string, only used if %time or %utc are specified in the
      // the format string.
      t := `2006-01-02 15:05:05.000 MST`

      // Create the message object.
      // Note that this is the same as this because I used the defaults.
      //     msg.NewMsg("MyPackage", "", "", []io.Writer{})
      l, e := msg.NewMsg(n, f, t, w)
      if err != nil { panic(err) }
      log = l
    }

    func test() {
      log.Debug("message of type %v", "debug")
      log.Info("info message")
      log.Warning("warning")

      // Now print messages to stdout and to a log while in this scope.
      fp := os.Create("log.txt")
      log.Writers = append(log.Writers, fp)

      // This stuff will go to stdout and the log file.
      log.Info("both")
      log.ErrNoExit("bad stuff happened but i can recover!")

      // Clean up by removing the file from the writers and then
      // closing it.
      log.Writers = log.Writers[:len(log.Writers)-1]
      fp.Close()
  }
*/
package msg

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"runtime"
	"strconv"
	"time"
)

// Interface defines logger functions.
type Interface interface {
	Debug(f string, a ...interface{})
	Info(f string, a ...interface{})
	Warn(f string, a ...interface{})
	Err(f string, a ...interface{})
	ErrNoExit(f string, a ...interface{})

	DebugWithLevel(l int, f string, a ...interface{})
	InfoWithLevel(l int, f string, a ...interface{})
	WarnWithLevel(l int, f string, a ...interface{})
	ErrWithLevel(l int, f string, a ...interface{})
	ErrNoExitWithLevel(l int, f string, a ...interface{})
}

// Object defines the logger.
type Object struct {
	// Name is the package name. It is accessed in the format string by %pkg.
	Name string

	// DebugEnabled enables debug messages if true.
	// It is true by default.
	DebugEnabled bool

	// InfoEnabled enables info messages if true.
	// It is true by default.
	InfoEnabled bool

	// WarningEnabled enables warning messages if true.
	// It is true by default.
	WarningEnabled bool

	// ErrorExitCode is the exit code to use for the Error function.
	// The default is 1.
	ErrorExitCode int

	// Writers for the message output.
	// If no writers are specified, messages go to os.Stdout.
	Writers []io.Writer

	// TimeFormat is the format of the prefix timestamp.
	// See time.Format for details.
	// The default format is: "2006-01-02 15:05:05.000 MST"
	TimeFormat string

	// Format is the template for the output. It has the following specifiers.
	//
	//   %file is the caller file name
	//   %func is the function name
	//   %line is the line number
	//   %msg  is the actual message
	//   %pkg  is the package name
	//   %time is the time format in the current locale
	//   %utc is the time format in the UTC locale
	//   %type is the msg type: DEBUG, INFO, WARNING, ERROR
	//   %% is a single % character
	//
	// You can explicitly format each field by specifying the formatting
	// options in parentheses.
	//
	//   %(-28)time
	//
	// Any other text is left verbatim.
	//
	// The default format is.
	//   `%(-27)time %(-7)type %file %line - %msg`
	Format string

	// outputFormat created by NewMsgObject and used to generate a message.
	outputFormat string

	// outputFlds created by NewMsgObject and used to specify the fields.
	outputFlds []string
}

// NewMsg makes a message object.
//   n - package name
//   f - format string, set to "" to get the default.
//   t - time stamp format, set to "" to get the default
//   w - the list of writers, if empty all messages go to stdout
func NewMsg(n string, f string, t string, w []io.Writer) (obj *Object, err error) {
	obj = new(Object)
	obj.Name = n
	obj.DebugEnabled = true
	obj.InfoEnabled = true
	obj.WarningEnabled = true
	obj.ErrorExitCode = 1

	if len(w) == 0 {
		obj.Writers = append(obj.Writers, os.Stdout)
	} else {
		obj.Writers = w
	}

	// Set the time format. If it is empty, set the default.
	if t == "" {
		obj.TimeFormat = "2006-01-02 15:05:05.000 MST"
	} else {
		obj.TimeFormat = t
	}

	// Set the format. If it is empty use the default.
	if f == "" {
		obj.Format = `%(-27)time %(-7)type %file %line - %msg`
	} else {
		obj.Format = f
	}

	// Parse the format.
	ofmt, oflds, err := ParseFormatString(f)
	obj.outputFormat = ofmt
	obj.outputFlds = oflds

	return
}

/*
Debug prints a debug message obtaining the callers filename, function and
line number.

Example:
      msg.Debug("%v = %v", key, value)
*/
func (o Object) Debug(f string, a ...interface{}) {
	if o.DebugEnabled {
		o.Print("DEBUG", 2, f, a...)
	}
}

/*
DebugWithLevel prints a debug message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.DebugWithLevel(2, "%v = %v", key, value)
*/
func (o Object) DebugWithLevel(l int, f string, a ...interface{}) {
	if o.DebugEnabled {
		o.Print("DEBUG", l, f, a...)
	}
}

/*
Info prints an info message obtaining the callers filename, function and
line number.

Example:
      msg.Info("%v = %v", key, value)
*/
func (o Object) Info(f string, a ...interface{}) {
	if o.InfoEnabled {
		o.Print("INFO", 2, f, a...)
	}
}

/*
InfoWithLevel prints an info message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.InfoWithLevel(2, "%v = %v", key, value)
*/
func (o Object) InfoWithLevel(l int, f string, a ...interface{}) {
	if o.InfoEnabled {
		o.Print("INFO", l, f, a...)
	}
}

/*
Warn prints a warning message obtaining the callers filename, function and
line number.

Example:
      msg.Warn("%v = %v", key, value)
*/
func (o Object) Warn(f string, a ...interface{}) {
	if o.WarningEnabled {
		o.Print("WARNING", 2, f, a...)
	}
}

/*
WarnWithLevel prints a warning message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.WarnWithLevel(2, "%v = %v", key, value)
*/
func (o Object) WarnWithLevel(l int, f string, a ...interface{}) {
	if o.WarningEnabled {
		o.Print("WARNING", 2, f, a...)
	}
}

/*
Err prints an error message obtaining the callers filename, function and
line number and exits. It cannot be disabled.

Example:
      msg.Err("%v = %v", key, value)
*/
func (o Object) Err(f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
	os.Exit(o.ErrorExitCode)
}

/*
ErrWithLevel prints an error message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug() and exits. It cannot be disabled.

Example:
      msg.ErrWithLevel(2, "%v = %v", key, value)
*/
func (o Object) ErrWithLevel(l int, f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
	os.Exit(o.ErrorExitCode)
}

/*
ErrNoExit prints an error message obtaining the callers filename, function and
line number. It does not exit and cannot be disabled.

Example:
      msg.ErrNoExit("%v = %v", key, value)
*/
func (o Object) ErrNoExit(f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
}

/*
ErrNoExitWithLevel prints an error message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug(). It does not exit and cannot be disabled.

Example:
      msg.ErrNoExitWithLevel(2, "%v = %v", key, value)
*/
func (o Object) ErrNoExitWithLevel(l int, f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
}

/*
Print is the basis of all message printers. It normally would not be called
directly.

      t - is the type, normally one of DEBUG, INFO, WARNING or ERROR
      l - is the caller level: 0 is this function, 1 is the caller, 2 is the callers caller and so on
      f - format string
      a - argument list
*/
func (o Object) Print(t string, l int, f string, a ...interface{}) {
	pc, fname, lineno, _ := runtime.Caller(l)
	fct := runtime.FuncForPC(pc).Name()
	fname = path.Base(fname[0 : len(fname)-3]) // strip off ".go"

	// The variables map for the format string.
	m := map[string]string{
		"file": fname,
		"func": fct,
		"line": strconv.Itoa(lineno),
		"msg":  fmt.Sprintf(f, a...),
		"pkg":  o.Name,
		"time": time.Now().Truncate(time.Millisecond).Format(o.TimeFormat),
		"utc":  time.Now().UTC().Truncate(time.Millisecond).Format(o.TimeFormat),
		"type": t,
	}

	// Collect the field values.
	var flds []interface{}
	for _, k := range o.outputFlds {
		if v, ok := m[k]; ok {
			flds = append(flds, v)
		} else {
			// This is, essentially, an assert. It should never happen.
			fmt.Fprintf(os.Stderr, "ERROR: unexpected condition, invalid specification id '%v'\n", k)
			os.Exit(1)
		}
	}

	// Create the formatted output string.
	s := fmt.Sprintf(o.outputFormat, flds...) + "\n"

	// Output it for each writer.
	for _, w := range o.Writers {
		fmt.Fprintf(w, s)
	}
}

/*
ParseFormatString transforms a format template to a format string
and the list of fields to print in each message.

It is meant to be used internally by NewMsg().

Here is an example transformation:

      input = "MYSTUFF %(-27)time %(-7)type %file %line - %msg"

      // TRANSFORM
      ofmt  = "MYSTUFF %-27v %-7v %v %v - %v"
      oids  = ["time", "type", "type", "file", "line", "msg"]
*/
func ParseFormatString(input string) (ofmt string, oids []string, err error) {
	ofmtb := []byte{}
	valid := []string{"file", "func", "line", "msg", "pkg", "time", "type", "utc"}
	ics := []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-$")

	// Define the parse states.
	//    normal - capture each byte
	//    spec   - capture a specification of the form %<id> or %(<fmt>)<id>.
	type state int
	const (
		normal state = iota
		spec
	)
	s := normal
	ib := []byte(input)
	for i := 0; i < len(ib); i++ {
		b := ib[i]
		switch s {
		case normal:
			// normal state, this is all of the stuff in the
			// template that is not part of a specification.
			if b == '%' {
				s = spec
			} else {
				ofmtb = append(ofmtb, b)
			}
		case spec:
			s = normal // after parsing the spec go back to the normal state

			// specification state, parse specifications of the form:
			//   %(<fmt>)<id>
			//   %<id>
			beg := i - 1
			if b == '(' {
				// This is a format specification. Capture it.
				// If ')' is not found, report an error.
				j := i // ib[j] == '('
				for ; i < len(ib) && ib[i] != ')'; i++ {
				}
				if i >= len(ib) {
					err = fmt.Errorf("missing ')' for '%v'", string(ib[beg:]))
					return
				}
				ofmtb = append(ofmtb, '%')
				ofmtb = append(ofmtb, ib[j+1:i]...)
				ofmtb = append(ofmtb, 'v')
				i++ // point past the ')'
			} else {
				ofmtb = append(ofmtb, []byte("%v")[:]...)
			}

			// Now parse out the id.
			id := ""
			for _, v := range valid {
				ba := []byte(v)
				if bytes.HasPrefix(ib[i:], ba) {
					// We MAY have a match.
					// for example '%line' matches but '%linex' does not.
					i += len(ba)
					if i < len(ib) {
						bs := []byte{ib[i]}
						if bytes.Contains([]byte(ics), bs) {
							ofmt = string(ofmtb)
							ba = append(ba, ib[i])
							err = fmt.Errorf("unrecognized specification id '%v'", string(ba))
							return
						}
					}
					id = string(ba)
					i--
					break
				}
			}
			if id == "" {
				ofmt = string(ofmtb)
				err = fmt.Errorf("specification syntax error '%v'", string(ib[beg:]))
				return
			}
			oids = append(oids, id)
		}
	}
	ofmt = string(ofmtb)
	return
}
<commit_msg>Fixed the example in the comments<commit_after>/*
Package msg creates a simple messaging package with flexible formatting that
can be used for stdout and/or logging.

Here is an example use.

    import (
      "jlinoff/utils/msg"
      "io"
      "os"
    )

    // My package logger.
    var log *msg.Object

    // Initialize it at startup.
    func init() {
      // Only write to stdout.
      w := []io.Writer{os.Stdout}

      // The name of my package. It is only used %pkg is specified in the
      // format string.
      n := "MyPackage"

      // Format string. Note that i could use %utc instead of %time to get
      // UTC time.
      f := `%pkg %(-27)time %(-7)type %file %line - %msg`

      // Time format string, only used if %time or %utc are specified in the
      // the format string.
      t := `2006-01-02 15:05:05.000 MST`

      // Create the message object.
      // Note that this is the same as this because I used the defaults.
      //     msg.NewMsg("MyPackage", "", "", []io.Writer{})
      l, e := msg.NewMsg(n, f, t, w)
      if e != nil { panic(e) }
      log = l
    }

    func test() {
      log.Debug("message of type %v", "debug")
      log.Info("info message")
      log.Warn("warning")

      // Now print messages to stdout and to a log while in this scope.
      fp, _ := os.Create("log.txt")
      log.Writers = append(log.Writers, fp)

      // This stuff will go to stdout and the log file.
      log.Info("both")
      log.ErrNoExit("bad stuff happened but i can recover!")

      // Clean up by removing the file from the writers and then
      // closing it.
      log.Writers = log.Writers[:len(log.Writers)-1]
      fp.Close()
  }
*/
package msg

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"runtime"
	"strconv"
	"time"
)

// Interface defines logger functions.
type Interface interface {
	Debug(f string, a ...interface{})
	Info(f string, a ...interface{})
	Warn(f string, a ...interface{})
	Err(f string, a ...interface{})
	ErrNoExit(f string, a ...interface{})

	DebugWithLevel(l int, f string, a ...interface{})
	InfoWithLevel(l int, f string, a ...interface{})
	WarnWithLevel(l int, f string, a ...interface{})
	ErrWithLevel(l int, f string, a ...interface{})
	ErrNoExitWithLevel(l int, f string, a ...interface{})
}

// Object defines the logger.
type Object struct {
	// Name is the package name. It is accessed in the format string by %pkg.
	Name string

	// DebugEnabled enables debug messages if true.
	// It is true by default.
	DebugEnabled bool

	// InfoEnabled enables info messages if true.
	// It is true by default.
	InfoEnabled bool

	// WarningEnabled enables warning messages if true.
	// It is true by default.
	WarningEnabled bool

	// ErrorExitCode is the exit code to use for the Error function.
	// The default is 1.
	ErrorExitCode int

	// Writers for the message output.
	// If no writers are specified, messages go to os.Stdout.
	Writers []io.Writer

	// TimeFormat is the format of the prefix timestamp.
	// See time.Format for details.
	// The default format is: "2006-01-02 15:05:05.000 MST"
	TimeFormat string

	// Format is the template for the output. It has the following specifiers.
	//
	//   %file is the caller file name
	//   %func is the function name
	//   %line is the line number
	//   %msg  is the actual message
	//   %pkg  is the package name
	//   %time is the time format in the current locale
	//   %utc is the time format in the UTC locale
	//   %type is the msg type: DEBUG, INFO, WARNING, ERROR
	//   %% is a single % character
	//
	// You can explicitly format each field by specifying the formatting
	// options in parentheses.
	//
	//   %(-28)time
	//
	// Any other text is left verbatim.
	//
	// The default format is.
	//   `%(-27)time %(-7)type %file %line - %msg`
	Format string

	// outputFormat created by NewMsgObject and used to generate a message.
	outputFormat string

	// outputFlds created by NewMsgObject and used to specify the fields.
	outputFlds []string
}

// NewMsg makes a message object.
//   n - package name
//   f - format string, set to "" to get the default.
//   t - time stamp format, set to "" to get the default
//   w - the list of writers, if empty all messages go to stdout
func NewMsg(n string, f string, t string, w []io.Writer) (obj *Object, err error) {
	obj = new(Object)
	obj.Name = n
	obj.DebugEnabled = true
	obj.InfoEnabled = true
	obj.WarningEnabled = true
	obj.ErrorExitCode = 1

	if len(w) == 0 {
		obj.Writers = append(obj.Writers, os.Stdout)
	} else {
		obj.Writers = w
	}

	// Set the time format. If it is empty, set the default.
	if t == "" {
		obj.TimeFormat = "2006-01-02 15:05:05.000 MST"
	} else {
		obj.TimeFormat = t
	}

	// Set the format. If it is empty use the default.
	if f == "" {
		obj.Format = `%(-27)time %(-7)type %file %line - %msg`
	} else {
		obj.Format = f
	}

	// Parse the format.
	ofmt, oflds, err := ParseFormatString(f)
	obj.outputFormat = ofmt
	obj.outputFlds = oflds

	return
}

/*
Debug prints a debug message obtaining the callers filename, function and
line number.

Example:
      msg.Debug("%v = %v", key, value)
*/
func (o Object) Debug(f string, a ...interface{}) {
	if o.DebugEnabled {
		o.Print("DEBUG", 2, f, a...)
	}
}

/*
DebugWithLevel prints a debug message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.DebugWithLevel(2, "%v = %v", key, value)
*/
func (o Object) DebugWithLevel(l int, f string, a ...interface{}) {
	if o.DebugEnabled {
		o.Print("DEBUG", l, f, a...)
	}
}

/*
Info prints an info message obtaining the callers filename, function and
line number.

Example:
      msg.Info("%v = %v", key, value)
*/
func (o Object) Info(f string, a ...interface{}) {
	if o.InfoEnabled {
		o.Print("INFO", 2, f, a...)
	}
}

/*
InfoWithLevel prints an info message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.InfoWithLevel(2, "%v = %v", key, value)
*/
func (o Object) InfoWithLevel(l int, f string, a ...interface{}) {
	if o.InfoEnabled {
		o.Print("INFO", l, f, a...)
	}
}

/*
Warn prints a warning message obtaining the callers filename, function and
line number.

Example:
      msg.Warn("%v = %v", key, value)
*/
func (o Object) Warn(f string, a ...interface{}) {
	if o.WarningEnabled {
		o.Print("WARNING", 2, f, a...)
	}
}

/*
WarnWithLevel prints a warning message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug().

Example:
      msg.WarnWithLevel(2, "%v = %v", key, value)
*/
func (o Object) WarnWithLevel(l int, f string, a ...interface{}) {
	if o.WarningEnabled {
		o.Print("WARNING", 2, f, a...)
	}
}

/*
Err prints an error message obtaining the callers filename, function and
line number and exits. It cannot be disabled.

Example:
      msg.Err("%v = %v", key, value)
*/
func (o Object) Err(f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
	os.Exit(o.ErrorExitCode)
}

/*
ErrWithLevel prints an error message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug() and exits. It cannot be disabled.

Example:
      msg.ErrWithLevel(2, "%v = %v", key, value)
*/
func (o Object) ErrWithLevel(l int, f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
	os.Exit(o.ErrorExitCode)
}

/*
ErrNoExit prints an error message obtaining the callers filename, function and
line number. It does not exit and cannot be disabled.

Example:
      msg.ErrNoExit("%v = %v", key, value)
*/
func (o Object) ErrNoExit(f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
}

/*
ErrNoExitWithLevel prints an error message obtaining the filename, function and
line number from the caller specified by level "l". l=2 is the same
as Debug(). It does not exit and cannot be disabled.

Example:
      msg.ErrNoExitWithLevel(2, "%v = %v", key, value)
*/
func (o Object) ErrNoExitWithLevel(l int, f string, a ...interface{}) {
	o.Print("ERROR", 2, f, a...)
}

/*
Print is the basis of all message printers. It normally would not be called
directly.

      t - is the type, normally one of DEBUG, INFO, WARNING or ERROR
      l - is the caller level: 0 is this function, 1 is the caller, 2 is the callers caller and so on
      f - format string
      a - argument list
*/
func (o Object) Print(t string, l int, f string, a ...interface{}) {
	pc, fname, lineno, _ := runtime.Caller(l)
	fct := runtime.FuncForPC(pc).Name()
	fname = path.Base(fname[0 : len(fname)-3]) // strip off ".go"

	// The variables map for the format string.
	m := map[string]string{
		"file": fname,
		"func": fct,
		"line": strconv.Itoa(lineno),
		"msg":  fmt.Sprintf(f, a...),
		"pkg":  o.Name,
		"time": time.Now().Truncate(time.Millisecond).Format(o.TimeFormat),
		"utc":  time.Now().UTC().Truncate(time.Millisecond).Format(o.TimeFormat),
		"type": t,
	}

	// Collect the field values.
	var flds []interface{}
	for _, k := range o.outputFlds {
		if v, ok := m[k]; ok {
			flds = append(flds, v)
		} else {
			// This is, essentially, an assert. It should never happen.
			fmt.Fprintf(os.Stderr, "ERROR: unexpected condition, invalid specification id '%v'\n", k)
			os.Exit(1)
		}
	}

	// Create the formatted output string.
	s := fmt.Sprintf(o.outputFormat, flds...) + "\n"

	// Output it for each writer.
	for _, w := range o.Writers {
		fmt.Fprintf(w, s)
	}
}

/*
ParseFormatString transforms a format template to a format string
and the list of fields to print in each message.

It is meant to be used internally by NewMsg().

Here is an example transformation:

      input = "MYSTUFF %(-27)time %(-7)type %file %line - %msg"

      // TRANSFORM
      ofmt  = "MYSTUFF %-27v %-7v %v %v - %v"
      oids  = ["time", "type", "type", "file", "line", "msg"]
*/
func ParseFormatString(input string) (ofmt string, oids []string, err error) {
	ofmtb := []byte{}
	valid := []string{"file", "func", "line", "msg", "pkg", "time", "type", "utc"}
	ics := []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-$")

	// Define the parse states.
	//    normal - capture each byte
	//    spec   - capture a specification of the form %<id> or %(<fmt>)<id>.
	type state int
	const (
		normal state = iota
		spec
	)
	s := normal
	ib := []byte(input)
	for i := 0; i < len(ib); i++ {
		b := ib[i]
		switch s {
		case normal:
			// normal state, this is all of the stuff in the
			// template that is not part of a specification.
			if b == '%' {
				s = spec
			} else {
				ofmtb = append(ofmtb, b)
			}
		case spec:
			s = normal // after parsing the spec go back to the normal state

			// specification state, parse specifications of the form:
			//   %(<fmt>)<id>
			//   %<id>
			beg := i - 1
			if b == '(' {
				// This is a format specification. Capture it.
				// If ')' is not found, report an error.
				j := i // ib[j] == '('
				for ; i < len(ib) && ib[i] != ')'; i++ {
				}
				if i >= len(ib) {
					err = fmt.Errorf("missing ')' for '%v'", string(ib[beg:]))
					return
				}
				ofmtb = append(ofmtb, '%')
				ofmtb = append(ofmtb, ib[j+1:i]...)
				ofmtb = append(ofmtb, 'v')
				i++ // point past the ')'
			} else {
				ofmtb = append(ofmtb, []byte("%v")[:]...)
			}

			// Now parse out the id.
			id := ""
			for _, v := range valid {
				ba := []byte(v)
				if bytes.HasPrefix(ib[i:], ba) {
					// We MAY have a match.
					// for example '%line' matches but '%linex' does not.
					i += len(ba)
					if i < len(ib) {
						bs := []byte{ib[i]}
						if bytes.Contains([]byte(ics), bs) {
							ofmt = string(ofmtb)
							ba = append(ba, ib[i])
							err = fmt.Errorf("unrecognized specification id '%v'", string(ba))
							return
						}
					}
					id = string(ba)
					i--
					break
				}
			}
			if id == "" {
				ofmt = string(ofmtb)
				err = fmt.Errorf("specification syntax error '%v'", string(ib[beg:]))
				return
			}
			oids = append(oids, id)
		}
	}
	ofmt = string(ofmtb)
	return
}
<|endoftext|>