<commit_before>package icmp

import (
    "golang.org/x/net/icmp"
    "close/stats"
    "close/config"
    "os"
    "log"
    "fmt"
    "time"
    "close/worker"
)

type PingConfig struct {
    Target      string              `json:"target" long:"target"`
    Proto       string              `json:"proto" long:"protocol" value-name:"ipv4|ipv6" default:"ipv4"`
    ID          int                 `json:"id" long:"id"`
    Interval    time.Duration       `json:"interval" long:"interval" value-name:"<count>(ns|us|ms|s|m|h)" default:"1s"`
}

func (self PingConfig) Worker() (worker.Worker, error) {
    return NewPinger(self)
}

type PingStats struct {
    ID              int
    Time            time.Time       // ping request was sent out

    RTT             time.Duration
}

func (self PingStats) StatsID() stats.ID {
    return stats.ID{
        Type:       "icmp_ping",
        Instance:   fmt.Sprintf("%d", self.ID),
    }
}

func (self PingStats) StatsTime() time.Time {
    return self.Time
}

func (self PingStats) StatsFields() map[string]interface{} {
    return map[string]interface{}{
        // timing
        "rtt": self.RTT.Seconds(),
    }
}

func (self PingStats) String() string {
    return fmt.Sprintf("rtt=%.2fms",
        self.RTT.Seconds() * 1000,
    )
}

type pingResult struct {
    ID    uint16
    Seq   uint16
    Time  time.Time
}

type Pinger struct {
    config      PingConfig

    log         *log.Logger

    conn        *Conn

    configC     chan  config.Config
    statsC      chan  stats.Stats
    receiverC   chan  pingResult
}

func NewPinger(config PingConfig) (*Pinger, error) {
    p := &Pinger{
        log:        log.New(os.Stderr, "ping: ", 0),
    }

    // start
    if err := p.apply(config); err != nil {
        return nil, err
    }

    return p, nil
}

func (p *Pinger) String() string {
    return fmt.Sprintf("Ping %v", p.config.Target)
}

func (p *Pinger) Config() config.Config {
    return &p.config
}

func (p *Pinger) StatsWriter(statsWriter *stats.Writer) error {
    p.statsC = statsWriter.StatsWriter()

    return nil
}

func (p *Pinger) ConfigSub(configSub *config.Sub) error {
    // copy for updates
    pingConfig := p.config

    if configChan, err := configSub.Start(&pingConfig); err != nil {
        return err
    } else {
        p.configC = configChan

        return nil
    }
}

// Apply configuration to state
// TODO: teardown old state?
func (p *Pinger) apply(config PingConfig) error {
    if config.ID == 0 {
        config.ID = os.Getpid()
    }

    if conn, err := NewConn(config); err != nil {
        return err
    } else {
        p.conn = conn
    }

    p.receiverC = make(chan pingResult)

    // good
    p.config = config

    go p.receiver(p.receiverC, p.conn)

    return nil
}

// mainloop
func (p *Pinger) Run() error {
    if p.statsC != nil {
        defer close(p.statsC)
    }
    defer p.log.Printf("stopped\n")

    // state
    var id = uint16(p.config.ID)
    var seq uint16
    timerChan := time.Tick(p.config.Interval)
    startTimes  := make(map[uint16]time.Time)

    for {
        select {
        case <-timerChan:
            seq++

            if err := p.send(id, seq); err != nil {
                return err
            } else {
                startTimes[seq] = time.Now()
            }

        case result, ok := <-p.receiverC:
            if !ok {
                return nil
            }
            if startTime, ok := startTimes[result.Seq]; ok {
                rtt := result.Time.Sub(startTime)

                if p.statsC != nil {
                    p.statsC <- PingStats{
                        ID:         p.config.ID,
                        Time:       startTime,
                        RTT:        rtt,
                    }
                }

                delete(startTimes, result.Seq)
            }

        case configConfig := <-p.configC:
            config := configConfig.(*PingConfig)

            p.log.Printf("config: %v\n", config)

            // TODO: apply()

//      case <-expiryTicker.C:
        }
    }
}

func (p *Pinger) Stop() {
    p.log.Printf("stopping...\n")

    // causes recevier() to close(receiverC)
    p.conn.IcmpConn.Close()
}

func (p *Pinger) send(id uint16, seq uint16) error {
    wm := p.conn.NewMessage(id, seq)

    if err := p.conn.Write(wm); err != nil {
        return fmt.Errorf("icmp.PacketConn %v: WriteTo %v: %v", p.conn.IcmpConn, p.conn.TargetAddr, err)
    }

    return nil
}

func (p *Pinger) receiver(receiverC chan pingResult, conn *Conn) {
    defer close(receiverC)

    icmpProto := conn.Proto.ianaProto
    icmpConn := conn.IcmpConn

    for {
        buf := make([]byte, 1500)
        if readSize, _, err := icmpConn.ReadFrom(buf); err != nil {
            p.log.Printf("icmp.PacketConn %v: ReadFrom: %v\n", icmpConn, err)

            // quit if the connection is closed
            return
        } else {
            buf = buf[:readSize]
        }

        recvTime := time.Now()

        if icmpMessage, err := icmp.ParseMessage(icmpProto, buf); err != nil {
            p.log.Printf("icmp.ParseMessage: %v\n", err)
            continue
        } else if icmpEcho, ok := icmpMessage.Body.(*icmp.Echo); ok {
            receiverC <- pingResult{
                ID:     uint16(icmpEcho.ID),
                Seq:    uint16(icmpEcho.Seq),
                Time:   recvTime,
            }
        }
    }
}
<commit_msg>icmp/ping: don't use ID as stats Instance, since running in docker means os.Getpid() == 1 for all workers<commit_after>package icmp

import (
    "golang.org/x/net/icmp"
    "close/stats"
    "close/config"
    "os"
    "log"
    "fmt"
    "time"
    "close/worker"
)

type PingConfig struct {
    Target      string              `json:"target" long:"target"`
    Proto       string              `json:"proto" long:"protocol" value-name:"ipv4|ipv6" default:"ipv4"`
    ID          int                 `json:"id" long:"id"`
    Interval    time.Duration       `json:"interval" long:"interval" value-name:"<count>(ns|us|ms|s|m|h)" default:"1s"`
}

func (self PingConfig) Worker() (worker.Worker, error) {
    return NewPinger(self)
}

type PingStats struct {
    ID              int
    Time            time.Time       // ping request was sent out

    RTT             time.Duration
}

func (self PingStats) StatsID() stats.ID {
    // use default Instance:
    return stats.ID{
        Type:       "icmp_ping",
    }
}

func (self PingStats) StatsTime() time.Time {
    return self.Time
}

func (self PingStats) StatsFields() map[string]interface{} {
    return map[string]interface{}{
        // timing
        "rtt": self.RTT.Seconds(),
    }
}

func (self PingStats) String() string {
    return fmt.Sprintf("rtt=%.2fms",
        self.RTT.Seconds() * 1000,
    )
}

type pingResult struct {
    ID    uint16
    Seq   uint16
    Time  time.Time
}

type Pinger struct {
    config      PingConfig

    log         *log.Logger

    conn        *Conn

    configC     chan  config.Config
    statsC      chan  stats.Stats
    receiverC   chan  pingResult
}

func NewPinger(config PingConfig) (*Pinger, error) {
    p := &Pinger{
        log:        log.New(os.Stderr, "ping: ", 0),
    }

    // start
    if err := p.apply(config); err != nil {
        return nil, err
    }

    return p, nil
}

func (p *Pinger) String() string {
    return fmt.Sprintf("Ping %v", p.config.Target)
}

func (p *Pinger) Config() config.Config {
    return &p.config
}

func (p *Pinger) StatsWriter(statsWriter *stats.Writer) error {
    p.statsC = statsWriter.StatsWriter()

    return nil
}

func (p *Pinger) ConfigSub(configSub *config.Sub) error {
    // copy for updates
    pingConfig := p.config

    if configChan, err := configSub.Start(&pingConfig); err != nil {
        return err
    } else {
        p.configC = configChan

        return nil
    }
}

// Apply configuration to state
// TODO: teardown old state?
func (p *Pinger) apply(config PingConfig) error {
    if config.ID == 0 {
        // XXX: this is going to be 1 when running within docker..
        config.ID = os.Getpid()
    }

    if conn, err := NewConn(config); err != nil {
        return err
    } else {
        p.conn = conn
    }

    p.receiverC = make(chan pingResult)

    // good
    p.config = config

    go p.receiver(p.receiverC, p.conn)

    return nil
}

// mainloop
func (p *Pinger) Run() error {
    if p.statsC != nil {
        defer close(p.statsC)
    }
    defer p.log.Printf("stopped\n")

    // state
    var id = uint16(p.config.ID)
    var seq uint16
    timerChan := time.Tick(p.config.Interval)
    startTimes  := make(map[uint16]time.Time)

    for {
        select {
        case <-timerChan:
            seq++

            if err := p.send(id, seq); err != nil {
                return err
            } else {
                startTimes[seq] = time.Now()
            }

        case result, ok := <-p.receiverC:
            if !ok {
                return nil
            }
            if startTime, ok := startTimes[result.Seq]; ok {
                rtt := result.Time.Sub(startTime)

                if p.statsC != nil {
                    p.statsC <- PingStats{
                        ID:         p.config.ID,
                        Time:       startTime,
                        RTT:        rtt,
                    }
                }

                delete(startTimes, result.Seq)
            }

        case configConfig := <-p.configC:
            config := configConfig.(*PingConfig)

            p.log.Printf("config: %v\n", config)

            // TODO: apply()

//      case <-expiryTicker.C:
        }
    }
}

func (p *Pinger) Stop() {
    p.log.Printf("stopping...\n")

    // causes recevier() to close(receiverC)
    p.conn.IcmpConn.Close()
}

func (p *Pinger) send(id uint16, seq uint16) error {
    wm := p.conn.NewMessage(id, seq)

    if err := p.conn.Write(wm); err != nil {
        return fmt.Errorf("icmp.PacketConn %v: WriteTo %v: %v", p.conn.IcmpConn, p.conn.TargetAddr, err)
    }

    return nil
}

