<commit_before>// Copyright 2016 Mender Software AS
//
//    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 (
	"bufio"
	"os"
	"strings"
	"time"

	"github.com/mendersoftware/log"
)

type Controller interface {
	Bootstrap() error
	TransitionState() MenderState
	GetCurrentImageID() string
	GetUpdatePollInterval() time.Duration
	LastError() error
	HasUpgrade() (bool, error)

	UInstallCommitRebooter
	Updater
}

const (
	defaultManifestFile = "/etc/build_mender"
	defaultKeyFile      = "mender-agent.pem"
	defaultDataStore    = "/var/lib/mender"
)

type MenderState int

// State transitions:
//
//            unknown
//               |
//               v
//             init
//               |
//               v
//          bootstrapped
//               |
//       +-------+-------------+
//       |                     |
//       v                     v
// fresh update         wait for update
//
// Any state can transition to MenderStateError, setting LastError()
// to the error that triggered the transition

const (
	MenderStateUnknown MenderState = iota
	// initial state
	MenderStateInit
	// client is bootstrapped, i.e. ready to go
	MenderStateBootstrapped
	// update applied, waiting for commit
	MenderStateRunningWithFreshUpdate
	// wait for new update
	MenderStateWaitForUpdate
	// error occurred, call Controller.LastError() to obtain the
	// error
	MenderStateError
)

type mender struct {
	UInstallCommitRebooter
	Updater
	env            BootEnvReadWriter
	state          MenderState
	config         menderConfig
	manifestFile   string
	deviceKey      *Keystore
	forceBootstrap bool
	lastError      error
}

type MenderPieces struct {
	updater Updater
	device  UInstallCommitRebooter
	env     BootEnvReadWriter
	store   Store
}

func NewMender(config menderConfig, pieces MenderPieces) *mender {

	ks := NewKeystore(pieces.store)
	if ks == nil {
		return nil
	}

	m := &mender{
		UInstallCommitRebooter: pieces.device,
		Updater:                pieces.updater,
		env:                    pieces.env,
		manifestFile:           defaultManifestFile,
		deviceKey:              ks,
		config:                 config,
	}

	if err := m.deviceKey.Load(m.config.DeviceKey); err != nil && IsNoKeys(err) == false {
		log.Errorf("failed to load device keys: %s", err)
		return nil
	}

	return m
}

func (m *mender) TransitionState() MenderState {
	m.updateState()
	return m.state
}

func (m mender) GetCurrentImageID() string {
	// This is where Yocto stores buid information
	manifest, err := os.Open(m.manifestFile)
	if err != nil {
		log.Error("Can not read current image id.")
		return ""
	}

	imageID := ""

	scanner := bufio.NewScanner(manifest)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		log.Debug("Read data from device manifest file: ", line)
		if strings.HasPrefix(line, "IMAGE_ID") {
			log.Debug("Found device id line: ", line)
			lineID := strings.Split(line, "=")
			if len(lineID) != 2 {
				log.Errorf("Broken device manifest file: (%v)", lineID)
				return ""
			}
			log.Debug("Current image id: ", strings.TrimSpace(lineID[1]))
			return strings.TrimSpace(lineID[1])
		}
	}
	if err := scanner.Err(); err != nil {
		log.Error(err)
	}
	return imageID
}

func (m *mender) changeState(state MenderState) {
	log.Infof("Mender state: %v -> %v", m.state, state)
	m.state = state
}

func (m *mender) HasUpgrade() (bool, error) {
	env, err := m.env.ReadEnv("upgrade_available")
	if err != nil {
		return false, err
	}
	upgradeAvailable := env["upgrade_available"]

	// we are after update
	if upgradeAvailable == "1" {
		return true, nil
	}
	return false, nil
}

func (m *mender) updateState() {

	newstate := MenderStateUnknown
	var merr error

	switch m.state {
	case MenderStateInit:
		if m.needsBootstrap() {
			if err := m.doBootstrap(); err != nil {
				newstate = MenderStateError
				merr = err
			} else {
				newstate = MenderStateBootstrapped
			}
		} else {
			newstate = MenderStateBootstrapped
		}
	case MenderStateBootstrapped:
		upg, err := m.HasUpgrade()
		if err != nil {
			newstate = MenderStateError
			merr = err
		} else {
			if upg {
				newstate = MenderStateRunningWithFreshUpdate
			} else {
				newstate = MenderStateWaitForUpdate
			}
		}
	}

	// record last errpr
	if newstate == MenderStateError {
		m.lastError = merr
	}
	m.changeState(newstate)
}

func (m *mender) ForceBootstrap() {
	m.forceBootstrap = true
}

func (m *mender) needsBootstrap() bool {
	if m.forceBootstrap {
		return true
	}

	if m.deviceKey.Private() == nil {
		log.Debugf("needs keys")
		return true
	}

	return false
}

func (m *mender) Bootstrap() error {
	if !m.needsBootstrap() {
		return nil
	}

	return m.doBootstrap()
}

func (m *mender) doBootstrap() error {
	if m.deviceKey.Private() == nil {
		log.Infof("device keys not present, generating")
		if err := m.deviceKey.Generate(); err != nil {
			return err
		}

		if err := m.deviceKey.Save(m.config.DeviceKey); err != nil {
			log.Errorf("faiiled to save keys to %s: %s",
				m.config.DeviceKey, err)
			return err
		}
	}

	m.forceBootstrap = false

	return nil
}

func (m *mender) LastError() error {
	return m.lastError
}

func (m mender) GetUpdatePollInterval() time.Duration {
	return time.Duration(m.config.PollIntervalSeconds) * time.Second
}
<commit_msg>mender: extern Controler interface with CheckUpdate() helper<commit_after>// Copyright 2016 Mender Software AS
//
//    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 (
	"bufio"
	"os"
	"strings"
	"time"

	"github.com/mendersoftware/log"
	"github.com/pkg/errors"
)

type Controller interface {
	Bootstrap() error
	TransitionState() MenderState
	GetCurrentImageID() string
	GetUpdatePollInterval() time.Duration
	LastError() error
	HasUpgrade() (bool, error)
	CheckUpdate() (*UpdateResponse, error)

	UInstallCommitRebooter
	Updater
}

const (
	defaultManifestFile = "/etc/build_mender"
	defaultKeyFile      = "mender-agent.pem"
	defaultDataStore    = "/var/lib/mender"
)

type MenderState int

// State transitions:
//
//            unknown
//               |
//               v
//             init
//               |
//               v
//          bootstrapped
//               |
//       +-------+-------------+
//       |                     |
//       v                     v
// fresh update         wait for update
//
// Any state can transition to MenderStateError, setting LastError()
// to the error that triggered the transition

const (
	MenderStateUnknown MenderState = iota
	// initial state
	MenderStateInit
	// client is bootstrapped, i.e. ready to go
	MenderStateBootstrapped
	// update applied, waiting for commit
	MenderStateRunningWithFreshUpdate
	// wait for new update
	MenderStateWaitForUpdate
	// error occurred, call Controller.LastError() to obtain the
	// error
	MenderStateError
)

type mender struct {
	UInstallCommitRebooter
	Updater
	env            BootEnvReadWriter
	state          MenderState
	config         menderConfig
	manifestFile   string
	deviceKey      *Keystore
	forceBootstrap bool
	lastError      error
}

type MenderPieces struct {
	updater Updater
	device  UInstallCommitRebooter
	env     BootEnvReadWriter
	store   Store
}

func NewMender(config menderConfig, pieces MenderPieces) *mender {

	ks := NewKeystore(pieces.store)
	if ks == nil {
		return nil
	}

	m := &mender{
		UInstallCommitRebooter: pieces.device,
		Updater:                pieces.updater,
		env:                    pieces.env,
		manifestFile:           defaultManifestFile,
		deviceKey:              ks,
		config:                 config,
	}

	if err := m.deviceKey.Load(m.config.DeviceKey); err != nil && IsNoKeys(err) == false {
		log.Errorf("failed to load device keys: %s", err)
		return nil
	}

	return m
}

func (m *mender) TransitionState() MenderState {
	m.updateState()
	return m.state
}

func (m mender) GetCurrentImageID() string {
	// This is where Yocto stores buid information
	manifest, err := os.Open(m.manifestFile)
	if err != nil {
		log.Error("Can not read current image id.")
		return ""
	}

	imageID := ""

	scanner := bufio.NewScanner(manifest)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		log.Debug("Read data from device manifest file: ", line)
		if strings.HasPrefix(line, "IMAGE_ID") {
			log.Debug("Found device id line: ", line)
			lineID := strings.Split(line, "=")
			if len(lineID) != 2 {
				log.Errorf("Broken device manifest file: (%v)", lineID)
				return ""
			}
			log.Debug("Current image id: ", strings.TrimSpace(lineID[1]))
			return strings.TrimSpace(lineID[1])
		}
	}
	if err := scanner.Err(); err != nil {
		log.Error(err)
	}
	return imageID
}

func (m *mender) changeState(state MenderState) {
	log.Infof("Mender state: %v -> %v", m.state, state)
	m.state = state
}

func (m *mender) HasUpgrade() (bool, error) {
	env, err := m.env.ReadEnv("upgrade_available")
	if err != nil {
		return false, err
	}
	upgradeAvailable := env["upgrade_available"]

	// we are after update
	if upgradeAvailable == "1" {
		return true, nil
	}
	return false, nil
}

func (m *mender) updateState() {

	newstate := MenderStateUnknown
	var merr error

	switch m.state {
	case MenderStateInit:
		if m.needsBootstrap() {
			if err := m.doBootstrap(); err != nil {
				newstate = MenderStateError
				merr = err
			} else {
				newstate = MenderStateBootstrapped
			}
		} else {
			newstate = MenderStateBootstrapped
		}
	case MenderStateBootstrapped:
		upg, err := m.HasUpgrade()
		if err != nil {
			newstate = MenderStateError
			merr = err
		} else {
			if upg {
				newstate = MenderStateRunningWithFreshUpdate
			} else {
				newstate = MenderStateWaitForUpdate
			}
		}
	}

	// record last errpr
	if newstate == MenderStateError {
		m.lastError = merr
	}
	m.changeState(newstate)
}

func (m *mender) ForceBootstrap() {
	m.forceBootstrap = true
}

func (m *mender) needsBootstrap() bool {
	if m.forceBootstrap {
		return true
	}

	if m.deviceKey.Private() == nil {
		log.Debugf("needs keys")
		return true
	}

	return false
}

func (m *mender) Bootstrap() error {
	if !m.needsBootstrap() {
		return nil
	}

	return m.doBootstrap()
}

func (m *mender) doBootstrap() error {
	if m.deviceKey.Private() == nil {
		log.Infof("device keys not present, generating")
		if err := m.deviceKey.Generate(); err != nil {
			return err
		}

		if err := m.deviceKey.Save(m.config.DeviceKey); err != nil {
			log.Errorf("faiiled to save keys to %s: %s",
				m.config.DeviceKey, err)
			return err
		}
	}

	m.forceBootstrap = false

	return nil
}

func (m *mender) LastError() error {
	return m.lastError
}

// Check if new update is available. In case of errors, returns nil and error
// that occurred. If no update is available *UpdateResponse is nil, otherwise it
// contains update information.
func (m *mender) CheckUpdate() (*UpdateResponse, error) {
	currentImageID := m.GetCurrentImageID()
	//TODO: if currentImageID == "" {
	// 	return errors.New("")
	// }

	haveUpdate, err := m.Updater.GetScheduledUpdate(m.config.ServerURL, m.config.DeviceID)
	if err != nil {
		log.Error("Error receiving scheduled update data: ", err)
		return nil, err
	}

	log.Debug("Received correct response for update request.")

	update, ok := haveUpdate.(UpdateResponse)
	if !ok {
		return nil, errors.Errorf("not an update response?")
	}

	if update.Image.YoctoID == currentImageID {
		return nil, nil
	}
	return &update, nil
}

