<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017, 2018 Red Hat, Inc.
 *
 */

package vnc

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"k8s.io/client-go/tools/clientcmd"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const LISTEN_TIMEOUT = 60 * time.Second

const FLAG = "vnc"

func NewCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "vnc (vm)",
		Short:   "Open a vnc connection to a virtual machine.",
		Example: usage(),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := VNC{clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

type VNC struct {
	clientConfig clientcmd.ClientConfig
}

func (o *VNC) Run(cmd *cobra.Command, args []string) error {
	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	vm := args[0]

	virtCli, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return err
	}

	//                                       -> pipeInWriter  -> pipeInReader
	// remote-viewer -> unix sock connection
	//                                       <- pipeOutReader <- pipeOutWriter
	pipeInReader, pipeInWriter := io.Pipe()
	pipeOutReader, pipeOutWriter := io.Pipe()

	k8ResChan := make(chan error)
	listenResChan := make(chan error)
	viewResChan := make(chan error)
	stopChan := make(chan struct{}, 1)
	doneChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)

	lnAddr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("Can't resolve the address: %s", err.Error())
	}

	// The local tcp server is used to proxy the podExec websock connection to remote-viewer
	ln, err := net.ListenTCP("tcp", lnAddr)
	if err != nil {
		return fmt.Errorf("Can't listen on unix socket: %s", err.Error())
	}

	port := ln.Addr().(*net.TCPAddr).Port

	// setup connection with VM
	go func() {
		err := virtCli.VM(namespace).VNC(vm, pipeInReader, pipeOutWriter)
		k8ResChan <- err
	}()

	// wait for remote-viewer to connect to our local proxy server
	go func() {
		start := time.Now()
		glog.Infof("connection timeout: %v", LISTEN_TIMEOUT)
		// exit early if spawning remote-viewer fails
		ln.SetDeadline(time.Now().Add(LISTEN_TIMEOUT))

		fd, err := ln.Accept()
		if err != nil {
			glog.V(2).Infof("Failed to accept unix sock connection. %s", err.Error())
			listenResChan <- err
		}
		defer fd.Close()

		glog.V(2).Infof("remote-viewer connected in %v", time.Now().Sub(start))

		// write to FD <- pipeOutReader
		go func() {
			_, err := io.Copy(fd, pipeOutReader)
			readStop <- err
		}()

		// read from FD -> pipeInWriter
		go func() {
			_, err := io.Copy(pipeInWriter, fd)
			writeStop <- err
		}()

		<-doneChan
		listenResChan <- err
	}()

	// execute remote viewer
	go func() {
		args := []string{fmt.Sprintf("vnc://127.0.0.1:%d", port)}
		if glog.V(4) {
			args = append(args, "--debug")
			glog.Infof("remote-viewer commandline: %v", args)
		}

		cmnd := exec.Command("remote-viewer", args...)

		output, err := cmnd.CombinedOutput()
		if err != nil {
			glog.Errorf("remote-viewer execution failed: %v, output: %v", err, string(output))
		} else {
			glog.V(2).Infof("remote-viewer output: %v", string(output))
		}
		viewResChan <- err
		close(doneChan)
	}()

	go func() {
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
		close(stopChan)
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-k8ResChan:
	case err = <-viewResChan:
	case err = <-listenResChan:
	}

	if err != nil {
		return fmt.Errorf("Error encountered: %s", err.Error())
	}
	return nil
}

func usage() string {
	usage := "# Connect to testvm via remote-viewer:\n"
	usage += "./virtctl vnc testvm"
	return usage
}
<commit_msg>vnc: dont' try to connect if VM is unreachable<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017, 2018 Red Hat, Inc.
 *
 */

package vnc

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/clientcmd"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const LISTEN_TIMEOUT = 60 * time.Second

const FLAG = "vnc"

func NewCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "vnc (vm)",
		Short:   "Open a vnc connection to a virtual machine.",
		Example: usage(),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := VNC{clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

type VNC struct {
	clientConfig clientcmd.ClientConfig
}

func (o *VNC) Run(cmd *cobra.Command, args []string) error {
	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	vm := args[0]

	virtCli, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return err
	}

	var opts k8smetav1.GetOptions
	vmObj, err := virtCli.VM(namespace).Get(vm, opts)
	if err != nil {
		return fmt.Errorf("Can't access VM %s: %s", vm, err.Error())
	}
	if !vmObj.IsRunning() {
		return fmt.Errorf("Can't connect to not-running VM %s (status=%s)", vm, vmObj.Status.Phase)
	}

	lnAddr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("Can't resolve the address: %s", err.Error())
	}

	// The local tcp server is used to proxy the podExec websock connection to remote-viewer
	ln, err := net.ListenTCP("tcp", lnAddr)
	if err != nil {
		return fmt.Errorf("Can't listen on unix socket: %s", err.Error())
	}
	// End of pre-flight checks. Everything looks good, we can start
	// the goroutines and let the data flow

	//                                       -> pipeInWriter  -> pipeInReader
	// remote-viewer -> unix sock connection
	//                                       <- pipeOutReader <- pipeOutWriter
	pipeInReader, pipeInWriter := io.Pipe()
	pipeOutReader, pipeOutWriter := io.Pipe()

	k8ResChan := make(chan error)
	listenResChan := make(chan error)
	viewResChan := make(chan error)
	stopChan := make(chan struct{}, 1)
	doneChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)

	// setup connection with VM
	go func() {
		err := virtCli.VM(namespace).VNC(vm, pipeInReader, pipeOutWriter)
		k8ResChan <- err
	}()

	// wait for remote-viewer to connect to our local proxy server
	go func() {
		start := time.Now()
		glog.Infof("connection timeout: %v", LISTEN_TIMEOUT)
		// exit early if spawning remote-viewer fails
		ln.SetDeadline(time.Now().Add(LISTEN_TIMEOUT))

		fd, err := ln.Accept()
		if err != nil {
			glog.V(2).Infof("Failed to accept unix sock connection. %s", err.Error())
			listenResChan <- err
		}
		defer fd.Close()

		glog.V(2).Infof("remote-viewer connected in %v", time.Now().Sub(start))

		// write to FD <- pipeOutReader
		go func() {
			_, err := io.Copy(fd, pipeOutReader)
			readStop <- err
		}()

		// read from FD -> pipeInWriter
		go func() {
			_, err := io.Copy(pipeInWriter, fd)
			writeStop <- err
		}()

		// don't terminate until remote-viewer is done
		<-doneChan
		listenResChan <- err
	}()

	// execute remote viewer
	go func() {
		port := ln.Addr().(*net.TCPAddr).Port
		args := []string{fmt.Sprintf("vnc://127.0.0.1:%d", port)}
		if glog.V(4) {
			args = append(args, "--debug")
			glog.Infof("remote-viewer commandline: %v", args)
		}

		cmnd := exec.Command("remote-viewer", args...)

		output, err := cmnd.CombinedOutput()
		if err != nil {
			glog.Errorf("remote-viewer execution failed: %v, output: %v", err, string(output))
		} else {
			glog.V(2).Infof("remote-viewer output: %v", string(output))
		}
		viewResChan <- err
		close(doneChan)
	}()

	go func() {
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
		close(stopChan)
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-k8ResChan:
	case err = <-viewResChan:
	case err = <-listenResChan:
	}

	if err != nil {
		return fmt.Errorf("Error encountered: %s", err.Error())
	}
	return nil
}

func usage() string {
	usage := "# Connect to testvm via remote-viewer:\n"
	usage += "./virtctl vnc testvm"
	return usage
}
<|endoftext|>


<commit_before>// +build appengine

package datastore

import (
	"fmt"
	"reflect"
	"strings"

	"gnd.la/config"
	"gnd.la/log"
	"gnd.la/orm/driver"
	"gnd.la/orm/operation"
	"gnd.la/orm/query"
	"gnd.la/util/types"

	"appengine"
	"appengine/datastore"
)

type Driver struct {
	c      appengine.Context
	logger *log.Logger
}

func (d *Driver) Check() error {
	return nil
}

func (d *Driver) Initialize(ms []driver.Model) error {
	// No need to create tables in the datastore. Instead,
	// check that the models can be stored.
	return nil
}

func (d *Driver) Query(m driver.Model, q query.Q, sort []driver.Sort, limit int, offset int) driver.Iter {
	dq, err := d.makeQuery(m, q, sort, limit, offset)
	if err != nil {
		return &Iter{err: err}
	}
	return &Iter{iter: dq.Run(d.c)}
}

func (d *Driver) Count(m driver.Model, q query.Q, limit int, offset int) (uint64, error) {
	dq, err := d.makeQuery(m, q, nil, limit, offset)
	if err != nil {
		return 0, err
	}
	c, err := dq.Count(d.c)
	return uint64(c), err
}

func (d *Driver) Exists(m driver.Model, q query.Q) (bool, error) {
	dq, err := d.makeQuery(m, q, nil, 1, -1)
	if err != nil {
		return false, err
	}
	c, err := dq.Count(d.c)
	return c != 0, err
}

func (d *Driver) Insert(m driver.Model, data interface{}) (driver.Result, error) {
	var id int64
	fields := m.Fields()
	var pkVal *reflect.Value
	// TODO: If the PK is supplied by the user rather than auto-assigned, it
	// might conflict with PKs generated by datastore.AllocateIDs().
	if fields.PrimaryKey >= 0 {
		p := d.primaryKey(fields, data)
		if p.IsValid() && types.Kind(p.Kind()) == types.Int {
			id = p.Int()
			if id == 0 {
				// Must assign PK field value after calling AllocateIDs
				pkVal = &p
			}
		}
	}
	name := m.Table()
	// Make all objects of a given kind ancestors of the same key. While
	// this hurts scalability, it makes all reads strongly consistent.
	parent := d.parentKey(m)
	var err error
	if id == 0 {
		id, _, err = datastore.AllocateIDs(d.c, name, parent, 1)
		if err != nil {
			return nil, err
		}
	}
	if fields.AutoincrementPk && pkVal != nil {
		pkVal.SetInt(int64(id))
	}
	key := datastore.NewKey(d.c, name, "", id, parent)
	log.Debugf("DATASTORE: put %s %v", key, data)
	_, err = datastore.Put(d.c, key, data)
	if err != nil {
		return nil, err
	}
	return &result{key: key, count: 1}, nil
}

func (d *Driver) Operate(m driver.Model, q query.Q, ops []*operation.Operation) (driver.Result, error) {
	return nil, fmt.Errorf("datastore driver does not support Operate")
}

func (d *Driver) Update(m driver.Model, q query.Q, data interface{}) (driver.Result, error) {
	keys, err := d.getKeys(m, q)
	if err != nil {
		return nil, err
	}
	src := make([]interface{}, len(keys))
	for ii := range src {
		src[ii] = data
	}
	// Multi variants need to be run in transactions, otherwise some
	// might fail while others succeed
	err = datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		_, e := datastore.PutMulti(c, keys, src)
		return e
	}, nil)
	if err != nil {
		return nil, err
	}
	return &result{count: len(keys)}, nil
}

func (d *Driver) Upsert(m driver.Model, q query.Q, data interface{}) (driver.Result, error) {
	return nil, nil
}

func (d *Driver) Delete(m driver.Model, q query.Q) (driver.Result, error) {
	keys, err := d.getKeys(m, q)
	if err != nil {
		return nil, err
	}
	// See comment around PutMulti
	err = datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		return datastore.DeleteMulti(c, keys)
	}, nil)
	if err != nil {
		return nil, err
	}
	return &result{count: len(keys)}, nil
}

func (d *Driver) getKeys(m driver.Model, q query.Q) ([]*datastore.Key, error) {
	dq, err := d.makeQuery(m, q, nil, -1, -1)
	if err != nil {
		return nil, err
	}
	iter := dq.KeysOnly().Run(d.c)
	var keys []*datastore.Key
	for {
		key, err := iter.Next(nil)
		if err != nil {
			if err == datastore.Done {
				break
			}
			return nil, err
		}
		keys = append(keys, key)
	}
	return keys, nil
}