func (p *Pinger) receiver(receiverC chan pingResult, conn *Conn) {
    defer close(receiverC)

    icmpProto := conn.Proto.ianaProto
    icmpConn := conn.IcmpConn

    for {
        buf := make([]byte, 1500)
        if readSize, _, err := icmpConn.ReadFrom(buf); err != nil {
            p.log.Printf("icmp.PacketConn %v: ReadFrom: %v\n", icmpConn, err)

            // quit if the connection is closed
            return
        } else {
            buf = buf[:readSize]
        }

        recvTime := time.Now()

        if icmpMessage, err := icmp.ParseMessage(icmpProto, buf); err != nil {
            p.log.Printf("icmp.ParseMessage: %v\n", err)
            continue
        } else if icmpEcho, ok := icmpMessage.Body.(*icmp.Echo); ok {
            receiverC <- pingResult{
                ID:     uint16(icmpEcho.ID),
                Seq:    uint16(icmpEcho.Seq),
                Time:   recvTime,
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 db defines the interface that fleetspeak expects from its persistence
// layer. Each installation will need to choose and configure a Store
// implementation. An example implementation meant for testing and small scale
// deployments is in the server/sqlite directory.
//
// It also includes some utility methods and types meant for use by Store
// implementations.
//
// SECURITY NOTE:
//
// The endpoints provide much of the data passed through this interface.
// Implementations are responsible for using safe coding practices to prevent
// SQL injection and similar attacks.
package db

import (
	"context"
	"fmt"
	"io"
	"math"
	"time"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/ids"

	tpb "github.com/golang/protobuf/ptypes/timestamp"
	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
	mpb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak_monitoring"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// A Store describes the full persistence mechanism required by the base
// fleetspeak system. These operations must be thread safe. These must also be
// all-or-nothing, fully committed on success, and are otherwise trusted to be
// individually transactional.
type Store interface {
	MessageStore
	ClientStore
	BroadcastStore
	FileStore

	// IsNotFound returns whether an error returned by the Datastore indicates that
	// a record was not found.
	IsNotFound(error) bool

	// Close shuts down the Store, releasing any held resources.
	Close() error
}

// ClientData contains basic data about a client.
type ClientData struct {
	Key    []byte        // The der encoded public key for the client.
	Labels []*fspb.Label // The client's labels.

	// Whether the client_id has been blacklisted. Once blacklisted any contact
	// from this client_id will result in an rekey request.
	Blacklisted bool
}

// A ContactID identifies a communication with a client. The form is determined
// by the Datastore implementation - it is treated as an opaque string by the
// rest of the FS system.
type ContactID string

// MessageStore provides methods to store and query messages.
//
// Notionally, a MessageStore is backed by a table where each row is a fspb.Message record,
// along with with one of the following:
//
// 1) due time and retry count - If the message is not processed or delivered
// before due time, it will be tried again. A count of the number of retries is
// maintained and used to compute the next due time.
//
// 2) completion time - When a record is processed or acknowledged by a client, the
// message is marked as completed by saving a completion time.
//
// Furthermore it is possible to register a MessageProcessor with each
// MessageStore which then receives notifications that server messages are ready
// for processing. In multi-server installations, the datastore should attempt
// to provide eactly one notification to some Fleetspeak server each time the
// message becomes overdue.
type MessageStore interface {
	// StoreMessages records msgs. If contact is not the empty string, it attaches
	// them to the associated contact.
	//
	// It is not an error for a message to already exist. In this case, the only
	// fields examined by the method are MessageId and Result, and Result will be
	// updated if it is a supported transition:
	//
	// No Result -> Success Result
	// No Result -> Failed Result
	// Failed Result -> Success Result
	//
	// All other transitions are silently ignored.
	//
	// A message is eligible to be returned by ClientMessagesForProcessing or the
	// registered MessageProcessor iff it does not yet have a Result. Also,
	// setting a Result will delete the message's Data field payload.
	StoreMessages(ctx context.Context, msgs []*fspb.Message, contact ContactID) error

	// ClientMessagesForProcessing returns up to lim messages that are due to be
	// processed by a client. It also increments the time at which the
	// messages will again become overdue using rp.
	//
	// Note that if an error occurs partway through the loading of messages,
	// the already loaded messages may be returned along with the error. In
	// particular, datastore implementations may want to do this if the ctx
	// times out before all messages are found and updated.
	ClientMessagesForProcessing(ctx context.Context, id common.ClientID, lim int) ([]*fspb.Message, error)

	// GetMessages retrieves specific messages.
	GetMessages(ctx context.Context, ids []common.MessageID, wantData bool) ([]*fspb.Message, error)

	// GetMessageStatus retrieves the current status of a message.
	GetMessageResult(ctx context.Context, id common.MessageID) (*fspb.MessageResult, error)

	// RegisterMessageProcessor installs a MessageProcessor which will be
	// called when a message is overdue for processing.
	RegisterMessageProcessor(mp MessageProcessor)

	// StopMessageProcessor causes the datastore to stop making calls to the
	// registered MessageProcessor. It only returns once all existing calls
	// to MessageProcessor have completed.
	StopMessageProcessor()
}

// A MessageProcessor receives messages that are overdue and should be reprocessing.
type MessageProcessor interface {
	// ProcessMessage is called by the Datastore to indicate that the
	// provided message is overdue and that processing should be attempted
	// again.
	//
	// This call will be repeated until MarkMessage(Processed|Failed) is
	// successfully called on m.
	ProcessMessages(msgs []*fspb.Message)
}

type ContactData struct {
	ClientID                 common.ClientID
	NonceSent, NonceReceived uint64
	Addr                     string
	ClientClock              *tpb.Timestamp
}

// ClientStore provides methods to store and retrieve information about clients.
type ClientStore interface {
	// ListClients returns basic information about clients. If ids is empty, it
	// returns all clients.
	ListClients(ctx context.Context, ids []common.ClientID) ([]*spb.Client, error)

	// GetClientData retrieves the current data about the client identified
	// by id.
	GetClientData(ctx context.Context, id common.ClientID) (*ClientData, error)

	// AddClient creates a new client.
	AddClient(ctx context.Context, id common.ClientID, data *ClientData) error

	// AddClientLabel records that a client now has a label.
	AddClientLabel(ctx context.Context, id common.ClientID, l *fspb.Label) error

	// RemoveLabel records that a client no longer has a label.
	RemoveClientLabel(ctx context.Context, id common.ClientID, l *fspb.Label) error

	// BlacklistClient records that a client_id is no longer trusted and should be
	// recreated.
	BlacklistClient(ctx context.Context, id common.ClientID) error

	// RecordClientContact records an authenticated contact with a
	// client. On success provides a contact id - an opaque string which can
	// be used to link messages to a contact.
	RecordClientContact(ctx context.Context, data ContactData) (ContactID, error)

	// ListClientContacts lists all of the contacts in the database for a given
	// client.
	//
	// NOTE: This method is explicitly permitted to return data up to 30 seconds
	// stale. Also, it is normal (and expected) for a datastore to delete contact
	// older than a few weeks.
	ListClientContacts(ctx context.Context, id common.ClientID) ([]*spb.ClientContact, error)

	// LinkMessagesToContact associates messages with a contact - it records
	// that they were sent or received during the given contact.
	LinkMessagesToContact(ctx context.Context, contact ContactID, msgs []common.MessageID) error

	// Writes resource-usage data received from a client to the data-store.
	RecordResourceUsageData(ctx context.Context, id common.ClientID, rud mpb.ResourceUsageData) error

	// Fetches at most 'limit' resource-usage records for a given client from the data-store.
	// TODO: Add more complex queries.
	FetchResourceUsageRecords(ctx context.Context, id common.ClientID, limit int) ([]*spb.ClientResourceUsageRecord, error)
}

// Broadcast limits with special meaning.
const (
	BroadcastDisabled  = uint64(0)
	BroadcastUnlimited = uint64(math.MaxInt64) // The sqlite datastore's uint64 doesn't support full uint64 range.
)

// A BroadcastInfo describes a broadcast and contains the static broadcast
// information, plus the current limit and count of messages sent.
type BroadcastInfo struct {
	Broadcast *spb.Broadcast
	Sent      uint64
	Limit     uint64
}

// An AllocationInfo describes an allocation. An allocation is the right to send
// some broadcast to up to Limit machines before Expiry.
type AllocationInfo struct {
	ID     ids.AllocationID
	Limit  uint64
	Expiry time.Time
}

// ComputeBroadcastAllocation computes how large a new allocation should be. It
// is meant to be used by implementations of Broadcaststore.
//
// It takes the allocation's current message limit, also the number already
// sent, the number already allocated, and the target fraction of the allocation
// to claim. It returns the number of messages that should be allocated to a new
// allocation, also the new total number of messages allocated.
func ComputeBroadcastAllocation(messageLimit, allocated, sent uint64, frac float32) (toAllocate, newAllocated uint64) {
	// Allocations for unlimited broadcasts don't count; such allocations are only used
	// to keep tract of the number sent.
	if messageLimit == BroadcastUnlimited {
		return BroadcastUnlimited, allocated
	}
	a := allocated
	if sent > a {
		a = sent
	}
	if a > messageLimit {
		return 0, allocated
	}
	toAllocate = uint64(float32(messageLimit-a) * frac)
	if toAllocate == 0 {
		toAllocate = 1
	}
	if toAllocate > messageLimit-a {
		toAllocate = messageLimit - a
	}
	newAllocated = toAllocate + allocated
	return
}

// ComputeBroadcastAllocationCleanup computes the new number of messages
// allocated when cleaning up an allocation. It takes the number of messages
// that were allocated to the allocation, also the current total number of messages
// allocated from the broadcast.
func ComputeBroadcastAllocationCleanup(allocationLimit, allocated uint64) (uint64, error) {
	if allocationLimit == BroadcastUnlimited {
		return allocated, nil
	}
	if allocationLimit > allocated {
		return 0, fmt.Errorf("allocationLimit = %v, which is larger than allocated = %v", allocationLimit, allocated)
	}
	return allocated - allocationLimit, nil
}

// BroadcastStore provides methods to store and retrieve information about broadcasts.
type BroadcastStore interface {
	// CreateBroadcast stores a new broadcast message.
	CreateBroadcast(ctx context.Context, b *spb.Broadcast, limit uint64) error

	// SetBroadcastLimit adjusts the limit of an existing broadcast.
	SetBroadcastLimit(ctx context.Context, id ids.BroadcastID, limit uint64) error

	// SaveBroadcastMessage saves a new broadcast message.
	SaveBroadcastMessage(ctx context.Context, msg *fspb.Message, bid ids.BroadcastID, cid common.ClientID, aid ids.AllocationID) error

	// ListActiveBroadcasts lists broadcasts which could be sent to some
	// client.
	ListActiveBroadcasts(ctx context.Context) ([]*BroadcastInfo, error)

	// ListSentBroadcasts returns identifiers for those broadcasts which have already been sent to a client.
	ListSentBroadcasts(ctx context.Context, id common.ClientID) ([]ids.BroadcastID, error)

	// CreateAllocation creates an allocation for a given Broadcast,
	// reserving frac of the unallocated broadcast limit until
	// expiry. Return nil if there is no message allocation available.
	CreateAllocation(ctx context.Context, id ids.BroadcastID, frac float32, expiry time.Time) (*AllocationInfo, error)

	// CleanupAllocation deletes the identified allocation record and
	// updates the broadcast sent count according to the number that were
	// actually sent under the given allocation.
	CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID) error
}

// ReadSeekerCloser groups io.ReadSeeker and io.Closer.
type ReadSeekerCloser interface {
	io.ReadSeeker
	io.Closer
}

// NOOPCloser wraps an io.ReadSeeker to trivially turn it into a ReadSeekerCloser.
type NOOPCloser struct {
	io.ReadSeeker
}

// Close implements io.Closer.
func (c NOOPCloser) Close() error {
	return nil
}

// FileStore provides methods to store and retrieve files. Files are keyed by an associated
// service and name.
//
// SECURITY NOTES:
//
// Fleetspeak doesn't provide any ACL support for files - all files are readable
// by any client.
//
// Implementations are responsible for validating and/or sanitizing the
// identifiers provided. For example, an implementation backed by a filesystem
// would need to protect against path traversal vulnerabilities.
type FileStore interface {
	// StoreFile stores data into the Filestore, organized by service and name.
	StoreFile(ctx context.Context, service, name string, data io.Reader) error

	// StatFile returns the modification time of a file previously stored by
	// StoreFile. Returns ErrNotFound if not found.
	StatFile(ctx context.Context, servce, name string) (time.Time, error)

	// ReadFile returns the data and modification time of file previously
	// stored by StoreFile.  Caller is responsible for closing data.
	//
	// Note: Calls to data are permitted to fail if ctx is canceled or expired.
	ReadFile(ctx context.Context, service, name string) (data ReadSeekerCloser, modtime time.Time, err error)
}
<commit_msg>Add comments for ContactData struct.<commit_after>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 db defines the interface that fleetspeak expects from its persistence
// layer. Each installation will need to choose and configure a Store
// implementation. An example implementation meant for testing and small scale
// deployments is in the server/sqlite directory.
//
// It also includes some utility methods and types meant for use by Store
// implementations.
//
// SECURITY NOTE:
//
// The endpoints provide much of the data passed through this interface.
// Implementations are responsible for using safe coding practices to prevent
// SQL injection and similar attacks.
package db

import (
	"context"
	"fmt"
	"io"
	"math"
	"time"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/ids"

	tpb "github.com/golang/protobuf/ptypes/timestamp"
	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
	mpb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak_monitoring"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// A Store describes the full persistence mechanism required by the base
// fleetspeak system. These operations must be thread safe. These must also be
// all-or-nothing, fully committed on success, and are otherwise trusted to be
// individually transactional.
type Store interface {
	MessageStore
	ClientStore
	BroadcastStore
	FileStore

	// IsNotFound returns whether an error returned by the Datastore indicates that
	// a record was not found.
	IsNotFound(error) bool

	// Close shuts down the Store, releasing any held resources.
	Close() error
}

// ClientData contains basic data about a client.
type ClientData struct {
	Key    []byte        // The der encoded public key for the client.
	Labels []*fspb.Label // The client's labels.

	// Whether the client_id has been blacklisted. Once blacklisted any contact
	// from this client_id will result in an rekey request.
	Blacklisted bool
}

// A ContactID identifies a communication with a client. The form is determined
// by the Datastore implementation - it is treated as an opaque string by the
// rest of the FS system.
type ContactID string

// MessageStore provides methods to store and query messages.
//
// Notionally, a MessageStore is backed by a table where each row is a fspb.Message record,
// along with with one of the following:
//
// 1) due time and retry count - If the message is not processed or delivered
// before due time, it will be tried again. A count of the number of retries is
// maintained and used to compute the next due time.
//
// 2) completion time - When a record is processed or acknowledged by a client, the
// message is marked as completed by saving a completion time.
//
// Furthermore it is possible to register a MessageProcessor with each
// MessageStore which then receives notifications that server messages are ready
// for processing. In multi-server installations, the datastore should attempt
// to provide eactly one notification to some Fleetspeak server each time the
// message becomes overdue.
type MessageStore interface {
	// StoreMessages records msgs. If contact is not the empty string, it attaches
	// them to the associated contact.
	//
	// It is not an error for a message to already exist. In this case, the only
	// fields examined by the method are MessageId and Result, and Result will be
	// updated if it is a supported transition:
	//
	// No Result -> Success Result
	// No Result -> Failed Result
	// Failed Result -> Success Result
	//
	// All other transitions are silently ignored.
	//
	// A message is eligible to be returned by ClientMessagesForProcessing or the
	// registered MessageProcessor iff it does not yet have a Result. Also,
	// setting a Result will delete the message's Data field payload.
	StoreMessages(ctx context.Context, msgs []*fspb.Message, contact ContactID) error

	// ClientMessagesForProcessing returns up to lim messages that are due to be
	// processed by a client. It also increments the time at which the
	// messages will again become overdue using rp.
	//
	// Note that if an error occurs partway through the loading of messages,
	// the already loaded messages may be returned along with the error. In
	// particular, datastore implementations may want to do this if the ctx
	// times out before all messages are found and updated.
	ClientMessagesForProcessing(ctx context.Context, id common.ClientID, lim int) ([]*fspb.Message, error)

	// GetMessages retrieves specific messages.
	GetMessages(ctx context.Context, ids []common.MessageID, wantData bool) ([]*fspb.Message, error)

	// GetMessageStatus retrieves the current status of a message.
	GetMessageResult(ctx context.Context, id common.MessageID) (*fspb.MessageResult, error)

	// RegisterMessageProcessor installs a MessageProcessor which will be
	// called when a message is overdue for processing.
	RegisterMessageProcessor(mp MessageProcessor)

	// StopMessageProcessor causes the datastore to stop making calls to the
	// registered MessageProcessor. It only returns once all existing calls
	// to MessageProcessor have completed.
	StopMessageProcessor()
}

// A MessageProcessor receives messages that are overdue and should be reprocessing.
type MessageProcessor interface {
	// ProcessMessage is called by the Datastore to indicate that the
	// provided message is overdue and that processing should be attempted
	// again.
	//
	// This call will be repeated until MarkMessage(Processed|Failed) is
	// successfully called on m.
	ProcessMessages(msgs []*fspb.Message)
}

// ContactData provides basic information about a client's contact with a FS
// server.
type ContactData struct {
	ClientID                 common.ClientID // ID of the client.
	NonceSent, NonceReceived uint64          // Nonce sent to the client and received from the client.
	Addr                     string          // Observed client network address.
	ClientClock              *tpb.Timestamp  // Client's report of its current clock setting.
}

// ClientStore provides methods to store and retrieve information about clients.
type ClientStore interface {
	// ListClients returns basic information about clients. If ids is empty, it
	// returns all clients.
	ListClients(ctx context.Context, ids []common.ClientID) ([]*spb.Client, error)

	// GetClientData retrieves the current data about the client identified
	// by id.
	GetClientData(ctx context.Context, id common.ClientID) (*ClientData, error)

	// AddClient creates a new client.
	AddClient(ctx context.Context, id common.ClientID, data *ClientData) error

	// AddClientLabel records that a client now has a label.
	AddClientLabel(ctx context.Context, id common.ClientID, l *fspb.Label) error

	// RemoveLabel records that a client no longer has a label.
	RemoveClientLabel(ctx context.Context, id common.ClientID, l *fspb.Label) error

	// BlacklistClient records that a client_id is no longer trusted and should be
	// recreated.
	BlacklistClient(ctx context.Context, id common.ClientID) error

	// RecordClientContact records an authenticated contact with a
	// client. On success provides a contact id - an opaque string which can
	// be used to link messages to a contact.
	RecordClientContact(ctx context.Context, data ContactData) (ContactID, error)

	// ListClientContacts lists all of the contacts in the database for a given
	// client.
	//
	// NOTE: This method is explicitly permitted to return data up to 30 seconds
	// stale. Also, it is normal (and expected) for a datastore to delete contact
	// older than a few weeks.
	ListClientContacts(ctx context.Context, id common.ClientID) ([]*spb.ClientContact, error)

	// LinkMessagesToContact associates messages with a contact - it records
	// that they were sent or received during the given contact.
	LinkMessagesToContact(ctx context.Context, contact ContactID, msgs []common.MessageID) error

	// Writes resource-usage data received from a client to the data-store.
	RecordResourceUsageData(ctx context.Context, id common.ClientID, rud mpb.ResourceUsageData) error

	// Fetches at most 'limit' resource-usage records for a given client from the data-store.
	// TODO: Add more complex queries.
	FetchResourceUsageRecords(ctx context.Context, id common.ClientID, limit int) ([]*spb.ClientResourceUsageRecord, error)
}

// Broadcast limits with special meaning.
const (
	BroadcastDisabled  = uint64(0)
	BroadcastUnlimited = uint64(math.MaxInt64) // The sqlite datastore's uint64 doesn't support full uint64 range.
)

// A BroadcastInfo describes a broadcast and contains the static broadcast
// information, plus the current limit and count of messages sent.
type BroadcastInfo struct {
	Broadcast *spb.Broadcast
	Sent      uint64
	Limit     uint64
}

// An AllocationInfo describes an allocation. An allocation is the right to send
// some broadcast to up to Limit machines before Expiry.
type AllocationInfo struct {
	ID     ids.AllocationID
	Limit  uint64
	Expiry time.Time
}

// ComputeBroadcastAllocation computes how large a new allocation should be. It
// is meant to be used by implementations of Broadcaststore.
//
// It takes the allocation's current message limit, also the number already
// sent, the number already allocated, and the target fraction of the allocation
// to claim. It returns the number of messages that should be allocated to a new
// allocation, also the new total number of messages allocated.
func ComputeBroadcastAllocation(messageLimit, allocated, sent uint64, frac float32) (toAllocate, newAllocated uint64) {
	// Allocations for unlimited broadcasts don't count; such allocations are only used
	// to keep tract of the number sent.
	if messageLimit == BroadcastUnlimited {
		return BroadcastUnlimited, allocated
	}
	a := allocated
	if sent > a {
		a = sent
	}
	if a > messageLimit {
		return 0, allocated
	}
	toAllocate = uint64(float32(messageLimit-a) * frac)
	if toAllocate == 0 {
		toAllocate = 1
	}
	if toAllocate > messageLimit-a {
		toAllocate = messageLimit - a
	}
	newAllocated = toAllocate + allocated
	return
}

// ComputeBroadcastAllocationCleanup computes the new number of messages
// allocated when cleaning up an allocation. It takes the number of messages
// that were allocated to the allocation, also the current total number of messages
// allocated from the broadcast.
func ComputeBroadcastAllocationCleanup(allocationLimit, allocated uint64) (uint64, error) {
	if allocationLimit == BroadcastUnlimited {
		return allocated, nil
	}
	if allocationLimit > allocated {
		return 0, fmt.Errorf("allocationLimit = %v, which is larger than allocated = %v", allocationLimit, allocated)
	}
	return allocated - allocationLimit, nil
}

// BroadcastStore provides methods to store and retrieve information about broadcasts.
type BroadcastStore interface {
	// CreateBroadcast stores a new broadcast message.
	CreateBroadcast(ctx context.Context, b *spb.Broadcast, limit uint64) error

	// SetBroadcastLimit adjusts the limit of an existing broadcast.
	SetBroadcastLimit(ctx context.Context, id ids.BroadcastID, limit uint64) error

	// SaveBroadcastMessage saves a new broadcast message.
	SaveBroadcastMessage(ctx context.Context, msg *fspb.Message, bid ids.BroadcastID, cid common.ClientID, aid ids.AllocationID) error

	// ListActiveBroadcasts lists broadcasts which could be sent to some
	// client.
	ListActiveBroadcasts(ctx context.Context) ([]*BroadcastInfo, error)

	// ListSentBroadcasts returns identifiers for those broadcasts which have already been sent to a client.
	ListSentBroadcasts(ctx context.Context, id common.ClientID) ([]ids.BroadcastID, error)

	// CreateAllocation creates an allocation for a given Broadcast,
	// reserving frac of the unallocated broadcast limit until
	// expiry. Return nil if there is no message allocation available.
	CreateAllocation(ctx context.Context, id ids.BroadcastID, frac float32, expiry time.Time) (*AllocationInfo, error)

	// CleanupAllocation deletes the identified allocation record and
	// updates the broadcast sent count according to the number that were
	// actually sent under the given allocation.
	CleanupAllocation(ctx context.Context, bid ids.BroadcastID, aid ids.AllocationID) error
}

// ReadSeekerCloser groups io.ReadSeeker and io.Closer.
type ReadSeekerCloser interface {
	io.ReadSeeker
	io.Closer
}

// NOOPCloser wraps an io.ReadSeeker to trivially turn it into a ReadSeekerCloser.
type NOOPCloser struct {
	io.ReadSeeker
}

// Close implements io.Closer.
func (c NOOPCloser) Close() error {
	return nil
}

// FileStore provides methods to store and retrieve files. Files are keyed by an associated
// service and name.
//
// SECURITY NOTES:
//
// Fleetspeak doesn't provide any ACL support for files - all files are readable
// by any client.
//
// Implementations are responsible for validating and/or sanitizing the
// identifiers provided. For example, an implementation backed by a filesystem
// would need to protect against path traversal vulnerabilities.
type FileStore interface {
	// StoreFile stores data into the Filestore, organized by service and name.
	StoreFile(ctx context.Context, service, name string, data io.Reader) error

	// StatFile returns the modification time of a file previously stored by
	// StoreFile. Returns ErrNotFound if not found.
	StatFile(ctx context.Context, servce, name string) (time.Time, error)

	// ReadFile returns the data and modification time of file previously
	// stored by StoreFile.  Caller is responsible for closing data.
	//
	// Note: Calls to data are permitted to fail if ctx is canceled or expired.
	ReadFile(ctx context.Context, service, name string) (data ReadSeekerCloser, modtime time.Time, err error)
}
<|endoftext|>


<commit_before>// Instruments allows you to collects metrics over discrete time intervals.
//
// Collected metrics will only reflect observations from last time window only,
// rather than including observations from prior windows, contrary to EWMA based metrics.
//
// 	timer := instruments.NewTimer(-1)
//
//	registry := instruments.NewRegistry()
//	registry.Register("processing-time", timer)
//
//	go reporter.Log("process", registry, time.Minute)
//
//	timer.Time(func() {
//	  ...
//	})
//
// Instruments support two types of instruments:
// Discrete instruments return a single value, and Sample instruments a sorted array of values.
//
// Theses base instruments are available:
//
// - Rate: tracks the rate of values per seconds.
//
// - Reservoir: randomly samples values.
//
// - Derive: tracks the rate of values based on the delta with previous value.
//
// - Gauge: tracks last value.
//
// - Timer: tracks durations.
//
// You can create custom instruments or compose new instruments form the built-in
// instruments as long as they implements the Sample or Discrete interfaces.
//
// Registry enforce the Discrete and Sample interfaces,
// creating a custom Reporter should be trivial, for example:
//
// 	for k, m := range r.Instruments() {
// 	 	switch i := m.(type) {
// 	 	case instruments.Discrete:
// 	 	 	s := i.Snapshot()
// 	 	 	report(k, s)
// 	 	case instruments.Sample:
// 	 	 	s := instruments.Quantile(i.Snapshot(), 0.95)
// 	 	 	report(k, s)
// 	 	}
//	}
//
package instruments

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

const rateScale = 1e-9

// Discrete represents a single value instrument.
type Discrete interface {
	Snapshot() int64
}

// Sample represents a sample instrument.
type Sample interface {
	Snapshot() []int64
}

// Rate tracks the rate of values per second.
type Rate struct {
	count int64
	time  int64
}

// NewRate creates a new rate instrument.
func NewRate() *Rate {
	return &Rate{
		time: time.Now().UnixNano(),
	}
}

// Update updates rate value.
func (r *Rate) Update(v int64) {
	atomic.AddInt64(&r.count, v)
}

// Snapshot returns the number of values per second since the last snapshot,
// and reset the count to zero.
func (r *Rate) Snapshot() int64 {
	now := time.Now().UnixNano()
	t := atomic.SwapInt64(&r.time, now)
	c := atomic.SwapInt64(&r.count, 0)
	s := float64(c) / rateScale / float64(now-t)
	return Ceil(s)
}

// Derive tracks the rate of deltas per seconds.
type Derive struct {
	rate  *Rate
	value int64
}

// NewDerive creates a new derive instruments.
func NewDerive(v int64) *Derive {
	return &Derive{
		value: v,
		rate:  NewRate(),
	}
}

// Update update rate value based on the stored previous value.
func (d *Derive) Update(v int64) {
	p := atomic.SwapInt64(&d.value, v)
	d.rate.Update(v - p)
}

// Snapshot returns the number of values per seconds since the last snapshot,
// and reset the count to zero.
func (d *Derive) Snapshot() int64 {
	return d.rate.Snapshot()
}

// Reservoir tracks a sample of values.
type Reservoir struct {
	size   int64
	values []int64
	m      sync.Mutex
}

const defaultReservoirSize = 1028

// NewReservoir creates a new reservoir of the given size.
// If size is negative, it will create a sample of DefaultReservoirSize size.
func NewReservoir(size int64) *Reservoir {
	if size <= 0 {
		size = defaultReservoirSize
	}
	return &Reservoir{
		values: make([]int64, size),
	}
}

// Update fills the sample randomly with given value,
// for reference, see: http://en.wikipedia.org/wiki/Reservoir_sampling
func (r *Reservoir) Update(v int64) {
	r.m.Lock()
	defer r.m.Unlock()
	s := atomic.AddInt64(&r.size, 1)
	if int(s) <= len(r.values) {
		// Not full
		r.values[s-1] = v
	} else {
		// Full
		l := rand.Int63n(s)
		if int(l) < len(r.values) {
			r.values[l] = v
		}
	}
}

// Snapshot returns sample as a sorted array.
func (r *Reservoir) Snapshot() []int64 {
	r.m.Lock()
	defer r.m.Unlock()
	s := atomic.SwapInt64(&r.size, 0)
	v := make([]int64, min(int(s), len(r.values)))
	copy(v, r.values)
	r.values = make([]int64, cap(r.values))
	sorted(v)
	return v
}

// Gauge tracks a value.
type Gauge struct {
	value int64
}

// NewGauge creates a new Gauge with the given value.
func NewGauge(v int64) *Gauge {
	return &Gauge{
		value: v,
	}
}

// Update updates the current stored value.
func (g *Gauge) Update(v int64) {
	atomic.StoreInt64(&g.value, v)
}

// Snapshot returns the current value.
func (g *Gauge) Snapshot() int64 {
	return atomic.LoadInt64(&g.value)
}

// Timer tracks durations.
type Timer struct {
	r *Reservoir
}

// NewTimer creates a new Timer with the given sample size.
func NewTimer(size int64) *Timer {
	return &Timer{
		r: NewReservoir(size),
	}
}

// Update adds duration to the sample in ms.
func (t *Timer) Update(d time.Duration) {
	v := Floor(d.Seconds() * 1000)
	t.r.Update(v)
}

// Snapshot returns durations sample as a sorted array.
func (t *Timer) Snapshot() []int64 {
	return t.r.Snapshot()
}

// Time records given function execution time.
func (t *Timer) Time(f func()) {
	ts := time.Now()
	f()
	t.Update(time.Since(ts))
}
<commit_msg>allow to change rate unit<commit_after>// Instruments allows you to collects metrics over discrete time intervals.
//
// Collected metrics will only reflect observations from last time window only,
// rather than including observations from prior windows, contrary to EWMA based metrics.
//
// 	timer := instruments.NewTimer(-1)
//
//	registry := instruments.NewRegistry()
//	registry.Register("processing-time", timer)
//
//	go reporter.Log("process", registry, time.Minute)
//
//	timer.Time(func() {
//	  ...
//	})
//
// Instruments support two types of instruments:
// Discrete instruments return a single value, and Sample instruments a sorted array of values.
//
// Theses base instruments are available:
//
// - Rate: tracks the rate of values per seconds.
//
// - Reservoir: randomly samples values.
//
// - Derive: tracks the rate of values based on the delta with previous value.
//
// - Gauge: tracks last value.
//
// - Timer: tracks durations.
//
// You can create custom instruments or compose new instruments form the built-in
// instruments as long as they implements the Sample or Discrete interfaces.
//
// Registry enforce the Discrete and Sample interfaces,
// creating a custom Reporter should be trivial, for example:
//
// 	for k, m := range r.Instruments() {
// 	 	switch i := m.(type) {
// 	 	case instruments.Discrete:
// 	 	 	s := i.Snapshot()
// 	 	 	report(k, s)
// 	 	case instruments.Sample:
// 	 	 	s := instruments.Quantile(i.Snapshot(), 0.95)
// 	 	 	report(k, s)
// 	 	}
//	}
//
package instruments

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

const rateScale = 1e-9

// Discrete represents a single value instrument.
type Discrete interface {
	Snapshot() int64
}

// Sample represents a sample instrument.
type Sample interface {
	Snapshot() []int64
}

// Scale returns a conversion factor from one unit to another.
func Scale(o, d time.Duration) float64 {
	return float64(o) / float64(d)
}

// Rate tracks the rate of values per second.
type Rate struct {
	count int64
	time  int64
	scale float64
}

// NewRate creates a new rate instrument.
func NewRate() *Rate {
	return NewRateScale(rateScale)
}

// NewRateScale creates a new rate instruments with the given conversion factor.
func NewRateScale(s float64) *Rate {
	return &Rate{
		time:  time.Now().UnixNano(),
		scale: s,
	}
}

// Update updates rate value.
func (r *Rate) Update(v int64) {
	atomic.AddInt64(&r.count, v)
}

// Snapshot returns the number of values per second since the last snapshot,
// and reset the count to zero.
func (r *Rate) Snapshot() int64 {
	now := time.Now().UnixNano()
	t := atomic.SwapInt64(&r.time, now)
	c := atomic.SwapInt64(&r.count, 0)
	s := float64(c) / r.scale / float64(now-t)
	return Ceil(s)
}

// Derive tracks the rate of deltas per seconds.
type Derive struct {
	rate  *Rate
	value int64
}

// NewDerive creates a new derive instruments.
func NewDerive(v int64) *Derive {
	return &Derive{
		value: v,
		rate:  NewRate(),
	}
}

// NewDeriveScale creates a new derive instruments with the given conversion factor.
func NewDeriveScale(v int64, s float64) *Derive {
	return &Derive{
		value: v,
		rate:  NewRateScale(s),
	}
}

// Update update rate value based on the stored previous value.
func (d *Derive) Update(v int64) {
	p := atomic.SwapInt64(&d.value, v)
	d.rate.Update(v - p)
}

// Snapshot returns the number of values per seconds since the last snapshot,
// and reset the count to zero.
func (d *Derive) Snapshot() int64 {
	return d.rate.Snapshot()
}

// Reservoir tracks a sample of values.
type Reservoir struct {
	size   int64
	values []int64
	m      sync.Mutex
}

const defaultReservoirSize = 1028

// NewReservoir creates a new reservoir of the given size.
// If size is negative, it will create a sample of DefaultReservoirSize size.
func NewReservoir(size int64) *Reservoir {
	if size <= 0 {
		size = defaultReservoirSize
	}
	return &Reservoir{
		values: make([]int64, size),
	}
}

// Update fills the sample randomly with given value,
// for reference, see: http://en.wikipedia.org/wiki/Reservoir_sampling
func (r *Reservoir) Update(v int64) {
	r.m.Lock()
	defer r.m.Unlock()
	s := atomic.AddInt64(&r.size, 1)
	if int(s) <= len(r.values) {
		// Not full
		r.values[s-1] = v
	} else {
		// Full
		l := rand.Int63n(s)
		if int(l) < len(r.values) {
			r.values[l] = v
		}
	}
}

// Snapshot returns sample as a sorted array.
func (r *Reservoir) Snapshot() []int64 {
	r.m.Lock()
	defer r.m.Unlock()
	s := atomic.SwapInt64(&r.size, 0)
	v := make([]int64, min(int(s), len(r.values)))
	copy(v, r.values)
	r.values = make([]int64, cap(r.values))
	sorted(v)
	return v
}

// Gauge tracks a value.
type Gauge struct {
	value int64
}

// NewGauge creates a new Gauge with the given value.
func NewGauge(v int64) *Gauge {
	return &Gauge{
		value: v,
	}
}

// Update updates the current stored value.
func (g *Gauge) Update(v int64) {
	atomic.StoreInt64(&g.value, v)
}

// Snapshot returns the current value.
func (g *Gauge) Snapshot() int64 {
	return atomic.LoadInt64(&g.value)
}

// Timer tracks durations.
type Timer struct {
	r *Reservoir
}

// NewTimer creates a new Timer with the given sample size.
func NewTimer(size int64) *Timer {
	return &Timer{
		r: NewReservoir(size),
	}
}

// Update adds duration to the sample in ms.
func (t *Timer) Update(d time.Duration) {
	v := Floor(d.Seconds() * 1000)
	t.r.Update(v)
}

// Snapshot returns durations sample as a sorted array.
func (t *Timer) Snapshot() []int64 {
	return t.r.Snapshot()
}

// Time records given function execution time.
func (t *Timer) Time(f func()) {
	ts := time.Now()
	f()
	t.Update(time.Since(ts))
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package codec

import (
	"encoding/binary"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/util/types"
	"github.com/pingcap/tidb/util/types/json"
)

// First byte in the encoded value which specifies the encoding type.
const (
	NilFlag          byte = 0
	bytesFlag        byte = 1
	compactBytesFlag byte = 2
	intFlag          byte = 3
	uintFlag         byte = 4
	floatFlag        byte = 5
	decimalFlag      byte = 6
	durationFlag     byte = 7
	varintFlag       byte = 8
	uvarintFlag      byte = 9
	jsonFlag         byte = 10
	maxFlag          byte = 250
)

// encode will encode a datum and append it to a byte slice. If comparable is true, the encoded bytes can be sorted as it's original order.
// If hash is true, the encoded bytes can be checked equal as it's original value.
func encode(b []byte, vals []types.Datum, comparable bool, hash bool) ([]byte, error) {
	for _, val := range vals {
		switch val.Kind() {
		case types.KindInt64:
			b = encodeSignedInt(b, val.GetInt64(), comparable)
		case types.KindUint64:
			if hash {
				int := val.GetInt64()
				if int < 0 {
					b = encodeUnsignedInt(b, uint64(int), comparable)
				} else {
					b = encodeSignedInt(b, int, comparable)
				}
			} else {
				b = encodeUnsignedInt(b, val.GetUint64(), comparable)
			}
		case types.KindFloat32, types.KindFloat64:
			b = append(b, floatFlag)
			b = EncodeFloat(b, val.GetFloat64())
		case types.KindString, types.KindBytes:
			b = encodeBytes(b, val.GetBytes(), comparable)
		case types.KindMysqlTime:
			b = append(b, uintFlag)
			t := val.GetMysqlTime()
			// Encoding timestamp need to consider timezone.
			// If it's not in UTC, transform to UTC first.
			if t.Type == mysql.TypeTimestamp && t.TimeZone != time.UTC {
				t.ConvertTimeZone(t.TimeZone, time.UTC)
			}
			v, err := t.ToPackedUint()
			if err != nil {
				return nil, errors.Trace(err)
			}
			b = EncodeUint(b, v)
		case types.KindMysqlDuration:
			// duration may have negative value, so we cannot use String to encode directly.
			b = append(b, durationFlag)
			b = EncodeInt(b, int64(val.GetMysqlDuration().Duration))
		case types.KindMysqlDecimal:
			b = append(b, decimalFlag)
			if hash {
				// If hash is true, we only consider the original value of this decimal and ignore it's precision.
				dec := val.GetMysqlDecimal()
				precision, frac := dec.PrecisionAndFrac()
				bin, err := dec.ToBin(precision, frac)
				if err != nil {
					return nil, errors.Trace(err)
				}
				b = append(b, bin...)
			} else {
				b = EncodeDecimal(b, val)
			}
		case types.KindMysqlHex:
			b = encodeSignedInt(b, int64(val.GetMysqlHex().ToNumber()), comparable)
		case types.KindMysqlBit:
			b = encodeUnsignedInt(b, uint64(val.GetMysqlBit().ToNumber()), comparable)
		case types.KindMysqlEnum:
			b = encodeUnsignedInt(b, uint64(val.GetMysqlEnum().ToNumber()), comparable)
		case types.KindMysqlSet:
			b = encodeUnsignedInt(b, uint64(val.GetMysqlSet().ToNumber()), comparable)
		case types.KindMysqlJSON:
			b = append(b, jsonFlag)
			b = append(b, json.Serialize(val.GetMysqlJSON())...)
		case types.KindNull:
			b = append(b, NilFlag)
		case types.KindMinNotNull:
			b = append(b, bytesFlag)
		case types.KindMaxValue:
			b = append(b, maxFlag)
		default:
			return nil, errors.Errorf("unsupport encode type %d", val.Kind())
		}
	}

	return b, nil
}

func encodeBytes(b []byte, v []byte, comparable bool) []byte {
	if comparable {
		b = append(b, bytesFlag)
		b = EncodeBytes(b, v)
	} else {
		b = append(b, compactBytesFlag)
		b = EncodeCompactBytes(b, v)
	}
	return b
}

func encodeSignedInt(b []byte, v int64, comparable bool) []byte {
	if comparable {
		b = append(b, intFlag)
		b = EncodeInt(b, v)
	} else {
		b = append(b, varintFlag)
		b = EncodeVarint(b, v)
	}
	return b
}

func encodeUnsignedInt(b []byte, v uint64, comparable bool) []byte {
	if comparable {
		b = append(b, uintFlag)
		b = EncodeUint(b, v)
	} else {
		b = append(b, uvarintFlag)
		b = EncodeUvarint(b, v)
	}
	return b
}

// EncodeKey appends the encoded values to byte slice b, returns the appended
// slice. It guarantees the encoded value is in ascending order for comparison.
// For Decimal type, datum must set datum's length and frac.
func EncodeKey(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, true, false)
}

// EncodeValue appends the encoded values to byte slice b, returning the appended
// slice. It does not guarantee the order for comparison.
func EncodeValue(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, false, false)
}

// HashValues appends the encoded values to byte slice b, returning the appended
// slice. If two datums are equal, they will generate the same bytes.
func HashValues(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, false, true)
}

// Decode decodes values from a byte slice generated with EncodeKey or EncodeValue
// before.
// size is the size of decoded datum slice.
func Decode(b []byte, size int) ([]types.Datum, error) {
	if len(b) < 1 {
		return nil, errors.New("invalid encoded key")
	}

	var (
		err    error
		values = make([]types.Datum, 0, size)
	)

	for len(b) > 0 {
		var d types.Datum
		b, d, err = DecodeOne(b)
		if err != nil {
			return nil, errors.Trace(err)
		}

		values = append(values, d)
	}

	return values, nil
}

// DecodeOne decodes on datum from a byte slice generated with EncodeKey or EncodeValue.
func DecodeOne(b []byte) (remain []byte, d types.Datum, err error) {
	if len(b) < 1 {
		return nil, d, errors.New("invalid encoded key")
	}
	flag := b[0]
	b = b[1:]
	switch flag {
	case intFlag:
		var v int64
		b, v, err = DecodeInt(b)
		d.SetInt64(v)
	case uintFlag:
		var v uint64
		b, v, err = DecodeUint(b)
		d.SetUint64(v)
	case varintFlag:
		var v int64
		b, v, err = DecodeVarint(b)
		d.SetInt64(v)
	case uvarintFlag:
		var v uint64
		b, v, err = DecodeUvarint(b)
		d.SetUint64(v)
	case floatFlag:
		var v float64
		b, v, err = DecodeFloat(b)
		d.SetFloat64(v)
	case bytesFlag:
		var v []byte
		b, v, err = DecodeBytes(b)
		d.SetBytes(v)
	case compactBytesFlag:
		var v []byte
		b, v, err = DecodeCompactBytes(b)
		d.SetBytes(v)
	case decimalFlag:
		b, d, err = DecodeDecimal(b)
	case durationFlag:
		var r int64
		b, r, err = DecodeInt(b)
		if err == nil {
			// use max fsp, let outer to do round manually.
			v := types.Duration{Duration: time.Duration(r), Fsp: types.MaxFsp}
			d.SetValue(v)
		}
	case jsonFlag:
		var j json.JSON
		j, err = json.Deserialize(b)
		if err == nil {
			d.SetMysqlJSON(j)
		}
	case NilFlag:
	default:
		return b, d, errors.Errorf("invalid encoded key flag %v", flag)
	}
	if err != nil {
		return b, d, errors.Trace(err)
	}
	return b, d, nil
}

// CutOne cuts the first encoded value from b.
// It will return the first encoded item and the remains as byte slice.
func CutOne(b []byte) (data []byte, remain []byte, err error) {
	l, err := peek(b)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	return b[:l], b[l:], nil
}

// SetRawValues set raw datum values from a row data.
func SetRawValues(data []byte, values []types.Datum) error {
	for i := 0; i < len(values); i++ {
		l, err := peek(data)
		if err != nil {
			return errors.Trace(err)
		}
		values[i].SetRaw(data[:l:l])
		data = data[l:]
	}
	return nil
}

// peek peeks the first encoded value from b and returns its length.
func peek(b []byte) (length int, err error) {
	if len(b) < 1 {
		return 0, errors.New("invalid encoded key")
	}
	flag := b[0]
	length++
	b = b[1:]
	var l int
	switch flag {
	case NilFlag:
	case intFlag, uintFlag, floatFlag, durationFlag:
		// Those types are stored in 8 bytes.
		l = 8
	case bytesFlag:
		l, err = peekBytes(b, false)
	case compactBytesFlag:
		l, err = peekCompactBytes(b)
	case decimalFlag:
		l, err = types.DecimalPeak(b)
	case varintFlag:
		l, err = peekVarint(b)
	case uvarintFlag:
		l, err = peekUvarint(b)
	case jsonFlag:
		l, err = json.PeekBytesAsJSON(b)
	default:
		return 0, errors.Errorf("invalid encoded key flag %v", flag)
	}
	if err != nil {
		return 0, errors.Trace(err)
	}
	length += l
	return
}

func peekBytes(b []byte, reverse bool) (int, error) {
	offset := 0
	for {
		if len(b) < offset+encGroupSize+1 {
			return 0, errors.New("insufficient bytes to decode value")
		}
		// The byte slice is encoded into many groups.
		// For each group, there are 8 bytes for data and 1 byte for marker.
		marker := b[offset+encGroupSize]
		var padCount byte
		if reverse {
			padCount = marker
		} else {
			padCount = encMarker - marker
		}
		offset += encGroupSize + 1
		// When padCount is not zero, it means we get the end of the byte slice.
		if padCount != 0 {
			break
		}
	}
	return offset, nil
}

func peekCompactBytes(b []byte) (int, error) {
	// Get length.
	v, n := binary.Varint(b)
	vi := int(v)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	} else if n == 0 {
		return 0, errors.New("insufficient bytes to decode value")
	}
	if len(b) < vi+n {
		return 0, errors.Errorf("insufficient bytes to decode value, expected length: %v", n)
	}
	return n + vi, nil
}

func peekVarint(b []byte) (int, error) {
	_, n := binary.Varint(b)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	}
	return n, nil
}

func peekUvarint(b []byte) (int, error) {
	_, n := binary.Uvarint(b)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	}
	return n, nil
}
<commit_msg>codec: use reference instead of value copy of "[]types.Datum" (#4408)<commit_after>// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package codec

import (
	"encoding/binary"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/util/types"
	"github.com/pingcap/tidb/util/types/json"
)

// First byte in the encoded value which specifies the encoding type.
const (
	NilFlag          byte = 0
	bytesFlag        byte = 1
	compactBytesFlag byte = 2
	intFlag          byte = 3
	uintFlag         byte = 4
	floatFlag        byte = 5
	decimalFlag      byte = 6
	durationFlag     byte = 7
	varintFlag       byte = 8
	uvarintFlag      byte = 9
	jsonFlag         byte = 10
	maxFlag          byte = 250
)

// encode will encode a datum and append it to a byte slice. If comparable is true, the encoded bytes can be sorted as it's original order.
// If hash is true, the encoded bytes can be checked equal as it's original value.
func encode(b []byte, vals []types.Datum, comparable bool, hash bool) ([]byte, error) {
	for i, length := 0, len(vals); i < length; i++ {
		switch vals[i].Kind() {
		case types.KindInt64:
			b = encodeSignedInt(b, vals[i].GetInt64(), comparable)
		case types.KindUint64:
			if hash {
				int := vals[i].GetInt64()
				if int < 0 {
					b = encodeUnsignedInt(b, uint64(int), comparable)
				} else {
					b = encodeSignedInt(b, int, comparable)
				}
			} else {
				b = encodeUnsignedInt(b, vals[i].GetUint64(), comparable)
			}
		case types.KindFloat32, types.KindFloat64:
			b = append(b, floatFlag)
			b = EncodeFloat(b, vals[i].GetFloat64())
		case types.KindString, types.KindBytes:
			b = encodeBytes(b, vals[i].GetBytes(), comparable)
		case types.KindMysqlTime:
			b = append(b, uintFlag)
			t := vals[i].GetMysqlTime()
			// Encoding timestamp need to consider timezone.
			// If it's not in UTC, transform to UTC first.
			if t.Type == mysql.TypeTimestamp && t.TimeZone != time.UTC {
				t.ConvertTimeZone(t.TimeZone, time.UTC)
			}
			v, err := t.ToPackedUint()
			if err != nil {
				return nil, errors.Trace(err)
			}
			b = EncodeUint(b, v)
		case types.KindMysqlDuration:
			// duration may have negative value, so we cannot use String to encode directly.
			b = append(b, durationFlag)
			b = EncodeInt(b, int64(vals[i].GetMysqlDuration().Duration))
		case types.KindMysqlDecimal:
			b = append(b, decimalFlag)
			if hash {
				// If hash is true, we only consider the original value of this decimal and ignore it's precision.
				dec := vals[i].GetMysqlDecimal()
				precision, frac := dec.PrecisionAndFrac()
				bin, err := dec.ToBin(precision, frac)
				if err != nil {
					return nil, errors.Trace(err)
				}
				b = append(b, bin...)
			} else {
				b = EncodeDecimal(b, vals[i])
			}
		case types.KindMysqlHex:
			b = encodeSignedInt(b, int64(vals[i].GetMysqlHex().ToNumber()), comparable)
		case types.KindMysqlBit:
			b = encodeUnsignedInt(b, uint64(vals[i].GetMysqlBit().ToNumber()), comparable)
		case types.KindMysqlEnum:
			b = encodeUnsignedInt(b, uint64(vals[i].GetMysqlEnum().ToNumber()), comparable)
		case types.KindMysqlSet:
			b = encodeUnsignedInt(b, uint64(vals[i].GetMysqlSet().ToNumber()), comparable)
		case types.KindMysqlJSON:
			b = append(b, jsonFlag)
			b = append(b, json.Serialize(vals[i].GetMysqlJSON())...)
		case types.KindNull:
			b = append(b, NilFlag)
		case types.KindMinNotNull:
			b = append(b, bytesFlag)
		case types.KindMaxValue:
			b = append(b, maxFlag)
		default:
			return nil, errors.Errorf("unsupport encode type %d", vals[i].Kind())
		}
	}

	return b, nil
}

func encodeBytes(b []byte, v []byte, comparable bool) []byte {
	if comparable {
		b = append(b, bytesFlag)
		b = EncodeBytes(b, v)
	} else {
		b = append(b, compactBytesFlag)
		b = EncodeCompactBytes(b, v)
	}
	return b
}

func encodeSignedInt(b []byte, v int64, comparable bool) []byte {
	if comparable {
		b = append(b, intFlag)
		b = EncodeInt(b, v)
	} else {
		b = append(b, varintFlag)
		b = EncodeVarint(b, v)
	}
	return b
}

func encodeUnsignedInt(b []byte, v uint64, comparable bool) []byte {
	if comparable {
		b = append(b, uintFlag)
		b = EncodeUint(b, v)
	} else {
		b = append(b, uvarintFlag)
		b = EncodeUvarint(b, v)
	}
	return b
}

// EncodeKey appends the encoded values to byte slice b, returns the appended
// slice. It guarantees the encoded value is in ascending order for comparison.
// For Decimal type, datum must set datum's length and frac.
func EncodeKey(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, true, false)
}