func (m mender) GetUpdatePollInterval() time.Duration {
	return time.Duration(m.config.PollIntervalSeconds) * time.Second
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

// Package cache implements dns cache feature with edns-client-subnet support.
package cache

// Cache that holds RRs.

import (
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
	"strconv"
)

// Elem hold an answer and additional section that returned from the cache.
// The signature is put in answer, extra is empty there. This wastes some memory.
type elem struct {
	expiration time.Time // time added + TTL, after this the elem is invalid
	m          *dns.Msg
}

// Cache is a cache that holds on the a number of RRs or DNS messages. The cache
// eviction is randomized.
type Cache struct {
	sync.RWMutex

	capacity int
	table    map[string]*elem
}

// New returns a new cache with the capacity and the ttl specified.
func New(capacity int) *Cache {
	if capacity <= 0 {
		return nil
	}
	c := new(Cache)
	c.table = make(map[string]*elem)
	c.capacity = capacity
	return c
}

func (c *Cache) Capacity() int { return c.capacity }

func (c *Cache) Remove(s string) {
	c.Lock()
	delete(c.table, s)
	c.Unlock()
}

// EvictRandom removes a random member a the cache.
// Must be called under a write lock.
func (c *Cache) EvictRandom() {
	clen := len(c.table)
	if clen < c.capacity {
		return
	}
	i := c.capacity - clen
	for k := range c.table {
		delete(c.table, k)
		i--
		if i == 0 {
			break
		}
	}
}

// InsertMessage inserts a message in the Cache. We will cache it for ttl seconds, which
// should be a small (60...300) integer.
func (c *Cache) InsertMessage(s string, m *dns.Msg) {
	if c.capacity <= 0 || len(m.Answer) == 0 {
		return
	}

	c.Lock()
	ttl := time.Duration(m.Answer[0].Header().Ttl) * time.Second
	if _, ok := c.table[s]; !ok {
		c.table[s] = &elem{time.Now().UTC().Add(ttl), m.Copy()}
	}
	log.Debug("Cached: " + s)
	c.EvictRandom()
	c.Unlock()
}

// Search returns a dns.Msg, the expiration time and a boolean indicating if we found something
// in the cache.
func (c *Cache) Search(s string) (*dns.Msg, time.Time, bool) {
	if c.capacity <= 0 {
		return nil, time.Time{}, false
	}
	c.RLock()
	if e, ok := c.table[s]; ok {
		e1 := e.m.Copy()
		c.RUnlock()
		return e1, e.expiration, true
	}
	c.RUnlock()
	return nil, time.Time{}, false
}

// Key creates a hash key from a question section.
func Key(q dns.Question, ednsIP string) string {
	return q.Name + " " + strconv.Itoa(int(q.Qtype)) + " " + ednsIP
}

// Hit returns a dns message from the cache. If the message's TTL is expired nil
// is returned and the message is removed from the cache.
func (c *Cache) Hit(key string, msgid uint16) *dns.Msg {
	m, exp, hit := c.Search(key)
	if hit {
		// Cache hit! \o/
		if time.Since(exp) < 0 {
			m.Id = msgid
			m.Compress = true
			// Even if something ended up with the TC bit *in* the cache, set it to off
			m.Truncated = false
			return m
		}
		// Expired! /o\
		c.Remove(key)
	}
	return nil
}
<commit_msg>fix null deference when dns.Msg is nil (#47)<commit_after>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

// Package cache implements dns cache feature with edns-client-subnet support.
package cache

// Cache that holds RRs.

import (
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
	"strconv"
)

// Elem hold an answer and additional section that returned from the cache.
// The signature is put in answer, extra is empty there. This wastes some memory.
type elem struct {
	expiration time.Time // time added + TTL, after this the elem is invalid
	m          *dns.Msg
}

// Cache is a cache that holds on the a number of RRs or DNS messages. The cache
// eviction is randomized.
type Cache struct {
	sync.RWMutex

	capacity int
	table    map[string]*elem
}

// New returns a new cache with the capacity and the ttl specified.
func New(capacity int) *Cache {
	if capacity <= 0 {
		return nil
	}
	c := new(Cache)
	c.table = make(map[string]*elem)
	c.capacity = capacity
	return c
}

func (c *Cache) Capacity() int { return c.capacity }

func (c *Cache) Remove(s string) {
	c.Lock()
	delete(c.table, s)
	c.Unlock()
}

// EvictRandom removes a random member a the cache.
// Must be called under a write lock.
func (c *Cache) EvictRandom() {
	clen := len(c.table)
	if clen < c.capacity {
		return
	}
	i := c.capacity - clen
	for k := range c.table {
		delete(c.table, k)
		i--
		if i == 0 {
			break
		}
	}
}

// InsertMessage inserts a message in the Cache. We will cache it for ttl seconds, which
// should be a small (60...300) integer.
func (c *Cache) InsertMessage(s string, m *dns.Msg) {
	if c.capacity <= 0 || m == nil || len(m.Answer) == 0 {
		return
	}

	c.Lock()
	ttl := time.Duration(m.Answer[0].Header().Ttl) * time.Second
	if _, ok := c.table[s]; !ok {
		c.table[s] = &elem{time.Now().UTC().Add(ttl), m.Copy()}
	}
	log.Debug("Cached: " + s)
	c.EvictRandom()
	c.Unlock()
}

// Search returns a dns.Msg, the expiration time and a boolean indicating if we found something
// in the cache.
func (c *Cache) Search(s string) (*dns.Msg, time.Time, bool) {
	if c.capacity <= 0 {
		return nil, time.Time{}, false
	}
	c.RLock()
	if e, ok := c.table[s]; ok {
		e1 := e.m.Copy()
		c.RUnlock()
		return e1, e.expiration, true
	}
	c.RUnlock()
	return nil, time.Time{}, false
}

// Key creates a hash key from a question section.
func Key(q dns.Question, ednsIP string) string {
	return q.Name + " " + strconv.Itoa(int(q.Qtype)) + " " + ednsIP
}

// Hit returns a dns message from the cache. If the message's TTL is expired nil
// is returned and the message is removed from the cache.
func (c *Cache) Hit(key string, msgid uint16) *dns.Msg {
	m, exp, hit := c.Search(key)
	if hit {
		// Cache hit! \o/
		if time.Since(exp) < 0 {
			m.Id = msgid
			m.Compress = true
			// Even if something ended up with the TC bit *in* the cache, set it to off
			m.Truncated = false
			return m
		}
		// Expired! /o\
		c.Remove(key)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ebml decodes EBML data.
//
// EBML is short for Extensible Binary Meta Language. EBML specifies a
// binary and octet (byte) aligned format inspired by the principle of
// XML. EBML itself is a generalized description of the technique of
// binary markup. Like XML, it is completely agnostic to any data that it
// might contain.
// For a specification, see http://ebml.sourceforge.net/specs/
package ebml

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"reflect"
	"strconv"
)

var Verbose bool = false

// ReachedPayloadError is generated when a field tagged with
// ebmlstop:"1" is reached.
type ReachedPayloadError struct {
	*Element
}

func (r ReachedPayloadError) Error() string {
	return "Reached payload"
}

// Element represents an EBML-encoded chunk of data.
type Element struct {
	*limitedReadSeeker
	Offset int64
	Id     uint
}

func (e *Element) String() string {
	return fmt.Sprintf("{ReadSeeker: %+v Offset: %+v Id: %x}", e.limitedReadSeeker, e.Offset, e.Id)
}

// Creates the root element corresponding to the data available in r.
func RootElement(rs io.ReadSeeker) (*Element, error) {
	e := &Element{newLimitedReadSeeker(rs, math.MaxInt64), 0, 0}
	return e, nil
}

func remaining(x int8) (rem int) {
	for x > 0 {
		rem++
		x += x
	}
	return
}

func parseVint(data []uint8) (val uint64) {
	for i, l := 0, len(data); i < l; i++ {
		val <<= 8
		val += uint64(data[i])
	}
	return
}

func readVintData(r io.Reader) (data []uint8, err error, rem int) {
	var v [16]uint8
	_, err = io.ReadFull(r, v[:1])
	if err == nil {
		rem = remaining(int8(v[0]))
		_, err = io.ReadFull(r, v[1:rem+1])
	}
	if err == nil {
		data = v[0 : rem+1]
	}
	return
}

func readVint(r io.Reader) (val uint64, err error, rem int) {
	var data []uint8
	data, err, rem = readVintData(r)
	if err == nil {
		val = parseVint(data)
	}
	return
}

func readSize(r io.Reader) (int64, error) {
	val, err, rem := readVint(r)
	return int64(val & ^(128 << uint(rem*8-rem))), err
}

// Next returns the next child element in an element.
func (e *Element) Next() (*Element, error) {
	if Verbose {
		log.Println("next", e)
	}
	off, err := e.Seek(0, 1)
	if err != nil {
		log.Panic(err)
	}
	id, err, _ := readVint(e)
	if err != nil {
		return nil, err
	}
	sz, err := readSize(e)
	if err != nil {
		return nil, err
	}
	ret := &Element{newLimitedReadSeeker(e, sz), off, uint(id)}
	if Verbose {
		log.Println("--->", ret)
	}
	return ret, err
}

func (e *Element) Size() int64 {
	return e.limitedReadSeeker.N
}

func (e *Element) readUint64() (uint64, error) {
	d, err := e.ReadData()
	var i int
	sz := len(d)
	var val uint64
	for i = 0; i < sz; i++ {
		val <<= 8
		val += uint64(d[i])
	}
	return val, err
}

func (e *Element) readUint() (uint, error) {
	val, err := e.readUint64()
	return uint(val), err
}

func (e *Element) readString() (string, error) {
	s, err := e.ReadData()
	return string(s), err
}

func (e *Element) ReadData() (d []byte, err error) {
	sz := e.Size()
	d = make([]uint8, sz, sz)
	_, err = io.ReadFull(e, d)
	return
}

func (e *Element) readFloat() (val float64, err error) {
	var uval uint64
	sz := e.Size()
	uval, err = e.readUint64()
	if sz == 8 {
		val = math.Float64frombits(uval)
	} else {
		val = float64(math.Float32frombits(uint32(uval)))
	}
	return
}

func (e *Element) skip() (err error) {
	_, err = e.Seek(e.Size(), 1)
	return
}

// Unmarshal reads EBML data from r into data. Data must be a pointer
// to a struct. Fields present in the struct but absent in the stream
// will just keep their zero value.
// Returns an error that can be an io.Error or a ReachedPayloadError.
func (e *Element) Unmarshal(val interface{}) error {
	return e.readStruct(reflect.Indirect(reflect.ValueOf(val)))
}

func getTag(f reflect.StructField, s string) uint {
	sid := f.Tag.Get(s)
	id, _ := strconv.ParseUint(sid, 16, 0)
	return uint(id)
}

func lookup(reqid uint, t reflect.Type) int {
	for i, l := 0, t.NumField(); i < l; i++ {
		f := t.Field(i)
		if getTag(f, "ebml") == reqid {
			return i - 1000000*int(getTag(f, "ebmlstop"))
		}
	}
	return -1
}

func setDefaults(v reflect.Value) {
	t := v.Type()
	for i, l := 0, t.NumField(); i < l; i++ {
		fv := v.Field(i)
		switch fv.Kind() {
		case reflect.Int, reflect.Uint,
			reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
			reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
			reflect.Float32, reflect.Float64,
			reflect.String:
			setFieldDefaults(fv, t.Field(i), v)
		case reflect.Array, reflect.Struct, reflect.Slice:
			break
		default:
			log.Panic("Unsupported type")
		}
	}
}

func setFieldDefaults(v reflect.Value, sf reflect.StructField, s reflect.Value) {
	if v.CanInterface() && reflect.DeepEqual(
		v.Interface(), reflect.Zero(v.Type()).Interface()) {
		tag := sf.Tag.Get("ebmldef")
		if tag != "" {
			switch v.Kind() {
			case reflect.Int, reflect.Int64:
				u, _ := strconv.ParseInt(tag, 10, 0)
				v.SetInt(int64(u))
			case reflect.Uint, reflect.Uint64:
				u, _ := strconv.ParseUint(tag, 10, 0)
				v.SetUint(u)
			case reflect.Float32, reflect.Float64:
				f, _ := strconv.ParseFloat(tag, 64)
				v.SetFloat(f)
			case reflect.String:
				v.SetString(tag)
			default:
				log.Panic("Unsupported default value")
			}
		}
		ltag := sf.Tag.Get("ebmldeflink")
		if ltag != "" {
			v.Set(s.FieldByName(ltag))
		}
	}
}

func (e *Element) readStruct(v reflect.Value) (err error) {
	t := v.Type()
	for err == nil {
		var ne *Element
		ne, err = e.Next()
		if err == io.EOF {
			err = nil
			break
		}
		i := lookup(ne.Id, t)
		if i >= 0 {
			err = ne.readField(v.Field(i))
		} else if i == -1 {
			err = ne.skip()
		} else {
			var curr int64
			curr, err = e.Seek(ne.Offset, 0)
			if err != nil || curr != ne.Offset {
				log.Panic(err, curr)
			}
			err = ReachedPayloadError{e}
		}
	}
	setDefaults(v)
	return
}

func (e *Element) readField(v reflect.Value) (err error) {
	switch v.Kind() {
	case reflect.Struct:
		err = e.readStruct(v)
	case reflect.Slice:
		err = e.readSlice(v)
	case reflect.Array:
		for i, l := 0, v.Len(); i < l && err == nil; i++ {
			err = e.readStruct(v.Index(i))
		}
	case reflect.String:
		var s string
		s, err = e.readString()
		v.SetString(s)
	case reflect.Int, reflect.Int64:
		var u uint64
		u, err = e.readUint64()
		v.SetInt(int64(u))
	case reflect.Uint, reflect.Uint64:
		var u uint64
		u, err = e.readUint64()
		v.SetUint(u)
	case reflect.Float32, reflect.Float64:
		var f float64
		f, err = e.readFloat()
		v.SetFloat(f)
	default:
		err = errors.New("Unknown type: " + v.String())
	}
	return
}

func (e *Element) readSlice(v reflect.Value) (err error) {
	switch v.Type().Elem().Kind() {
	case reflect.Uint8:
		var sl []uint8
		sl, err = e.ReadData()
		if err == nil {
			if !v.CanSet() {
				log.Panic("can't set ", v, e)
			}
			v.Set(reflect.ValueOf(sl))
		}
	case reflect.Struct:
		vl := v.Len()
		ne := reflect.New(v.Type().Elem())
		nsl := reflect.Append(v, reflect.Indirect(ne))
		v.Set(nsl)
		err = e.readStruct(v.Index(vl))
	default:
		err = errors.New("Unknown slice type: " + v.String())
	}
	return
}
<commit_msg>EBML spec allows zero-padded strings<commit_after>// Copyright (c) 2012, Jorge Acereda Maciá. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

// Package ebml decodes EBML data.
//
// EBML is short for Extensible Binary Meta Language. EBML specifies a
// binary and octet (byte) aligned format inspired by the principle of
// XML. EBML itself is a generalized description of the technique of
// binary markup. Like XML, it is completely agnostic to any data that it
// might contain.
// For a specification, see http://ebml.sourceforge.net/specs/
package ebml

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"reflect"
	"strconv"
)

var Verbose bool = false

// ReachedPayloadError is generated when a field tagged with
// ebmlstop:"1" is reached.
type ReachedPayloadError struct {
	*Element
}

func (r ReachedPayloadError) Error() string {
	return "Reached payload"
}

// Element represents an EBML-encoded chunk of data.
type Element struct {
	*limitedReadSeeker
	Offset int64
	Id     uint
}

func (e *Element) String() string {
	return fmt.Sprintf("{ReadSeeker: %+v Offset: %+v Id: %x}", e.limitedReadSeeker, e.Offset, e.Id)
}

// Creates the root element corresponding to the data available in r.
func RootElement(rs io.ReadSeeker) (*Element, error) {
	e := &Element{newLimitedReadSeeker(rs, math.MaxInt64), 0, 0}
	return e, nil
}

func remaining(x int8) (rem int) {
	for x > 0 {
		rem++
		x += x
	}
	return
}

func parseVint(data []uint8) (val uint64) {
	for i, l := 0, len(data); i < l; i++ {
		val <<= 8
		val += uint64(data[i])
	}
	return
}

func readVintData(r io.Reader) (data []uint8, err error, rem int) {
	var v [16]uint8
	_, err = io.ReadFull(r, v[:1])
	if err == nil {
		rem = remaining(int8(v[0]))
		_, err = io.ReadFull(r, v[1:rem+1])
	}
	if err == nil {
		data = v[0 : rem+1]
	}
	return
}

func readVint(r io.Reader) (val uint64, err error, rem int) {
	var data []uint8
	data, err, rem = readVintData(r)
	if err == nil {
		val = parseVint(data)
	}
	return
}

func readSize(r io.Reader) (int64, error) {
	val, err, rem := readVint(r)
	return int64(val & ^(128 << uint(rem*8-rem))), err
}

// Next returns the next child element in an element.
func (e *Element) Next() (*Element, error) {
	if Verbose {
		log.Println("next", e)
	}
	off, err := e.Seek(0, 1)
	if err != nil {
		log.Panic(err)
	}
	id, err, _ := readVint(e)
	if err != nil {
		return nil, err
	}
	sz, err := readSize(e)
	if err != nil {
		return nil, err
	}
	ret := &Element{newLimitedReadSeeker(e, sz), off, uint(id)}
	if Verbose {
		log.Println("--->", ret)
	}
	return ret, err
}

func (e *Element) Size() int64 {
	return e.limitedReadSeeker.N
}

func (e *Element) readUint64() (uint64, error) {
	d, err := e.ReadData()
	var i int
	sz := len(d)
	var val uint64
	for i = 0; i < sz; i++ {
		val <<= 8
		val += uint64(d[i])
	}
	return val, err
}

func (e *Element) readUint() (uint, error) {
	val, err := e.readUint64()
	return uint(val), err
}

func (e *Element) readString() (string, error) {
	s, err := e.ReadData()
	sl := len(s)
	for sl > 0 && s[sl-1] == 0 {
		sl--
	}
	return string(s[:sl]), err
}

func (e *Element) ReadData() (d []byte, err error) {
	sz := e.Size()
	d = make([]uint8, sz, sz)
	_, err = io.ReadFull(e, d)
	return
}

func (e *Element) readFloat() (val float64, err error) {
	var uval uint64
	sz := e.Size()
	uval, err = e.readUint64()
	if sz == 8 {
		val = math.Float64frombits(uval)
	} else {
		val = float64(math.Float32frombits(uint32(uval)))
	}
	return
}

func (e *Element) skip() (err error) {
	_, err = e.Seek(e.Size(), 1)
	return
}

// Unmarshal reads EBML data from r into data. Data must be a pointer
// to a struct. Fields present in the struct but absent in the stream
// will just keep their zero value.
// Returns an error that can be an io.Error or a ReachedPayloadError.
func (e *Element) Unmarshal(val interface{}) error {
	return e.readStruct(reflect.Indirect(reflect.ValueOf(val)))
}

func getTag(f reflect.StructField, s string) uint {
	sid := f.Tag.Get(s)
	id, _ := strconv.ParseUint(sid, 16, 0)
	return uint(id)
}

func lookup(reqid uint, t reflect.Type) int {
	for i, l := 0, t.NumField(); i < l; i++ {
		f := t.Field(i)
		if getTag(f, "ebml") == reqid {
			return i - 1000000*int(getTag(f, "ebmlstop"))
		}
	}
	return -1
}

func setDefaults(v reflect.Value) {
	t := v.Type()
	for i, l := 0, t.NumField(); i < l; i++ {
		fv := v.Field(i)
		switch fv.Kind() {
		case reflect.Int, reflect.Uint,
			reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
			reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
			reflect.Float32, reflect.Float64,
			reflect.String:
			setFieldDefaults(fv, t.Field(i), v)
		case reflect.Array, reflect.Struct, reflect.Slice:
			break
		default:
			log.Panic("Unsupported type")
		}
	}
}

func setFieldDefaults(v reflect.Value, sf reflect.StructField, s reflect.Value) {
	if v.CanInterface() && reflect.DeepEqual(
		v.Interface(), reflect.Zero(v.Type()).Interface()) {
		tag := sf.Tag.Get("ebmldef")
		if tag != "" {
			switch v.Kind() {
			case reflect.Int, reflect.Int64:
				u, _ := strconv.ParseInt(tag, 10, 0)
				v.SetInt(int64(u))
			case reflect.Uint, reflect.Uint64:
				u, _ := strconv.ParseUint(tag, 10, 0)
				v.SetUint(u)
			case reflect.Float32, reflect.Float64:
				f, _ := strconv.ParseFloat(tag, 64)
				v.SetFloat(f)
			case reflect.String:
				v.SetString(tag)
			default:
				log.Panic("Unsupported default value")
			}
		}
		ltag := sf.Tag.Get("ebmldeflink")
		if ltag != "" {
			v.Set(s.FieldByName(ltag))
		}
	}
}

func (e *Element) readStruct(v reflect.Value) (err error) {
	t := v.Type()
	for err == nil {
		var ne *Element
		ne, err = e.Next()
		if err == io.EOF {
			err = nil
			break
		}
		i := lookup(ne.Id, t)
		if i >= 0 {
			err = ne.readField(v.Field(i))
		} else if i == -1 {
			err = ne.skip()
		} else {
			var curr int64
			curr, err = e.Seek(ne.Offset, 0)
			if err != nil || curr != ne.Offset {
				log.Panic(err, curr)
			}
			err = ReachedPayloadError{e}
		}
	}
	setDefaults(v)
	return
}

func (e *Element) readField(v reflect.Value) (err error) {
	switch v.Kind() {
	case reflect.Struct:
		err = e.readStruct(v)
	case reflect.Slice:
		err = e.readSlice(v)
	case reflect.Array:
		for i, l := 0, v.Len(); i < l && err == nil; i++ {
			err = e.readStruct(v.Index(i))
		}
	case reflect.String:
		var s string
		s, err = e.readString()
		v.SetString(s)
	case reflect.Int, reflect.Int64:
		var u uint64
		u, err = e.readUint64()
		v.SetInt(int64(u))
	case reflect.Uint, reflect.Uint64:
		var u uint64
		u, err = e.readUint64()
		v.SetUint(u)
	case reflect.Float32, reflect.Float64:
		var f float64
		f, err = e.readFloat()
		v.SetFloat(f)
	default:
		err = errors.New("Unknown type: " + v.String())
	}
	return
}

func (e *Element) readSlice(v reflect.Value) (err error) {
	switch v.Type().Elem().Kind() {
	case reflect.Uint8:
		var sl []uint8
		sl, err = e.ReadData()
		if err == nil {
			if !v.CanSet() {
				log.Panic("can't set ", v, e)
			}
			v.Set(reflect.ValueOf(sl))
		}
	case reflect.Struct:
		vl := v.Len()
		ne := reflect.New(v.Type().Elem())
		nsl := reflect.Append(v, reflect.Indirect(ne))
		v.Set(nsl)
		err = e.readStruct(v.Index(vl))
	default:
		err = errors.New("Unknown slice type: " + v.String())
	}
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package priorities

import (
	"context"
	"sync"
	"sync/atomic"

	"k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/kubernetes/pkg/scheduler/algorithm"
	"k8s.io/kubernetes/pkg/scheduler/algorithm/predicates"
	priorityutil "k8s.io/kubernetes/pkg/scheduler/algorithm/priorities/util"
	schedulerapi "k8s.io/kubernetes/pkg/scheduler/api"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"

	"k8s.io/klog"
)

// InterPodAffinity contains information to calculate inter pod affinity.
type InterPodAffinity struct {
	info                  predicates.NodeInfo
	nodeLister            algorithm.NodeLister
	podLister             algorithm.PodLister
	hardPodAffinityWeight int32
}

// NewInterPodAffinityPriority creates an InterPodAffinity.
func NewInterPodAffinityPriority(
	info predicates.NodeInfo,
	nodeLister algorithm.NodeLister,
	podLister algorithm.PodLister,
	hardPodAffinityWeight int32) PriorityFunction {
	interPodAffinity := &InterPodAffinity{
		info:                  info,
		nodeLister:            nodeLister,
		podLister:             podLister,
		hardPodAffinityWeight: hardPodAffinityWeight,
	}
	return interPodAffinity.CalculateInterPodAffinityPriority
}

type podAffinityPriorityMap struct {
	sync.Mutex

	// nodes contain all nodes that should be considered
	nodes []*v1.Node
	// counts store the mapping from node name to so-far computed score of
	// the node.
	counts map[string]*int64
	// The first error that we faced.
	firstError error
}

func newPodAffinityPriorityMap(nodes []*v1.Node) *podAffinityPriorityMap {
	return &podAffinityPriorityMap{
		nodes:  nodes,
		counts: make(map[string]*int64, len(nodes)),
	}
}

func (p *podAffinityPriorityMap) setError(err error) {
	p.Lock()
	defer p.Unlock()
	if p.firstError == nil {
		p.firstError = err
	}
}

func (p *podAffinityPriorityMap) processTerm(term *v1.PodAffinityTerm, podDefiningAffinityTerm, podToCheck *v1.Pod, fixedNode *v1.Node, weight int64) {
	namespaces := priorityutil.GetNamespacesFromPodAffinityTerm(podDefiningAffinityTerm, term)
	selector, err := metav1.LabelSelectorAsSelector(term.LabelSelector)
	if err != nil {
		p.setError(err)
		return
	}
	match := priorityutil.PodMatchesTermsNamespaceAndSelector(podToCheck, namespaces, selector)
	if match {
		for _, node := range p.nodes {
			if priorityutil.NodesHaveSameTopologyKey(node, fixedNode, term.TopologyKey) {
				atomic.AddInt64(p.counts[node.Name], weight)
			}
		}
	}
}

func (p *podAffinityPriorityMap) processTerms(terms []v1.WeightedPodAffinityTerm, podDefiningAffinityTerm, podToCheck *v1.Pod, fixedNode *v1.Node, multiplier int) {
	for i := range terms {
		term := &terms[i]
		p.processTerm(&term.PodAffinityTerm, podDefiningAffinityTerm, podToCheck, fixedNode, int64(term.Weight*int32(multiplier)))
	}
}

// CalculateInterPodAffinityPriority compute a sum by iterating through the elements of weightedPodAffinityTerm and adding
// "weight" to the sum if the corresponding PodAffinityTerm is satisfied for
// that node; the node(s) with the highest sum are the most preferred.
// Symmetry need to be considered for preferredDuringSchedulingIgnoredDuringExecution from podAffinity & podAntiAffinity,
// symmetry need to be considered for hard requirements from podAffinity
func (ipa *InterPodAffinity) CalculateInterPodAffinityPriority(pod *v1.Pod, nodeNameToInfo map[string]*schedulernodeinfo.NodeInfo, nodes []*v1.Node) (schedulerapi.HostPriorityList, error) {
	affinity := pod.Spec.Affinity
	hasAffinityConstraints := affinity != nil && affinity.PodAffinity != nil
	hasAntiAffinityConstraints := affinity != nil && affinity.PodAntiAffinity != nil

	// priorityMap stores the mapping from node name to so-far computed score of
	// the node.
	pm := newPodAffinityPriorityMap(nodes)
	allNodeNames := make([]string, 0, len(nodeNameToInfo))
	for name := range nodeNameToInfo {
		allNodeNames = append(allNodeNames, name)
		pm.counts[name] = new(int64)
	}

	// convert the topology key based weights to the node name based weights
	var maxCount, minCount int64

	processPod := func(existingPod *v1.Pod) error {
		existingPodNode, err := ipa.info.GetNodeInfo(existingPod.Spec.NodeName)
		if err != nil {
			if apierrors.IsNotFound(err) {
				klog.Errorf("Node not found, %v", existingPod.Spec.NodeName)
				return nil
			}
			return err
		}
		existingPodAffinity := existingPod.Spec.Affinity
		existingHasAffinityConstraints := existingPodAffinity != nil && existingPodAffinity.PodAffinity != nil
		existingHasAntiAffinityConstraints := existingPodAffinity != nil && existingPodAffinity.PodAntiAffinity != nil

		if hasAffinityConstraints {
			// For every soft pod affinity term of <pod>, if <existingPod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPods>`s node by the term`s weight.
			terms := affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, pod, existingPod, existingPodNode, 1)
		}
		if hasAntiAffinityConstraints {
			// For every soft pod anti-affinity term of <pod>, if <existingPod> matches the term,
			// decrement <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>`s node by the term`s weight.
			terms := affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, pod, existingPod, existingPodNode, -1)
		}

		if existingHasAffinityConstraints {
			// For every hard pod affinity term of <existingPod>, if <pod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the constant <ipa.hardPodAffinityWeight>
			if ipa.hardPodAffinityWeight > 0 {
				terms := existingPodAffinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution
				// TODO: Uncomment this block when implement RequiredDuringSchedulingRequiredDuringExecution.
				//if len(existingPodAffinity.PodAffinity.RequiredDuringSchedulingRequiredDuringExecution) != 0 {
				//	terms = append(terms, existingPodAffinity.PodAffinity.RequiredDuringSchedulingRequiredDuringExecution...)
				//}
				for _, term := range terms {
					pm.processTerm(&term, existingPod, pod, existingPodNode, int64(ipa.hardPodAffinityWeight))
				}
			}
			// For every soft pod affinity term of <existingPod>, if <pod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the term's weight.
			terms := existingPodAffinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, existingPod, pod, existingPodNode, 1)
		}
		if existingHasAntiAffinityConstraints {
			// For every soft pod anti-affinity term of <existingPod>, if <pod> matches the term,
			// decrement <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the term's weight.
			terms := existingPodAffinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, existingPod, pod, existingPodNode, -1)
		}
		return nil
	}
	processNode := func(i int) {
		nodeInfo := nodeNameToInfo[allNodeNames[i]]
		if nodeInfo.Node() != nil {
			if hasAffinityConstraints || hasAntiAffinityConstraints {
				// We need to process all the nodes.
				for _, existingPod := range nodeInfo.Pods() {
					if err := processPod(existingPod); err != nil {
						pm.setError(err)
					}
				}
			} else {
				// The pod doesn't have any constraints - we need to check only existing
				// ones that have some.
				for _, existingPod := range nodeInfo.PodsWithAffinity() {
					if err := processPod(existingPod); err != nil {
						pm.setError(err)
					}
				}
			}
		}
	}
	workqueue.ParallelizeUntil(context.TODO(), 16, len(allNodeNames), processNode)
	if pm.firstError != nil {
		return nil, pm.firstError
	}

	for _, node := range nodes {
		if *pm.counts[node.Name] > maxCount {
			maxCount = *pm.counts[node.Name]
		}
		if *pm.counts[node.Name] < minCount {
			minCount = *pm.counts[node.Name]
		}
	}

	// calculate final priority score for each node
	result := make(schedulerapi.HostPriorityList, 0, len(nodes))
	for _, node := range nodes {
		fScore := float64(0)
		if (maxCount - minCount) > 0 {
			fScore = float64(schedulerapi.MaxPriority) * (float64(*pm.counts[node.Name]-minCount) / float64(maxCount-minCount))
		}
		result = append(result, schedulerapi.HostPriority{Host: node.Name, Score: int(fScore)})
		if klog.V(10) {
			klog.Infof("%v -> %v: InterPodAffinityPriority, Score: (%d)", pod.Name, node.Name, int(fScore))
		}
	}
	return result, nil
}
<commit_msg>Revert "scheduler: performance improvement on PodAffinity"<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package priorities