func (d *Driver) makeQuery(m driver.Model, q query.Q, sort []driver.Sort, limit int, offset int) (*datastore.Query, error) {
	if m.Join() != nil {
		return nil, errJoinNotSupported
	}
	dq := datastore.NewQuery(m.Table()).Ancestor(d.parentKey(m))
	var err error
	if dq, err = d.applyQuery(m, dq, q); err != nil {
		return nil, err
	}
	for _, v := range sort {
		field := v.Field()
		if v.Direction() == driver.DESC {
			field = "-" + field
		}
		dq = dq.Order(field)
	}
	if limit >= 0 {
		dq = dq.Limit(limit)
	}
	if offset > 0 {
		dq = dq.Offset(limit)
	}
	return dq, nil
}

func (d *Driver) parentKey(m driver.Model) *datastore.Key {
	return datastore.NewKey(d.c, m.Table(), "", -1, nil)
}

func (d *Driver) applyQuery(m driver.Model, dq *datastore.Query, q query.Q) (*datastore.Query, error) {
	var field *query.Field
	var op string
	switch x := q.(type) {
	case *query.Eq:
		field = &x.Field
		op = " ="
	case *query.Lt:
		field = &x.Field
		op = " <"
	case *query.Lte:
		field = &x.Field
		op = " <="
	case *query.Gt:
		field = &x.Field
		op = " >"
	case *query.Gte:
		field = &x.Field
		op = " >="
	case *query.And:
		var err error
		for _, v := range x.Conditions {
			dq, err = d.applyQuery(m, dq, v)
			if err != nil {
				return nil, err
			}
		}
	case nil:
	default:
		return nil, fmt.Errorf("datastore does not support %T queries", q)
	}
	if field != nil {
		if _, ok := field.Value.(query.F); ok {
			return nil, fmt.Errorf("datastore queries can't reference other properties (%v)", field.Value)
		}
		name := field.Field
		if strings.IndexByte(name, '.') >= 0 {
			// GAE flattens embedded fields, so we must remove
			// the parts of the field which refer to a flattened
			// field.
			fields := m.Fields()
			if idx, ok := fields.QNameMap[name]; ok {
				indexes := fields.Indexes[idx]
				parts := strings.Split(name, ".")
				if len(indexes) == len(parts) {
					var final []string
					typ := fields.Type
					for ii, v := range indexes {
						f := typ.Field(v)
						if !f.Anonymous {
							final = append(final, parts[ii])
						}
						typ = f.Type
					}
					name = strings.Join(final, ".")
				}
			}
		}
		log.Debugf("DATASTORE: filter %s %s %v", m, name+op, field.Value)
		dq = dq.Filter(name+op, field.Value)
	}
	return dq, nil
}

func (d *Driver) Close() error {
	return nil
}

func (d *Driver) Upserts() bool {
	return false
}

func (d *Driver) Tags() []string {
	return []string{"datastore"}
}

func (d *Driver) SetLogger(logger *log.Logger) {
	d.logger = logger
}

func (d *Driver) SetContext(ctx appengine.Context) {
	d.c = ctx
}

func (d *Driver) Begin() (driver.Tx, error) {
	return nil, errTransactionNotSupported
}

func (d *Driver) Commit() error {
	return driver.ErrNotInTransaction
}

func (d *Driver) Rollback() error {
	return driver.ErrNotInTransaction
}

func (d *Driver) Transaction(f func(driver.Driver) error) error {
	return datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		drv := *d
		drv.c = c
		return f(&drv)
	}, nil)
}

func (d *Driver) Capabilities() driver.Capability {
	return driver.CAP_TRANSACTION | driver.CAP_AUTO_ID | driver.CAP_EVENTUAL | driver.CAP_PK
}

func (d *Driver) primaryKey(f *driver.Fields, data interface{}) reflect.Value {
	return fieldByIndex(reflect.ValueOf(data), f.Indexes[f.PrimaryKey])
}

func (d *Driver) HasFunc(fname string, retType reflect.Type) bool {
	return false
}

func (d *Driver) Connection() interface{} {
	return d.c
}

func fieldByIndex(val reflect.Value, indexes []int) reflect.Value {
	for _, v := range indexes {
		if val.Type().Kind() == reflect.Ptr {
			if val.IsNil() {
				return reflect.Value{}
			}
			val = val.Elem()
		}
		val = val.Field(v)
	}
	return val
}

func datastoreOpener(url *config.URL) (driver.Driver, error) {
	return &Driver{}, nil
}

func init() {
	driver.Register("datastore", datastoreOpener)
}
<commit_msg>Always fail when a datastore field can't be mapped with the model<commit_after>// +build appengine

package datastore

import (
	"fmt"
	"reflect"
	"strings"

	"gnd.la/config"
	"gnd.la/log"
	"gnd.la/orm/driver"
	"gnd.la/orm/operation"
	"gnd.la/orm/query"
	"gnd.la/util/types"

	"appengine"
	"appengine/datastore"
)

type Driver struct {
	c      appengine.Context
	logger *log.Logger
}

func (d *Driver) Check() error {
	return nil
}

func (d *Driver) Initialize(ms []driver.Model) error {
	// No need to create tables in the datastore. Instead,
	// check that the models can be stored.
	return nil
}

func (d *Driver) Query(m driver.Model, q query.Q, sort []driver.Sort, limit int, offset int) driver.Iter {
	dq, err := d.makeQuery(m, q, sort, limit, offset)
	if err != nil {
		return &Iter{err: err}
	}
	return &Iter{iter: dq.Run(d.c)}
}

func (d *Driver) Count(m driver.Model, q query.Q, limit int, offset int) (uint64, error) {
	dq, err := d.makeQuery(m, q, nil, limit, offset)
	if err != nil {
		return 0, err
	}
	c, err := dq.Count(d.c)
	return uint64(c), err
}

func (d *Driver) Exists(m driver.Model, q query.Q) (bool, error) {
	dq, err := d.makeQuery(m, q, nil, 1, -1)
	if err != nil {
		return false, err
	}
	c, err := dq.Count(d.c)
	return c != 0, err
}

func (d *Driver) Insert(m driver.Model, data interface{}) (driver.Result, error) {
	var id int64
	fields := m.Fields()
	var pkVal *reflect.Value
	// TODO: If the PK is supplied by the user rather than auto-assigned, it
	// might conflict with PKs generated by datastore.AllocateIDs().
	if fields.PrimaryKey >= 0 {
		p := d.primaryKey(fields, data)
		if p.IsValid() && types.Kind(p.Kind()) == types.Int {
			id = p.Int()
			if id == 0 {
				// Must assign PK field value after calling AllocateIDs
				pkVal = &p
			}
		}
	}
	name := m.Table()
	// Make all objects of a given kind ancestors of the same key. While
	// this hurts scalability, it makes all reads strongly consistent.
	parent := d.parentKey(m)
	var err error
	if id == 0 {
		id, _, err = datastore.AllocateIDs(d.c, name, parent, 1)
		if err != nil {
			return nil, err
		}
	}
	if fields.AutoincrementPk && pkVal != nil {
		pkVal.SetInt(int64(id))
	}
	key := datastore.NewKey(d.c, name, "", id, parent)
	log.Debugf("DATASTORE: put %s %v", key, data)
	_, err = datastore.Put(d.c, key, data)
	if err != nil {
		return nil, err
	}
	return &result{key: key, count: 1}, nil
}

func (d *Driver) Operate(m driver.Model, q query.Q, ops []*operation.Operation) (driver.Result, error) {
	return nil, fmt.Errorf("datastore driver does not support Operate")
}

func (d *Driver) Update(m driver.Model, q query.Q, data interface{}) (driver.Result, error) {
	keys, err := d.getKeys(m, q)
	if err != nil {
		return nil, err
	}
	src := make([]interface{}, len(keys))
	for ii := range src {
		src[ii] = data
	}
	// Multi variants need to be run in transactions, otherwise some
	// might fail while others succeed
	err = datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		_, e := datastore.PutMulti(c, keys, src)
		return e
	}, nil)
	if err != nil {
		return nil, err
	}
	return &result{count: len(keys)}, nil
}

func (d *Driver) Upsert(m driver.Model, q query.Q, data interface{}) (driver.Result, error) {
	return nil, nil
}

func (d *Driver) Delete(m driver.Model, q query.Q) (driver.Result, error) {
	keys, err := d.getKeys(m, q)
	if err != nil {
		return nil, err
	}
	// See comment around PutMulti
	err = datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		return datastore.DeleteMulti(c, keys)
	}, nil)
	if err != nil {
		return nil, err
	}
	return &result{count: len(keys)}, nil
}

func (d *Driver) getKeys(m driver.Model, q query.Q) ([]*datastore.Key, error) {
	dq, err := d.makeQuery(m, q, nil, -1, -1)
	if err != nil {
		return nil, err
	}
	iter := dq.KeysOnly().Run(d.c)
	var keys []*datastore.Key
	for {
		key, err := iter.Next(nil)
		if err != nil {
			if err == datastore.Done {
				break
			}
			return nil, err
		}
		keys = append(keys, key)
	}
	return keys, nil
}

func (d *Driver) makeQuery(m driver.Model, q query.Q, sort []driver.Sort, limit int, offset int) (*datastore.Query, error) {
	if m.Join() != nil {
		return nil, errJoinNotSupported
	}
	dq := datastore.NewQuery(m.Table()).Ancestor(d.parentKey(m))
	var err error
	if dq, err = d.applyQuery(m, dq, q); err != nil {
		return nil, err
	}
	for _, v := range sort {
		field := v.Field()
		if v.Direction() == driver.DESC {
			field = "-" + field
		}
		dq = dq.Order(field)
	}
	if limit >= 0 {
		dq = dq.Limit(limit)
	}
	if offset > 0 {
		dq = dq.Offset(limit)
	}
	return dq, nil
}

func (d *Driver) parentKey(m driver.Model) *datastore.Key {
	return datastore.NewKey(d.c, m.Table(), "", -1, nil)
}

func (d *Driver) applyQuery(m driver.Model, dq *datastore.Query, q query.Q) (*datastore.Query, error) {
	var field *query.Field
	var op string
	switch x := q.(type) {
	case *query.Eq:
		field = &x.Field
		op = " ="
	case *query.Lt:
		field = &x.Field
		op = " <"
	case *query.Lte:
		field = &x.Field
		op = " <="
	case *query.Gt:
		field = &x.Field
		op = " >"
	case *query.Gte:
		field = &x.Field
		op = " >="
	case *query.And:
		var err error
		for _, v := range x.Conditions {
			dq, err = d.applyQuery(m, dq, v)
			if err != nil {
				return nil, err
			}
		}
	case nil:
	default:
		return nil, fmt.Errorf("datastore does not support %T queries", q)
	}
	if field != nil {
		if _, ok := field.Value.(query.F); ok {
			return nil, fmt.Errorf("datastore queries can't reference other properties (%v)", field.Value)
		}
		name := field.Field
		fields := m.Fields()
		idx, ok := fields.QNameMap[name]
		if !ok {
			return nil, fmt.Errorf("can't map field %q to a datastore name", name)
		}
		if strings.IndexByte(name, '.') >= 0 {
			// GAE flattens embedded fields, so we must remove
			// the parts of the field which refer to a flattened
			// field.
			indexes := fields.Indexes[idx]
			parts := strings.Split(name, ".")
			if len(indexes) == len(parts) {
				var final []string
				typ := fields.Type
				for ii, v := range indexes {
					f := typ.Field(v)
					if !f.Anonymous {
						final = append(final, parts[ii])
					}
					typ = f.Type
				}
				name = strings.Join(final, ".")
			}
		}
		log.Debugf("DATASTORE: filter %s %s %v", m, name+op, field.Value)
		dq = dq.Filter(name+op, field.Value)
	}
	return dq, nil
}