// EncodeValue appends the encoded values to byte slice b, returning the appended
// slice. It does not guarantee the order for comparison.
func EncodeValue(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, false, false)
}

// HashValues appends the encoded values to byte slice b, returning the appended
// slice. If two datums are equal, they will generate the same bytes.
func HashValues(b []byte, v ...types.Datum) ([]byte, error) {
	return encode(b, v, false, true)
}

// Decode decodes values from a byte slice generated with EncodeKey or EncodeValue
// before.
// size is the size of decoded datum slice.
func Decode(b []byte, size int) ([]types.Datum, error) {
	if len(b) < 1 {
		return nil, errors.New("invalid encoded key")
	}

	var (
		err    error
		values = make([]types.Datum, 0, size)
	)

	for len(b) > 0 {
		var d types.Datum
		b, d, err = DecodeOne(b)
		if err != nil {
			return nil, errors.Trace(err)
		}

		values = append(values, d)
	}

	return values, nil
}

// DecodeOne decodes on datum from a byte slice generated with EncodeKey or EncodeValue.
func DecodeOne(b []byte) (remain []byte, d types.Datum, err error) {
	if len(b) < 1 {
		return nil, d, errors.New("invalid encoded key")
	}
	flag := b[0]
	b = b[1:]
	switch flag {
	case intFlag:
		var v int64
		b, v, err = DecodeInt(b)
		d.SetInt64(v)
	case uintFlag:
		var v uint64
		b, v, err = DecodeUint(b)
		d.SetUint64(v)
	case varintFlag:
		var v int64
		b, v, err = DecodeVarint(b)
		d.SetInt64(v)
	case uvarintFlag:
		var v uint64
		b, v, err = DecodeUvarint(b)
		d.SetUint64(v)
	case floatFlag:
		var v float64
		b, v, err = DecodeFloat(b)
		d.SetFloat64(v)
	case bytesFlag:
		var v []byte
		b, v, err = DecodeBytes(b)
		d.SetBytes(v)
	case compactBytesFlag:
		var v []byte
		b, v, err = DecodeCompactBytes(b)
		d.SetBytes(v)
	case decimalFlag:
		b, d, err = DecodeDecimal(b)
	case durationFlag:
		var r int64
		b, r, err = DecodeInt(b)
		if err == nil {
			// use max fsp, let outer to do round manually.
			v := types.Duration{Duration: time.Duration(r), Fsp: types.MaxFsp}
			d.SetValue(v)
		}
	case jsonFlag:
		var j json.JSON
		j, err = json.Deserialize(b)
		if err == nil {
			d.SetMysqlJSON(j)
		}
	case NilFlag:
	default:
		return b, d, errors.Errorf("invalid encoded key flag %v", flag)
	}
	if err != nil {
		return b, d, errors.Trace(err)
	}
	return b, d, nil
}