import (
	"context"
	"sync"

	"k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/kubernetes/pkg/scheduler/algorithm"
	"k8s.io/kubernetes/pkg/scheduler/algorithm/predicates"
	priorityutil "k8s.io/kubernetes/pkg/scheduler/algorithm/priorities/util"
	schedulerapi "k8s.io/kubernetes/pkg/scheduler/api"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"

	"k8s.io/klog"
)

// InterPodAffinity contains information to calculate inter pod affinity.
type InterPodAffinity struct {
	info                  predicates.NodeInfo
	nodeLister            algorithm.NodeLister
	podLister             algorithm.PodLister
	hardPodAffinityWeight int32
}

// NewInterPodAffinityPriority creates an InterPodAffinity.
func NewInterPodAffinityPriority(
	info predicates.NodeInfo,
	nodeLister algorithm.NodeLister,
	podLister algorithm.PodLister,
	hardPodAffinityWeight int32) PriorityFunction {
	interPodAffinity := &InterPodAffinity{
		info:                  info,
		nodeLister:            nodeLister,
		podLister:             podLister,
		hardPodAffinityWeight: hardPodAffinityWeight,
	}
	return interPodAffinity.CalculateInterPodAffinityPriority
}

type podAffinityPriorityMap struct {
	sync.Mutex

	// nodes contain all nodes that should be considered
	nodes []*v1.Node
	// counts store the mapping from node name to so-far computed score of
	// the node.
	counts map[string]float64
	// The first error that we faced.
	firstError error
}