func (d *Driver) Close() error {
	return nil
}

func (d *Driver) Upserts() bool {
	return false
}

func (d *Driver) Tags() []string {
	return []string{"datastore"}
}

func (d *Driver) SetLogger(logger *log.Logger) {
	d.logger = logger
}

func (d *Driver) SetContext(ctx appengine.Context) {
	d.c = ctx
}

func (d *Driver) Begin() (driver.Tx, error) {
	return nil, errTransactionNotSupported
}

func (d *Driver) Commit() error {
	return driver.ErrNotInTransaction
}

func (d *Driver) Rollback() error {
	return driver.ErrNotInTransaction
}

func (d *Driver) Transaction(f func(driver.Driver) error) error {
	return datastore.RunInTransaction(d.c, func(c appengine.Context) error {
		drv := *d
		drv.c = c
		return f(&drv)
	}, nil)
}

func (d *Driver) Capabilities() driver.Capability {
	return driver.CAP_TRANSACTION | driver.CAP_AUTO_ID | driver.CAP_EVENTUAL | driver.CAP_PK
}

func (d *Driver) primaryKey(f *driver.Fields, data interface{}) reflect.Value {
	return fieldByIndex(reflect.ValueOf(data), f.Indexes[f.PrimaryKey])
}

func (d *Driver) HasFunc(fname string, retType reflect.Type) bool {
	return false
}

func (d *Driver) Connection() interface{} {
	return d.c
}

func fieldByIndex(val reflect.Value, indexes []int) reflect.Value {
	for _, v := range indexes {
		if val.Type().Kind() == reflect.Ptr {
			if val.IsNil() {
				return reflect.Value{}
			}
			val = val.Elem()
		}
		val = val.Field(v)
	}
	return val
}

func datastoreOpener(url *config.URL) (driver.Driver, error) {
	return &Driver{}, nil
}

func init() {
	driver.Register("datastore", datastoreOpener)
}
<|endoftext|>


<commit_before>package statserver

import (
	"math"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/getlantern/geolookup"
)

var (
	publishInterval   = 10 * time.Second
	retryWaitTime     = 500 * time.Millisecond
	maxGeolocateTries = 10
)

// Peer represents information about a peer
type Peer struct {
	IP              string    `json:"peerid"`
	LastConnected   time.Time `json:"lastConnected"`
	BytesDn         int64     `json:"bytesDn"`
	BytesUp         int64     `json:"bytesUp"`
	BytesUpDn       int64     `json:"bytesUpDn"`
	BPSDn           int64     `json:"bpsDn"`
	BPSUp           int64     `json:"bpsUp"`
	BPSUpDn         int64     `json:"bpsUpDn"`
	Country         string    `json:"country"`
	Latitude        float64   `json:"lat"`
	Longitude       float64   `json:"lon"`
	pub             publish
	atLastReporting *Peer
	lastReported    time.Time
	reportedFinal   bool
}

// publish is a function to which a peer can publish itself
type publish func(peer *Peer)

func newPeer(ip string, pub publish) *Peer {
	peer := &Peer{
		IP:              ip,
		pub:             pub,
		lastReported:    time.Now(),
		atLastReporting: &Peer{},
	}
	*peer.atLastReporting = *peer
	go peer.run()
	return peer
}

func (peer *Peer) run() {
	err := peer.geolocate()

	if err != nil {
		log.Errorf("Unable to geolocate peer after %d attempts, stopping reporting. Last error was: %s", maxGeolocateTries, err)
		return
	}

	for {
		newActivity := peer.LastConnected != peer.atLastReporting.LastConnected
		if newActivity {
			// We have new activity, meaning that we will eventually need to
			// report a final update
			peer.reportedFinal = false
		}

		// Only report if there's been activity or we need to make our final report
		shouldReport := newActivity || !peer.reportedFinal
		if shouldReport {
			log.Tracef("%v reporting", peer.IP)

			// Calculate stats
			now := time.Now()
			peer.lastReported = now
			delta := peer.lastReported.Sub(peer.atLastReporting.lastReported).Seconds()
			peer.BytesUpDn = peer.BytesUp + peer.BytesDn
			peer.BPSDn = int64(float64(peer.BytesDn-peer.atLastReporting.BytesDn) / delta)
			peer.BPSUp = int64(float64(peer.BytesUp-peer.atLastReporting.BytesUp) / delta)
			peer.BPSUpDn = peer.BPSDn + peer.BPSUp

			// Remember copy of peer as last reported
			*peer.atLastReporting = *peer

			// Publish copy of peer
			peer.pub(peer.atLastReporting)

			if shouldReport && !newActivity {
				log.Tracef("%v just reported its final update", peer.IP)
				peer.reportedFinal = true
			}
		}

		time.Sleep(publishInterval)
	}
}

func (peer *Peer) geolocate() error {
	var err error

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

		if i > 0 {
			// Maximum sleep time: 2^(maxGeolocateTries - 1) * retryWaitTime
			retryWait := time.Duration(math.Pow(2, float64(i)) * float64(retryWaitTime))
			log.Errorf("Failed geolocation attempt %d/%d,  waiting %v before retrying: %s", i, maxGeolocateTries-1, retryWait, err)
			time.Sleep(retryWait)
		}

		err = peer.doGeolocate()

		if err == nil {
			return nil
		}

	}

	return err
}

func (peer *Peer) doGeolocate() error {
	geodata, _, err := geolookup.LookupIPWithClient(peer.IP, geoClient.Load().(*http.Client))

	if err != nil {
		return err
	}

	peer.Country = geodata.Country.IsoCode
	peer.Latitude = geodata.Location.Latitude
	peer.Longitude = geodata.Location.Longitude

	return nil
}

func (peer *Peer) onBytesReceived(bytes int64) {
	peer.LastConnected = time.Now()
	atomic.AddInt64(&peer.BytesUp, bytes)
}

func (peer *Peer) onBytesSent(bytes int64) {
	peer.LastConnected = time.Now()
	atomic.AddInt64(&peer.BytesDn, bytes)
}
<commit_msg>Cleaning up data races in peer stats closes #2980<commit_after>package statserver

import (
	"math"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/getlantern/geolookup"
)

var (
	publishInterval   = 10 * time.Second
	retryWaitTime     = 500 * time.Millisecond
	maxGeolocateTries = 10
)

// Peer represents information about a peer
type Peer struct {
	IP              string    `json:"peerid"`
	LastConnected   time.Time `json:"lastConnected"`
	BytesDn         int64     `json:"bytesDn"`
	BytesUp         int64     `json:"bytesUp"`
	BytesUpDn       int64     `json:"bytesUpDn"`
	BPSDn           int64     `json:"bpsDn"`
	BPSUp           int64     `json:"bpsUp"`
	BPSUpDn         int64     `json:"bpsUpDn"`
	Country         string    `json:"country"`
	Latitude        float64   `json:"lat"`
	Longitude       float64   `json:"lon"`
	pub             publish
	atLastReporting *Peer
	lastReported    time.Time
	reportedFinal   bool
	sync.Mutex
}

// publish is a function to which a peer can publish itself
type publish func(peer *Peer)

func newPeer(ip string, pub publish) *Peer {
	peer := &Peer{
		IP:              ip,
		pub:             pub,
		lastReported:    time.Now(),
		atLastReporting: &Peer{},
	}
	*peer.atLastReporting = *peer
	go peer.run()
	return peer
}

func (peer *Peer) run() {
	err := peer.geolocate()

	if err != nil {
		log.Errorf("Unable to geolocate peer after %d attempts, stopping reporting. Last error was: %s", maxGeolocateTries, err)
		return
	}

	for {
		newActivity := peer.lastConnected() != peer.atLastReporting.lastConnected()
		if newActivity {
			// We have new activity, meaning that we will eventually need to
			// report a final update
			peer.reportedFinal = false
		}

		// Only report if there's been activity or we need to make our final report
		shouldReport := newActivity || !peer.reportedFinal
		if shouldReport {
			log.Tracef("%v reporting", peer.IP)

			// Calculate stats
			now := time.Now()
			peer.lastReported = now
			delta := peer.lastReported.Sub(peer.atLastReporting.lastReported).Seconds()
			peer.BytesUpDn = peer.bytesUp() + peer.bytesDn()
			peer.BPSDn = int64(float64(peer.bytesDn()-peer.atLastReporting.bytesDn()) / delta)
			peer.BPSUp = int64(float64(peer.bytesUp()-peer.atLastReporting.bytesUp()) / delta)
			peer.BPSUpDn = peer.BPSDn + peer.BPSUp

			// Remember copy of peer as last reported
			peer.Lock()
			*peer.atLastReporting = *peer
			peer.Unlock()

			// Publish copy of peer
			peer.pub(peer.atLastReporting)

			if shouldReport && !newActivity {
				log.Tracef("%v just reported its final update", peer.IP)
				peer.reportedFinal = true
			}
		}

		time.Sleep(publishInterval)
	}
}

func (peer *Peer) geolocate() error {
	var err error

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

		if i > 0 {
			// Maximum sleep time: 2^(maxGeolocateTries - 1) * retryWaitTime
			retryWait := time.Duration(math.Pow(2, float64(i)) * float64(retryWaitTime))
			log.Errorf("Failed geolocation attempt %d/%d,  waiting %v before retrying: %s", i, maxGeolocateTries-1, retryWait, err)
			time.Sleep(retryWait)
		}

		err = peer.doGeolocate()

		if err == nil {
			return nil
		}

	}

	return err
}

func (peer *Peer) doGeolocate() error {
	geodata, _, err := geolookup.LookupIPWithClient(peer.IP, geoClient.Load().(*http.Client))

	if err != nil {
		return err
	}

	peer.Country = geodata.Country.IsoCode
	peer.Latitude = geodata.Location.Latitude
	peer.Longitude = geodata.Location.Longitude

	return nil
}

func (peer *Peer) lastConnected() time.Time {
	peer.Lock()
	defer peer.Unlock()
	return peer.LastConnected
}

func (peer *Peer) setLastConnected() {
	peer.Lock()
	defer peer.Unlock()
	peer.LastConnected = time.Now()
}

func (peer *Peer) bytesDn() int64 {
	return atomic.LoadInt64(&peer.BytesDn)
}

func (peer *Peer) bytesUp() int64 {
	return atomic.LoadInt64(&peer.BytesUp)
}

func (peer *Peer) onBytesReceived(bytes int64) {
	peer.setLastConnected()
	atomic.AddInt64(&peer.BytesUp, bytes)
}

func (peer *Peer) onBytesSent(bytes int64) {
	peer.setLastConnected()
	atomic.AddInt64(&peer.BytesDn, bytes)
}
<|endoftext|>


<commit_before>package forward

import (
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"

	"github.com/getlantern/http-proxy/utils"
	"github.com/getlantern/idletiming"
)

type Forwarder struct {
	log          utils.Logger
	errHandler   utils.ErrorHandler
	roundTripper http.RoundTripper
	rewriter     RequestRewriter
	next         http.Handler

	idleTimeout time.Duration
}

type optSetter func(f *Forwarder) error

func RoundTripper(r http.RoundTripper) optSetter {
	return func(f *Forwarder) error {
		f.roundTripper = r
		return nil
	}
}

type RequestRewriter interface {
	Rewrite(r *http.Request)
}

func Rewriter(r RequestRewriter) optSetter {
	return func(f *Forwarder) error {
		f.rewriter = r
		return nil
	}
}

func Logger(l utils.Logger) optSetter {
	return func(f *Forwarder) error {
		f.log = l
		return nil
	}
}

func IdleTimeoutSetter(i time.Duration) optSetter {
	return func(f *Forwarder) error {
		f.idleTimeout = i
		return nil
	}
}