// CutOne cuts the first encoded value from b.
// It will return the first encoded item and the remains as byte slice.
func CutOne(b []byte) (data []byte, remain []byte, err error) {
	l, err := peek(b)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	return b[:l], b[l:], nil
}

// SetRawValues set raw datum values from a row data.
func SetRawValues(data []byte, values []types.Datum) error {
	for i := 0; i < len(values); i++ {
		l, err := peek(data)
		if err != nil {
			return errors.Trace(err)
		}
		values[i].SetRaw(data[:l:l])
		data = data[l:]
	}
	return nil
}

// peek peeks the first encoded value from b and returns its length.
func peek(b []byte) (length int, err error) {
	if len(b) < 1 {
		return 0, errors.New("invalid encoded key")
	}
	flag := b[0]
	length++
	b = b[1:]
	var l int
	switch flag {
	case NilFlag:
	case intFlag, uintFlag, floatFlag, durationFlag:
		// Those types are stored in 8 bytes.
		l = 8
	case bytesFlag:
		l, err = peekBytes(b, false)
	case compactBytesFlag:
		l, err = peekCompactBytes(b)
	case decimalFlag:
		l, err = types.DecimalPeak(b)
	case varintFlag:
		l, err = peekVarint(b)
	case uvarintFlag:
		l, err = peekUvarint(b)
	case jsonFlag:
		l, err = json.PeekBytesAsJSON(b)
	default:
		return 0, errors.Errorf("invalid encoded key flag %v", flag)
	}
	if err != nil {
		return 0, errors.Trace(err)
	}
	length += l
	return
}