func newPodAffinityPriorityMap(nodes []*v1.Node) *podAffinityPriorityMap {
	return &podAffinityPriorityMap{
		nodes:  nodes,
		counts: make(map[string]float64, len(nodes)),
	}
}

func (p *podAffinityPriorityMap) setError(err error) {
	p.Lock()
	defer p.Unlock()
	if p.firstError == nil {
		p.firstError = err
	}
}

func (p *podAffinityPriorityMap) processTerm(term *v1.PodAffinityTerm, podDefiningAffinityTerm, podToCheck *v1.Pod, fixedNode *v1.Node, weight float64) {
	namespaces := priorityutil.GetNamespacesFromPodAffinityTerm(podDefiningAffinityTerm, term)
	selector, err := metav1.LabelSelectorAsSelector(term.LabelSelector)
	if err != nil {
		p.setError(err)
		return
	}
	match := priorityutil.PodMatchesTermsNamespaceAndSelector(podToCheck, namespaces, selector)
	if match {
		func() {
			p.Lock()
			defer p.Unlock()
			for _, node := range p.nodes {
				if priorityutil.NodesHaveSameTopologyKey(node, fixedNode, term.TopologyKey) {
					p.counts[node.Name] += weight
				}
			}
		}()
	}
}

func (p *podAffinityPriorityMap) processTerms(terms []v1.WeightedPodAffinityTerm, podDefiningAffinityTerm, podToCheck *v1.Pod, fixedNode *v1.Node, multiplier int) {
	for i := range terms {
		term := &terms[i]
		p.processTerm(&term.PodAffinityTerm, podDefiningAffinityTerm, podToCheck, fixedNode, float64(term.Weight*int32(multiplier)))
	}
}

// CalculateInterPodAffinityPriority compute a sum by iterating through the elements of weightedPodAffinityTerm and adding
// "weight" to the sum if the corresponding PodAffinityTerm is satisfied for
// that node; the node(s) with the highest sum are the most preferred.
// Symmetry need to be considered for preferredDuringSchedulingIgnoredDuringExecution from podAffinity & podAntiAffinity,
// symmetry need to be considered for hard requirements from podAffinity
func (ipa *InterPodAffinity) CalculateInterPodAffinityPriority(pod *v1.Pod, nodeNameToInfo map[string]*schedulernodeinfo.NodeInfo, nodes []*v1.Node) (schedulerapi.HostPriorityList, error) {
	affinity := pod.Spec.Affinity
	hasAffinityConstraints := affinity != nil && affinity.PodAffinity != nil
	hasAntiAffinityConstraints := affinity != nil && affinity.PodAntiAffinity != nil

	allNodeNames := make([]string, 0, len(nodeNameToInfo))
	for name := range nodeNameToInfo {
		allNodeNames = append(allNodeNames, name)
	}

	// convert the topology key based weights to the node name based weights
	var maxCount float64
	var minCount float64
	// priorityMap stores the mapping from node name to so-far computed score of
	// the node.
	pm := newPodAffinityPriorityMap(nodes)

	processPod := func(existingPod *v1.Pod) error {
		existingPodNode, err := ipa.info.GetNodeInfo(existingPod.Spec.NodeName)
		if err != nil {
			if apierrors.IsNotFound(err) {
				klog.Errorf("Node not found, %v", existingPod.Spec.NodeName)
				return nil
			}
			return err
		}
		existingPodAffinity := existingPod.Spec.Affinity
		existingHasAffinityConstraints := existingPodAffinity != nil && existingPodAffinity.PodAffinity != nil
		existingHasAntiAffinityConstraints := existingPodAffinity != nil && existingPodAffinity.PodAntiAffinity != nil

		if hasAffinityConstraints {
			// For every soft pod affinity term of <pod>, if <existingPod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPods>`s node by the term`s weight.
			terms := affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, pod, existingPod, existingPodNode, 1)
		}
		if hasAntiAffinityConstraints {
			// For every soft pod anti-affinity term of <pod>, if <existingPod> matches the term,
			// decrement <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>`s node by the term`s weight.
			terms := affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, pod, existingPod, existingPodNode, -1)
		}

		if existingHasAffinityConstraints {
			// For every hard pod affinity term of <existingPod>, if <pod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the constant <ipa.hardPodAffinityWeight>
			if ipa.hardPodAffinityWeight > 0 {
				terms := existingPodAffinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution
				// TODO: Uncomment this block when implement RequiredDuringSchedulingRequiredDuringExecution.
				//if len(existingPodAffinity.PodAffinity.RequiredDuringSchedulingRequiredDuringExecution) != 0 {
				//	terms = append(terms, existingPodAffinity.PodAffinity.RequiredDuringSchedulingRequiredDuringExecution...)
				//}
				for _, term := range terms {
					pm.processTerm(&term, existingPod, pod, existingPodNode, float64(ipa.hardPodAffinityWeight))
				}
			}
			// For every soft pod affinity term of <existingPod>, if <pod> matches the term,
			// increment <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the term's weight.
			terms := existingPodAffinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, existingPod, pod, existingPodNode, 1)
		}
		if existingHasAntiAffinityConstraints {
			// For every soft pod anti-affinity term of <existingPod>, if <pod> matches the term,
			// decrement <pm.counts> for every node in the cluster with the same <term.TopologyKey>
			// value as that of <existingPod>'s node by the term's weight.
			terms := existingPodAffinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution
			pm.processTerms(terms, existingPod, pod, existingPodNode, -1)
		}
		return nil
	}
	processNode := func(i int) {
		nodeInfo := nodeNameToInfo[allNodeNames[i]]
		if nodeInfo.Node() != nil {
			if hasAffinityConstraints || hasAntiAffinityConstraints {
				// We need to process all the nodes.
				for _, existingPod := range nodeInfo.Pods() {
					if err := processPod(existingPod); err != nil {
						pm.setError(err)
					}
				}
			} else {
				// The pod doesn't have any constraints - we need to check only existing
				// ones that have some.
				for _, existingPod := range nodeInfo.PodsWithAffinity() {
					if err := processPod(existingPod); err != nil {
						pm.setError(err)
					}
				}
			}
		}
	}
	workqueue.ParallelizeUntil(context.TODO(), 16, len(allNodeNames), processNode)
	if pm.firstError != nil {
		return nil, pm.firstError
	}

	for _, node := range nodes {
		if pm.counts[node.Name] > maxCount {
			maxCount = pm.counts[node.Name]
		}
		if pm.counts[node.Name] < minCount {
			minCount = pm.counts[node.Name]
		}
	}

	// calculate final priority score for each node
	result := make(schedulerapi.HostPriorityList, 0, len(nodes))
	for _, node := range nodes {
		fScore := float64(0)
		if (maxCount - minCount) > 0 {
			fScore = float64(schedulerapi.MaxPriority) * ((pm.counts[node.Name] - minCount) / (maxCount - minCount))
		}
		result = append(result, schedulerapi.HostPriority{Host: node.Name, Score: int(fScore)})
		if klog.V(10) {
			klog.Infof("%v -> %v: InterPodAffinityPriority, Score: (%d)", pod.Name, node.Name, int(fScore))
		}
	}
	return result, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package nodelabel

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/apis/config/validation"
	"k8s.io/kubernetes/pkg/scheduler/framework"
)