func New(next http.Handler, setters ...optSetter) (*Forwarder, error) {
	var dialerFunc func(string, string) (net.Conn, error)

	var timeoutTransport http.RoundTripper = &http.Transport{
		Dial:                dialerFunc,
		TLSHandshakeTimeout: 10 * time.Second,
	}
	f := &Forwarder{
		log:          utils.NullLogger,
		errHandler:   utils.DefaultHandler,
		roundTripper: timeoutTransport,
		next:         next,
		idleTimeout:  30,
	}
	for _, s := range setters {
		if err := s(f); err != nil {
			return nil, err
		}
	}
	if f.rewriter == nil {
		f.rewriter = &HeaderRewriter{
			TrustForwardHeader: true,
			Hostname:           "",
		}
	}

	dialerFunc = func(network, addr string) (conn net.Conn, err error) {
		conn, err = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial(network, addr)
		if err != nil {
			return
		}

		idleConn := idletiming.Conn(conn, f.idleTimeout, func() {
			if conn != nil {
				conn.Close()
			}
		})
		return idleConn, err
	}

	return f, nil
}

func (f *Forwarder) ServeHTTP(w http.ResponseWriter, req *http.Request) {

	// Create a copy of the request suitable for our needs
	reqClone, err := f.cloneRequest(req, req.URL)
	if err != nil {
		f.log.Errorf("Error forwarding to %v, error: %v", req.Host, err)
		f.errHandler.ServeHTTP(w, req, err)
		return
	}
	f.rewriter.Rewrite(reqClone)

	// Forward the request and get a response
	start := time.Now().UTC()
	response, err := f.roundTripper.RoundTrip(reqClone)
	if err != nil {
		f.log.Errorf("Error forwarding to %v, error: %v", req.Host, err)
		f.errHandler.ServeHTTP(w, req, err)
		return
	}
	f.log.Infof("Round trip: %v, code: %v, duration: %v\n",
		req.URL, response.StatusCode, time.Now().UTC().Sub(start))

	if f.log.IsLevel(utils.DEBUG) {
		respStr, _ := httputil.DumpResponse(response, true)
		f.log.Debugf("Forward Middleware received response:\n%s", respStr)
	}

	// Forward the response to the origin
	copyHeadersForForwarding(w.Header(), response.Header)
	w.WriteHeader(response.StatusCode)

	// It became nil in a Co-Advisor test though the doc says it will never be nil
	if response.Body != nil {
		_, err = io.Copy(w, response.Body)
		if err != nil {
			f.log.Errorf("%v\n", err)
		}

		response.Body.Close()
	}
}

func (f *Forwarder) cloneRequest(req *http.Request, u *url.URL) (*http.Request, error) {
	outReq := new(http.Request)
	// Beware, this will make a shallow copy. We have to copy all maps
	*outReq = *req

	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	outReq.Close = false

	// Request Header
	outReq.Header = make(http.Header)
	copyHeadersForForwarding(outReq.Header, req.Header)

	// Request URL
	scheme := "http"
	outReq.URL = cloneURL(req.URL)
	outReq.URL.Scheme = scheme
	outReq.URL.Host = req.Host
	outReq.URL.Opaque = req.RequestURI
	// raw query is already included in RequestURI, so ignore it to avoid dupes
	outReq.URL.RawQuery = ""
	// Do not pass client Host header unless optsetter PassHostHeader is set.

	// Trailer support
	// We are forced to do this because Go's server won't allow us to read the trailers otherwise
	_, err := httputil.DumpRequestOut(req, true)
	if err != nil {
		f.log.Errorf("Error: %v", err)
		return nil, err
	}

	//f.log.Errorf("%v\n", string(buf))
	rcloser := ioutil.NopCloser(req.Body)
	outReq.Body = rcloser

	chunkedTrasfer := false
	for _, enc := range req.TransferEncoding {
		if enc == "chunked" {
			chunkedTrasfer = true
			break
		}
	}

	// Append Trailer
	if chunkedTrasfer && len(req.Trailer) > 0 && req.Header.Get("Content-Length") == "" {
		outReq.Trailer = http.Header{}
		for k, vv := range req.Trailer {
			for _, v := range vv {
				outReq.Trailer.Add(k, v)
			}
		}
	}

	return outReq, nil
}
<commit_msg>Check for Content-Length meaningless in this scenario<commit_after>package forward

import (
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"

	"github.com/getlantern/http-proxy/utils"
	"github.com/getlantern/idletiming"
)

type Forwarder struct {
	log          utils.Logger
	errHandler   utils.ErrorHandler
	roundTripper http.RoundTripper
	rewriter     RequestRewriter
	next         http.Handler

	idleTimeout time.Duration
}

type optSetter func(f *Forwarder) error

func RoundTripper(r http.RoundTripper) optSetter {
	return func(f *Forwarder) error {
		f.roundTripper = r
		return nil
	}
}

type RequestRewriter interface {
	Rewrite(r *http.Request)
}

func Rewriter(r RequestRewriter) optSetter {
	return func(f *Forwarder) error {
		f.rewriter = r
		return nil
	}
}

func Logger(l utils.Logger) optSetter {
	return func(f *Forwarder) error {
		f.log = l
		return nil
	}
}

func IdleTimeoutSetter(i time.Duration) optSetter {
	return func(f *Forwarder) error {
		f.idleTimeout = i
		return nil
	}
}

func New(next http.Handler, setters ...optSetter) (*Forwarder, error) {
	var dialerFunc func(string, string) (net.Conn, error)

	var timeoutTransport http.RoundTripper = &http.Transport{
		Dial:                dialerFunc,
		TLSHandshakeTimeout: 10 * time.Second,
	}
	f := &Forwarder{
		log:          utils.NullLogger,
		errHandler:   utils.DefaultHandler,
		roundTripper: timeoutTransport,
		next:         next,
		idleTimeout:  30,
	}
	for _, s := range setters {
		if err := s(f); err != nil {
			return nil, err
		}
	}
	if f.rewriter == nil {
		f.rewriter = &HeaderRewriter{
			TrustForwardHeader: true,
			Hostname:           "",
		}
	}

	dialerFunc = func(network, addr string) (conn net.Conn, err error) {
		conn, err = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial(network, addr)
		if err != nil {
			return
		}

		idleConn := idletiming.Conn(conn, f.idleTimeout, func() {
			if conn != nil {
				conn.Close()
			}
		})
		return idleConn, err
	}

	return f, nil
}

func (f *Forwarder) ServeHTTP(w http.ResponseWriter, req *http.Request) {

	// Create a copy of the request suitable for our needs
	reqClone, err := f.cloneRequest(req, req.URL)
	if err != nil {
		f.log.Errorf("Error forwarding to %v, error: %v", req.Host, err)
		f.errHandler.ServeHTTP(w, req, err)
		return
	}
	f.rewriter.Rewrite(reqClone)

	// Forward the request and get a response
	start := time.Now().UTC()
	response, err := f.roundTripper.RoundTrip(reqClone)
	if err != nil {
		f.log.Errorf("Error forwarding to %v, error: %v", req.Host, err)
		f.errHandler.ServeHTTP(w, req, err)
		return
	}
	f.log.Infof("Round trip: %v, code: %v, duration: %v\n",
		req.URL, response.StatusCode, time.Now().UTC().Sub(start))

	if f.log.IsLevel(utils.DEBUG) {
		respStr, _ := httputil.DumpResponse(response, true)
		f.log.Debugf("Forward Middleware received response:\n%s", respStr)
	}

	// Forward the response to the origin
	copyHeadersForForwarding(w.Header(), response.Header)
	w.WriteHeader(response.StatusCode)

	// It became nil in a Co-Advisor test though the doc says it will never be nil
	if response.Body != nil {
		_, err = io.Copy(w, response.Body)
		if err != nil {
			f.log.Errorf("%v\n", err)
		}

		response.Body.Close()
	}
}

func (f *Forwarder) cloneRequest(req *http.Request, u *url.URL) (*http.Request, error) {
	outReq := new(http.Request)
	// Beware, this will make a shallow copy. We have to copy all maps
	*outReq = *req

	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	outReq.Close = false

	// Request Header
	outReq.Header = make(http.Header)
	copyHeadersForForwarding(outReq.Header, req.Header)

	// Request URL
	scheme := "http"
	outReq.URL = cloneURL(req.URL)
	outReq.URL.Scheme = scheme
	outReq.URL.Host = req.Host
	outReq.URL.Opaque = req.RequestURI
	// raw query is already included in RequestURI, so ignore it to avoid dupes
	outReq.URL.RawQuery = ""
	// Do not pass client Host header unless optsetter PassHostHeader is set.

	// Trailer support
	// We are forced to do this because Go's server won't allow us to read the trailers otherwise
	_, err := httputil.DumpRequestOut(req, true)
	if err != nil {
		f.log.Errorf("Error: %v", err)
		return nil, err
	}

	rcloser := ioutil.NopCloser(req.Body)
	outReq.Body = rcloser

	chunkedTrasfer := false
	for _, enc := range req.TransferEncoding {
		if enc == "chunked" {
			chunkedTrasfer = true
			break
		}
	}

	// Append Trailer
	if chunkedTrasfer && len(req.Trailer) > 0 {
		outReq.Trailer = http.Header{}
		for k, vv := range req.Trailer {
			for _, v := range vv {
				outReq.Trailer.Add(k, v)
			}
		}
	}

	return outReq, nil
}
<|endoftext|>


<commit_before>// Package base provides a base implementation of the storage driver that can
// be used to implement common checks. The goal is to increase the amount of
// code sharing.
//
// The canonical approach to use this class is to embed in the exported driver
// struct such that calls are proxied through this implementation. First,
// declare the internal driver, as follows:
//
// 	type driver struct { ... internal ...}
//
// The resulting type should implement StorageDriver such that it can be the
// target of a Base struct. The exported type can then be declared as follows:
//
// 	type Driver struct {
// 		Base
// 	}
//
// Because Driver embeds Base, it effectively implements Base. If the driver
// needs to intercept a call, before going to base, Driver should implement
// that method. Effectively, Driver can intercept calls before coming in and
// driver implements the actual logic.
//
// To further shield the embed from other packages, it is recommended to
// employ a private embed struct:
//
// 	type baseEmbed struct {
// 		base.Base
// 	}
//
// Then, declare driver to embed baseEmbed, rather than Base directly:
//
// 	type Driver struct {
// 		baseEmbed
// 	}
//
// The type now implements StorageDriver, proxying through Base, without
// exporting an unnessecary field.
package base

import (
	"io"

	storagedriver "github.com/docker/distribution/registry/storage/driver"
)

// Base provides a wrapper around a storagedriver implementation that provides
// common path and bounds checking.
type Base struct {
	storagedriver.StorageDriver
}

// GetContent wraps GetContent of underlying storage driver.
func (base *Base) GetContent(path string) ([]byte, error) {
	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.GetContent(path)
}