func peekBytes(b []byte, reverse bool) (int, error) {
	offset := 0
	for {
		if len(b) < offset+encGroupSize+1 {
			return 0, errors.New("insufficient bytes to decode value")
		}
		// The byte slice is encoded into many groups.
		// For each group, there are 8 bytes for data and 1 byte for marker.
		marker := b[offset+encGroupSize]
		var padCount byte
		if reverse {
			padCount = marker
		} else {
			padCount = encMarker - marker
		}
		offset += encGroupSize + 1
		// When padCount is not zero, it means we get the end of the byte slice.
		if padCount != 0 {
			break
		}
	}
	return offset, nil
}

func peekCompactBytes(b []byte) (int, error) {
	// Get length.
	v, n := binary.Varint(b)
	vi := int(v)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	} else if n == 0 {
		return 0, errors.New("insufficient bytes to decode value")
	}
	if len(b) < vi+n {
		return 0, errors.Errorf("insufficient bytes to decode value, expected length: %v", n)
	}
	return n + vi, nil
}

func peekVarint(b []byte) (int, error) {
	_, n := binary.Varint(b)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	}
	return n, nil
}

func peekUvarint(b []byte) (int, error) {
	_, n := binary.Uvarint(b)
	if n < 0 {
		return 0, errors.New("value larger than 64 bits")
	}
	return n, nil
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

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

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

package log

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"sync"
	"time"
)