// Name of this plugin.
const Name = "NodeLabel"

const (
	// ErrReasonPresenceViolated is used for CheckNodeLabelPresence predicate error.
	ErrReasonPresenceViolated = "node(s) didn't have the requested labels"
)

// New initializes a new plugin and returns it.
func New(plArgs runtime.Object, handle framework.Handle) (framework.Plugin, error) {
	args, err := getArgs(plArgs)
	if err != nil {
		return nil, err
	}
	if err := validation.ValidateNodeLabelArgs(nil, &args); err != nil {
		return nil, err
	}

	return &NodeLabel{
		handle: handle,
		args:   args,
	}, nil
}

func getArgs(obj runtime.Object) (config.NodeLabelArgs, error) {
	ptr, ok := obj.(*config.NodeLabelArgs)
	if !ok {
		return config.NodeLabelArgs{}, fmt.Errorf("want args to be of type NodeLabelArgs, got %T", obj)
	}
	return *ptr, nil
}

// NodeLabel checks whether a pod can fit based on the node labels which match a filter that it requests.
type NodeLabel struct {
	handle framework.Handle
	args   config.NodeLabelArgs
}

var _ framework.FilterPlugin = &NodeLabel{}
var _ framework.ScorePlugin = &NodeLabel{}
var _ framework.EnqueueExtensions = &NodeLabel{}

// Name returns name of the plugin. It is used in logs, etc.
func (pl *NodeLabel) Name() string {
	return Name
}

// Filter invoked at the filter extension point.
// It checks whether all of the specified labels exists on a node or not, regardless of their value
//
// Consider the cases where the nodes are placed in regions/zones/racks and these are identified by labels
// In some cases, it is required that only nodes that are part of ANY of the defined regions/zones/racks be selected
//
// Alternately, eliminating nodes that have a certain label, regardless of value, is also useful
// A node may have a label with "retiring" as key and the date as the value
// and it may be desirable to avoid scheduling new pods on this node.
func (pl *NodeLabel) Filter(ctx context.Context, _ *framework.CycleState, pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {
	node := nodeInfo.Node()
	if node == nil {
		return framework.NewStatus(framework.Error, "node not found")
	}

	size := int64(len(pl.args.PresentLabels) + len(pl.args.AbsentLabels))
	if size == 0 {
		return nil
	}

	nodeLabels := labels.Set(node.Labels)
	check := func(labels []string, presence bool) bool {
		for _, label := range labels {
			exists := nodeLabels.Has(label)
			if exists != presence {
				return false
			}
		}
		return true
	}
	if check(pl.args.PresentLabels, true) && check(pl.args.AbsentLabels, false) {
		return nil
	}

	return framework.NewStatus(framework.UnschedulableAndUnresolvable, ErrReasonPresenceViolated)
}

// Score invoked at the score extension point.
func (pl *NodeLabel) Score(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) (int64, *framework.Status) {
	nodeInfo, err := pl.handle.SnapshotSharedLister().NodeInfos().Get(nodeName)
	if err != nil {
		return 0, framework.AsStatus(fmt.Errorf("getting node %q from Snapshot: %w", nodeName, err))
	}

	node := nodeInfo.Node()

	size := int64(len(pl.args.PresentLabelsPreference) + len(pl.args.AbsentLabelsPreference))
	if size == 0 {
		return 0, nil
	}

	score := int64(0)
	for _, label := range pl.args.PresentLabelsPreference {
		if labels.Set(node.Labels).Has(label) {
			score += framework.MaxNodeScore
		}
	}
	for _, label := range pl.args.AbsentLabelsPreference {
		if !labels.Set(node.Labels).Has(label) {
			score += framework.MaxNodeScore
		}
	}

	// Take average score for each label to ensure the score doesn't exceed MaxNodeScore.
	score /= size

	return score, nil
}

// ScoreExtensions of the Score plugin.
func (pl *NodeLabel) ScoreExtensions() framework.ScoreExtensions {
	return nil
}

// EventsToRegister returns the possible events that may make a Pod
// failed by this plugin schedulable.
func (pl *NodeLabel) EventsToRegister() []framework.ClusterEvent {
	return []framework.ClusterEvent{
		{Resource: framework.Node, ActionType: framework.Add | framework.UpdateNodeLabel},
	}
}
<commit_msg>use '_' in nodelabel plugin for unused parameters<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package nodelabel

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/apis/config/validation"
	"k8s.io/kubernetes/pkg/scheduler/framework"
)

// Name of this plugin.
const Name = "NodeLabel"

const (
	// ErrReasonPresenceViolated is used for CheckNodeLabelPresence predicate error.
	ErrReasonPresenceViolated = "node(s) didn't have the requested labels"
)

// New initializes a new plugin and returns it.
func New(plArgs runtime.Object, handle framework.Handle) (framework.Plugin, error) {
	args, err := getArgs(plArgs)
	if err != nil {
		return nil, err
	}
	if err := validation.ValidateNodeLabelArgs(nil, &args); err != nil {
		return nil, err
	}

	return &NodeLabel{
		handle: handle,
		args:   args,
	}, nil
}

func getArgs(obj runtime.Object) (config.NodeLabelArgs, error) {
	ptr, ok := obj.(*config.NodeLabelArgs)
	if !ok {
		return config.NodeLabelArgs{}, fmt.Errorf("want args to be of type NodeLabelArgs, got %T", obj)
	}
	return *ptr, nil
}

// NodeLabel checks whether a pod can fit based on the node labels which match a filter that it requests.
type NodeLabel struct {
	handle framework.Handle
	args   config.NodeLabelArgs
}

var _ framework.FilterPlugin = &NodeLabel{}
var _ framework.ScorePlugin = &NodeLabel{}
var _ framework.EnqueueExtensions = &NodeLabel{}

// Name returns name of the plugin. It is used in logs, etc.
func (pl *NodeLabel) Name() string {
	return Name
}

// Filter invoked at the filter extension point.
// It checks whether all of the specified labels exists on a node or not, regardless of their value
//
// Consider the cases where the nodes are placed in regions/zones/racks and these are identified by labels
// In some cases, it is required that only nodes that are part of ANY of the defined regions/zones/racks be selected
//
// Alternately, eliminating nodes that have a certain label, regardless of value, is also useful
// A node may have a label with "retiring" as key and the date as the value
// and it may be desirable to avoid scheduling new pods on this node.
func (pl *NodeLabel) Filter(ctx context.Context, _ *framework.CycleState, _ *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {
	node := nodeInfo.Node()
	if node == nil {
		return framework.NewStatus(framework.Error, "node not found")
	}

	size := int64(len(pl.args.PresentLabels) + len(pl.args.AbsentLabels))
	if size == 0 {
		return nil
	}

	nodeLabels := labels.Set(node.Labels)
	check := func(labels []string, presence bool) bool {
		for _, label := range labels {
			exists := nodeLabels.Has(label)
			if exists != presence {
				return false
			}
		}
		return true
	}
	if check(pl.args.PresentLabels, true) && check(pl.args.AbsentLabels, false) {
		return nil
	}

	return framework.NewStatus(framework.UnschedulableAndUnresolvable, ErrReasonPresenceViolated)
}

// Score invoked at the score extension point.
func (pl *NodeLabel) Score(ctx context.Context, _ *framework.CycleState, _ *v1.Pod, nodeName string) (int64, *framework.Status) {
	nodeInfo, err := pl.handle.SnapshotSharedLister().NodeInfos().Get(nodeName)
	if err != nil {
		return 0, framework.AsStatus(fmt.Errorf("getting node %q from Snapshot: %w", nodeName, err))
	}

	node := nodeInfo.Node()

	size := int64(len(pl.args.PresentLabelsPreference) + len(pl.args.AbsentLabelsPreference))
	if size == 0 {
		return 0, nil
	}

	score := int64(0)
	for _, label := range pl.args.PresentLabelsPreference {
		if labels.Set(node.Labels).Has(label) {
			score += framework.MaxNodeScore
		}
	}
	for _, label := range pl.args.AbsentLabelsPreference {
		if !labels.Set(node.Labels).Has(label) {
			score += framework.MaxNodeScore
		}
	}

	// Take average score for each label to ensure the score doesn't exceed MaxNodeScore.
	score /= size

	return score, nil
}

// ScoreExtensions of the Score plugin.
func (pl *NodeLabel) ScoreExtensions() framework.ScoreExtensions {
	return nil
}

// EventsToRegister returns the possible events that may make a Pod
// failed by this plugin schedulable.
func (pl *NodeLabel) EventsToRegister() []framework.ClusterEvent {
	return []framework.ClusterEvent{
		{Resource: framework.Node, ActionType: framework.Add | framework.UpdateNodeLabel},
	}
}
<|endoftext|>


<commit_before>package p2p

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/p2p/discover"
	"github.com/ethereum/go-ethereum/p2p/nat"
	"github.com/ethereum/go-ethereum/rlp"
)

const (
	defaultDialTimeout   = 10 * time.Second
	refreshPeersInterval = 30 * time.Second

	// total timeout for encryption handshake and protocol
	// handshake in both directions.
	handshakeTimeout = 5 * time.Second
	// maximum time allowed for reading a complete message.
	// this is effectively the amount of time a connection can be idle.
	frameReadTimeout = 1 * time.Minute
	// maximum amount of time allowed for writing a complete message.
	frameWriteTimeout = 5 * time.Second
)

var srvjslog = logger.NewJsonLogger()

// Server manages all peer connections.
//
// The fields of Server are used as configuration parameters.
// You should set them before starting the Server. Fields may not be
// modified while the server is running.
type Server struct {
	// This field must be set to a valid secp256k1 private key.
	PrivateKey *ecdsa.PrivateKey

	// MaxPeers is the maximum number of peers that can be
	// connected. It must be greater than zero.
	MaxPeers int

	// Name sets the node name of this server.
	// Use common.MakeName to create a name that follows existing conventions.
	Name string

	// Bootstrap nodes are used to establish connectivity
	// with the rest of the network.
	BootstrapNodes []*discover.Node

	// Protocols should contain the protocols supported
	// by the server. Matching protocols are launched for
	// each peer.
	Protocols []Protocol

	// If ListenAddr is set to a non-nil address, the server
	// will listen for incoming connections.
	//
	// If the port is zero, the operating system will pick a port. The
	// ListenAddr field will be updated with the actual address when
	// the server is started.
	ListenAddr string

	// If set to a non-nil value, the given NAT port mapper
	// is used to make the listening port available to the
	// Internet.
	NAT nat.Interface

	// If Dialer is set to a non-nil value, the given Dialer
	// is used to dial outbound peer connections.
	Dialer *net.Dialer

	// If NoDial is true, the server will not dial any peers.
	NoDial bool

	// Hooks for testing. These are useful because we can inhibit
	// the whole protocol stack.
	setupFunc
	newPeerHook

	ourHandshake *protoHandshake

	lock    sync.RWMutex // protects running and peers
	running bool
	peers   map[discover.NodeID]*Peer

	ntab     *discover.Table
	listener net.Listener

	quit        chan struct{}
	loopWG      sync.WaitGroup // {dial,listen,nat}Loop
	peerWG      sync.WaitGroup // active peer goroutines
	peerConnect chan *discover.Node
}