// PutContent wraps PutContent of underlying storage driver.
func (base *Base) PutContent(path string, content []byte) error {
	if !storagedriver.PathRegexp.MatchString(path) {
		return storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.PutContent(path, content)
}

// ReadStream wraps ReadStream of underlying storage driver.
func (base *Base) ReadStream(path string, offset int64) (io.ReadCloser, error) {
	if offset < 0 {
		return nil, storagedriver.InvalidOffsetError{Path: path, Offset: offset}
	}

	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.ReadStream(path, offset)
}

// WriteStream wraps WriteStream of underlying storage driver.
func (base *Base) WriteStream(path string, offset int64, reader io.Reader) (nn int64, err error) {
	if offset < 0 {
		return 0, storagedriver.InvalidOffsetError{Path: path, Offset: offset}
	}

	if !storagedriver.PathRegexp.MatchString(path) {
		return 0, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.WriteStream(path, offset, reader)
}

// Stat wraps Stat of underlying storage driver.
func (base *Base) Stat(path string) (storagedriver.FileInfo, error) {
	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.Stat(path)
}

// List wraps List of underlying storage driver.
func (base *Base) List(path string) ([]string, error) {
	if !storagedriver.PathRegexp.MatchString(path) && path != "/" {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.List(path)
}

// Move wraps Move of underlying storage driver.
func (base *Base) Move(sourcePath string, destPath string) error {
	if !storagedriver.PathRegexp.MatchString(sourcePath) {
		return storagedriver.InvalidPathError{Path: sourcePath}
	} else if !storagedriver.PathRegexp.MatchString(destPath) {
		return storagedriver.InvalidPathError{Path: destPath}
	}

	return base.StorageDriver.Move(sourcePath, destPath)
}

// Delete wraps Delete of underlying storage driver.
func (base *Base) Delete(path string) error {
	if !storagedriver.PathRegexp.MatchString(path) {
		return storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.Delete(path)
}

// URLFor wraps URLFor of underlying storage driver.
func (base *Base) URLFor(path string, options map[string]interface{}) (string, error) {
	if !storagedriver.PathRegexp.MatchString(path) {
		return "", storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.URLFor(path, options)
}
<commit_msg>Trace function calls to Base storage driver<commit_after>// Package base provides a base implementation of the storage driver that can
// be used to implement common checks. The goal is to increase the amount of
// code sharing.
//
// The canonical approach to use this class is to embed in the exported driver
// struct such that calls are proxied through this implementation. First,
// declare the internal driver, as follows:
//
// 	type driver struct { ... internal ...}
//
// The resulting type should implement StorageDriver such that it can be the
// target of a Base struct. The exported type can then be declared as follows:
//
// 	type Driver struct {
// 		Base
// 	}
//
// Because Driver embeds Base, it effectively implements Base. If the driver
// needs to intercept a call, before going to base, Driver should implement
// that method. Effectively, Driver can intercept calls before coming in and
// driver implements the actual logic.
//
// To further shield the embed from other packages, it is recommended to
// employ a private embed struct:
//
// 	type baseEmbed struct {
// 		base.Base
// 	}
//
// Then, declare driver to embed baseEmbed, rather than Base directly:
//
// 	type Driver struct {
// 		baseEmbed
// 	}
//
// The type now implements StorageDriver, proxying through Base, without
// exporting an unnessecary field.
package base

import (
	"io"

	"github.com/docker/distribution/context"
	storagedriver "github.com/docker/distribution/registry/storage/driver"
)

// Base provides a wrapper around a storagedriver implementation that provides
// common path and bounds checking.
type Base struct {
	storagedriver.StorageDriver
}

// GetContent wraps GetContent of underlying storage driver.
func (base *Base) GetContent(path string) ([]byte, error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.GetContent")

	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.GetContent(path)
}

// PutContent wraps PutContent of underlying storage driver.
func (base *Base) PutContent(path string, content []byte) error {
	_, done := context.WithTrace(context.Background())
	defer done("Base.PutContent")

	if !storagedriver.PathRegexp.MatchString(path) {
		return storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.PutContent(path, content)
}

// ReadStream wraps ReadStream of underlying storage driver.
func (base *Base) ReadStream(path string, offset int64) (io.ReadCloser, error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.ReadStream")

	if offset < 0 {
		return nil, storagedriver.InvalidOffsetError{Path: path, Offset: offset}
	}

	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.ReadStream(path, offset)
}

// WriteStream wraps WriteStream of underlying storage driver.
func (base *Base) WriteStream(path string, offset int64, reader io.Reader) (nn int64, err error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.WriteStream")

	if offset < 0 {
		return 0, storagedriver.InvalidOffsetError{Path: path, Offset: offset}
	}

	if !storagedriver.PathRegexp.MatchString(path) {
		return 0, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.WriteStream(path, offset, reader)
}

// Stat wraps Stat of underlying storage driver.
func (base *Base) Stat(path string) (storagedriver.FileInfo, error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.Stat")

	if !storagedriver.PathRegexp.MatchString(path) {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.Stat(path)
}

// List wraps List of underlying storage driver.
func (base *Base) List(path string) ([]string, error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.List")

	if !storagedriver.PathRegexp.MatchString(path) && path != "/" {
		return nil, storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.List(path)
}

// Move wraps Move of underlying storage driver.
func (base *Base) Move(sourcePath string, destPath string) error {
	_, done := context.WithTrace(context.Background())
	defer done("Base.Move")

	if !storagedriver.PathRegexp.MatchString(sourcePath) {
		return storagedriver.InvalidPathError{Path: sourcePath}
	} else if !storagedriver.PathRegexp.MatchString(destPath) {
		return storagedriver.InvalidPathError{Path: destPath}
	}

	return base.StorageDriver.Move(sourcePath, destPath)
}

// Delete wraps Delete of underlying storage driver.
func (base *Base) Delete(path string) error {
	_, done := context.WithTrace(context.Background())
	defer done("Base.Move")

	if !storagedriver.PathRegexp.MatchString(path) {
		return storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.Delete(path)
}

// URLFor wraps URLFor of underlying storage driver.
func (base *Base) URLFor(path string, options map[string]interface{}) (string, error) {
	_, done := context.WithTrace(context.Background())
	defer done("Base.URLFor")

	if !storagedriver.PathRegexp.MatchString(path) {
		return "", storagedriver.InvalidPathError{Path: path}
	}

	return base.StorageDriver.URLFor(path, options)
}
<|endoftext|>


<commit_before>package relationships

import (
	"errors"
	"fmt"
	"github.com/bor3ham/reja/schema"
)

type GenericForeignKey struct {
	RelationshipStub
	Key            string
	TypeColumnName string
	IDColumnName   string
	Nullable       bool
	Default        func(schema.Context, interface{}) Pointer
	ValidTypes     []string
}

func (gfk GenericForeignKey) GetKey() string {
	return gfk.Key
}
func (gfk GenericForeignKey) GetType() string {
	return ""
}

func (gfk GenericForeignKey) GetSelectExtraColumns() []string {
	return []string{
		gfk.TypeColumnName,
		gfk.IDColumnName,
	}
}
func (gfk GenericForeignKey) GetSelectExtraVariables() []interface{} {
	var typeDest *string
	var idDest *string
	return []interface{}{
		&typeDest,
		&idDest,
	}
}

func (gfk GenericForeignKey) GetDefaultValue() interface{} {
	return schema.Result{}
}
func (gfk GenericForeignKey) GetValues(
	c schema.Context,
	m *schema.Model,
	ids []string,
	extra [][]interface{},
) (
	map[string]interface{},
	map[string]map[string][]string,
) {
	values := map[string]interface{}{}
	maps := map[string]map[string][]string{}
	for index, result := range extra {
		myId := ids[index]

		// parse extra columns
		modelType, ok := result[0].(**string)
		if !ok {
			panic("Unable to convert extra type")
		}
		stringId, ok := result[1].(**string)
		if !ok {
			panic("Unable to convert extra fk id")
		}

		// check value does not already exist
		// a foreign key can only have one value
		_, exists := values[myId]
		if exists {
			existingValue, ok := values[myId].(Pointer)
			if !ok {
				panic("Unable to convert previous value")
			}
			if *stringId == nil {
				if existingValue.Data != nil {
					panic("Contradictory values in query results")
				}
			} else {
				if existingValue.Data == nil ||
					*existingValue.Data.ID != **stringId ||
					existingValue.Data.Type != **modelType {
					panic("Contradictory values in query results")
				}
			}

			continue
		}

		selfLink := relationLink(c, m.Type, myId, gfk.Key)
		newValue := schema.Result{
			Links: map[string]*string{
				"self": &selfLink,
			},
		}
		if *stringId != nil {
			newValue.Data = schema.InstancePointer{
				Type: **modelType,
				ID:   *stringId,
			}
		}
		// update the value
		values[myId] = newValue

		// add to relation map
		if *stringId != nil {
			_, exists = maps[myId]
			if !exists {
				maps[myId] = map[string][]string{}
			}
			_, exists = maps[myId][**modelType]
			if !exists {
				maps[myId][**modelType] = []string{}
			}
			maps[myId][**modelType] = append(maps[myId][**modelType], **stringId)
		}
	}

	return values, maps
}

func (gfk *GenericForeignKey) DefaultFallback(
	c schema.Context,
	val interface{},
	instance interface{},
) (
	interface{},
	error,
) {
	var gfkVal Pointer
	if val == nil {
		gfkVal = Pointer{Provided: false}
	} else {
		var err error
		gfkVal, err = ParseResultPointer(val)
		if err != nil {
			return nil, err
		}
	}

	if !gfkVal.Provided {
		if gfk.Default != nil {
			return gfk.Default(c, instance), nil
		}
		return nil, nil
	}
	return gfkVal, nil
}
func (gfk *GenericForeignKey) Validate(c schema.Context, val interface{}) (interface{}, error) {
	gfkVal := AssertPointer(val)

	if gfkVal.Data == nil {
		if !gfk.Nullable {
			return nil, errors.New(fmt.Sprintf(
				"Relationship '%s' invalid: Cannot be null.",
				gfk.Key,
			))
		}
		return gfkVal, nil
	}

	valType := gfkVal.Data.Type
	if gfkVal.Data.ID == nil {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: Missing ID.",
			gfk.Key,
		))
	}
	valID := *gfkVal.Data.ID

	// check that the type is valid
	if len(gfk.ValidTypes) > 0 {
		valid := false
		for _, relationType := range gfk.ValidTypes {
			if relationType == valType {
				valid = true
			}
		}
		if !valid {
			return nil, errors.New(fmt.Sprintf(
				"Relationship '%s' invalid: Bad type for relation.",
				gfk.Key,
			))
		}
	}
	// check that the object exists
	model := c.GetServer().GetModel(valType)
	// validate the type exists
	if model == nil {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: Non existent type.",
			gfk.Key,
		))
	}
	include := schema.Include{
		Children: map[string]*schema.Include{},
	}
	instances, _, err := c.GetObjectsByIDs(model, []string{valID}, &include)
	if err != nil {
		panic(err)
	}
	if len(instances) == 0 {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: %s ID '%s' does not exist.",
			gfk.Key,
			valType,
			valID,
		))
	}
	return gfkVal, nil
}

func (gfk *GenericForeignKey) GetInsertColumns(val interface{}) []string {
	return []string{
		gfk.TypeColumnName,
		gfk.IDColumnName,
	}
}
func (gfk *GenericForeignKey) GetInsertValues(val interface{}) []interface{} {
	resultVal := AssertPointer(val)
	if resultVal.Data == nil {
		return []interface{}{
			nil,
			nil,
		}
	}
	return []interface{}{
		resultVal.Data.Type,
		resultVal.Data.ID,
	}
}

func AssertGenericForeignKey(val interface{}) schema.Result {
	gfkVal, ok := val.(schema.Result)
	if !ok {
		panic("Bad generic foreign key value")
	}
	return gfkVal
}
<commit_msg>Add gfk patching<commit_after>package relationships

import (
	"errors"
	"fmt"
	"github.com/bor3ham/reja/schema"
)

type GenericForeignKey struct {
	RelationshipStub
	Key            string
	TypeColumnName string
	IDColumnName   string
	Nullable       bool
	Default        func(schema.Context, interface{}) Pointer
	ValidTypes     []string
}

func (gfk GenericForeignKey) GetKey() string {
	return gfk.Key
}
func (gfk GenericForeignKey) GetType() string {
	return ""
}

func (gfk GenericForeignKey) GetSelectExtraColumns() []string {
	return []string{
		gfk.TypeColumnName,
		gfk.IDColumnName,
	}
}
func (gfk GenericForeignKey) GetSelectExtraVariables() []interface{} {
	var typeDest *string
	var idDest *string
	return []interface{}{
		&typeDest,
		&idDest,
	}
}

func (gfk GenericForeignKey) GetDefaultValue() interface{} {
	return schema.Result{}
}
func (gfk GenericForeignKey) GetValues(
	c schema.Context,
	m *schema.Model,
	ids []string,
	extra [][]interface{},
) (
	map[string]interface{},
	map[string]map[string][]string,
) {
	values := map[string]interface{}{}
	maps := map[string]map[string][]string{}
	for index, result := range extra {
		myId := ids[index]

		// parse extra columns
		modelType, ok := result[0].(**string)
		if !ok {
			panic("Unable to convert extra type")
		}
		stringId, ok := result[1].(**string)
		if !ok {
			panic("Unable to convert extra fk id")
		}

		// check value does not already exist
		// a foreign key can only have one value
		_, exists := values[myId]
		if exists {
			existingValue, ok := values[myId].(Pointer)
			if !ok {
				panic("Unable to convert previous value")
			}
			if *stringId == nil {
				if existingValue.Data != nil {
					panic("Contradictory values in query results")
				}
			} else {
				if existingValue.Data == nil ||
					*existingValue.Data.ID != **stringId ||
					existingValue.Data.Type != **modelType {
					panic("Contradictory values in query results")
				}
			}

			continue
		}

		selfLink := relationLink(c, m.Type, myId, gfk.Key)
		newValue := schema.Result{
			Links: map[string]*string{
				"self": &selfLink,
			},
		}
		if *stringId != nil {
			newValue.Data = schema.InstancePointer{
				Type: **modelType,
				ID:   *stringId,
			}
		}
		// update the value
		values[myId] = newValue

		// add to relation map
		if *stringId != nil {
			_, exists = maps[myId]
			if !exists {
				maps[myId] = map[string][]string{}
			}
			_, exists = maps[myId][**modelType]
			if !exists {
				maps[myId][**modelType] = []string{}
			}
			maps[myId][**modelType] = append(maps[myId][**modelType], **stringId)
		}
	}

	return values, maps
}

func (gfk *GenericForeignKey) DefaultFallback(
	c schema.Context,
	val interface{},
	instance interface{},
) (
	interface{},
	error,
) {
	var gfkVal Pointer
	if val == nil {
		gfkVal = Pointer{Provided: false}
	} else {
		var err error
		gfkVal, err = ParseResultPointer(val)
		if err != nil {
			return nil, err
		}
	}

	if !gfkVal.Provided {
		if gfk.Default != nil {
			return gfk.Default(c, instance), nil
		}
		return nil, nil
	}
	return gfkVal, nil
}
func (gfk *GenericForeignKey) Validate(c schema.Context, val interface{}) (interface{}, error) {
	gfkVal := AssertPointer(val)

	if gfkVal.Data == nil {
		if !gfk.Nullable {
			return nil, errors.New(fmt.Sprintf(
				"Relationship '%s' invalid: Cannot be null.",
				gfk.Key,
			))
		}
		return gfkVal, nil
	}

	valType := gfkVal.Data.Type
	if gfkVal.Data.ID == nil {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: Missing ID.",
			gfk.Key,
		))
	}
	valID := *gfkVal.Data.ID

	// check that the type is valid
	if len(gfk.ValidTypes) > 0 {
		valid := false
		for _, relationType := range gfk.ValidTypes {
			if relationType == valType {
				valid = true
			}
		}
		if !valid {
			return nil, errors.New(fmt.Sprintf(
				"Relationship '%s' invalid: Bad type for relation.",
				gfk.Key,
			))
		}
	}
	// check that the object exists
	model := c.GetServer().GetModel(valType)
	// validate the type exists
	if model == nil {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: Non existent type.",
			gfk.Key,
		))
	}
	include := schema.Include{
		Children: map[string]*schema.Include{},
	}
	instances, _, err := c.GetObjectsByIDs(model, []string{valID}, &include)
	if err != nil {
		panic(err)
	}
	if len(instances) == 0 {
		return nil, errors.New(fmt.Sprintf(
			"Relationship '%s' invalid: %s ID '%s' does not exist.",
			gfk.Key,
			valType,
			valID,
		))
	}
	return gfkVal, nil
}
func (gfk *GenericForeignKey) ValidateUpdate(
	c schema.Context,
	newVal interface{},
	oldVal interface{},
) (
	interface{},
	error,
) {
	// extract new value
	newPointer, err := ParseResultPointer(newVal)
	if err != nil {
		return nil, err
	}
	// if not provided, return nothing
	if !newPointer.Provided {
		return nil, nil
	}
	// clean and check validity of new value
	valid, err := gfk.Validate(c, newPointer)
	if err != nil {
		return nil, err
	}
	validNewPointer := AssertPointer(valid)

	// extract old value
	oldResult, ok := oldVal.(schema.Result)
	if !ok {
		panic("Bad old generic foreign key value")
	}
	var oldValue Pointer
	if oldResult.Data == nil {
		oldValue = Pointer{}
	} else {
		oldPointer, ok := oldResult.Data.(schema.InstancePointer)
		if !ok {
			panic("Bad old generic foreign key value")
		}
		oldValue = Pointer{Data: &oldPointer}
	}

	// return nothing if no changes
	if validNewPointer.Equal(oldValue) {
		return nil, nil
	}
	// otherwise return new validated value
	return validNewPointer, nil
}