var logger = New(os.Stdout, NewTextFormatter(), WarningLevel)

func init() {
	logger.callDepth = 4

	// TODO add item in configuaration file
	lvl := os.Getenv("LOG_LEVEL")
	if len(lvl) == 0 {
		logger.SetLevel(InfoLevel)
		return
	}

	level, err := parseLevel(lvl)
	if err != nil {
		logger.SetLevel(InfoLevel)
		return
	}

	logger.SetLevel(level)

}

// Logger provides a struct with fields that describe the details of logger.
type Logger struct {
	out       io.Writer
	fmtter    Formatter
	lvl       Level
	callDepth int
	skipLine  bool
	mu        sync.Mutex
}

// New returns a customized Logger
func New(out io.Writer, fmtter Formatter, lvl Level) *Logger {
	return &Logger{
		out:       out,
		fmtter:    fmtter,
		lvl:       lvl,
		callDepth: 2,
	}
}

//SetOutput sets the output of Logger l
func (l *Logger) SetOutput(out io.Writer) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.out = out
}

//SetFormatter sets the formatter of Logger l
func (l *Logger) SetFormatter(fmtter Formatter) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.fmtter = fmtter
}

//SetLevel sets the level of Logger l
func (l *Logger) SetLevel(lvl Level) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.lvl = lvl
}