type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, bool) (*conn, error)
type newPeerHook func(*Peer)

// Peers returns all connected peers.
func (srv *Server) Peers() (peers []*Peer) {
	srv.lock.RLock()
	defer srv.lock.RUnlock()
	for _, peer := range srv.peers {
		if peer != nil {
			peers = append(peers, peer)
		}
	}
	return
}

// PeerCount returns the number of connected peers.
func (srv *Server) PeerCount() int {
	srv.lock.RLock()
	n := len(srv.peers)
	srv.lock.RUnlock()
	return n
}

// SuggestPeer creates a connection to the given Node if it
// is not already connected.
func (srv *Server) SuggestPeer(n *discover.Node) {
	srv.peerConnect <- n
}

// Broadcast sends an RLP-encoded message to all connected peers.
// This method is deprecated and will be removed later.
func (srv *Server) Broadcast(protocol string, code uint64, data interface{}) error {
	var payload []byte
	if data != nil {
		var err error
		payload, err = rlp.EncodeToBytes(data)
		if err != nil {
			return err
		}
	}
	srv.lock.RLock()
	defer srv.lock.RUnlock()
	for _, peer := range srv.peers {
		if peer != nil {
			var msg = Msg{Code: code}
			if data != nil {
				msg.Payload = bytes.NewReader(payload)
				msg.Size = uint32(len(payload))
			}
			peer.writeProtoMsg(protocol, msg)
		}
	}
	return nil
}

// Start starts running the server.
// Servers can be re-used and started again after stopping.
func (srv *Server) Start() (err error) {
	srv.lock.Lock()
	defer srv.lock.Unlock()
	if srv.running {
		return errors.New("server already running")
	}
	glog.V(logger.Info).Infoln("Starting Server")

	// static fields
	if srv.PrivateKey == nil {
		return fmt.Errorf("Server.PrivateKey must be set to a non-nil key")
	}
	if srv.MaxPeers <= 0 {
		return fmt.Errorf("Server.MaxPeers must be > 0")
	}
	srv.quit = make(chan struct{})
	srv.peers = make(map[discover.NodeID]*Peer)
	srv.peerConnect = make(chan *discover.Node)
	if srv.setupFunc == nil {
		srv.setupFunc = setupConn
	}

	// node table
	ntab, err := discover.ListenUDP(srv.PrivateKey, srv.ListenAddr, srv.NAT)
	if err != nil {
		return err
	}
	srv.ntab = ntab

	// handshake
	srv.ourHandshake = &protoHandshake{Version: baseProtocolVersion, Name: srv.Name, ID: ntab.Self().ID}
	for _, p := range srv.Protocols {
		srv.ourHandshake.Caps = append(srv.ourHandshake.Caps, p.cap())
	}

	// listen/dial
	if srv.ListenAddr != "" {
		if err := srv.startListening(); err != nil {
			return err
		}
	}
	if srv.Dialer == nil {
		srv.Dialer = &net.Dialer{Timeout: defaultDialTimeout}
	}
	if !srv.NoDial {
		srv.loopWG.Add(1)
		go srv.dialLoop()
	}
	if srv.NoDial && srv.ListenAddr == "" {
		glog.V(logger.Warn).Infoln("I will be kind-of useless, neither dialing nor listening.")
	}

	srv.running = true
	return nil
}

func (srv *Server) startListening() error {
	listener, err := net.Listen("tcp", srv.ListenAddr)
	if err != nil {
		return err
	}
	laddr := listener.Addr().(*net.TCPAddr)
	srv.ListenAddr = laddr.String()
	srv.listener = listener
	srv.loopWG.Add(1)
	go srv.listenLoop()
	if !laddr.IP.IsLoopback() && srv.NAT != nil {
		srv.loopWG.Add(1)
		go func() {
			nat.Map(srv.NAT, srv.quit, "tcp", laddr.Port, laddr.Port, "ethereum p2p")
			srv.loopWG.Done()
		}()
	}
	return nil
}

// Stop terminates the server and all active peer connections.
// It blocks until all active connections have been closed.
func (srv *Server) Stop() {
	srv.lock.Lock()
	if !srv.running {
		srv.lock.Unlock()
		return
	}
	srv.running = false
	srv.lock.Unlock()

	glog.V(logger.Info).Infoln("Stopping Server")
	srv.ntab.Close()
	if srv.listener != nil {
		// this unblocks listener Accept
		srv.listener.Close()
	}
	close(srv.quit)
	srv.loopWG.Wait()

	// No new peers can be added at this point because dialLoop and
	// listenLoop are down. It is safe to call peerWG.Wait because
	// peerWG.Add is not called outside of those loops.
	for _, peer := range srv.peers {
		peer.Disconnect(DiscQuitting)
	}
	srv.peerWG.Wait()
}

// Self returns the local node's endpoint information.
func (srv *Server) Self() *discover.Node {
	return srv.ntab.Self()
}

// main loop for adding connections via listening
func (srv *Server) listenLoop() {
	defer srv.loopWG.Done()
	glog.V(logger.Info).Infoln("Listening on", srv.listener.Addr())
	for {
		conn, err := srv.listener.Accept()
		if err != nil {
			return
		}
		glog.V(logger.Debug).Infof("Accepted conn %v\n", conn.RemoteAddr())
		srv.peerWG.Add(1)
		go srv.startPeer(conn, nil)
	}
}

func (srv *Server) dialLoop() {
	var (
		dialed      = make(chan *discover.Node)
		dialing     = make(map[discover.NodeID]bool)
		findresults = make(chan []*discover.Node)
		refresh     = time.NewTimer(0)
	)
	defer srv.loopWG.Done()
	defer refresh.Stop()

	// TODO: maybe limit number of active dials
	dial := func(dest *discover.Node) {
		// Don't dial nodes that would fail the checks in addPeer.
		// This is important because the connection handshake is a lot
		// of work and we'd rather avoid doing that work for peers
		// that can't be added.
		srv.lock.RLock()
		ok, _ := srv.checkPeer(dest.ID)
		srv.lock.RUnlock()
		if !ok || dialing[dest.ID] {
			return
		}

		dialing[dest.ID] = true
		srv.peerWG.Add(1)
		go func() {
			srv.dialNode(dest)
			dialed <- dest
		}()
	}

	srv.ntab.Bootstrap(srv.BootstrapNodes)
	for {
		select {
		case <-refresh.C:
			// Grab some nodes to connect to if we're not at capacity.
			srv.lock.RLock()
			needpeers := len(srv.peers) < srv.MaxPeers
			srv.lock.RUnlock()
			if needpeers {
				go func() {
					var target discover.NodeID
					rand.Read(target[:])
					findresults <- srv.ntab.Lookup(target)
				}()
				refresh.Stop()
			}

		case dest := <-srv.peerConnect:
			dial(dest)
		case dests := <-findresults:
			for _, dest := range dests {
				dial(dest)
			}
			refresh.Reset(refreshPeersInterval)
		case dest := <-dialed:
			delete(dialing, dest.ID)

		case <-srv.quit:
			// TODO: maybe wait for active dials
			return
		}
	}
}

func (srv *Server) dialNode(dest *discover.Node) {
	addr := &net.TCPAddr{IP: dest.IP, Port: dest.TCPPort}
	glog.V(logger.Debug).Infof("Dialing %v\n", dest)
	conn, err := srv.Dialer.Dial("tcp", addr.String())
	if err != nil {
		// dialLoop adds to the wait group counter when launching
		// dialNode, so we need to count it down again. startPeer also
		// does that when an error occurs.
		srv.peerWG.Done()
		glog.V(logger.Detail).Infof("dial error: %v", err)
		return
	}
	srv.startPeer(conn, dest)
}

func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
	// TODO: handle/store session token

	// Run setupFunc, which should create an authenticated connection
	// and run the capability exchange. Note that any early error
	// returns during that exchange need to call peerWG.Done because
	// the callers of startPeer added the peer to the wait group already.
	fd.SetDeadline(time.Now().Add(handshakeTimeout))
	srv.lock.RLock()
	atcap := len(srv.peers) == srv.MaxPeers
	srv.lock.RUnlock()
	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, atcap)
	if err != nil {
		fd.Close()
		glog.V(logger.Debug).Infof("Handshake with %v failed: %v", fd.RemoteAddr(), err)
		srv.peerWG.Done()
		return
	}
	conn.MsgReadWriter = &netWrapper{
		wrapped: conn.MsgReadWriter,
		conn:    fd, rtimeout: frameReadTimeout, wtimeout: frameWriteTimeout,
	}
	p := newPeer(fd, conn, srv.Protocols)
	if ok, reason := srv.addPeer(conn.ID, p); !ok {
		glog.V(logger.Detail).Infof("Not adding %v (%v)\n", p, reason)
		p.politeDisconnect(reason)
		srv.peerWG.Done()
		return
	}
	// The handshakes are done and it passed all checks.
	// Spawn the Peer loops.
	go srv.runPeer(p)
}

func (srv *Server) runPeer(p *Peer) {
	glog.V(logger.Debug).Infof("Added %v\n", p)
	srvjslog.LogJson(&logger.P2PConnected{
		RemoteId:            p.ID().String(),
		RemoteAddress:       p.RemoteAddr().String(),
		RemoteVersionString: p.Name(),
		NumConnections:      srv.PeerCount(),
	})
	if srv.newPeerHook != nil {
		srv.newPeerHook(p)
	}
	discreason := p.run()
	srv.removePeer(p)
	glog.V(logger.Debug).Infof("Removed %v (%v)\n", p, discreason)
	srvjslog.LogJson(&logger.P2PDisconnected{
		RemoteId:       p.ID().String(),
		NumConnections: srv.PeerCount(),
	})
}

func (srv *Server) addPeer(id discover.NodeID, p *Peer) (bool, DiscReason) {
	srv.lock.Lock()
	defer srv.lock.Unlock()
	if ok, reason := srv.checkPeer(id); !ok {
		return false, reason
	}
	srv.peers[id] = p
	return true, 0
}

func (srv *Server) checkPeer(id discover.NodeID) (bool, DiscReason) {
	switch {
	case !srv.running:
		return false, DiscQuitting
	case len(srv.peers) >= srv.MaxPeers:
		return false, DiscTooManyPeers
	case srv.peers[id] != nil:
		return false, DiscAlreadyConnected
	case id == srv.Self().ID:
		return false, DiscSelf
	default:
		return true, 0
	}
}

func (srv *Server) removePeer(p *Peer) {
	srv.lock.Lock()
	delete(srv.peers, p.ID())
	srv.lock.Unlock()
	srv.peerWG.Done()
}
<commit_msg>p2p: limit number of lingering inbound pre-handshake connections<commit_after>package p2p

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/p2p/discover"
	"github.com/ethereum/go-ethereum/p2p/nat"
	"github.com/ethereum/go-ethereum/rlp"
)

const (
	defaultDialTimeout   = 10 * time.Second
	refreshPeersInterval = 30 * time.Second

	// This is the maximum number of inbound connection
	// that are allowed to linger between 'accepted' and
	// 'added as peer'.
	maxAcceptConns = 50

	// total timeout for encryption handshake and protocol
	// handshake in both directions.
	handshakeTimeout = 5 * time.Second
	// maximum time allowed for reading a complete message.
	// this is effectively the amount of time a connection can be idle.
	frameReadTimeout = 1 * time.Minute
	// maximum amount of time allowed for writing a complete message.
	frameWriteTimeout = 5 * time.Second
)

var srvjslog = logger.NewJsonLogger()