func (gfk *GenericForeignKey) GetInsertColumns(val interface{}) []string {
	return []string{
		gfk.TypeColumnName,
		gfk.IDColumnName,
	}
}
func (gfk *GenericForeignKey) GetInsertValues(val interface{}) []interface{} {
	resultVal := AssertPointer(val)
	if resultVal.Data == nil {
		return []interface{}{
			nil,
			nil,
		}
	}
	return []interface{}{
		resultVal.Data.Type,
		resultVal.Data.ID,
	}
}

func AssertGenericForeignKey(val interface{}) schema.Result {
	gfkVal, ok := val.(schema.Result)
	if !ok {
		panic("Bad generic foreign key value")
	}
	return gfkVal
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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.
*/

// Multitenant autoscaler executable.
package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/controller"
	"github.com/knative/pkg/injection"
	"github.com/knative/pkg/injection/clients/kubeclient"
	endpointsinformer "github.com/knative/pkg/injection/informers/kubeinformers/corev1/endpoints"
	"github.com/knative/pkg/logging"
	"github.com/knative/pkg/metrics"
	pkgmetrics "github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/pkg/system"
	"github.com/knative/serving/pkg/apis/serving"
	"github.com/knative/serving/pkg/autoscaler"
	"github.com/knative/serving/pkg/autoscaler/statserver"
	"github.com/knative/serving/pkg/reconciler/autoscaling/hpa"
	"github.com/knative/serving/pkg/reconciler/autoscaling/kpa"
	"github.com/knative/serving/pkg/resources"
	basecmd "github.com/kubernetes-incubator/custom-metrics-apiserver/pkg/cmd"
	"github.com/spf13/pflag"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	corev1informers "k8s.io/client-go/informers/core/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

const (
	statsServerAddr = ":8080"
	statsBufferLen  = 1000
	component       = "autoscaler"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func main() {
	// Initialize early to get access to flags and merge them with the autoscaler flags.
	customMetricsAdapter := &basecmd.AdapterBase{}
	customMetricsAdapter.Flags().AddGoFlagSet(flag.CommandLine)
	pflag.Parse()

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	// Set up our logger.
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	logger, atomicLevel := logging.NewLoggerFromConfig(loggingConfig, component)
	defer flush(logger)
	ctx = logging.WithLogger(ctx, logger)

	cfg, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatalw("Error building kubeconfig", zap.Error(err))
	}

	logger.Infof("Registering %d clients", len(injection.Default.GetClients()))
	logger.Infof("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	logger.Infof("Registering %d informers", len(injection.Default.GetInformers()))
	logger.Infof("Registering %d controllers", 2)

	// Adjust our client's rate limits based on the number of controller's we are running.
	cfg.QPS = 2 * rest.DefaultQPS
	cfg.Burst = 2 * rest.DefaultBurst

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// statsCh is the main communication channel between the stats channel and multiscaler.
	statsCh := make(chan *autoscaler.StatMessage, statsBufferLen)
	defer close(statsCh)

	cmw := configmap.NewInformedWatcher(kubeclient.Get(ctx), system.Namespace())
	// Watch the logging config map and dynamically update logging levels.
	cmw.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	cmw.Watch(metrics.ConfigMapName(), metrics.UpdateExporterFromConfigMap(component, logger))

	endpointsInformer := endpointsinformer.Get(ctx)

	collector := autoscaler.NewMetricCollector(statsScraperFactoryFunc(endpointsInformer.Lister()), logger)
	customMetricsAdapter.WithCustomMetrics(autoscaler.NewMetricProvider(collector))

	// Set up scalers.
	// uniScalerFactory depends endpointsInformer to be set.
	multiScaler := autoscaler.NewMultiScaler(ctx.Done(), uniScalerFactoryFunc(endpointsInformer, collector), logger)

	psInformerFactory := resources.NewPodScalableInformerFactory(ctx)
	controllers := []*controller.Impl{
		kpa.NewController(ctx, cmw, multiScaler, collector, psInformerFactory),
		hpa.NewController(ctx, cmw, collector, psInformerFactory),
	}

	// Set up a statserver.
	statsServer := statserver.New(statsServerAddr, statsCh, logger)
	defer statsServer.Shutdown(time.Second * 5)

	// Start watching the configs.
	if err := cmw.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start watching configs", zap.Error(err))
	}

	// Start all of the informers and wait for them to sync.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", err)
	}

	go controller.StartAll(ctx.Done(), controllers...)

	// Run the controllers and the statserver in a group.
	var eg errgroup.Group
	eg.Go(func() error {
		return customMetricsAdapter.Run(ctx.Done())
	})

	eg.Go(func() error {
		return statsServer.ListenAndServe()
	})

	go func() {
		for {
			sm, ok := <-statsCh
			if !ok {
				break
			}
			collector.Record(sm.Key, sm.Stat)
			multiScaler.Poke(sm.Key, sm.Stat)
		}
	}()

	egCh := make(chan struct{})

	go func() {
		if err := eg.Wait(); err != nil {
			logger.Errorw("Group error.", zap.Error(err))
		}
		close(egCh)
	}()

	select {
	case <-egCh:
	case <-ctx.Done():
	}
}

func uniScalerFactoryFunc(endpointsInformer corev1informers.EndpointsInformer, metricClient autoscaler.MetricClient) func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
	return func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
		if v, ok := decider.Labels[serving.ConfigurationLabelKey]; !ok || v == "" {
			return nil, fmt.Errorf("label %q not found or empty in Decider %s", serving.ConfigurationLabelKey, decider.Name)
		}
		if decider.Spec.ServiceName == "" {
			return nil, fmt.Errorf("%s decider has empty ServiceName", decider.Name)
		}

		serviceName := decider.Labels[serving.ServiceLabelKey] // This can be empty.
		configName := decider.Labels[serving.ConfigurationLabelKey]

		// Create a stats reporter which tags statistics by PA namespace, configuration name, and PA name.
		reporter, err := autoscaler.NewStatsReporter(decider.Namespace, serviceName, configName, decider.Name)
		if err != nil {
			return nil, err
		}

		podCounter := resources.NewScopedEndpointsCounter(endpointsInformer.Lister(), decider.Namespace, decider.Name)
		return autoscaler.New(decider.Namespace, decider.Name, metricClient, podCounter, decider.Spec, reporter)
	}
}