//SetOutput sets the output of default Logger
func SetOutput(out io.Writer) {
	logger.SetOutput(out)
}

//SetFormatter sets the formatter of default Logger
func SetFormatter(fmtter Formatter) {
	logger.SetFormatter(fmtter)
}

//SetLevel sets the level of default Logger
func SetLevel(lvl Level) {
	logger.SetLevel(lvl)
}

func (l *Logger) output(record *Record) (err error) {
	b, err := l.fmtter.Format(record)
	if err != nil {
		return
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	_, err = l.out.Write(b)

	return
}

// Debug ...
func (l *Logger) Debug(v ...interface{}) {
	if l.lvl <= DebugLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), DebugLevel)
		l.output(record)
	}
}

// Debugf ...
func (l *Logger) Debugf(format string, v ...interface{}) {
	if l.lvl <= DebugLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), DebugLevel)
		l.output(record)
	}
}

// Info ...
func (l *Logger) Info(v ...interface{}) {
	if l.lvl <= InfoLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), "", InfoLevel)
		l.output(record)
	}
}

// Infof ...
func (l *Logger) Infof(format string, v ...interface{}) {
	if l.lvl <= InfoLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), "", InfoLevel)
		l.output(record)
	}
}

// Warning ...
func (l *Logger) Warning(v ...interface{}) {
	if l.lvl <= WarningLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), "", WarningLevel)
		l.output(record)
	}
}

// Warningf ...
func (l *Logger) Warningf(format string, v ...interface{}) {
	if l.lvl <= WarningLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), "", WarningLevel)
		l.output(record)
	}
}

// Error ...
func (l *Logger) Error(v ...interface{}) {
	if l.lvl <= ErrorLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), ErrorLevel)
		l.output(record)
	}
}

// Errorf ...
func (l *Logger) Errorf(format string, v ...interface{}) {
	if l.lvl <= ErrorLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), ErrorLevel)
		l.output(record)
	}
}

// Fatal ...
func (l *Logger) Fatal(v ...interface{}) {
	if l.lvl <= FatalLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), FatalLevel)
		l.output(record)
	}
	os.Exit(1)
}

// Fatalf ...
func (l *Logger) Fatalf(format string, v ...interface{}) {
	if l.lvl <= FatalLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), FatalLevel)
		l.output(record)
	}
	os.Exit(1)
}

func (l *Logger) getLine() string {
	if l.skipLine {
		return ""
	}
	return line(l.callDepth)
}

// Debug ...
func Debug(v ...interface{}) {
	logger.Debug(v...)
}

// Debugf ...
func Debugf(format string, v ...interface{}) {
	logger.Debugf(format, v...)
}

// Info ...
func Info(v ...interface{}) {
	logger.Info(v...)
}

// Infof ...
func Infof(format string, v ...interface{}) {
	logger.Infof(format, v...)
}

// Warning  ...
func Warning(v ...interface{}) {
	logger.Warning(v...)
}

// Warningf ...
func Warningf(format string, v ...interface{}) {
	logger.Warningf(format, v...)
}

// Error ...
func Error(v ...interface{}) {
	logger.Error(v...)
}

// Errorf ...
func Errorf(format string, v ...interface{}) {
	logger.Errorf(format, v...)
}

// Fatal ...
func Fatal(v ...interface{}) {
	logger.Fatal(v...)
}

// Fatalf ...
func Fatalf(format string, v ...interface{}) {
	logger.Fatalf(format, v...)
}

func line(calldepth int) string {
	_, file, line, ok := runtime.Caller(calldepth)
	if !ok {
		file = "???"
		line = 0
	}

	for i := len(file) - 2; i > 0; i-- {
		if file[i] == os.PathSeparator {
			file = file[i+1:]
			break
		}
	}

	return fmt.Sprintf("[%s:%d]:", file, line)
}
<commit_msg>update the depth in New func<commit_after>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

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

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

package log

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"sync"
	"time"
)

var logger = New(os.Stdout, NewTextFormatter(), WarningLevel)

func init() {
	logger.callDepth = 4

	// TODO add item in configuaration file
	lvl := os.Getenv("LOG_LEVEL")
	if len(lvl) == 0 {
		logger.SetLevel(InfoLevel)
		return
	}

	level, err := parseLevel(lvl)
	if err != nil {
		logger.SetLevel(InfoLevel)
		return
	}

	logger.SetLevel(level)

}

// Logger provides a struct with fields that describe the details of logger.
type Logger struct {
	out       io.Writer
	fmtter    Formatter
	lvl       Level
	callDepth int
	skipLine  bool
	mu        sync.Mutex
}

// New returns a customized Logger
func New(out io.Writer, fmtter Formatter, lvl Level) *Logger {
	return &Logger{
		out:       out,
		fmtter:    fmtter,
		lvl:       lvl,
		callDepth: 3,
	}
}

//SetOutput sets the output of Logger l
func (l *Logger) SetOutput(out io.Writer) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.out = out
}

//SetFormatter sets the formatter of Logger l
func (l *Logger) SetFormatter(fmtter Formatter) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.fmtter = fmtter
}

//SetLevel sets the level of Logger l
func (l *Logger) SetLevel(lvl Level) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.lvl = lvl
}

//SetOutput sets the output of default Logger
func SetOutput(out io.Writer) {
	logger.SetOutput(out)
}

//SetFormatter sets the formatter of default Logger
func SetFormatter(fmtter Formatter) {
	logger.SetFormatter(fmtter)
}

//SetLevel sets the level of default Logger
func SetLevel(lvl Level) {
	logger.SetLevel(lvl)
}

func (l *Logger) output(record *Record) (err error) {
	b, err := l.fmtter.Format(record)
	if err != nil {
		return
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	_, err = l.out.Write(b)

	return
}

// Debug ...
func (l *Logger) Debug(v ...interface{}) {
	if l.lvl <= DebugLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), DebugLevel)
		l.output(record)
	}
}

// Debugf ...
func (l *Logger) Debugf(format string, v ...interface{}) {
	if l.lvl <= DebugLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), DebugLevel)
		l.output(record)
	}
}

// Info ...
func (l *Logger) Info(v ...interface{}) {
	if l.lvl <= InfoLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), "", InfoLevel)
		l.output(record)
	}
}

// Infof ...
func (l *Logger) Infof(format string, v ...interface{}) {
	if l.lvl <= InfoLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), "", InfoLevel)
		l.output(record)
	}
}

// Warning ...
func (l *Logger) Warning(v ...interface{}) {
	if l.lvl <= WarningLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), "", WarningLevel)
		l.output(record)
	}
}

// Warningf ...
func (l *Logger) Warningf(format string, v ...interface{}) {
	if l.lvl <= WarningLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), "", WarningLevel)
		l.output(record)
	}
}

// Error ...
func (l *Logger) Error(v ...interface{}) {
	if l.lvl <= ErrorLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), ErrorLevel)
		l.output(record)
	}
}

// Errorf ...
func (l *Logger) Errorf(format string, v ...interface{}) {
	if l.lvl <= ErrorLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), ErrorLevel)
		l.output(record)
	}
}

// Fatal ...
func (l *Logger) Fatal(v ...interface{}) {
	if l.lvl <= FatalLevel {
		record := NewRecord(time.Now(), fmt.Sprint(v...), l.getLine(), FatalLevel)
		l.output(record)
	}
	os.Exit(1)
}

// Fatalf ...
func (l *Logger) Fatalf(format string, v ...interface{}) {
	if l.lvl <= FatalLevel {
		record := NewRecord(time.Now(), fmt.Sprintf(format, v...), l.getLine(), FatalLevel)
		l.output(record)
	}
	os.Exit(1)
}

func (l *Logger) getLine() string {
	if l.skipLine {
		return ""
	}
	return line(l.callDepth)
}

// Debug ...
func Debug(v ...interface{}) {
	logger.Debug(v...)
}

// Debugf ...
func Debugf(format string, v ...interface{}) {
	logger.Debugf(format, v...)
}

// Info ...
func Info(v ...interface{}) {
	logger.Info(v...)
}

// Infof ...
func Infof(format string, v ...interface{}) {
	logger.Infof(format, v...)
}

// Warning  ...
func Warning(v ...interface{}) {
	logger.Warning(v...)
}

// Warningf ...
func Warningf(format string, v ...interface{}) {
	logger.Warningf(format, v...)
}

// Error ...
func Error(v ...interface{}) {
	logger.Error(v...)
}

// Errorf ...
func Errorf(format string, v ...interface{}) {
	logger.Errorf(format, v...)
}

// Fatal ...
func Fatal(v ...interface{}) {
	logger.Fatal(v...)
}

// Fatalf ...
func Fatalf(format string, v ...interface{}) {
	logger.Fatalf(format, v...)
}

func line(calldepth int) string {
	_, file, line, ok := runtime.Caller(calldepth)
	if !ok {
		file = "???"
		line = 0
	}

	for i := len(file) - 2; i > 0; i-- {
		if file[i] == os.PathSeparator {
			file = file[i+1:]
			break
		}
	}

	return fmt.Sprintf("[%s:%d]:", file, line)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"strconv"
	"time"

	"github.com/Sirupsen/logrus"
	log "github.com/cihub/seelog"
	"github.com/stampzilla/gozwave"
	"github.com/stampzilla/gozwave/commands"
	"github.com/stampzilla/gozwave/events"
	"github.com/stampzilla/gozwave/nodes"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/pkg/notifier"
	"github.com/stampzilla/stampzilla-go/protocol"
	"github.com/stampzilla/stampzilla-go/protocol/devices"
)

var VERSION string = "dev"
var BUILD_DATE string = ""

// MAIN - This is run when the init function is done

var notify *notifier.Notify