// Server manages all peer connections.
//
// The fields of Server are used as configuration parameters.
// You should set them before starting the Server. Fields may not be
// modified while the server is running.
type Server struct {
	// This field must be set to a valid secp256k1 private key.
	PrivateKey *ecdsa.PrivateKey

	// MaxPeers is the maximum number of peers that can be
	// connected. It must be greater than zero.
	MaxPeers int

	// Name sets the node name of this server.
	// Use common.MakeName to create a name that follows existing conventions.
	Name string

	// Bootstrap nodes are used to establish connectivity
	// with the rest of the network.
	BootstrapNodes []*discover.Node

	// Protocols should contain the protocols supported
	// by the server. Matching protocols are launched for
	// each peer.
	Protocols []Protocol

	// If ListenAddr is set to a non-nil address, the server
	// will listen for incoming connections.
	//
	// If the port is zero, the operating system will pick a port. The
	// ListenAddr field will be updated with the actual address when
	// the server is started.
	ListenAddr string

	// If set to a non-nil value, the given NAT port mapper
	// is used to make the listening port available to the
	// Internet.
	NAT nat.Interface

	// If Dialer is set to a non-nil value, the given Dialer
	// is used to dial outbound peer connections.
	Dialer *net.Dialer

	// If NoDial is true, the server will not dial any peers.
	NoDial bool

	// Hooks for testing. These are useful because we can inhibit
	// the whole protocol stack.
	setupFunc
	newPeerHook

	ourHandshake *protoHandshake

	lock    sync.RWMutex // protects running and peers
	running bool
	peers   map[discover.NodeID]*Peer

	ntab     *discover.Table
	listener net.Listener

	quit        chan struct{}
	loopWG      sync.WaitGroup // {dial,listen,nat}Loop
	peerWG      sync.WaitGroup // active peer goroutines
	peerConnect chan *discover.Node
}

type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, bool) (*conn, error)
type newPeerHook func(*Peer)

// Peers returns all connected peers.
func (srv *Server) Peers() (peers []*Peer) {
	srv.lock.RLock()
	defer srv.lock.RUnlock()
	for _, peer := range srv.peers {
		if peer != nil {
			peers = append(peers, peer)
		}
	}
	return
}

// PeerCount returns the number of connected peers.
func (srv *Server) PeerCount() int {
	srv.lock.RLock()
	n := len(srv.peers)
	srv.lock.RUnlock()
	return n
}

// SuggestPeer creates a connection to the given Node if it
// is not already connected.
func (srv *Server) SuggestPeer(n *discover.Node) {
	srv.peerConnect <- n
}

// Broadcast sends an RLP-encoded message to all connected peers.
// This method is deprecated and will be removed later.
func (srv *Server) Broadcast(protocol string, code uint64, data interface{}) error {
	var payload []byte
	if data != nil {
		var err error
		payload, err = rlp.EncodeToBytes(data)
		if err != nil {
			return err
		}
	}
	srv.lock.RLock()
	defer srv.lock.RUnlock()
	for _, peer := range srv.peers {
		if peer != nil {
			var msg = Msg{Code: code}
			if data != nil {
				msg.Payload = bytes.NewReader(payload)
				msg.Size = uint32(len(payload))
			}
			peer.writeProtoMsg(protocol, msg)
		}
	}
	return nil
}

// Start starts running the server.
// Servers can be re-used and started again after stopping.
func (srv *Server) Start() (err error) {
	srv.lock.Lock()
	defer srv.lock.Unlock()
	if srv.running {
		return errors.New("server already running")
	}
	glog.V(logger.Info).Infoln("Starting Server")

	// static fields
	if srv.PrivateKey == nil {
		return fmt.Errorf("Server.PrivateKey must be set to a non-nil key")
	}
	if srv.MaxPeers <= 0 {
		return fmt.Errorf("Server.MaxPeers must be > 0")
	}
	srv.quit = make(chan struct{})
	srv.peers = make(map[discover.NodeID]*Peer)
	srv.peerConnect = make(chan *discover.Node)
	if srv.setupFunc == nil {
		srv.setupFunc = setupConn
	}

	// node table
	ntab, err := discover.ListenUDP(srv.PrivateKey, srv.ListenAddr, srv.NAT)
	if err != nil {
		return err
	}
	srv.ntab = ntab

	// handshake
	srv.ourHandshake = &protoHandshake{Version: baseProtocolVersion, Name: srv.Name, ID: ntab.Self().ID}
	for _, p := range srv.Protocols {
		srv.ourHandshake.Caps = append(srv.ourHandshake.Caps, p.cap())
	}

	// listen/dial
	if srv.ListenAddr != "" {
		if err := srv.startListening(); err != nil {
			return err
		}
	}
	if srv.Dialer == nil {
		srv.Dialer = &net.Dialer{Timeout: defaultDialTimeout}
	}
	if !srv.NoDial {
		srv.loopWG.Add(1)
		go srv.dialLoop()
	}
	if srv.NoDial && srv.ListenAddr == "" {
		glog.V(logger.Warn).Infoln("I will be kind-of useless, neither dialing nor listening.")
	}

	srv.running = true
	return nil
}

func (srv *Server) startListening() error {
	listener, err := net.Listen("tcp", srv.ListenAddr)
	if err != nil {
		return err
	}
	laddr := listener.Addr().(*net.TCPAddr)
	srv.ListenAddr = laddr.String()
	srv.listener = listener
	srv.loopWG.Add(1)
	go srv.listenLoop()
	if !laddr.IP.IsLoopback() && srv.NAT != nil {
		srv.loopWG.Add(1)
		go func() {
			nat.Map(srv.NAT, srv.quit, "tcp", laddr.Port, laddr.Port, "ethereum p2p")
			srv.loopWG.Done()
		}()
	}
	return nil
}

// Stop terminates the server and all active peer connections.
// It blocks until all active connections have been closed.
func (srv *Server) Stop() {
	srv.lock.Lock()
	if !srv.running {
		srv.lock.Unlock()
		return
	}
	srv.running = false
	srv.lock.Unlock()

	glog.V(logger.Info).Infoln("Stopping Server")
	srv.ntab.Close()
	if srv.listener != nil {
		// this unblocks listener Accept
		srv.listener.Close()
	}
	close(srv.quit)
	srv.loopWG.Wait()

	// No new peers can be added at this point because dialLoop and
	// listenLoop are down. It is safe to call peerWG.Wait because
	// peerWG.Add is not called outside of those loops.
	for _, peer := range srv.peers {
		peer.Disconnect(DiscQuitting)
	}
	srv.peerWG.Wait()
}

// Self returns the local node's endpoint information.
func (srv *Server) Self() *discover.Node {
	return srv.ntab.Self()
}

// main loop for adding connections via listening
func (srv *Server) listenLoop() {
	defer srv.loopWG.Done()

	// This channel acts as a semaphore limiting
	// active inbound connections that are lingering pre-handshake.
	// If all slots are taken, no further connections are accepted.
	slots := make(chan struct{}, maxAcceptConns)
	for i := 0; i < maxAcceptConns; i++ {
		slots <- struct{}{}
	}

	glog.V(logger.Info).Infoln("Listening on", srv.listener.Addr())
	for {
		<-slots
		conn, err := srv.listener.Accept()
		if err != nil {
			return
		}
		glog.V(logger.Debug).Infof("Accepted conn %v\n", conn.RemoteAddr())
		srv.peerWG.Add(1)
		go func() {
			srv.startPeer(conn, nil)
			slots <- struct{}{}
		}()
	}
}

func (srv *Server) dialLoop() {
	var (
		dialed      = make(chan *discover.Node)
		dialing     = make(map[discover.NodeID]bool)
		findresults = make(chan []*discover.Node)
		refresh     = time.NewTimer(0)
	)
	defer srv.loopWG.Done()
	defer refresh.Stop()

	// TODO: maybe limit number of active dials
	dial := func(dest *discover.Node) {
		// Don't dial nodes that would fail the checks in addPeer.
		// This is important because the connection handshake is a lot
		// of work and we'd rather avoid doing that work for peers
		// that can't be added.
		srv.lock.RLock()
		ok, _ := srv.checkPeer(dest.ID)
		srv.lock.RUnlock()
		if !ok || dialing[dest.ID] {
			return
		}

		dialing[dest.ID] = true
		srv.peerWG.Add(1)
		go func() {
			srv.dialNode(dest)
			dialed <- dest
		}()
	}

	srv.ntab.Bootstrap(srv.BootstrapNodes)
	for {
		select {
		case <-refresh.C:
			// Grab some nodes to connect to if we're not at capacity.
			srv.lock.RLock()
			needpeers := len(srv.peers) < srv.MaxPeers
			srv.lock.RUnlock()
			if needpeers {
				go func() {
					var target discover.NodeID
					rand.Read(target[:])
					findresults <- srv.ntab.Lookup(target)
				}()
				refresh.Stop()
			}

		case dest := <-srv.peerConnect:
			dial(dest)
		case dests := <-findresults:
			for _, dest := range dests {
				dial(dest)
			}
			refresh.Reset(refreshPeersInterval)
		case dest := <-dialed:
			delete(dialing, dest.ID)

		case <-srv.quit:
			// TODO: maybe wait for active dials
			return
		}
	}
}

func (srv *Server) dialNode(dest *discover.Node) {
	addr := &net.TCPAddr{IP: dest.IP, Port: dest.TCPPort}
	glog.V(logger.Debug).Infof("Dialing %v\n", dest)
	conn, err := srv.Dialer.Dial("tcp", addr.String())
	if err != nil {
		// dialLoop adds to the wait group counter when launching
		// dialNode, so we need to count it down again. startPeer also
		// does that when an error occurs.
		srv.peerWG.Done()
		glog.V(logger.Detail).Infof("dial error: %v", err)
		return
	}
	srv.startPeer(conn, dest)
}

func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
	// TODO: handle/store session token

	// Run setupFunc, which should create an authenticated connection
	// and run the capability exchange. Note that any early error
	// returns during that exchange need to call peerWG.Done because
	// the callers of startPeer added the peer to the wait group already.
	fd.SetDeadline(time.Now().Add(handshakeTimeout))
	srv.lock.RLock()
	atcap := len(srv.peers) == srv.MaxPeers
	srv.lock.RUnlock()
	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, atcap)
	if err != nil {
		fd.Close()
		glog.V(logger.Debug).Infof("Handshake with %v failed: %v", fd.RemoteAddr(), err)
		srv.peerWG.Done()
		return
	}
	conn.MsgReadWriter = &netWrapper{
		wrapped: conn.MsgReadWriter,
		conn:    fd, rtimeout: frameReadTimeout, wtimeout: frameWriteTimeout,
	}
	p := newPeer(fd, conn, srv.Protocols)
	if ok, reason := srv.addPeer(conn.ID, p); !ok {
		glog.V(logger.Detail).Infof("Not adding %v (%v)\n", p, reason)
		p.politeDisconnect(reason)
		srv.peerWG.Done()
		return
	}
	// The handshakes are done and it passed all checks.
	// Spawn the Peer loops.
	go srv.runPeer(p)
}

func (srv *Server) runPeer(p *Peer) {
	glog.V(logger.Debug).Infof("Added %v\n", p)
	srvjslog.LogJson(&logger.P2PConnected{
		RemoteId:            p.ID().String(),
		RemoteAddress:       p.RemoteAddr().String(),
		RemoteVersionString: p.Name(),
		NumConnections:      srv.PeerCount(),
	})
	if srv.newPeerHook != nil {
		srv.newPeerHook(p)
	}
	discreason := p.run()
	srv.removePeer(p)
	glog.V(logger.Debug).Infof("Removed %v (%v)\n", p, discreason)
	srvjslog.LogJson(&logger.P2PDisconnected{
		RemoteId:       p.ID().String(),
		NumConnections: srv.PeerCount(),
	})
}

func (srv *Server) addPeer(id discover.NodeID, p *Peer) (bool, DiscReason) {
	srv.lock.Lock()
	defer srv.lock.Unlock()
	if ok, reason := srv.checkPeer(id); !ok {
		return false, reason
	}
	srv.peers[id] = p
	return true, 0
}

func (srv *Server) checkPeer(id discover.NodeID) (bool, DiscReason) {
	switch {
	case !srv.running:
		return false, DiscQuitting
	case len(srv.peers) >= srv.MaxPeers:
		return false, DiscTooManyPeers
	case srv.peers[id] != nil:
		return false, DiscAlreadyConnected
	case id == srv.Self().ID:
		return false, DiscSelf
	default:
		return true, 0
	}
}

func (srv *Server) removePeer(p *Peer) {
	srv.lock.Lock()
	delete(srv.peers, p.ID())
	srv.lock.Unlock()
	srv.peerWG.Done()
}
<|endoftext|>


<commit_before>package tracer

import (
	"math"
	"math/rand"

	"github.com/peterstace/grayt/ray"
	"github.com/peterstace/grayt/vect"
)