func statsScraperFactoryFunc(endpointsLister corev1listers.EndpointsLister) func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
	return func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
		podCounter := resources.NewScopedEndpointsCounter(endpointsLister, metric.Namespace, metric.Spec.ScrapeTarget)
		return autoscaler.NewServiceScraper(metric, podCounter)
	}
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	pkgmetrics.FlushExporter()
}
<commit_msg>Fix autoscaler shutdown error reporting. (#4425)<commit_after>/*
Copyright 2018 The Knative 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.
*/

// Multitenant autoscaler executable.
package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/controller"
	"github.com/knative/pkg/injection"
	"github.com/knative/pkg/injection/clients/kubeclient"
	endpointsinformer "github.com/knative/pkg/injection/informers/kubeinformers/corev1/endpoints"
	"github.com/knative/pkg/logging"
	"github.com/knative/pkg/metrics"
	pkgmetrics "github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/pkg/system"
	"github.com/knative/serving/pkg/apis/serving"
	"github.com/knative/serving/pkg/autoscaler"
	"github.com/knative/serving/pkg/autoscaler/statserver"
	"github.com/knative/serving/pkg/reconciler/autoscaling/hpa"
	"github.com/knative/serving/pkg/reconciler/autoscaling/kpa"
	"github.com/knative/serving/pkg/resources"
	basecmd "github.com/kubernetes-incubator/custom-metrics-apiserver/pkg/cmd"
	"github.com/spf13/pflag"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	corev1informers "k8s.io/client-go/informers/core/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

const (
	statsServerAddr = ":8080"
	statsBufferLen  = 1000
	component       = "autoscaler"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func main() {
	// Initialize early to get access to flags and merge them with the autoscaler flags.
	customMetricsAdapter := &basecmd.AdapterBase{}
	customMetricsAdapter.Flags().AddGoFlagSet(flag.CommandLine)
	pflag.Parse()

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	// Set up our logger.
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	logger, atomicLevel := logging.NewLoggerFromConfig(loggingConfig, component)
	defer flush(logger)
	ctx = logging.WithLogger(ctx, logger)

	cfg, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatalw("Error building kubeconfig", zap.Error(err))
	}

	logger.Infof("Registering %d clients", len(injection.Default.GetClients()))
	logger.Infof("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	logger.Infof("Registering %d informers", len(injection.Default.GetInformers()))
	logger.Infof("Registering %d controllers", 2)

	// Adjust our client's rate limits based on the number of controller's we are running.
	cfg.QPS = 2 * rest.DefaultQPS
	cfg.Burst = 2 * rest.DefaultBurst

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// statsCh is the main communication channel between the stats channel and multiscaler.
	statsCh := make(chan *autoscaler.StatMessage, statsBufferLen)
	defer close(statsCh)

	cmw := configmap.NewInformedWatcher(kubeclient.Get(ctx), system.Namespace())
	// Watch the logging config map and dynamically update logging levels.
	cmw.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	cmw.Watch(metrics.ConfigMapName(), metrics.UpdateExporterFromConfigMap(component, logger))

	endpointsInformer := endpointsinformer.Get(ctx)

	collector := autoscaler.NewMetricCollector(statsScraperFactoryFunc(endpointsInformer.Lister()), logger)
	customMetricsAdapter.WithCustomMetrics(autoscaler.NewMetricProvider(collector))

	// Set up scalers.
	// uniScalerFactory depends endpointsInformer to be set.
	multiScaler := autoscaler.NewMultiScaler(ctx.Done(), uniScalerFactoryFunc(endpointsInformer, collector), logger)

	psInformerFactory := resources.NewPodScalableInformerFactory(ctx)
	controllers := []*controller.Impl{
		kpa.NewController(ctx, cmw, multiScaler, collector, psInformerFactory),
		hpa.NewController(ctx, cmw, collector, psInformerFactory),
	}

	// Set up a statserver.
	statsServer := statserver.New(statsServerAddr, statsCh, logger)

	// Start watching the configs.
	if err := cmw.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start watching configs", zap.Error(err))
	}

	// Start all of the informers and wait for them to sync.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", err)
	}

	go controller.StartAll(ctx.Done(), controllers...)

	go func() {
		for {
			sm, ok := <-statsCh
			if !ok {
				break
			}
			collector.Record(sm.Key, sm.Stat)
			multiScaler.Poke(sm.Key, sm.Stat)
		}
	}()

	eg, egCtx := errgroup.WithContext(ctx)
	eg.Go(func() error {
		return customMetricsAdapter.Run(ctx.Done())
	})
	eg.Go(statsServer.ListenAndServe)

	// This will block until either a signal arrives or one of the grouped functions
	// returns an error.
	<-egCtx.Done()

	statsServer.Shutdown(5 * time.Second)
	if err := eg.Wait(); err != nil {
		logger.Errorw("Error while shutting down", zap.Error(err))
	}
}

func uniScalerFactoryFunc(endpointsInformer corev1informers.EndpointsInformer, metricClient autoscaler.MetricClient) func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
	return func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
		if v, ok := decider.Labels[serving.ConfigurationLabelKey]; !ok || v == "" {
			return nil, fmt.Errorf("label %q not found or empty in Decider %s", serving.ConfigurationLabelKey, decider.Name)
		}
		if decider.Spec.ServiceName == "" {
			return nil, fmt.Errorf("%s decider has empty ServiceName", decider.Name)
		}

		serviceName := decider.Labels[serving.ServiceLabelKey] // This can be empty.
		configName := decider.Labels[serving.ConfigurationLabelKey]

		// Create a stats reporter which tags statistics by PA namespace, configuration name, and PA name.
		reporter, err := autoscaler.NewStatsReporter(decider.Namespace, serviceName, configName, decider.Name)
		if err != nil {
			return nil, err
		}

		podCounter := resources.NewScopedEndpointsCounter(endpointsInformer.Lister(), decider.Namespace, decider.Name)
		return autoscaler.New(decider.Namespace, decider.Name, metricClient, podCounter, decider.Spec, reporter)
	}
}

func statsScraperFactoryFunc(endpointsLister corev1listers.EndpointsLister) func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
	return func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
		podCounter := resources.NewScopedEndpointsCounter(endpointsLister, metric.Namespace, metric.Spec.ScrapeTarget)
		return autoscaler.NewServiceScraper(metric, podCounter)
	}
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	pkgmetrics.FlushExporter()
}
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/alecthomas/kingpin"
	"github.com/couchbaselabs/logg"
	"github.com/spf13/cobra"
	"github.com/tleyden/todolite-appserver/libtodolite"
)

var (
	url        *string
	openOcrUrl *string
	since      *string
)

func init() {
	logg.LogKeys["CLI"] = true
	logg.LogKeys["TODOLITE"] = true
}

// changesfollowerCmd respresents the changesfollower command
var changesfollowerCmd = &cobra.Command{
	Use:   "changesfollower",
	Short: "Follows the todolite changes feed and performs actions",
	Long:  `This has the ability to run images through OCR.  In the future, it will be able to send push notifications, etc.`,
	Run: func(cmd *cobra.Command, args []string) {

		// kingpin.Parse()
		if *url == "" {
			kingpin.Errorf("URL is empty")
			return
		}
		if *openOcrUrl == "" {
			kingpin.Errorf("OpenOcr URL is empty")
			return
		}

		logg.LogTo("CLI", "url: %v openOcrUrl: %v", *url, *openOcrUrl)
		todoliteApp := libtodolite.NewTodoLiteApp(*url, *openOcrUrl)
		err := todoliteApp.InitApp()
		if err != nil {
			logg.LogPanic("Error initializing todo lite app: %v", err)
		}
		go todoliteApp.FollowChangesFeed(*since)
		select {}

	},
}

func init() {
	RootCmd.AddCommand(changesfollowerCmd)

	url = changesfollowerCmd.PersistentFlags().String("url", "", "Sync Gateway URL")

	openOcrUrl = changesfollowerCmd.PersistentFlags().String("openOcrUrl", "", "OpenOCR Server URL")

	since = changesfollowerCmd.PersistentFlags().String("since", "", "Since value to start changes feed at (last sequence)")

}
<commit_msg>remove dependency on kingpin<commit_after>package cmd

import (
	"fmt"

	"github.com/couchbaselabs/logg"
	"github.com/spf13/cobra"
	"github.com/tleyden/todolite-appserver/libtodolite"
)

var (
	url        *string
	openOcrUrl *string
	since      *string
)

func init() {
	logg.LogKeys["CLI"] = true
	logg.LogKeys["TODOLITE"] = true
}

// changesfollowerCmd respresents the changesfollower command
var changesfollowerCmd = &cobra.Command{
	Use:   "changesfollower",
	Short: "Follows the todolite changes feed and performs actions",
	Long:  `This has the ability to run images through OCR.  In the future, it will be able to send push notifications, etc.`,
	Run: func(cmd *cobra.Command, args []string) {

		if *url == "" {
			logg.LogError(fmt.Errorf("URL is empty"))
			return
		}
		if *openOcrUrl == "" {
			logg.LogError(fmt.Errorf("OpenOcr URL is empty"))
			return
		}

		logg.LogTo("CLI", "url: %v openOcrUrl: %v", *url, *openOcrUrl)
		todoliteApp := libtodolite.NewTodoLiteApp(*url, *openOcrUrl)
		err := todoliteApp.InitApp()
		if err != nil {
			logg.LogPanic("Error initializing todo lite app: %v", err)
		}
		go todoliteApp.FollowChangesFeed(*since)
		select {}

	},
}

func init() {
	RootCmd.AddCommand(changesfollowerCmd)

	url = changesfollowerCmd.PersistentFlags().String("url", "", "Sync Gateway URL")

	openOcrUrl = changesfollowerCmd.PersistentFlags().String("openOcrUrl", "", "OpenOCR Server URL")

	since = changesfollowerCmd.PersistentFlags().String("since", "", "Since value to start changes feed at (last sequence)")

}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative 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 main

import (
	// The set of controllers this controller process runs.
	"knative.dev/net-contour/pkg/reconciler/contour"

	// This defines the shared main for injected controllers.
	"knative.dev/pkg/injection/sharedmain"
)

func main() {
	sharedmain.Main("controller",
		contour.NewController,
	)
}
<commit_msg>Rename the component to avoid collisions with leaderelection. (#115)<commit_after>/*
Copyright 2019 The Knative 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 main

import (
	// The set of controllers this controller process runs.
	"knative.dev/net-contour/pkg/reconciler/contour"

	// This defines the shared main for injected controllers.
	"knative.dev/pkg/injection/sharedmain"
)

func main() {
	sharedmain.Main("contour-ingress-controller",
		contour.NewController,
	)
}
<|endoftext|>


<commit_before>package cmd

import (
	"log"
	"os"
	"path/filepath"

	"github.com/drud/ddev/pkg/ddevapp"
	"github.com/spf13/cobra"
)

// ConfigCommand represents the `ddev config` command
var ConfigCommand = &cobra.Command{
	Use:   "config",
	Short: "Create or modify a ddev application config in the current directory",
	Run: func(cmd *cobra.Command, args []string) {
		appRoot, err := filepath.Abs(filepath.Dir(os.Args[0]))
		if err != nil {
			log.Fatalf("Could not determine current working directory: %v\n", err)
		}

		c, err := ddevapp.NewConfig(appRoot)
		if err != nil {
			// If there is an error reading the config and the file exists, we're not sure
			// how to proceed.
			if c.ConfigExists() {
				log.Fatalf("Could not read config: %v", err)
			}
		}

		err = c.Config()
		if err != nil {
			log.Fatalf("There was a problem configuring your application: %v\n", err)
		}

		err = c.Write()
		if err != nil {
			log.Fatalf("Could not write ddev config file: %v\n", err)
		}
	},
}

func init() {
	RootCmd.AddCommand(ConfigCommand)
}
<commit_msg>Fixes #149 by using os.Getcwd() (#152)<commit_after>package cmd

import (
	"log"
	"os"

	"github.com/drud/ddev/pkg/ddevapp"
	"github.com/spf13/cobra"
)

// ConfigCommand represents the `ddev config` command
var ConfigCommand = &cobra.Command{
	Use:   "config",
	Short: "Create or modify a ddev application config in the current directory",
	Run: func(cmd *cobra.Command, args []string) {
		appRoot, err := os.Getwd()
		if err != nil {
			log.Fatalf("Could not determine current working directory: %v\n", err)
		}

		c, err := ddevapp.NewConfig(appRoot)
		if err != nil {
			// If there is an error reading the config and the file exists, we're not sure
			// how to proceed.
			if c.ConfigExists() {
				log.Fatalf("Could not read config: %v", err)
			}
		}

		err = c.Config()
		if err != nil {
			log.Fatalf("There was a problem configuring your application: %v\n", err)
		}

		err = c.Write()
		if err != nil {
			log.Fatalf("Could not write ddev config file: %v\n", err)
		}
	},
}

func init() {
	RootCmd.AddCommand(ConfigCommand)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/belogik/goes"
	flags "github.com/jessevdk/go-flags"
	"github.com/miku/dupsquash"
)

type Work struct {
	Indices    []string
	Connection dupsquash.SearchConnection

	Fields        []string
	LikeText      string
	MinTermFreq   int
	MaxQueryTerms int
	Size          int

	Values []string
}

func Worker(in chan *Work, out chan [][]string, wg *sync.WaitGroup) {
	defer wg.Done()
	for work := range in {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          work.Fields,
					"like_text":       work.LikeText,
					"min_term_freq":   work.MinTermFreq,
					"max_query_terms": work.MaxQueryTerms,
				},
			},
			"size": work.Size,
		}
		queryResults := Query(work.Connection, &work.Indices, &query)
		var results [][]string
		for _, result := range queryResults {
			parts := append(work.Values, result...)
			results = append(results, parts)
		}
		out <- results
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan [][]string, done chan bool) {
	for results := range in {
		for _, result := range results {
			writer.Write([]byte(strings.Join(result, "\t")))
			writer.Write([]byte("\n"))
		}
	}
	done <- true
}