func main() {
	log.Info("Starting ZWAVE node")

	debug := flag.Bool("v", false, "Verbose - show more debuging info")
	port := flag.String("controllerport", "/dev/ttyACM0", "SerialAPI communication port (to controller)")

	// Parse all commandline arguments, host and port parameters are added in the basenode init function
	flag.Parse()
	logrus.SetLevel(logrus.WarnLevel)
	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	//Get a config with the correct parameters
	config := basenode.NewConfig()

	//Activate the config
	basenode.SetConfig(config)

	z, err := gozwave.Connect(*port, "zwave-networkmap.json")
	if err != nil {
		log.Error(err)
		return
	}

	node := protocol.NewNode("zwave")
	node.Version = VERSION
	node.BuildDate = BUILD_DATE

	//Start communication with the server
	connection := basenode.Connect()
	notify = notifier.New(connection)
	notify.SetSource(node)

	// Thit worker keeps track on our connection state, if we are connected or not
	go monitorState(node, connection)

	state := NewState()
	node.SetState(state)
	state.zwave = z

	// This worker recives all incomming commands
	go serverRecv(node, connection)

	<-time.After(time.Second) // TODO: Wait for node.Uuid_ to be populated

	// Add all existing nodes to the state / device list
	for _, znode := range z.Nodes.All() {
		if znode.Id == 1 {
			continue
		}

		//state.Nodes = append(state.Nodes, newZwavenode(znode))
		state.Nodes[strconv.Itoa(znode.Id)] = newZwavenode(znode)
		n := state.GetNode(znode.Id)
		n.sync(znode)

		addOrUpdateDevice(node, znode)
	}
	connection.Send(node.Node())

	// Listen from events from the zwave-controller
	for {
		select {
		case event := <-z.GetNextEvent():
			log.Infof("Event: %#v", event)
			switch e := event.(type) {
			case events.NodeDiscoverd:
				znode := z.Nodes.Get(e.Address)
				log.Infof("%#v", znode)
				if znode != nil {
					state.Nodes[strconv.Itoa(znode.Id)] = newZwavenode(znode)
				}

			case events.NodeUpdated:
				n := state.GetNode(e.Address)
				if n != nil {
					znode := z.Nodes.Get(e.Address)

					addOrUpdateDevice(node, znode)
					n.sync(znode)
				}
			}

			connection.Send(node.Node())
		}
	}
}

func addOrUpdateDevice(node *protocol.Node, znode *nodes.Node) {
	if znode.Device == nil {
		return
	}

	log.Errorf("Endpoints: %#v", znode.Endpoints)

	for i := 0; i < len(znode.Endpoints); i++ {
		devid := strconv.Itoa(int(znode.Id) + (i * 1000))
		endpoint := ""
		if i > 0 {
			endpoint = strconv.Itoa(i)
		}

		//Dont add if it already exists
		if node.Devices().Exists(devid) {
			return
		}

		switch {
		case znode.HasCommand(commands.SwitchMultilevel):
			node.Devices().Add(&devices.Device{
				Type:   "dimmableLamp",
				Name:   znode.Device.Brand + " - " + znode.Device.Product + " (Address: " + devid + ")",
				Id:     devid,
				Online: true,
				Node:   node.Uuid(),
				StateMap: map[string]string{
					"on":    "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateBool.on" + endpoint,
					"level": "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateFloat.level" + endpoint,
				},
			})
		case znode.HasCommand(commands.SwitchBinary):
			node.Devices().Add(&devices.Device{
				Type:   "lamp",
				Name:   znode.Device.Brand + " - " + znode.Device.Product + " (Address: " + devid + ")",
				Id:     devid,
				Online: true,
				Node:   node.Uuid(),
				StateMap: map[string]string{
					"on": "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateBool.on" + endpoint,
				},
			})
		}
	}

}

// WORKER that monitors the current connection state
func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node.Node())
		case basenode.ConnectionStateDisconnected:
		}
	}
}

// WORKER that recives all incomming commands
func serverRecv(node *protocol.Node, connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(node, connection, d)
	}
}

// THis is called on each incomming command
func processCommand(node *protocol.Node, connection basenode.Connection, cmd protocol.Command) {
	if s, ok := node.State().(*State); ok {
		log.Infof("Incoming command from server: %#v \n", cmd, s)
		if len(cmd.Args) == 0 {
			return
		}

		id, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			log.Error(err)
			return
		}

		var device gozwave.Controllable

		endpoint := int(id / 1000)
		id = id - (endpoint * 1000)

		znode := s.zwave.Nodes.Get(id)
		if znode == nil {
			log.Error("Node not found")
			return
		}

		if id < 1000 && len(znode.Endpoints) < 2 {
			device = znode
		} else {
			device = znode.Endpoint(endpoint)
		}

		switch cmd.Cmd {
		case "on":
			device.On()
		case "off":
			device.Off()
		case "level":
			level, err := strconv.ParseFloat(cmd.Args[1], 64)
			if err != nil {
				log.Error(err)
				return
			}
			device.Level(level)
		default:
			log.Warnf("Unknown command '%s'", cmd.Cmd)
		}
	}
}
<commit_msg>Fixed restore state from zwave-configuration file<commit_after>package main

import (
	"flag"
	"strconv"
	"time"

	"github.com/Sirupsen/logrus"
	log "github.com/cihub/seelog"
	"github.com/stampzilla/gozwave"
	"github.com/stampzilla/gozwave/events"
	"github.com/stampzilla/gozwave/nodes"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/pkg/notifier"
	"github.com/stampzilla/stampzilla-go/protocol"
	"github.com/stampzilla/stampzilla-go/protocol/devices"
)

var VERSION string = "dev"
var BUILD_DATE string = ""

// MAIN - This is run when the init function is done

var notify *notifier.Notify

func main() {
	log.Info("Starting ZWAVE node")

	debug := flag.Bool("v", false, "Verbose - show more debuging info")
	port := flag.String("controllerport", "/dev/ttyACM0", "SerialAPI communication port (to controller)")

	// Parse all commandline arguments, host and port parameters are added in the basenode init function
	flag.Parse()
	logrus.SetLevel(logrus.WarnLevel)
	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	//Get a config with the correct parameters
	config := basenode.NewConfig()

	//Activate the config
	basenode.SetConfig(config)

	z, err := gozwave.Connect(*port, "zwave-networkmap.json")
	if err != nil {
		log.Error(err)
		return
	}

	node := protocol.NewNode("zwave")
	node.Version = VERSION
	node.BuildDate = BUILD_DATE

	//Start communication with the server
	connection := basenode.Connect()
	notify = notifier.New(connection)
	notify.SetSource(node)

	// Thit worker keeps track on our connection state, if we are connected or not
	go monitorState(node, connection)

	state := NewState()
	node.SetState(state)
	state.zwave = z

	// This worker recives all incomming commands
	go serverRecv(node, connection)

	<-time.After(time.Second) // TODO: Wait for node.Uuid_ to be populated

	// Add all existing nodes to the state / device list
	for _, znode := range z.Nodes.All() {
		if znode.Id == 1 {
			continue
		}

		//state.Nodes = append(state.Nodes, newZwavenode(znode))
		state.Nodes[strconv.Itoa(znode.Id)] = newZwavenode(znode)
		n := state.GetNode(znode.Id)
		n.sync(znode)

		addOrUpdateDevice(node, znode)
	}
	connection.Send(node.Node())

	// Listen from events from the zwave-controller
	for {
		select {
		case event := <-z.GetNextEvent():
			log.Infof("Event: %#v", event)
			switch e := event.(type) {
			case events.NodeDiscoverd:
				znode := z.Nodes.Get(e.Address)
				//spew.Dump(znode)
				if znode != nil {
					n := newZwavenode(znode)
					state.Nodes[strconv.Itoa(znode.Id)] = n

					addOrUpdateDevice(node, znode) // Device management
					n.sync(znode)                  // State management
				}

			case events.NodeUpdated:
				n := state.GetNode(e.Address)
				if n != nil {
					znode := z.Nodes.Get(e.Address)

					addOrUpdateDevice(node, znode) // Device management
					n.sync(znode)                  // State management
				}
			}

			connection.Send(node.Node())
		}
	}
}

func addOrUpdateDevice(node *protocol.Node, znode *nodes.Node) {
	if znode.Device == nil {
		return
	}

	log.Errorf("Endpoints: %#v", znode.Endpoints)

	for i := 0; i < len(znode.Endpoints); i++ {
		devid := strconv.Itoa(int(znode.Id) + (i * 1000))
		endpoint := ""
		if i > 0 {
			endpoint = strconv.Itoa(i)
		}

		//Dont add if it already exists
		if node.Devices().Exists(devid) {
			return
		}

		switch {
		case znode.IsDeviceClass(gozwave.GENERIC_TYPE_SWITCH_MULTILEVEL,
			gozwave.SPECIFIC_TYPE_POWER_SWITCH_MULTILEVEL):
			//znode.HasCommand(commands.SwitchMultilevel):
			node.Devices().Add(&devices.Device{
				Type:   "dimmableLamp",
				Name:   znode.Device.Brand + " - " + znode.Device.Product + " (Address: " + devid + ")",
				Id:     devid,
				Online: true,
				Node:   node.Uuid(),
				StateMap: map[string]string{
					"on":    "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateBool.on" + endpoint,
					"level": "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateFloat.level" + endpoint,
				},
			})
		//case znode.HasCommand(commands.SwitchBinary):
		case znode.IsDeviceClass(gozwave.GENERIC_TYPE_SWITCH_BINARY,
			gozwave.SPECIFIC_TYPE_POWER_SWITCH_BINARY):
			node.Devices().Add(&devices.Device{
				Type:   "lamp",
				Name:   znode.Device.Brand + " - " + znode.Device.Product + " (Address: " + devid + ")",
				Id:     devid,
				Online: true,
				Node:   node.Uuid(),
				StateMap: map[string]string{
					"on": "Nodes[" + strconv.Itoa(int(znode.Id)) + "]" + ".stateBool.on" + endpoint,
				},
			})
		}
	}

}

// WORKER that monitors the current connection state
func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node.Node())
		case basenode.ConnectionStateDisconnected:
		}
	}
}

// WORKER that recives all incomming commands
func serverRecv(node *protocol.Node, connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(node, connection, d)
	}
}

// THis is called on each incomming command
func processCommand(node *protocol.Node, connection basenode.Connection, cmd protocol.Command) {
	if s, ok := node.State().(*State); ok {
		log.Infof("Incoming command from server: %#v \n", cmd, s)
		if len(cmd.Args) == 0 {
			return
		}

		id, err := strconv.Atoi(cmd.Args[0])
		if err != nil {
			log.Error(err)
			return
		}

		var device gozwave.Controllable

		endpoint := int(id / 1000)
		id = id - (endpoint * 1000)

		znode := s.zwave.Nodes.Get(id)
		if znode == nil {
			log.Error("Node not found")
			return
		}

		if id < 1000 && len(znode.Endpoints) < 2 {
			device = znode
		} else {
			device = znode.Endpoint(endpoint)
		}

		switch cmd.Cmd {
		case "on":
			device.On()
		case "off":
			device.Off()
		case "level":
			level, err := strconv.ParseFloat(cmd.Args[1], 64)
			if err != nil {
				log.Error(err)
				return
			}
			device.Level(level)
		default:
			log.Warnf("Unknown command '%s'", cmd.Cmd)
		}
	}
}
<|endoftext|>