// Cameras produce rays that go from an eye to a virtual screen. The rays
// produced are specified via a coordiate system on the virtual screen.  The
// left side of the virtual screen has x coordinate -1, the right side of the
// virtual screen has coordinate +1. The top of the virtual screen has y
// coordinate v and the bottom of the virtual screen has y coordinate -v (where
// the value of v depends on the aspect ratio of the screen).
type Camera interface {
	MakeRay(x, y float64) ray.Ray
}

type RectilinearCamera struct {
	screen, eye struct {
		// X vectors go from the center of the screen or eye to the right of
		// the screen or eye.  Y vectors go from the center of the screen or
		// eye towards the top of the screen or eye.  Loc is the location of
		// the center of the screen or eye.
		loc, x, y vect.Vect
	}
}

// CameraConfig gives configuration options that are common to all camera
// types. This struct is a parameter to the camera factory functions, primarily
// so in the calling context, it's clear (from the name) whach each
// configuration option is (compared to just passing in Vects and float64s to
// the factory function).
type CameraConfig struct {
	Location      vect.Vect
	ViewDirection vect.Vect
	UpDirection   vect.Vect
	FieldOfView   float64
	FocalLength   float64
	FocalRatio    float64
}

func NewRectilinearCamera(conf CameraConfig) Camera {

	cam := &RectilinearCamera{}

	conf.UpDirection = conf.UpDirection.Unit()
	conf.ViewDirection = conf.ViewDirection.Unit()

	cam.screen.x = vect.Cross(conf.ViewDirection, conf.UpDirection)
	cam.screen.y = vect.Cross(cam.screen.x, conf.ViewDirection)

	cam.eye.x = cam.screen.x.Extended(conf.FocalLength / conf.FocalRatio)
	cam.eye.y = cam.screen.y.Extended(conf.FocalLength / conf.FocalRatio)
	cam.eye.loc = conf.Location

	halfScreenWidth := math.Tan(conf.FieldOfView/2) * conf.FocalLength
	cam.screen.x = cam.screen.x.Extended(halfScreenWidth)
	cam.screen.y = cam.screen.y.Extended(halfScreenWidth)
	cam.screen.loc = vect.Add(cam.eye.loc, conf.ViewDirection.Extended(conf.FocalLength))

	return cam
}

func (c *RectilinearCamera) MakeRay(x, y float64) ray.Ray {
	start := vect.Add(
		c.eye.loc,
		vect.Add(
			c.eye.x.Extended(2*rand.Float64()-1.0),
			c.eye.y.Extended(2*rand.Float64()-1.0),
		),
	)
	end := vect.Add(
		c.screen.loc,
		vect.Add(
			c.screen.x.Extended(x),
			c.screen.y.Extended(y),
		),
	)
	return ray.Ray{
		Start: start,
		Dir:   vect.Sub(end, start).Unit(),
	}
}
<commit_msg>Unexport RectilinearCamera<commit_after>package tracer

import (
	"math"
	"math/rand"

	"github.com/peterstace/grayt/ray"
	"github.com/peterstace/grayt/vect"
)

// Camera implementations produce rays that go from an eye to a virtual screen.
// The rays produced are specified via a coordiate system on the virtual
// screen.  The left side of the virtual screen has x coordinate -1, the right
// side of the virtual screen has coordinate +1. The top of the virtual screen
// has y coordinate v and the bottom of the virtual screen has y coordinate -v
// (where the value of v depends on the aspect ratio of the screen).
type Camera interface {
	MakeRay(x, y float64) ray.Ray
}

type rectilinearCamera struct {
	screen, eye struct {
		// X vectors go from the center of the screen or eye to the right of
		// the screen or eye.  Y vectors go from the center of the screen or
		// eye towards the top of the screen or eye.  Loc is the location of
		// the center of the screen or eye.
		loc, x, y vect.Vect
	}
}

// CameraConfig gives configuration options that are common to all camera
// types. This struct is a parameter to the camera factory functions, primarily
// so in the calling context, it's clear (from the name) whach each
// configuration option is (compared to just passing in Vects and float64s to
// the factory function).
type CameraConfig struct {
	Location      vect.Vect
	ViewDirection vect.Vect
	UpDirection   vect.Vect
	FieldOfView   float64
	FocalLength   float64
	FocalRatio    float64
}

// NewRectilinearCamera creates a rectilinear camera from a camera config.
func NewRectilinearCamera(conf CameraConfig) Camera {

	cam := &rectilinearCamera{}

	conf.UpDirection = conf.UpDirection.Unit()
	conf.ViewDirection = conf.ViewDirection.Unit()

	cam.screen.x = vect.Cross(conf.ViewDirection, conf.UpDirection)
	cam.screen.y = vect.Cross(cam.screen.x, conf.ViewDirection)

	cam.eye.x = cam.screen.x.Extended(conf.FocalLength / conf.FocalRatio)
	cam.eye.y = cam.screen.y.Extended(conf.FocalLength / conf.FocalRatio)
	cam.eye.loc = conf.Location

	halfScreenWidth := math.Tan(conf.FieldOfView/2) * conf.FocalLength
	cam.screen.x = cam.screen.x.Extended(halfScreenWidth)
	cam.screen.y = cam.screen.y.Extended(halfScreenWidth)
	cam.screen.loc = vect.Add(cam.eye.loc, conf.ViewDirection.Extended(conf.FocalLength))

	return cam
}

func (c *rectilinearCamera) MakeRay(x, y float64) ray.Ray {
	start := vect.Add(
		c.eye.loc,
		vect.Add(
			c.eye.x.Extended(2*rand.Float64()-1.0),
			c.eye.y.Extended(2*rand.Float64()-1.0),
		),
	)
	end := vect.Add(
		c.screen.loc,
		vect.Add(
			c.screen.x.Extended(x),
			c.screen.y.Extended(y),
		),
	)
	return ray.Ray{
		Start: start,
		Dir:   vect.Sub(end, start).Unit(),
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package tracing

import (
	"context"
	"net/http"

	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.opentelemetry.io/otel/exporters/otlp"
	"go.opentelemetry.io/otel/exporters/otlp/otlpgrpc"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	oteltrace "go.opentelemetry.io/otel/trace"

	"k8s.io/client-go/transport"
	"k8s.io/component-base/tracing/api/v1"
)

// NewProvider creates a TracerProvider in a component, and enforces recommended tracing behavior
func NewProvider(ctx context.Context,
	tracingConfig *v1.TracingConfiguration,
	addedOpts []otlpgrpc.Option,
	resourceOpts []resource.Option,
) (oteltrace.TracerProvider, error) {
	if tracingConfig == nil {
		return oteltrace.NewNoopTracerProvider(), nil
	}
	opts := append([]otlpgrpc.Option{}, addedOpts...)
	if tracingConfig.Endpoint != nil {
		opts = append(opts, otlpgrpc.WithEndpoint(*tracingConfig.Endpoint))
	}
	opts = append(opts, otlpgrpc.WithInsecure())
	driver := otlpgrpc.NewDriver(opts...)
	exporter, err := otlp.NewExporter(ctx, driver)
	if err != nil {
		return nil, err
	}
	res, err := resource.New(ctx, resourceOpts...)
	if err != nil {
		return nil, err
	}

	// sampler respects parent span's sampling rate or
	// otherwise never samples.
	sampler := sdktrace.NeverSample()
	// Or, emit spans for a fraction of transactions
	if tracingConfig.SamplingRatePerMillion != nil && *tracingConfig.SamplingRatePerMillion > 0 {
		sampler = sdktrace.TraceIDRatioBased(float64(*tracingConfig.SamplingRatePerMillion) / float64(1000000))
	}
	// batch span processor to aggregate spans before export.
	bsp := sdktrace.NewBatchSpanProcessor(exporter)
	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.ParentBased(sampler)),
		sdktrace.WithSpanProcessor(bsp),
		sdktrace.WithResource(res),
	)
	return tp, nil
}

// WithTracing adds tracing to requests if the incoming request is sampled
func WithTracing(handler http.Handler, tp oteltrace.TracerProvider, serviceName string) http.Handler {
	opts := []otelhttp.Option{
		otelhttp.WithPropagators(Propagators()),
		otelhttp.WithTracerProvider(tp),
	}
	// With Noop TracerProvider, the otelhttp still handles context propagation.
	// See https://github.com/open-telemetry/opentelemetry-go/tree/main/example/passthrough
	return otelhttp.NewHandler(handler, serviceName, opts...)
}

// WrapperFor can be used to add tracing to a *rest.Config.
// Example usage:
// tp := NewProvider(...)
// config, _ := rest.InClusterConfig()
// config.Wrap(WrapperFor(&tp))
// kubeclient, _ := clientset.NewForConfig(config)
func WrapperFor(tp oteltrace.TracerProvider) transport.WrapperFunc {
	return func(rt http.RoundTripper) http.RoundTripper {
		opts := []otelhttp.Option{
			otelhttp.WithPropagators(Propagators()),
			otelhttp.WithTracerProvider(tp),
		}
		// With Noop TracerProvider, the otelhttp still handles context propagation.
		// See https://github.com/open-telemetry/opentelemetry-go/tree/main/example/passthrough
		return otelhttp.NewTransport(rt, opts...)
	}
}

// Propagators returns the recommended set of propagators.
func Propagators() propagation.TextMapPropagator {
	return propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
}
<commit_msg>update code to use newer  otel api<commit_after>/*
Copyright 2021 The Kubernetes Authors.

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

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

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

package tracing

import (
	"context"
	"net/http"

	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	oteltrace "go.opentelemetry.io/otel/trace"

	"k8s.io/client-go/transport"
	"k8s.io/component-base/tracing/api/v1"
)

// NewProvider creates a TracerProvider in a component, and enforces recommended tracing behavior
func NewProvider(ctx context.Context,
	tracingConfig *v1.TracingConfiguration,
	addedOpts []otlptracegrpc.Option,
	resourceOpts []resource.Option,
) (oteltrace.TracerProvider, error) {
	if tracingConfig == nil {
		return oteltrace.NewNoopTracerProvider(), nil
	}
	opts := append([]otlptracegrpc.Option{}, addedOpts...)
	if tracingConfig.Endpoint != nil {
		opts = append(opts, otlptracegrpc.WithEndpoint(*tracingConfig.Endpoint))
	}
	opts = append(opts, otlptracegrpc.WithInsecure())
	exporter, err := otlptracegrpc.New(ctx, opts...)
	if err != nil {
		return nil, err
	}
	res, err := resource.New(ctx, resourceOpts...)
	if err != nil {
		return nil, err
	}

	// sampler respects parent span's sampling rate or
	// otherwise never samples.
	sampler := sdktrace.NeverSample()
	// Or, emit spans for a fraction of transactions
	if tracingConfig.SamplingRatePerMillion != nil && *tracingConfig.SamplingRatePerMillion > 0 {
		sampler = sdktrace.TraceIDRatioBased(float64(*tracingConfig.SamplingRatePerMillion) / float64(1000000))
	}
	// batch span processor to aggregate spans before export.
	bsp := sdktrace.NewBatchSpanProcessor(exporter)
	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.ParentBased(sampler)),
		sdktrace.WithSpanProcessor(bsp),
		sdktrace.WithResource(res),
	)
	return tp, nil
}

// WithTracing adds tracing to requests if the incoming request is sampled
func WithTracing(handler http.Handler, tp oteltrace.TracerProvider, serviceName string) http.Handler {
	opts := []otelhttp.Option{
		otelhttp.WithPropagators(Propagators()),
		otelhttp.WithTracerProvider(tp),
	}
	// With Noop TracerProvider, the otelhttp still handles context propagation.
	// See https://github.com/open-telemetry/opentelemetry-go/tree/main/example/passthrough
	return otelhttp.NewHandler(handler, serviceName, opts...)
}

// WrapperFor can be used to add tracing to a *rest.Config.
// Example usage:
// tp := NewProvider(...)
// config, _ := rest.InClusterConfig()
// config.Wrap(WrapperFor(&tp))
// kubeclient, _ := clientset.NewForConfig(config)
func WrapperFor(tp oteltrace.TracerProvider) transport.WrapperFunc {
	return func(rt http.RoundTripper) http.RoundTripper {
		opts := []otelhttp.Option{
			otelhttp.WithPropagators(Propagators()),
			otelhttp.WithTracerProvider(tp),
		}
		// With Noop TracerProvider, the otelhttp still handles context propagation.
		// See https://github.com/open-telemetry/opentelemetry-go/tree/main/example/passthrough
		return otelhttp.NewTransport(rt, opts...)
	}
}

// Propagators returns the recommended set of propagators.
func Propagators() propagation.TextMapPropagator {
	return propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
}
<|endoftext|>