// Query runs `query` over connection `conn` on `indices` and returns a slice of string slices
func Query(conn dupsquash.SearchConnection, indices *[]string, query *map[string]interface{}) [][]string {
	extraArgs := make(url.Values, 1)
	searchResults, err := conn.Search(*query, *indices, []string{""}, extraArgs)
	if err != nil {
		log.Fatalln(err)
	}
	results := make([][]string, len(searchResults.Hits.Hits))
	for i, hit := range searchResults.Hits.Hits {
		results[i] = []string{hit.Index, hit.Type, hit.Id, strconv.FormatFloat(hit.Score, 'f', 3, 64)}
	}
	return results
}

func main() {
	var opts struct {
		ElasticSearchHost string `long:"host" default:"localhost" description:"elasticsearch host" value-name:"HOST"`
		ElasticSearchPort string `long:"port" default:"9200" description:"elasticsearch port" value-name:"PORT"`

		Like string `short:"l" long:"like" description:"string to compare" value-name:"STRING"`

		LikeFile      string `short:"i" long:"file" description:"input file (TSV) with strings to compare" value-name:"FILENAME"`
		FileColumn    string `short:"f" long:"column" default:"1" description:"which column(s) to pick for the comparison" value-name:"COLUMN[S]"`
		FileDelimiter string `long:"delimiter" default:"\t" description:"column delimiter of the file" value-name:"DELIM"`
		FileNullValue string `long:"null-value" default:"<NULL>" description:"value that indicates empty value in input file" value-name:"STRING"`

		Index         string `long:"indices" description:"index or indices (space separated)" value-name:"NAME[S]"`
		IndexFields   string `short:"x" default:"content.245.a content.245.b" long:"index-fields"description:"which index fields to use for comparison" value-name:"NAME[S]"`
		MinTermFreq   int    `long:"min-term-freq" description:"passed on lucene option" default:"1" value-name:"N"`
		MaxQueryTerms int    `long:"max-query-terms" description:"passed on lucene option" default:"25" value-name:"N"`
		Size          int    `short:"s" long:"size" description:"number of results per query" default:"5" value-name:"N"`

		NumWorkers  int    `short:"w" long:"workers" default:"0" description:"number of workers, 0 means number of available cpus" value-name:"N"`
		ShowVersion bool   `short:"V" default:"false" long:"version" description:"show version and exit"`
		ShowHelp    bool   `short:"h" default:"false" long:"help" description:"show this help message"`
		CpuProfile  string `long:"cpuprofile" description:"write pprof file" value-name:"FILE"`
	}

	argparser := flags.NewParser(&opts, flags.PrintErrors|flags.PassDoubleDash)
	_, err := argparser.Parse()
	if err != nil {
		log.Fatal(err)
	}

	if opts.ShowVersion {
		fmt.Printf("%s\n", dupsquash.AppVersion)
		return
	}

	argparser.Usage = fmt.Sprintf("[OPTIONS]")
	if opts.ShowHelp {
		argparser.WriteHelp(os.Stdout)
		return
	}

	if opts.CpuProfile != "" {
		f, err := os.Create(opts.CpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if opts.NumWorkers == 0 {
		opts.NumWorkers = runtime.NumCPU()
		runtime.GOMAXPROCS(opts.NumWorkers)
	}

	if opts.NumWorkers < 0 {
		log.Fatal("value for --workers must be non-negative")
	}

	conn := goes.NewConnection(opts.ElasticSearchHost, opts.ElasticSearchPort)
	fields := strings.Fields(opts.IndexFields)
	indices := strings.Fields(opts.Index)

	if opts.LikeFile != "" {
		if _, err := os.Stat(opts.LikeFile); os.IsNotExist(err) {
			log.Fatalf("no such file or directory: %s\n", opts.LikeFile)
		}

		file, err := os.Open(opts.LikeFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		projector, err := dupsquash.ParseIndices(opts.FileColumn)
		if err != nil {
			log.Fatalf("could not parse column indices: %s\n", opts.FileColumn)
		}

		queue := make(chan *Work)
		results := make(chan [][]string)
		done := make(chan bool)

		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		go FanInWriter(writer, results, done)

		var wg sync.WaitGroup
		for i := 0; i < opts.NumWorkers; i++ {
			wg.Add(1)
			go Worker(queue, results, &wg)
		}

		for scanner.Scan() {
			values := strings.Split(scanner.Text(), opts.FileDelimiter)
			likeText, err := dupsquash.ConcatenateValuesNull(values, projector, opts.FileNullValue)
			if err != nil {
				log.Fatal(err)
			}

			work := Work{
				Indices:       indices,
				Connection:    conn,
				Fields:        fields,
				LikeText:      likeText,
				MinTermFreq:   opts.MinTermFreq,
				MaxQueryTerms: opts.MaxQueryTerms,
				Size:          opts.Size,
				Values:        values,
			}
			queue <- &work
		}

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		close(queue)
		wg.Wait()
		close(results)
		select {
		case <-time.After(1e9):
			break
		case <-done:
			break
		}
		return
	}

	if opts.Like != "" {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          fields,
					"like_text":       opts.Like,
					"min_term_freq":   opts.MinTermFreq,
					"max_query_terms": opts.MaxQueryTerms,
				},
			},
			"size": opts.Size,
		}

		results := Query(conn, &indices, &query)
		for _, result := range results {
			fmt.Println(strings.Join(result, "\t"))
		}
		return
	}
}
<commit_msg>switch to pkg/flag<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/belogik/goes"
	"github.com/miku/dupsquash"
)

type Work struct {
	Indices    []string
	Connection dupsquash.SearchConnection

	Fields        []string
	LikeText      string
	MinTermFreq   int
	MaxQueryTerms int
	Size          int

	Values []string
}

func Worker(in chan *Work, out chan [][]string, wg *sync.WaitGroup) {
	defer wg.Done()
	for work := range in {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          work.Fields,
					"like_text":       work.LikeText,
					"min_term_freq":   work.MinTermFreq,
					"max_query_terms": work.MaxQueryTerms,
				},
			},
			"size": work.Size,
		}
		queryResults := Query(work.Connection, &work.Indices, &query)
		var results [][]string
		for _, result := range queryResults {
			parts := append(work.Values, result...)
			results = append(results, parts)
		}
		out <- results
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan [][]string, done chan bool) {
	for results := range in {
		for _, result := range results {
			writer.Write([]byte(strings.Join(result, "\t")))
			writer.Write([]byte("\n"))
		}
	}
	done <- true
}

// Query runs `query` over connection `conn` on `indices` and returns a slice of string slices
func Query(conn dupsquash.SearchConnection, indices *[]string, query *map[string]interface{}) [][]string {
	extraArgs := make(url.Values, 1)
	searchResults, err := conn.Search(*query, *indices, []string{""}, extraArgs)
	if err != nil {
		log.Fatalln(err)
	}
	results := make([][]string, len(searchResults.Hits.Hits))
	for i, hit := range searchResults.Hits.Hits {
		results[i] = []string{hit.Index, hit.Type, hit.Id, strconv.FormatFloat(hit.Score, 'f', 3, 64)}
	}
	return results
}

func main() {

	esHost := flag.String("host", "localhost", "elasticsearch host")
	esPort := flag.String("port", "9200", "elasticsearch port")
	likeText := flag.String("like", "", "more like this queries like-text")
	likeFile := flag.String("file", "", "input file")
	fileColumn := flag.String("columns", "1", "which column to use as like-text")
	columnDelimiter := flag.String("delimiter", "\t", "column delimiter of the input file")
	columnNull := flag.String("null", "<NULL>", "column value to ignore")
	indicesString := flag.String("indices", "", "index or indices to query")
	indexFields := flag.String("fields", "content.245.a content.245.b", "index fields to query")
	minTermFreq := flag.Int("min-term-freq", 1, "min term frequency")
	maxQueryTerms := flag.Int("max-query-terms", 25, "max query terms")
	size := flag.Int("size", 5, "maximum number of similar records to report")
	numWorkers := flag.Int("workers", runtime.NumCPU(), "number of workers to use")
	version := flag.Bool("v", false, "prints current program version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")

	var PrintUsage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()

	if *version {
		fmt.Printf("%s\n", dupsquash.AppVersion)
		return
	}

	if *likeText == "" && *likeFile == "" {
		PrintUsage()
		os.Exit(1)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	runtime.GOMAXPROCS(*numWorkers)

	conn := goes.NewConnection(*esHost, *esPort)
	fields := strings.Fields(*indexFields)
	indices := strings.Fields(*indicesString)

	if *likeFile != "" {
		if _, err := os.Stat(*likeFile); os.IsNotExist(err) {
			log.Fatalf("no such file or directory: %s\n", *likeFile)
		}

		file, err := os.Open(*likeFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		projector, err := dupsquash.ParseIndices(*fileColumn)
		if err != nil {
			log.Fatalf("could not parse column indices: %s\n", *fileColumn)
		}

		queue := make(chan *Work)
		results := make(chan [][]string)
		done := make(chan bool)

		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		go FanInWriter(writer, results, done)

		var wg sync.WaitGroup
		for i := 0; i < *numWorkers; i++ {
			wg.Add(1)
			go Worker(queue, results, &wg)
		}

		for scanner.Scan() {
			values := strings.Split(scanner.Text(), *columnDelimiter)
			likeText, err := dupsquash.ConcatenateValuesNull(values, projector, *columnNull)
			if err != nil {
				log.Fatal(err)
			}

			work := Work{
				Indices:       indices,
				Connection:    conn,
				Fields:        fields,
				LikeText:      likeText,
				MinTermFreq:   *minTermFreq,
				MaxQueryTerms: *maxQueryTerms,
				Size:          *size,
				Values:        values,
			}
			queue <- &work
		}

		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}

		close(queue)
		wg.Wait()
		close(results)
		select {
		case <-time.After(1e9):
			break
		case <-done:
			break
		}
		return
	}

	if *likeText != "" {
		var query = map[string]interface{}{
			"query": map[string]interface{}{
				"more_like_this": map[string]interface{}{
					"fields":          fields,
					"like_text":       *likeText,
					"min_term_freq":   *minTermFreq,
					"max_query_terms": *maxQueryTerms,
				},
			},
			"size": *size,
		}

		results := Query(conn, &indices, &query)
		for _, result := range results {
			fmt.Println(strings.Join(result, "\t"))
		}
		return
	}
}
<|endoftext|>