<commit_before>// Copyright 2020 Asim Aslam
//
// 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.
//
// Original source: github.com/micro/go-micro/v3/api/handler/rpc/stream.go

package rpc

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/micro/micro/v3/service/api"
	"github.com/micro/micro/v3/service/client"
	"github.com/micro/micro/v3/service/errors"
	"github.com/micro/micro/v3/service/logger"
	raw "github.com/micro/micro/v3/util/codec/bytes"
	"github.com/micro/micro/v3/util/router"
)

const (
	// Time allowed to write a message to the client.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the client.
	pongWait = 60 * time.Second

	// Send pings to client with this period. Must be less than pongWait.
	pingPeriod = 15 * time.Second

	// Maximum message size allowed from client.
	maxMessageSize = 512
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func serveStream(ctx context.Context, w http.ResponseWriter, r *http.Request, service *api.Service, c client.Client) {
	// serve as websocket if thats the case
	if isWebSocket(r) {
		serveWebsocket(ctx, w, r, service, c)
		return
	}

	ct := r.Header.Get("Content-Type")
	// Strip charset from Content-Type (like `application/json; charset=UTF-8`)
	if idx := strings.IndexRune(ct, ';'); idx >= 0 {
		ct = ct[:idx]
	}

	payload, err := api.RequestPayload(r)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}
	if len(payload) == 0 {
		// make it valid json
		payload = []byte("{}")
	}

	var request interface{}
	if !bytes.Equal(payload, []byte(`{}`)) {
		switch ct {
		case "application/json", "":
			m := json.RawMessage(payload)
			request = &m
		default:
			request = &raw.Frame{Data: payload}
		}
	}

	// we always need to set content type for message
	if ct == "" {
		ct = "application/json"
	}
	req := c.NewRequest(
		service.Name,
		service.Endpoint.Name,
		request,
		client.WithContentType(ct),
		client.StreamingRequest(),
	)

	w.Header().Set("Content-Type", ct)

	// create custom router
	callOpt := client.WithRouter(router.New(service.Services))

	// create a new stream
	stream, err := c.Stream(ctx, req, callOpt)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		merr, ok := err.(*errors.Error)
		if ok {
			w.WriteHeader(int(merr.Code))
			w.Write([]byte(merr.Error()))
		}
		return
	}
	defer stream.Close()

	// send request even if nil because it triggers the call in case server expects no input
	// without this, we establish a connection but don't kick off the stream of communication
	if err = stream.Send(request); err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		merr, ok := err.(*errors.Error)
		if ok {
			w.WriteHeader(int(merr.Code))
			w.Write([]byte(merr.Error()))
		} else {
			w.WriteHeader(500)
			w.Write([]byte(err.Error()))
		}
		return
	}

	rsp := stream.Response()

	// receive from stream and send to client
	for {
		select {
		case <-ctx.Done():
			return
		case <-stream.Context().Done():
			return
		default:
			// read backend response body
			buf, err := rsp.Read()
			if err != nil {
				// clean exit
				if err == io.EOF {
					return
				}
				// wants to avoid import  grpc/status.Status
				if strings.Contains(err.Error(), "context canceled") {
					return
				}
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
				merr, ok := err.(*errors.Error)
				if ok {
					w.WriteHeader(int(merr.Code))
					w.Write([]byte(merr.Error()))
				}
				return
			}
			// send the buffer
			_, err = fmt.Fprint(w, string(buf))
			if err != nil {
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
			}

			// flush it
			flusher, ok := w.(http.Flusher)
			if ok {
				flusher.Flush()
			}
		}
	}
}

type stream struct {
	// message type requested (binary or text)
	messageType int
	// request context
	ctx context.Context
	// the websocket connection.
	conn *websocket.Conn
	// the downstream connection.
	stream client.Stream
}

func (s *stream) processWSReadsAndWrites() {
	defer func() {
		s.conn.Close()
	}()

	msgs := make(chan []byte)

	stopCtx, cancel := context.WithCancel(context.Background())
	wg := sync.WaitGroup{}
	wg.Add(3)
	go s.rspToBufLoop(cancel, &wg, stopCtx, msgs)
	go s.bufToClientLoop(cancel, &wg, stopCtx, msgs)
	go s.clientToServerLoop(cancel, &wg, stopCtx)
	wg.Wait()
}

func (s *stream) clientToServerLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context) {
	defer func() {
		s.conn.Close()
		cancel()
		wg.Done()
	}()
	s.conn.SetReadLimit(maxMessageSize)
	s.conn.SetReadDeadline(time.Now().Add(pongWait))
	s.conn.SetPongHandler(func(string) error { s.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })

	for {
		select {
		case <-stopCtx.Done():
			return
		default:
		}

		_, msg, err := s.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
			}
			return
		}

		var request interface{}
		switch s.messageType {
		case websocket.TextMessage:
			m := json.RawMessage(msg)
			request = &m
		default:
			request = &raw.Frame{Data: msg}
		}

		if err := s.stream.Send(request); err != nil {
			if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
				logger.Error(err)
			}
			return
		}
	}

}

func (s *stream) rspToBufLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context, msgs chan []byte) {
	defer func() {
		cancel()
		wg.Done()
	}()
	rsp := s.stream.Response()
	for {
		select {
		case <-stopCtx.Done():
			return
		default:
		}
		bytes, err := rsp.Read()
		if err != nil {
			if err == io.EOF {
				// clean exit
				return
			}
			s.conn.WriteMessage(websocket.CloseAbnormalClosure, []byte{})
			return
		}
		select {
		case <-stopCtx.Done():
			return
		case msgs <- bytes:
		}

	}

}

func (s *stream) bufToClientLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context, msgs chan []byte) {
	defer func() {
		cancel()
		wg.Done()
		s.stream.Close()
	}()
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		case <-stopCtx.Done():
			return
		case <-s.ctx.Done():
			return
		case <-s.stream.Context().Done():
			s.conn.WriteMessage(websocket.CloseMessage, []byte{})
			return
		case <-ticker.C:
			s.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := s.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		case msg := <-msgs:
			// read response body
			s.conn.SetWriteDeadline(time.Now().Add(writeWait))
			w, err := s.conn.NextWriter(s.messageType)
			if err != nil {
				return
			}
			if _, err := w.Write(msg); err != nil {
				return
			}
			if err := w.Close(); err != nil {
				return
			}
		}
	}

}

// serveWebsocket will stream rpc back over websockets assuming json
func serveWebsocket(ctx context.Context, w http.ResponseWriter, r *http.Request, service *api.Service, c client.Client) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}

	// determine the content type
	ct := r.Header.Get("Content-Type")
	// strip charset from Content-Type (like `application/json; charset=UTF-8`)
	if idx := strings.IndexRune(ct, ';'); idx >= 0 {
		ct = ct[:idx]
	}
	if len(ct) == 0 {
		ct = "application/json"
	}

	// create stream
	req := c.NewRequest(service.Name, service.Endpoint.Name, nil, client.WithContentType(ct), client.StreamingRequest())
	str, err := c.Stream(ctx, req, client.WithRouter(router.New(service.Services)))
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}

	// determine the message type
	msgType := websocket.BinaryMessage
	if ct == "application/json" {
		msgType = websocket.TextMessage
	}

	s := stream{ctx: ctx, conn: conn, stream: str, messageType: msgType}
	s.processWSReadsAndWrites()
}

func isStream(r *http.Request, srv *api.Service) bool {
	// check if the endpoint supports streaming
	for _, service := range srv.Services {
		for _, ep := range service.Endpoints {
			// skip if it doesn't match the name
			if ep.Name != srv.Endpoint.Name {
				continue
			}
			// matched if the name
			if v := ep.Metadata["stream"]; v == "true" {
				return true
			}
		}
	}

	return false
}

func isWebSocket(r *http.Request) bool {
	contains := func(key, val string) bool {
		vv := strings.Split(r.Header.Get(key), ",")
		for _, v := range vv {
			if val == strings.ToLower(strings.TrimSpace(v)) {
				return true
			}
		}
		return false
	}

	if contains("Connection", "upgrade") && contains("Upgrade", "websocket") {
		return true
	}

	return false
}
<commit_msg>return error json to client and close conns correctly (#1867)<commit_after>// Copyright 2020 Asim Aslam
//
// 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.
//
// Original source: github.com/micro/go-micro/v3/api/handler/rpc/stream.go

package rpc

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/micro/micro/v3/service/api"
	"github.com/micro/micro/v3/service/client"
	"github.com/micro/micro/v3/service/errors"
	"github.com/micro/micro/v3/service/logger"
	raw "github.com/micro/micro/v3/util/codec/bytes"
	"github.com/micro/micro/v3/util/router"
)

const (
	// Time allowed to write a message to the client.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the client.
	pongWait = 60 * time.Second

	// Send pings to client with this period. Must be less than pongWait.
	pingPeriod = 15 * time.Second

	// Maximum message size allowed from client.
	maxMessageSize = 512
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func serveStream(ctx context.Context, w http.ResponseWriter, r *http.Request, service *api.Service, c client.Client) {
	// serve as websocket if thats the case
	if isWebSocket(r) {
		serveWebsocket(ctx, w, r, service, c)
		return
	}

	ct := r.Header.Get("Content-Type")
	// Strip charset from Content-Type (like `application/json; charset=UTF-8`)
	if idx := strings.IndexRune(ct, ';'); idx >= 0 {
		ct = ct[:idx]
	}

	payload, err := api.RequestPayload(r)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}
	if len(payload) == 0 {
		// make it valid json
		payload = []byte("{}")
	}

	var request interface{}
	if !bytes.Equal(payload, []byte(`{}`)) {
		switch ct {
		case "application/json", "":
			m := json.RawMessage(payload)
			request = &m
		default:
			request = &raw.Frame{Data: payload}
		}
	}

	// we always need to set content type for message
	if ct == "" {
		ct = "application/json"
	}
	req := c.NewRequest(
		service.Name,
		service.Endpoint.Name,
		request,
		client.WithContentType(ct),
		client.StreamingRequest(),
	)

	w.Header().Set("Content-Type", ct)

	// create custom router
	callOpt := client.WithRouter(router.New(service.Services))

	// create a new stream
	stream, err := c.Stream(ctx, req, callOpt)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		merr, ok := err.(*errors.Error)
		if ok {
			w.WriteHeader(int(merr.Code))
			w.Write([]byte(merr.Error()))
		}
		return
	}
	defer stream.Close()

	// send request even if nil because it triggers the call in case server expects no input
	// without this, we establish a connection but don't kick off the stream of communication
	if err = stream.Send(request); err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		merr, ok := err.(*errors.Error)
		if ok {
			w.WriteHeader(int(merr.Code))
			w.Write([]byte(merr.Error()))
		} else {
			w.WriteHeader(500)
			w.Write([]byte(err.Error()))
		}
		return
	}

	rsp := stream.Response()

	// receive from stream and send to client
	for {
		select {
		case <-ctx.Done():
			return
		case <-stream.Context().Done():
			return
		default:
			// read backend response body
			buf, err := rsp.Read()
			if err != nil {
				// clean exit
				if err == io.EOF {
					return
				}
				// wants to avoid import  grpc/status.Status
				if strings.Contains(err.Error(), "context canceled") {
					return
				}
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
				merr, ok := err.(*errors.Error)
				if ok {
					w.WriteHeader(int(merr.Code))
					w.Write([]byte(merr.Error()))
				}
				return
			}
			// send the buffer
			_, err = fmt.Fprint(w, string(buf))
			if err != nil {
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
			}

			// flush it
			flusher, ok := w.(http.Flusher)
			if ok {
				flusher.Flush()
			}
		}
	}
}

type stream struct {
	// message type requested (binary or text)
	messageType int
	// request context
	ctx context.Context
	// the websocket connection.
	conn *websocket.Conn
	// the downstream connection.
	stream client.Stream
}

func (s *stream) processWSReadsAndWrites() {
	defer func() {
		s.conn.Close()
	}()

	msgs := make(chan []byte)

	stopCtx, cancel := context.WithCancel(context.Background())
	wg := sync.WaitGroup{}
	wg.Add(3)
	go s.rspToBufLoop(cancel, &wg, stopCtx, msgs)
	go s.bufToClientLoop(cancel, &wg, stopCtx, msgs)
	go s.clientToServerLoop(cancel, &wg, stopCtx)
	wg.Wait()
}

func (s *stream) clientToServerLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context) {
	defer func() {
		s.stream.Close()
		cancel()
		wg.Done()
	}()
	s.conn.SetReadLimit(maxMessageSize)
	s.conn.SetReadDeadline(time.Now().Add(pongWait))
	s.conn.SetPongHandler(func(string) error { s.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })

	for {
		select {
		case <-stopCtx.Done():
			return
		default:
		}

		_, msg, err := s.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
					logger.Error(err)
				}
			}
			return
		}

		var request interface{}
		switch s.messageType {
		case websocket.TextMessage:
			m := json.RawMessage(msg)
			request = &m
		default:
			request = &raw.Frame{Data: msg}
		}

		if err := s.stream.Send(request); err != nil {
			if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
				logger.Error(err)
			}
			return
		}
	}

}

func (s *stream) rspToBufLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context, msgs chan []byte) {
	defer func() {
		cancel()
		wg.Done()
	}()
	rsp := s.stream.Response()
	for {
		select {
		case <-stopCtx.Done():
			return
		default:
		}
		bytes, err := rsp.Read()
		if err != nil {
			if err == io.EOF {
				// clean exit
				return
			}
			// write error then close the connection
			b, _ := json.Marshal(err)
			s.conn.WriteMessage(s.messageType, b)
			s.conn.WriteMessage(websocket.CloseAbnormalClosure, []byte{})
			return
		}
		select {
		case <-stopCtx.Done():
			return
		case msgs <- bytes:
		}

	}

}

func (s *stream) bufToClientLoop(cancel context.CancelFunc, wg *sync.WaitGroup, stopCtx context.Context, msgs chan []byte) {
	defer func() {
		s.conn.Close()
		cancel()
		wg.Done()

	}()
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		case <-stopCtx.Done():
			return
		case <-s.ctx.Done():
			return
		case <-s.stream.Context().Done():
			s.conn.WriteMessage(websocket.CloseMessage, []byte{})
			return
		case <-ticker.C:
			s.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := s.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		case msg := <-msgs:
			// read response body
			s.conn.SetWriteDeadline(time.Now().Add(writeWait))
			w, err := s.conn.NextWriter(s.messageType)
			if err != nil {
				return
			}
			if _, err := w.Write(msg); err != nil {
				return
			}
			if err := w.Close(); err != nil {
				return
			}
		}
	}

}

// serveWebsocket will stream rpc back over websockets assuming json
func serveWebsocket(ctx context.Context, w http.ResponseWriter, r *http.Request, service *api.Service, c client.Client) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}

	// determine the content type
	ct := r.Header.Get("Content-Type")
	// strip charset from Content-Type (like `application/json; charset=UTF-8`)
	if idx := strings.IndexRune(ct, ';'); idx >= 0 {
		ct = ct[:idx]
	}
	if len(ct) == 0 {
		ct = "application/json"
	}

	// create stream
	req := c.NewRequest(service.Name, service.Endpoint.Name, nil, client.WithContentType(ct), client.StreamingRequest())
	str, err := c.Stream(ctx, req, client.WithRouter(router.New(service.Services)))
	if err != nil {
		if logger.V(logger.ErrorLevel, logger.DefaultLogger) {
			logger.Error(err)
		}
		return
	}

	// determine the message type
	msgType := websocket.BinaryMessage
	if ct == "application/json" {
		msgType = websocket.TextMessage
	}

	s := stream{ctx: ctx, conn: conn, stream: str, messageType: msgType}
	s.processWSReadsAndWrites()
}

func isStream(r *http.Request, srv *api.Service) bool {
	// check if the endpoint supports streaming
	for _, service := range srv.Services {
		for _, ep := range service.Endpoints {
			// skip if it doesn't match the name
			if ep.Name != srv.Endpoint.Name {
				continue
			}
			// matched if the name
			if v := ep.Metadata["stream"]; v == "true" {
				return true
			}
		}
	}

	return false
}

func isWebSocket(r *http.Request) bool {
	contains := func(key, val string) bool {
		vv := strings.Split(r.Header.Get(key), ",")
		for _, v := range vv {
			if val == strings.ToLower(strings.TrimSpace(v)) {
				return true
			}
		}
		return false
	}

	if contains("Connection", "upgrade") && contains("Upgrade", "websocket") {
		return true
	}

	return false
}
<|endoftext|>


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

package parserprovider

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/knadh/koanf"
	"github.com/knadh/koanf/providers/confmap"
	"github.com/magiconair/properties"

	"go.opentelemetry.io/collector/config/configparser"
)

type setFlagProvider struct {
	base ParserProvider
}

// NewSetFlag returns a config.ParserProvider, that wraps a "base" config.ParserProvider, then
// overrides properties from set flag(s) in the loaded Parser.
//
// The implementation reads set flag(s) from the cmd and concatenates them as a "properties" file.
// Then the properties file is read and properties are set to the loaded Parser.
func NewSetFlag(base ParserProvider) ParserProvider {
	return &setFlagProvider{
		base: base,
	}
}

func (sfl *setFlagProvider) Get() (*configparser.Parser, error) {
	flagProperties := getSetFlag()
	if len(flagProperties) == 0 {
		return sfl.base.Get()
	}

	b := &bytes.Buffer{}
	for _, property := range flagProperties {
		property = strings.TrimSpace(property)
		if _, err := fmt.Fprintf(b, "%s\n", property); err != nil {
			return nil, err
		}
	}

	var props *properties.Properties
	var err error
	if props, err = properties.Load(b.Bytes(), properties.UTF8); err != nil {
		return nil, err
	}

	// Create a map manually instead of using props.Map() to allow env var expansion
	// as used by original Viper-based configparser.Parser.
	parsed := map[string]interface{}{}
	for _, key := range props.Keys() {
		value, _ := props.Get(key)
		parsed[key] = value
	}

	propertyKoanf := koanf.New(".")
	if err = propertyKoanf.Load(confmap.Provider(parsed, "."), nil); err != nil {
		return nil, fmt.Errorf("failed to read set flag config: %v", err)
	}

	var cp *configparser.Parser
	if cp, err = sfl.base.Get(); err != nil {
		return nil, err
	}

	cp.MergeStringMap(propertyKoanf.Raw())

	return cp, nil
}
<commit_msg>Small nit in setflag ParserProvider (#3807)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package parserprovider

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/knadh/koanf/maps"
	"github.com/magiconair/properties"

	"go.opentelemetry.io/collector/config/configparser"
)

type setFlagProvider struct {
	base ParserProvider
}

// NewSetFlag returns a config.ParserProvider, that wraps a "base" config.ParserProvider, then
// overrides properties from set flag(s) in the loaded Parser.
//
// The implementation reads set flag(s) from the cmd and concatenates them as a "properties" file.
// Then the properties file is read and properties are set to the loaded Parser.
func NewSetFlag(base ParserProvider) ParserProvider {
	return &setFlagProvider{
		base: base,
	}
}

func (sfl *setFlagProvider) Get() (*configparser.Parser, error) {
	flagProperties := getSetFlag()
	if len(flagProperties) == 0 {
		return sfl.base.Get()
	}

	b := &bytes.Buffer{}
	for _, property := range flagProperties {
		property = strings.TrimSpace(property)
		if _, err := fmt.Fprintf(b, "%s\n", property); err != nil {
			return nil, err
		}
	}

	var props *properties.Properties
	var err error
	if props, err = properties.Load(b.Bytes(), properties.UTF8); err != nil {
		return nil, err
	}

	// Create a map manually instead of using props.Map() to allow env var expansion
	// as used by original Viper-based configparser.Parser.
	parsed := make(map[string]interface{}, props.Len())
	for _, key := range props.Keys() {
		value, _ := props.Get(key)
		parsed[key] = value
	}
	prop := maps.Unflatten(parsed, ".")

	var cp *configparser.Parser
	if cp, err = sfl.base.Get(); err != nil {
		return nil, err
	}
	return cp, cp.MergeStringMap(prop)
}
<|endoftext|>


<commit_before>package gorange

import (
	"fmt"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/karrick/goswarm"
)

// CachingClient memoizes responses from a Querier.
type CachingClient struct {
	config           cachingClientConfig
	cache            *goswarm.Simple
	lastRequestTimes *goswarm.Simple
	version          int64

	// handle safe shutdowns
	closeError chan error
	halt       chan struct{}
}

type cachingClientConfig struct {
	querier Querier
	stale   time.Duration // prune periodicity
	expiry  time.Duration // drop keys older than

	// when non-zero, check %version and drop expired items or refresh stale items
	checkVersionPeriodicity time.Duration

	// when non-zero, periodically garbage collect expired items
	gcPeriodicity time.Duration
}

// NewCachingClient returns CachingClient that attempts to respond to Query methods by consulting
// its TTL cache, then directing the call to the underlying Querier if a valid response is not
// stored.
func newCachingClient(config *cachingClientConfig) (*CachingClient, error) {
	if config == nil {
		config = &cachingClientConfig{}
	}
	if config.querier == nil {
		return nil, fmt.Errorf("cannot create CachingClient without querier")
	}
	if config.stale < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative stale duration: %v", config.stale)
	}
	if config.expiry < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative expiry duration: %v", config.expiry)
	}
	if config.checkVersionPeriodicity < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative checkVersionPeriodicity duration: %v", config.checkVersionPeriodicity)
	}
	if config.gcPeriodicity < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative gcPeriodicity duration: %v", config.gcPeriodicity)
	}
	// when good config, go ahead and create instance
	badStaleDuration := 1 * time.Minute
	badExpiryDuration := 5 * time.Minute

	// nil config implies treat like a conventional map used for concurrent access: values never go stale, never expire
	lastRequestTimes, err := goswarm.NewSimple(nil)
	if err != nil {
		return nil, err
	}
	cache, err := goswarm.NewSimple(&goswarm.Config{
		GoodStaleDuration:  config.stale,
		GoodExpiryDuration: config.expiry,
		BadStaleDuration:   badStaleDuration,
		BadExpiryDuration:  badExpiryDuration,
		GCPeriodicity:      config.gcPeriodicity,
		Lookup: func(url string) (interface{}, error) {
			results, err := config.querier.Query(url)
			if err != nil {
				if _, ok := err.(ErrRangeException); ok {
					now := time.Now()
					// RoundRobin stops looking for a non-error value, because
					// we send it an error as the value.
					return goswarm.TimedValue{
						Value:  nil,
						Err:    err,
						Stale:  now.Add(badStaleDuration),
						Expiry: now.Add(badExpiryDuration),
					}, nil
				}
				// Return the non-RangeException error so When more than one server,
				// RoundRobin will continue looking for a non-error value
				return nil, err
			}
			return results, nil
		},
	})
	if err != nil {
		return nil, err
	}
	c := &CachingClient{
		cache:            cache,
		config:           *config,
		lastRequestTimes: lastRequestTimes,
	}
	c.halt = make(chan struct{})
	c.closeError = make(chan error)
	go c.run()
	return c, nil
}

// Close releases all memory and go-routines used by the Simple swarm. If during instantiation,
// checkVersionPeriodicty was greater than the zero-value for time.Duration, this method may block
// while completing any in progress updates due to `%version` changes.
func (c *CachingClient) Close() error {
	close(c.halt)

	// Wait for run() loop to acknowledge signal that it's complete
	<-c.closeError // current run method always returns nil error

	return c.cache.Close()
}

// Query returns the response of the query, first checking in the TTL cache, then by actually
// invoking the Query method on the underlying Querier.
func (c *CachingClient) Query(query string) ([]string, error) {
	c.lastRequestTimes.Store(query, time.Now())
	raw, err := c.cache.Query(query)
	if err != nil {
		return nil, err
	}
	results, ok := raw.([]string)
	if !ok {
		return nil, fmt.Errorf("cannot convert %T to []string", raw)
	}
	return results, nil
}

func (c CachingClient) getLastRequestTime(key string) time.Time {
	lrt, ok := c.lastRequestTimes.Load(key)
	if !ok {
		panic(fmt.Errorf("SHOULD NEVER FIND KEY IN cache BUT NOT IN lastResponseTimes: %q", key))
	}
	return lrt.(time.Time)
}

func (c *CachingClient) refreshBasedOnVersion() error {
	results, err := c.config.querier.Query("%version")
	if err != nil {
		return err
	}
	if len(results) != 1 {
		return fmt.Errorf("%%version returned %d output lines; expected 1 line", len(results))
	}
	// version is an epoch timestamp
	version, err := strconv.ParseInt(results[0], 10, 64)
	if err != nil {
		return err
	}
	if version > c.version {
		c.refreshBefore(time.Unix(version, 0).Add(-c.config.stale))
		c.version = version
	}
	return nil
}

func (c *CachingClient) refreshBefore(cutoff time.Time) {
	// log.Printf("refreshBefore(%d)", cutoff.Unix())

	// To prevent overloading the range server with refresh requests for lots of keys at once,
	// trickle them in one-by-one.
	toRefresh := make(chan string, 64) // WARNING: must be at least 1 to prevent Range callback from dead locking
	var refresher sync.WaitGroup
	refresher.Add(1)
	go func() {
		for key := range toRefresh {
			c.cache.Update(key)
		}
		refresher.Done()
	}()

	// Go maps and goswarm.Sample's Range method allows deleting keys while iterating over the
	// map's key-value pairs. We'll use that to our advantage below.
	c.cache.Range(func(key string, tv *goswarm.TimedValue) {
		if tv.Err != nil {
			log.Printf("deleting result that is an error: %q", key)
			c.cache.Delete(key)
		} else if c.getLastRequestTime(key).Before(cutoff) {
			log.Printf("dropping because last requested quite a while ago: %q", key)
			c.cache.Delete(key)
		} else {
			log.Printf("enqueue request to update: %q", key)
			toRefresh <- key
		}
	})
	close(toRefresh)
	refresher.Wait()
}

func (c *CachingClient) run() {
	// If param is 0, client does not want to use the feature, so make it a very long
	// periodicity, and when the select case is chosen, skip calling the feature.
	checkVersionPeriodicity := c.config.checkVersionPeriodicity
	if checkVersionPeriodicity == 0 {
		checkVersionPeriodicity = 24 * time.Hour
	}
	stale := c.config.stale
	if stale == 0 {
		stale = 24 * time.Hour
	}

	for {
		select {
		case <-time.After(checkVersionPeriodicity):
			if c.config.checkVersionPeriodicity > 0 {
				_ = c.refreshBasedOnVersion() // ignoring error return value
			}
		case <-time.After(stale):
			if c.config.stale > 0 {
				c.refreshBefore(time.Now().Add(-c.config.expiry))
			}
		case <-c.halt:
			c.closeError <- nil
			// there is no cleanup required, so we just return
			return
		}
	}
}
<commit_msg>commented out debugging lines, doh!<commit_after>package gorange

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/karrick/goswarm"
)

// CachingClient memoizes responses from a Querier.
type CachingClient struct {
	config           cachingClientConfig
	cache            *goswarm.Simple
	lastRequestTimes *goswarm.Simple
	version          int64

	// handle safe shutdowns
	closeError chan error
	halt       chan struct{}
}

type cachingClientConfig struct {
	querier Querier
	stale   time.Duration // prune periodicity
	expiry  time.Duration // drop keys older than

	// when non-zero, check %version and drop expired items or refresh stale items
	checkVersionPeriodicity time.Duration

	// when non-zero, periodically garbage collect expired items
	gcPeriodicity time.Duration
}

// NewCachingClient returns CachingClient that attempts to respond to Query methods by consulting
// its TTL cache, then directing the call to the underlying Querier if a valid response is not
// stored.
func newCachingClient(config *cachingClientConfig) (*CachingClient, error) {
	if config == nil {
		config = &cachingClientConfig{}
	}
	if config.querier == nil {
		return nil, fmt.Errorf("cannot create CachingClient without querier")
	}
	if config.stale < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative stale duration: %v", config.stale)
	}
	if config.expiry < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative expiry duration: %v", config.expiry)
	}
	if config.checkVersionPeriodicity < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative checkVersionPeriodicity duration: %v", config.checkVersionPeriodicity)
	}
	if config.gcPeriodicity < 0 {
		return nil, fmt.Errorf("cannot create CachingClient with negative gcPeriodicity duration: %v", config.gcPeriodicity)
	}
	// when good config, go ahead and create instance
	badStaleDuration := 1 * time.Minute
	badExpiryDuration := 5 * time.Minute

	// nil config implies treat like a conventional map used for concurrent access: values never go stale, never expire
	lastRequestTimes, err := goswarm.NewSimple(nil)
	if err != nil {
		return nil, err
	}
	cache, err := goswarm.NewSimple(&goswarm.Config{
		GoodStaleDuration:  config.stale,
		GoodExpiryDuration: config.expiry,
		BadStaleDuration:   badStaleDuration,
		BadExpiryDuration:  badExpiryDuration,
		GCPeriodicity:      config.gcPeriodicity,
		Lookup: func(url string) (interface{}, error) {
			results, err := config.querier.Query(url)
			if err != nil {
				if _, ok := err.(ErrRangeException); ok {
					now := time.Now()
					// RoundRobin stops looking for a non-error value, because
					// we send it an error as the value.
					return goswarm.TimedValue{
						Value:  nil,
						Err:    err,
						Stale:  now.Add(badStaleDuration),
						Expiry: now.Add(badExpiryDuration),
					}, nil
				}
				// Return the non-RangeException error so When more than one server,
				// RoundRobin will continue looking for a non-error value
				return nil, err
			}
			return results, nil
		},
	})
	if err != nil {
		return nil, err
	}
	c := &CachingClient{
		cache:            cache,
		config:           *config,
		lastRequestTimes: lastRequestTimes,
	}
	c.halt = make(chan struct{})
	c.closeError = make(chan error)
	go c.run()
	return c, nil
}

// Close releases all memory and go-routines used by the Simple swarm. If during instantiation,
// checkVersionPeriodicty was greater than the zero-value for time.Duration, this method may block
// while completing any in progress updates due to `%version` changes.
func (c *CachingClient) Close() error {
	close(c.halt)

	// Wait for run() loop to acknowledge signal that it's complete
	<-c.closeError // current run method always returns nil error

	return c.cache.Close()
}

// Query returns the response of the query, first checking in the TTL cache, then by actually
// invoking the Query method on the underlying Querier.
func (c *CachingClient) Query(query string) ([]string, error) {
	c.lastRequestTimes.Store(query, time.Now())
	raw, err := c.cache.Query(query)
	if err != nil {
		return nil, err
	}
	results, ok := raw.([]string)
	if !ok {
		return nil, fmt.Errorf("cannot convert %T to []string", raw)
	}
	return results, nil
}

func (c CachingClient) getLastRequestTime(key string) time.Time {
	lrt, ok := c.lastRequestTimes.Load(key)
	if !ok {
		panic(fmt.Errorf("SHOULD NEVER FIND KEY IN cache BUT NOT IN lastResponseTimes: %q", key))
	}
	return lrt.(time.Time)
}

func (c *CachingClient) refreshBasedOnVersion() error {
	results, err := c.config.querier.Query("%version")
	if err != nil {
		return err
	}
	if len(results) != 1 {
		return fmt.Errorf("%%version returned %d output lines; expected 1 line", len(results))
	}
	// version is an epoch timestamp
	version, err := strconv.ParseInt(results[0], 10, 64)
	if err != nil {
		return err
	}
	if version > c.version {
		c.refreshBefore(time.Unix(version, 0).Add(-c.config.stale))
		c.version = version
	}
	return nil
}

func (c *CachingClient) refreshBefore(cutoff time.Time) {
	// log.Printf("refreshBefore(%d)", cutoff.Unix())

	// To prevent overloading the range server with refresh requests for lots of keys at once,
	// trickle them in one-by-one.
	toRefresh := make(chan string, 64) // WARNING: must be at least 1 to prevent Range callback from dead locking
	var refresher sync.WaitGroup
	refresher.Add(1)
	go func() {
		for key := range toRefresh {
			c.cache.Update(key)
		}
		refresher.Done()
	}()

	// Go maps and goswarm.Sample's Range method allows deleting keys while iterating over the
	// map's key-value pairs. We'll use that to our advantage below.
	c.cache.Range(func(key string, tv *goswarm.TimedValue) {
		if tv.Err != nil {
			// log.Printf("deleting result that is an error: %q", key)
			c.cache.Delete(key)
		} else if c.getLastRequestTime(key).Before(cutoff) {
			// log.Printf("dropping because last requested quite a while ago: %q", key)
			c.cache.Delete(key)
		} else {
			// log.Printf("enqueue request to update: %q", key)
			toRefresh <- key
		}
	})
	close(toRefresh)
	refresher.Wait()
}

func (c *CachingClient) run() {
	// If param is 0, client does not want to use the feature, so make it a very long
	// periodicity, and when the select case is chosen, skip calling the feature.
	checkVersionPeriodicity := c.config.checkVersionPeriodicity
	if checkVersionPeriodicity == 0 {
		checkVersionPeriodicity = 24 * time.Hour
	}
	stale := c.config.stale
	if stale == 0 {
		stale = 24 * time.Hour
	}

	for {
		select {
		case <-time.After(checkVersionPeriodicity):
			if c.config.checkVersionPeriodicity > 0 {
				_ = c.refreshBasedOnVersion() // ignoring error return value
			}
		case <-time.After(stale):
			if c.config.stale > 0 {
				c.refreshBefore(time.Now().Add(-c.config.expiry))
			}
		case <-c.halt:
			c.closeError <- nil
			// there is no cleanup required, so we just return
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build none

package main

import (
	"fmt"
	"log"
	"os"

	"github.com/ethereum/go-ethereum/crypto/secp256k1"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/whisper"
)

func main() {
	logger.AddLogSystem(logger.NewStdLogSystem(os.Stdout, log.LstdFlags, logger.InfoLevel))

	pub, _ := secp256k1.GenerateKeyPair()

	whisper := whisper.New()

	srv := p2p.Server{
		MaxPeers:   10,
		Identity:   p2p.NewSimpleClientIdentity("whisper-go", "1.0", "", string(pub)),
		ListenAddr: ":30300",
		NAT:        p2p.UPNP(),

		Protocols: []p2p.Protocol{whisper.Protocol()},
	}
	if err := srv.Start(); err != nil {
		fmt.Println("could not start server:", err)
		os.Exit(1)
	}

	select {}
}
<commit_msg>whisper: make the test app runnable & do something inside<commit_after>// +build none

// Contains a simple whisper peer setup and self messaging to allow playing
// around with the protocol and API without a fancy client implementation.

package main

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/p2p"
	"github.com/ethereum/go-ethereum/p2p/nat"
	"github.com/ethereum/go-ethereum/whisper"
)

func main() {
	logger.AddLogSystem(logger.NewStdLogSystem(os.Stdout, log.LstdFlags, logger.InfoLevel))

	// Generate the peer identity
	key, err := crypto.GenerateKey()
	if err != nil {
		fmt.Printf("Failed to generate peer key: %v.\n", err)
		os.Exit(-1)
	}
	name := common.MakeName("whisper-go", "1.0")
	shh := whisper.New()

	// Create an Ethereum peer to communicate through
	server := p2p.Server{
		PrivateKey: key,
		MaxPeers:   10,
		Name:       name,
		Protocols:  []p2p.Protocol{shh.Protocol()},
		ListenAddr: ":30300",
		NAT:        nat.Any(),
	}
	fmt.Println("Starting Ethereum peer...")
	if err := server.Start(); err != nil {
		fmt.Printf("Failed to start Ethereum peer: %v.\n", err)
		os.Exit(1)
	}

	// Send a message to self to check that something works
	payload := fmt.Sprintf("Hello world, this is %v. In case you're wondering, the time is %v", name, time.Now())
	if err := selfSend(shh, []byte(payload)); err != nil {
		fmt.Printf("Failed to self message: %v.\n", err)
		os.Exit(-1)
	}
}

// SendSelf wraps a payload into a Whisper envelope and forwards it to itself.
func selfSend(shh *whisper.Whisper, payload []byte) error {
	ok := make(chan struct{})

	// Start watching for self messages, output any arrivals
	id := shh.NewIdentity()
	shh.Watch(whisper.Filter{
		To: &id.PublicKey,
		Fn: func(msg *whisper.Message) {
			fmt.Printf("Message received: %s, signed with 0x%x.\n", string(msg.Payload), msg.Signature)
			close(ok)
		},
	})
	// Wrap the payload and encrypt it
	msg := whisper.NewMessage(payload)
	envelope, err := msg.Seal(whisper.DefaultPow, whisper.Opts{
		Ttl:  whisper.DefaultTtl,
		From: id,
		To:   &id.PublicKey,
	})
	if err != nil {
		return fmt.Errorf("failed to seal message: %v", err)
	}
	// Dump the message into the system and wait for it to pop back out
	if err := shh.Send(envelope); err != nil {
		return fmt.Errorf("failed to send self-message: %v", err)
	}
	select {
	case <-ok:
	case <-time.After(time.Second):
		return fmt.Errorf("failed to receive message in time")
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Brady Catherman
//
// 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.

// build (linux OR darwin)

package testlib

import (
	"syscall"
)

func testSyscallExit() {
	syscall.ExitProcess(1)
}
<commit_msg>Remove windows_test.go which didnt work anyway.<commit_after><|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"bytes"
	"sort"
	"time"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

var (
	_ currentStakerChainState = &currentStakerChainStateImpl{}
)

type currentStakerChainState interface {
	GetNextStaker() (addStakerTx *Tx, potentialReward uint64, err error)
	GetValidator(nodeID ids.ShortID) (currentValidator, error)

	UpdateStakers(
		addValidators []*validatorReward,
		addDelegators []*validatorReward,
		addSubnetValidators []*Tx,
		numTxsToRemove int,
	) (currentStakerChainState, error)
	DeleteNextStaker() (currentStakerChainState, error)

	Stakers() []*Tx // Sorted in removal order

	Apply(internalState)

	ValidatorSet(subnetID ids.ID) (validators.Set, error)
}

type currentStakerChainStateImpl struct {
	nextStaker *validatorReward

	// nodeID -> validator
	validatorsByNodeID map[ids.ShortID]*currentValidatorImpl

	// txID -> tx
	validatorsByTxID map[ids.ID]*validatorReward

	// list of current validators in order of their removal
	validators []*Tx

	addedStakers   []*validatorReward
	deletedStakers []*Tx
}

type validatorReward struct {
	addStakerTx     *Tx
	potentialReward uint64
}

func (cs *currentStakerChainStateImpl) GetNextStaker() (addStakerTx *Tx, potentialReward uint64, err error) {
	if cs.nextStaker == nil {
		return nil, 0, database.ErrNotFound
	}
	return cs.nextStaker.addStakerTx, cs.nextStaker.potentialReward, nil
}

func (cs *currentStakerChainStateImpl) GetValidator(nodeID ids.ShortID) (currentValidator, error) {
	vdr, exists := cs.validatorsByNodeID[nodeID]
	if !exists {
		return nil, database.ErrNotFound
	}
	return vdr, nil
}

func (cs *currentStakerChainStateImpl) UpdateStakers(
	addValidatorTxs []*validatorReward,
	addDelegatorTxs []*validatorReward,
	addSubnetValidatorTxs []*Tx,
	numTxsToRemove int,
) (currentStakerChainState, error) {
	newCS := &currentStakerChainStateImpl{
		validatorsByNodeID: make(map[ids.ShortID]*currentValidatorImpl, len(cs.validatorsByNodeID)+len(addValidatorTxs)),
		validatorsByTxID:   make(map[ids.ID]*validatorReward, len(cs.validatorsByTxID)+len(addValidatorTxs)+len(addDelegatorTxs)+len(addSubnetValidatorTxs)),
		validators:         cs.validators[numTxsToRemove:], // sorted in order of removal

		addedStakers:   append(addValidatorTxs, addDelegatorTxs...),
		deletedStakers: cs.validators[:numTxsToRemove],
	}

	for nodeID, vdr := range cs.validatorsByNodeID {
		newCS.validatorsByNodeID[nodeID] = vdr
	}

	for txID, vdr := range cs.validatorsByTxID {
		newCS.validatorsByTxID[txID] = vdr
	}

	if numAdded := len(addValidatorTxs) + len(addDelegatorTxs) + len(addSubnetValidatorTxs); numAdded != 0 {
		numCurrent := len(newCS.validators)
		newSize := numCurrent + numAdded
		newValidators := make([]*Tx, newSize)
		copy(newValidators, newCS.validators)
		copy(newValidators[numCurrent:], addSubnetValidatorTxs)

		numStart := numCurrent + len(addSubnetValidatorTxs)
		for i, tx := range addValidatorTxs {
			newValidators[numStart+i] = tx.addStakerTx
		}

		numStart = numCurrent + len(addSubnetValidatorTxs) + len(addValidatorTxs)
		for i, tx := range addDelegatorTxs {
			newValidators[numStart+i] = tx.addStakerTx
		}

		sortValidatorsByRemoval(newValidators)
		newCS.validators = newValidators

		for _, vdr := range addValidatorTxs {
			switch tx := vdr.addStakerTx.UnsignedTx.(type) {
			case *UnsignedAddValidatorTx:
				newCS.validatorsByNodeID[tx.Validator.NodeID] = &currentValidatorImpl{
					addValidatorTx:  tx,
					potentialReward: vdr.potentialReward,
				}
				newCS.validatorsByTxID[vdr.addStakerTx.ID()] = vdr
			default:
				return nil, errWrongTxType
			}
		}

		for _, vdr := range addDelegatorTxs {
			switch tx := vdr.addStakerTx.UnsignedTx.(type) {
			case *UnsignedAddDelegatorTx:
				oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
				newVdr := *oldVdr
				newVdr.delegators = make([]*UnsignedAddDelegatorTx, len(oldVdr.delegators)+1)
				copy(newVdr.delegators, oldVdr.delegators)
				newVdr.delegators[len(oldVdr.delegators)] = tx
				sortDelegatorsByRemoval(newVdr.delegators)
				newVdr.delegatorWeight += tx.Validator.Wght
				newCS.validatorsByNodeID[tx.Validator.NodeID] = &newVdr
				newCS.validatorsByTxID[vdr.addStakerTx.ID()] = vdr
			default:
				return nil, errWrongTxType
			}
		}

		for _, vdr := range addSubnetValidatorTxs {
			switch tx := vdr.UnsignedTx.(type) {
			case *UnsignedAddSubnetValidatorTx:
				oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
				newVdr := *oldVdr
				newVdr.subnets = make(map[ids.ID]*UnsignedAddSubnetValidatorTx, len(oldVdr.subnets)-1)
				for subnetID, addTx := range oldVdr.subnets {
					newVdr.subnets[subnetID] = addTx
				}
				newVdr.subnets[tx.Validator.Subnet] = tx
			default:
				return nil, errWrongTxType
			}

			wrappedTx := &validatorReward{
				addStakerTx: vdr,
			}
			newCS.validatorsByTxID[vdr.ID()] = wrappedTx
			newCS.addedStakers = append(newCS.addedStakers, wrappedTx)
		}
	}

	for i := 0; i < numTxsToRemove; i++ {
		removed := cs.validators[i]
		removedID := removed.ID()
		delete(newCS.validatorsByTxID, removedID)

		switch tx := removed.UnsignedTx.(type) {
		case *UnsignedAddSubnetValidatorTx:
			oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
			newVdr := *oldVdr
			newVdr.subnets = make(map[ids.ID]*UnsignedAddSubnetValidatorTx, len(oldVdr.subnets)-1)
			for subnetID, addTx := range oldVdr.subnets {
				if removedID != addTx.ID() {
					newVdr.subnets[subnetID] = addTx
				}
			}
			newCS.validatorsByNodeID[tx.Validator.NodeID] = &newVdr
		default:
			return nil, errWrongTxType
		}
	}

	newCS.setNextStaker()
	return newCS, nil
}

func (cs *currentStakerChainStateImpl) DeleteNextStaker() (currentStakerChainState, error) {
	removedTx, _, err := cs.GetNextStaker()
	if err != nil {
		return nil, err
	}
	removedTxID := removedTx.ID()

	newCS := &currentStakerChainStateImpl{
		validatorsByNodeID: make(map[ids.ShortID]*currentValidatorImpl, len(cs.validatorsByNodeID)),
		validatorsByTxID:   make(map[ids.ID]*validatorReward, len(cs.validatorsByTxID)-1),
		validators:         cs.validators[1:], // sorted in order of removal

		deletedStakers: []*Tx{removedTx},
	}

	switch tx := removedTx.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		for nodeID, vdr := range cs.validatorsByNodeID {
			if nodeID != tx.Validator.NodeID {
				newCS.validatorsByNodeID[nodeID] = vdr
			}
		}
	case *UnsignedAddDelegatorTx:
		for nodeID, vdr := range cs.validatorsByNodeID {
			if nodeID != tx.Validator.NodeID {
				newCS.validatorsByNodeID[nodeID] = vdr
			} else {
				newCS.validatorsByNodeID[nodeID] = &currentValidatorImpl{
					validatorImpl: validatorImpl{
						delegators: vdr.delegators[1:], // sorted in order of removal
						subnets:    vdr.subnets,
					},
					addValidatorTx:  vdr.addValidatorTx,
					delegatorWeight: vdr.delegatorWeight - tx.Validator.Wght,
					potentialReward: vdr.potentialReward,
				}
			}
		}
	default:
		return nil, errWrongTxType
	}

	for txID, vdr := range cs.validatorsByTxID {
		if txID != removedTxID {
			newCS.validatorsByTxID[txID] = vdr
		}
	}

	newCS.setNextStaker()
	return newCS, nil
}

func (cs *currentStakerChainStateImpl) Stakers() []*Tx {
	return cs.validators
}

func (cs *currentStakerChainStateImpl) Apply(is internalState) {
	for _, added := range cs.addedStakers {
		is.AddCurrentStaker(added.addStakerTx, added.potentialReward)
	}
	for _, deleted := range cs.deletedStakers {
		is.DeleteCurrentStaker(deleted)
	}
	is.SetCurrentStakerChainState(cs)

	// Validator changes should only be applied once.
	cs.addedStakers = nil
	cs.deletedStakers = nil
}

func (cs *currentStakerChainStateImpl) ValidatorSet(subnetID ids.ID) (validators.Set, error) {
	if subnetID == constants.PrimaryNetworkID {
		return cs.primaryValidatorSet()
	}
	return cs.subnetValidatorSet(subnetID)
}

func (cs *currentStakerChainStateImpl) primaryValidatorSet() (validators.Set, error) {
	vdrs := validators.NewSet()

	var err error
	for nodeID, vdr := range cs.validatorsByNodeID {
		vdrWeight := vdr.addValidatorTx.Validator.Wght
		vdrWeight, err = safemath.Add64(vdrWeight, vdr.delegatorWeight)
		if err != nil {
			return nil, err
		}
		if err := vdrs.AddWeight(nodeID, vdrWeight); err != nil {
			return nil, err
		}
	}

	return vdrs, nil
}

func (cs *currentStakerChainStateImpl) subnetValidatorSet(subnetID ids.ID) (validators.Set, error) {
	vdrs := validators.NewSet()

	for nodeID, vdr := range cs.validatorsByNodeID {
		subnetVDR, exists := vdr.subnets[subnetID]
		if !exists {
			continue
		}
		if err := vdrs.AddWeight(nodeID, subnetVDR.Validator.Wght); err != nil {
			return nil, err
		}
	}

	return vdrs, nil
}

func (cs *currentStakerChainStateImpl) setNextStaker() {
	for _, tx := range cs.validators {
		switch tx.UnsignedTx.(type) {
		case *UnsignedAddValidatorTx, *UnsignedAddDelegatorTx:
			cs.nextStaker = cs.validatorsByTxID[tx.ID()]
			return
		}
	}
}

type innerSortValidatorsByRemoval []*Tx

func (s innerSortValidatorsByRemoval) Less(i, j int) bool {
	iDel := s[i]
	jDel := s[j]

	var (
		iEndTime  time.Time
		iPriority byte
	)
	switch tx := iDel.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		iEndTime = tx.EndTime()
		iPriority = lowPriority
	case *UnsignedAddDelegatorTx:
		iEndTime = tx.EndTime()
		iPriority = mediumPriority
	case *UnsignedAddSubnetValidatorTx:
		iEndTime = tx.EndTime()
		iPriority = topPriority
	default:
		panic(errWrongTxType)
	}

	var (
		jEndTime  time.Time
		jPriority byte
	)
	switch tx := jDel.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		jEndTime = tx.EndTime()
		jPriority = lowPriority
	case *UnsignedAddDelegatorTx:
		jEndTime = tx.EndTime()
		jPriority = mediumPriority
	case *UnsignedAddSubnetValidatorTx:
		jEndTime = tx.EndTime()
		jPriority = topPriority
	default:
		panic(errWrongTxType)
	}

	if iEndTime.Before(jEndTime) {
		return true
	}
	if jEndTime.Before(iEndTime) {
		return false
	}

	// If the end times are the same, then we sort by the tx type. First we
	// remove UnsignedAddSubnetValidatorTxs, then UnsignedAddDelegatorTx, then
	// UnsignedAddValidatorTx.
	if iPriority > jPriority {
		return true
	}
	if iPriority < jPriority {
		return false
	}

	// If the end times are the same, and the tx types are the same, then we
	// sort by the txID.
	iTxID := iDel.ID()
	jTxID := jDel.ID()
	return bytes.Compare(iTxID[:], jTxID[:]) == -1
}

func (s innerSortValidatorsByRemoval) Len() int {
	return len(s)
}

func (s innerSortValidatorsByRemoval) Swap(i, j int) {
	s[j], s[i] = s[i], s[j]
}

func sortValidatorsByRemoval(s []*Tx) {
	sort.Sort(innerSortValidatorsByRemoval(s))
}

type innerSortDelegatorsByRemoval []*UnsignedAddDelegatorTx

func (s innerSortDelegatorsByRemoval) Less(i, j int) bool {
	iDel := s[i]
	jDel := s[j]

	iEndTime := iDel.EndTime()
	jEndTime := jDel.EndTime()
	if iEndTime.Before(jEndTime) {
		return true
	}
	if jEndTime.Before(iEndTime) {
		return false
	}

	// If the end times are the same, then we sort by the txID
	iTxID := iDel.ID()
	jTxID := jDel.ID()
	return bytes.Compare(iTxID[:], jTxID[:]) == -1
}

func (s innerSortDelegatorsByRemoval) Len() int {
	return len(s)
}

func (s innerSortDelegatorsByRemoval) Swap(i, j int) {
	s[j], s[i] = s[i], s[j]
}

func sortDelegatorsByRemoval(s []*UnsignedAddDelegatorTx) {
	sort.Sort(innerSortDelegatorsByRemoval(s))
}
<commit_msg>fixed map size estimate<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"bytes"
	"sort"
	"time"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

var (
	_ currentStakerChainState = &currentStakerChainStateImpl{}
)

type currentStakerChainState interface {
	GetNextStaker() (addStakerTx *Tx, potentialReward uint64, err error)
	GetValidator(nodeID ids.ShortID) (currentValidator, error)

	UpdateStakers(
		addValidators []*validatorReward,
		addDelegators []*validatorReward,
		addSubnetValidators []*Tx,
		numTxsToRemove int,
	) (currentStakerChainState, error)
	DeleteNextStaker() (currentStakerChainState, error)

	Stakers() []*Tx // Sorted in removal order

	Apply(internalState)

	ValidatorSet(subnetID ids.ID) (validators.Set, error)
}

type currentStakerChainStateImpl struct {
	nextStaker *validatorReward

	// nodeID -> validator
	validatorsByNodeID map[ids.ShortID]*currentValidatorImpl

	// txID -> tx
	validatorsByTxID map[ids.ID]*validatorReward

	// list of current validators in order of their removal
	validators []*Tx

	addedStakers   []*validatorReward
	deletedStakers []*Tx
}

type validatorReward struct {
	addStakerTx     *Tx
	potentialReward uint64
}

func (cs *currentStakerChainStateImpl) GetNextStaker() (addStakerTx *Tx, potentialReward uint64, err error) {
	if cs.nextStaker == nil {
		return nil, 0, database.ErrNotFound
	}
	return cs.nextStaker.addStakerTx, cs.nextStaker.potentialReward, nil
}

func (cs *currentStakerChainStateImpl) GetValidator(nodeID ids.ShortID) (currentValidator, error) {
	vdr, exists := cs.validatorsByNodeID[nodeID]
	if !exists {
		return nil, database.ErrNotFound
	}
	return vdr, nil
}

func (cs *currentStakerChainStateImpl) UpdateStakers(
	addValidatorTxs []*validatorReward,
	addDelegatorTxs []*validatorReward,
	addSubnetValidatorTxs []*Tx,
	numTxsToRemove int,
) (currentStakerChainState, error) {
	newCS := &currentStakerChainStateImpl{
		validatorsByNodeID: make(map[ids.ShortID]*currentValidatorImpl, len(cs.validatorsByNodeID)+len(addValidatorTxs)),
		validatorsByTxID:   make(map[ids.ID]*validatorReward, len(cs.validatorsByTxID)+len(addValidatorTxs)+len(addDelegatorTxs)+len(addSubnetValidatorTxs)),
		validators:         cs.validators[numTxsToRemove:], // sorted in order of removal

		addedStakers:   append(addValidatorTxs, addDelegatorTxs...),
		deletedStakers: cs.validators[:numTxsToRemove],
	}

	for nodeID, vdr := range cs.validatorsByNodeID {
		newCS.validatorsByNodeID[nodeID] = vdr
	}

	for txID, vdr := range cs.validatorsByTxID {
		newCS.validatorsByTxID[txID] = vdr
	}

	if numAdded := len(addValidatorTxs) + len(addDelegatorTxs) + len(addSubnetValidatorTxs); numAdded != 0 {
		numCurrent := len(newCS.validators)
		newSize := numCurrent + numAdded
		newValidators := make([]*Tx, newSize)
		copy(newValidators, newCS.validators)
		copy(newValidators[numCurrent:], addSubnetValidatorTxs)

		numStart := numCurrent + len(addSubnetValidatorTxs)
		for i, tx := range addValidatorTxs {
			newValidators[numStart+i] = tx.addStakerTx
		}

		numStart = numCurrent + len(addSubnetValidatorTxs) + len(addValidatorTxs)
		for i, tx := range addDelegatorTxs {
			newValidators[numStart+i] = tx.addStakerTx
		}

		sortValidatorsByRemoval(newValidators)
		newCS.validators = newValidators

		for _, vdr := range addValidatorTxs {
			switch tx := vdr.addStakerTx.UnsignedTx.(type) {
			case *UnsignedAddValidatorTx:
				newCS.validatorsByNodeID[tx.Validator.NodeID] = &currentValidatorImpl{
					addValidatorTx:  tx,
					potentialReward: vdr.potentialReward,
				}
				newCS.validatorsByTxID[vdr.addStakerTx.ID()] = vdr
			default:
				return nil, errWrongTxType
			}
		}

		for _, vdr := range addDelegatorTxs {
			switch tx := vdr.addStakerTx.UnsignedTx.(type) {
			case *UnsignedAddDelegatorTx:
				oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
				newVdr := *oldVdr
				newVdr.delegators = make([]*UnsignedAddDelegatorTx, len(oldVdr.delegators)+1)
				copy(newVdr.delegators, oldVdr.delegators)
				newVdr.delegators[len(oldVdr.delegators)] = tx
				sortDelegatorsByRemoval(newVdr.delegators)
				newVdr.delegatorWeight += tx.Validator.Wght
				newCS.validatorsByNodeID[tx.Validator.NodeID] = &newVdr
				newCS.validatorsByTxID[vdr.addStakerTx.ID()] = vdr
			default:
				return nil, errWrongTxType
			}
		}

		for _, vdr := range addSubnetValidatorTxs {
			switch tx := vdr.UnsignedTx.(type) {
			case *UnsignedAddSubnetValidatorTx:
				oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
				newVdr := *oldVdr
				newVdr.subnets = make(map[ids.ID]*UnsignedAddSubnetValidatorTx, len(oldVdr.subnets)+1)
				for subnetID, addTx := range oldVdr.subnets {
					newVdr.subnets[subnetID] = addTx
				}
				newVdr.subnets[tx.Validator.Subnet] = tx
			default:
				return nil, errWrongTxType
			}

			wrappedTx := &validatorReward{
				addStakerTx: vdr,
			}
			newCS.validatorsByTxID[vdr.ID()] = wrappedTx
			newCS.addedStakers = append(newCS.addedStakers, wrappedTx)
		}
	}

	for i := 0; i < numTxsToRemove; i++ {
		removed := cs.validators[i]
		removedID := removed.ID()
		delete(newCS.validatorsByTxID, removedID)

		switch tx := removed.UnsignedTx.(type) {
		case *UnsignedAddSubnetValidatorTx:
			oldVdr := newCS.validatorsByNodeID[tx.Validator.NodeID]
			newVdr := *oldVdr
			newVdr.subnets = make(map[ids.ID]*UnsignedAddSubnetValidatorTx, len(oldVdr.subnets)-1)
			for subnetID, addTx := range oldVdr.subnets {
				if removedID != addTx.ID() {
					newVdr.subnets[subnetID] = addTx
				}
			}
			newCS.validatorsByNodeID[tx.Validator.NodeID] = &newVdr
		default:
			return nil, errWrongTxType
		}
	}

	newCS.setNextStaker()
	return newCS, nil
}

func (cs *currentStakerChainStateImpl) DeleteNextStaker() (currentStakerChainState, error) {
	removedTx, _, err := cs.GetNextStaker()
	if err != nil {
		return nil, err
	}
	removedTxID := removedTx.ID()

	newCS := &currentStakerChainStateImpl{
		validatorsByNodeID: make(map[ids.ShortID]*currentValidatorImpl, len(cs.validatorsByNodeID)),
		validatorsByTxID:   make(map[ids.ID]*validatorReward, len(cs.validatorsByTxID)-1),
		validators:         cs.validators[1:], // sorted in order of removal

		deletedStakers: []*Tx{removedTx},
	}

	switch tx := removedTx.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		for nodeID, vdr := range cs.validatorsByNodeID {
			if nodeID != tx.Validator.NodeID {
				newCS.validatorsByNodeID[nodeID] = vdr
			}
		}
	case *UnsignedAddDelegatorTx:
		for nodeID, vdr := range cs.validatorsByNodeID {
			if nodeID != tx.Validator.NodeID {
				newCS.validatorsByNodeID[nodeID] = vdr
			} else {
				newCS.validatorsByNodeID[nodeID] = &currentValidatorImpl{
					validatorImpl: validatorImpl{
						delegators: vdr.delegators[1:], // sorted in order of removal
						subnets:    vdr.subnets,
					},
					addValidatorTx:  vdr.addValidatorTx,
					delegatorWeight: vdr.delegatorWeight - tx.Validator.Wght,
					potentialReward: vdr.potentialReward,
				}
			}
		}
	default:
		return nil, errWrongTxType
	}

	for txID, vdr := range cs.validatorsByTxID {
		if txID != removedTxID {
			newCS.validatorsByTxID[txID] = vdr
		}
	}

	newCS.setNextStaker()
	return newCS, nil
}

func (cs *currentStakerChainStateImpl) Stakers() []*Tx {
	return cs.validators
}

func (cs *currentStakerChainStateImpl) Apply(is internalState) {
	for _, added := range cs.addedStakers {
		is.AddCurrentStaker(added.addStakerTx, added.potentialReward)
	}
	for _, deleted := range cs.deletedStakers {
		is.DeleteCurrentStaker(deleted)
	}
	is.SetCurrentStakerChainState(cs)

	// Validator changes should only be applied once.
	cs.addedStakers = nil
	cs.deletedStakers = nil
}

func (cs *currentStakerChainStateImpl) ValidatorSet(subnetID ids.ID) (validators.Set, error) {
	if subnetID == constants.PrimaryNetworkID {
		return cs.primaryValidatorSet()
	}
	return cs.subnetValidatorSet(subnetID)
}

func (cs *currentStakerChainStateImpl) primaryValidatorSet() (validators.Set, error) {
	vdrs := validators.NewSet()

	var err error
	for nodeID, vdr := range cs.validatorsByNodeID {
		vdrWeight := vdr.addValidatorTx.Validator.Wght
		vdrWeight, err = safemath.Add64(vdrWeight, vdr.delegatorWeight)
		if err != nil {
			return nil, err
		}
		if err := vdrs.AddWeight(nodeID, vdrWeight); err != nil {
			return nil, err
		}
	}

	return vdrs, nil
}

func (cs *currentStakerChainStateImpl) subnetValidatorSet(subnetID ids.ID) (validators.Set, error) {
	vdrs := validators.NewSet()

	for nodeID, vdr := range cs.validatorsByNodeID {
		subnetVDR, exists := vdr.subnets[subnetID]
		if !exists {
			continue
		}
		if err := vdrs.AddWeight(nodeID, subnetVDR.Validator.Wght); err != nil {
			return nil, err
		}
	}

	return vdrs, nil
}

func (cs *currentStakerChainStateImpl) setNextStaker() {
	for _, tx := range cs.validators {
		switch tx.UnsignedTx.(type) {
		case *UnsignedAddValidatorTx, *UnsignedAddDelegatorTx:
			cs.nextStaker = cs.validatorsByTxID[tx.ID()]
			return
		}
	}
}

type innerSortValidatorsByRemoval []*Tx

func (s innerSortValidatorsByRemoval) Less(i, j int) bool {
	iDel := s[i]
	jDel := s[j]

	var (
		iEndTime  time.Time
		iPriority byte
	)
	switch tx := iDel.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		iEndTime = tx.EndTime()
		iPriority = lowPriority
	case *UnsignedAddDelegatorTx:
		iEndTime = tx.EndTime()
		iPriority = mediumPriority
	case *UnsignedAddSubnetValidatorTx:
		iEndTime = tx.EndTime()
		iPriority = topPriority
	default:
		panic(errWrongTxType)
	}

	var (
		jEndTime  time.Time
		jPriority byte
	)
	switch tx := jDel.UnsignedTx.(type) {
	case *UnsignedAddValidatorTx:
		jEndTime = tx.EndTime()
		jPriority = lowPriority
	case *UnsignedAddDelegatorTx:
		jEndTime = tx.EndTime()
		jPriority = mediumPriority
	case *UnsignedAddSubnetValidatorTx:
		jEndTime = tx.EndTime()
		jPriority = topPriority
	default:
		panic(errWrongTxType)
	}

	if iEndTime.Before(jEndTime) {
		return true
	}
	if jEndTime.Before(iEndTime) {
		return false
	}

	// If the end times are the same, then we sort by the tx type. First we
	// remove UnsignedAddSubnetValidatorTxs, then UnsignedAddDelegatorTx, then
	// UnsignedAddValidatorTx.
	if iPriority > jPriority {
		return true
	}
	if iPriority < jPriority {
		return false
	}

	// If the end times are the same, and the tx types are the same, then we
	// sort by the txID.
	iTxID := iDel.ID()
	jTxID := jDel.ID()
	return bytes.Compare(iTxID[:], jTxID[:]) == -1
}

func (s innerSortValidatorsByRemoval) Len() int {
	return len(s)
}

func (s innerSortValidatorsByRemoval) Swap(i, j int) {
	s[j], s[i] = s[i], s[j]
}

func sortValidatorsByRemoval(s []*Tx) {
	sort.Sort(innerSortValidatorsByRemoval(s))
}

type innerSortDelegatorsByRemoval []*UnsignedAddDelegatorTx

func (s innerSortDelegatorsByRemoval) Less(i, j int) bool {
	iDel := s[i]
	jDel := s[j]

	iEndTime := iDel.EndTime()
	jEndTime := jDel.EndTime()
	if iEndTime.Before(jEndTime) {
		return true
	}
	if jEndTime.Before(iEndTime) {
		return false
	}

	// If the end times are the same, then we sort by the txID
	iTxID := iDel.ID()
	jTxID := jDel.ID()
	return bytes.Compare(iTxID[:], jTxID[:]) == -1
}

func (s innerSortDelegatorsByRemoval) Len() int {
	return len(s)
}

func (s innerSortDelegatorsByRemoval) Swap(i, j int) {
	s[j], s[i] = s[i], s[j]
}

func sortDelegatorsByRemoval(s []*UnsignedAddDelegatorTx) {
	sort.Sort(innerSortDelegatorsByRemoval(s))
}
<|endoftext|>


<commit_before>package ss

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

// JSON field prefixes.
const (
	pre2D       = "media_box2d_"
	pre3D       = "media_box3d_"
	preFlyer    = "media_flyer_"
	preDate     = "date_"
	preGenre    = "genres_"
	preSynopsis = "synopsis_"
)

const (
	baseURL      = "http://www.screenscraper.fr/"
	gameInfoPath = "api/jeuInfos.php"
	userInfoPath = "api/ssuserInfos.php"
)

// ErrNotFound is the error returned when a ROM isn't found.
var ErrNotFound = errors.New("not found")

// DevInfo is the information about the developer and used across APIs.
type DevInfo struct {
	ID       string
	Password string
	Name     string
}

// UserInfo is information about the user making the call.
type UserInfo struct {
	ID       string
	Password string
}

type UserInfoResp struct {
	ID              string `xml:"ssuser>id"`
	Level           int    `xml:"ssuser>niveau"`
	Contribution    int    `xml:"ssuser>contribution"`
	UploadedSystems int    `xml:"ssuser>uploadsysteme"`
	UploadedInfo    int    `xml:"ssuser>uploadinfos"`
	ROMsAssociated  int    `xml:"ssuser>romasso"`
	UpdatedMedia    int    `xml:"ssuser>uploadmedia"`
	FavoriteRegion  string `xml:"ssuser>favregion"`
	MaxThreads      int    `xml:"ssuser>maxthreads"`
}

// GameInfoReq is the information we use in the GameInfo command.
type GameInfoReq struct {
	Name    string
	SHA1    string
	RomType string
}

type BoxArt struct {
	Box2D map[string]string `json:"media_boxs2d"`
	Box3D map[string]string `json:"media_boxs3d"`
}

type Media struct {
	Screenshot    string            `json:"media_screenshot"`
	ScreenMarquee string            `json:"media_screenmarquee"`
	Marquee       string            `json:"media_marquee"`
	Video         string            `json:"media_video"`
	Flyers        map[string]string `json:"media_flyers"`
	BoxArt        BoxArt            `json:"media_boxs"`
}

func getPrefix(m map[string]string, pre string) (string, bool) {
	for k, v := range m {
		if strings.HasPrefix(k, pre) {
			return v, true
		}
	}
	return "", false
}

func getSuffix(m map[string]string, pre string, suf []string) (string, bool) {
	if m == nil {
		return "", false
	}
	for _, x := range suf {
		if i, ok := m[pre+x]; ok {
			return i, true
		}
		if x == "xx" {
			if i, ok := getPrefix(m, pre); ok {
				return i, true
			}
		}
	}
	return "", false
}

func (m Media) Box2D(r []string) (string, bool) {
	return getSuffix(m.BoxArt.Box2D, pre2D, r)
}

func (m Media) Box3D(r []string) (string, bool) {
	return getSuffix(m.BoxArt.Box3D, pre3D, r)
}

func (m Media) Flyer(r []string) (string, bool) {
	return getSuffix(m.Flyers, preFlyer, r)
}

type ROM struct {
	FileName   string `json:"romfilename"`
	SHA1       string `json:"romsha1"`
	RegionsRaw string `json:"romregions"`
}

func (r ROM) Regions() []string {
	var x []string
	for _, p := range strings.Split(r.RegionsRaw, ",") {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		x = append(x, p)
	}
	return x
}

type Game struct {
	Synopsis  map[string]string          `json:"synopsis"`
	ID        string                     `json:"id"`
	Name      string                     `json:"nom"`
	Names     map[string]string          `json:"noms"`
	Regions   []string                   `json:"regionshortnames"`
	Publisher string                     `json:"editeur"`
	Developer string                     `json:"developpeur"`
	Players   string                     `json:"joueurs"`
	Rating    string                     `json:"note"`
	Dates     map[string]string          `json:"dates"`
	Genres    map[string]json.RawMessage `json:"genres:`
	Media     Media                      `json:"medias"`
	ROMs      []ROM                      `json:"roms"`
	genres    map[string]string
}

func (g Game) Date(r []string) (string, bool) {
	return getSuffix(g.Dates, preDate, r)
}

func (g *Game) decodeGenre() {
	g.genres = make(map[string]string)
	for k, v := range g.Genres {
		if strings.HasSuffix(k, "_medias") || strings.HasSuffix(k, "_id") {
			continue
		}
		s := []string{}
		if err := json.Unmarshal(v, &s); err == nil {
			g.genres[k] = strings.Join(s, " / ")
		}
	}
}

func (g Game) Genre(l []string) (string, bool) {
	if g.genres == nil {
		g.decodeGenre()
	}
	return getSuffix(g.genres, preGenre, l)
}

func (g Game) Desc(l []string) (string, bool) {
	return getSuffix(g.Synopsis, preSynopsis, l)
}

func (g Game) ROM(req GameInfoReq) ROM {
	for _, x := range g.ROMs {
		if x.SHA1 == req.SHA1 {
			return x
		}
	}
	for _, x := range g.ROMs {
		if x.FileName == req.Name {
			return x
		}
	}
	return ROM{}
}

type Response struct {
	Game Game `json:"jeu"`
}

type GameInfoResp struct {
	Response Response `json:"response"`
}

func SanitizeURL(s string) string {
	u, err := url.Parse(s)
	if err != nil {
		return s
	}
	v := u.Query()
	v.Set("devid", "xxx")
	v.Set("devpassword", "yyy")
	v.Set("softname", "zzz")
	v.Del("ssid")
	v.Del("sspassword")
	u.RawQuery = v.Encode()
	return u.String()
}

func User(dev DevInfo, user UserInfo) (*UserInfoResp, error) {
	u, err := url.Parse(baseURL)
	u.Path = userInfoPath
	q := url.Values{}
	q.Set("output", "xml")
	q.Set("devid", dev.ID)
	q.Set("devpassword", dev.Password)
	if dev.Name != "" {
		q.Set("softname", dev.Name)
	}
	q.Set("ssid", user.ID)
	q.Set("sspassword", user.Password)
	u.RawQuery = q.Encode()
	resp, err := http.Get(u.String())
	if err != nil {
		v := u.Query()
		v.Set("devid", "xxx")
		v.Set("devpassword", "yyy")
		v.Set("softname", "zzz")
		u.RawQuery = v.Encode()
		return nil, fmt.Errorf("getting game url:%s, error:%s", u, err)
	}
	defer resp.Body.Close()
	r := &UserInfoResp{}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if err := xml.Unmarshal(b, r); err != nil {
		return nil, fmt.Errorf("ss: cannot parse response: %q", err)
	}
	return r, nil
}

func Threads(dev DevInfo, user UserInfo) int {
	if user.ID == "" || user.Password == "" {
		return 1
	}
	i, err := User(dev, user)
	if err != nil {
		log.Print("error getting allowed threads defaulting to 1")
		return 1
	}
	if i.MaxThreads < 1 {
		return 1
	}
	return i.MaxThreads
}

// GameInfo is the call to get game info.
func GameInfo(dev DevInfo, user UserInfo, req GameInfoReq) (*GameInfoResp, error) {
	u, err := url.Parse(baseURL)
	u.Path = gameInfoPath
	q := url.Values{}
	q.Set("output", "json")
	q.Set("devid", dev.ID)
	q.Set("devpassword", dev.Password)
	if dev.Name != "" {
		q.Set("softname", dev.Name)
	}
	if user.ID != "" {
		q.Set("ssid", user.ID)
	}
	if user.Password != "" {
		q.Set("sspassword", user.Password)
	}
	if req.SHA1 != "" {
		q.Set("sha1", req.SHA1)
	}
	if req.RomType == "" {
		q.Set("romtype", "rom")
	} else {
		q.Set("romtype", req.RomType)
	}
	if req.Name != "" {
		q.Set("romnom", req.Name)
	}
	u.RawQuery = q.Encode()
	resp, err := http.Get(u.String())
	if err != nil {
		if uerr, ok := err.(*url.Error); ok {
			uerr.URL = SanitizeURL(uerr.URL)
			return nil, uerr
		}
		return nil, err
	}
	defer resp.Body.Close()
	r := &GameInfoResp{}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if bytes.HasPrefix(b, []byte("Erreur : Rom/Iso/Dossier non trouv")) {
		return nil, ErrNotFound
	}
	if err := json.Unmarshal(b, r); err != nil {
		if err.Error() == "invalid character 'm' looking for beginning of value" {
			return nil, fmt.Errorf("ss: cannot parse response(%s): %q", string(b), err)
		}
		if err.Error() == "invalid character 'A' looking for beginning of value" {
			return nil, fmt.Errorf("ss: cannot parse response(%s): %q", string(b), err)
		}
		return nil, fmt.Errorf("ss: cannot parse response: %q", err)
	}
	return r, nil
}
<commit_msg>Harden code to future changes in ScreenScraper API.<commit_after>package ss

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

// JSON field prefixes.
const (
	pre2D       = "media_box2d_"
	pre3D       = "media_box3d_"
	preFlyer    = "media_flyer_"
	preDate     = "date_"
	preGenre    = "genres_"
	preSynopsis = "synopsis_"
)

const (
	baseURL      = "http://www.screenscraper.fr/"
	gameInfoPath = "api/jeuInfos.php"
	userInfoPath = "api/ssuserInfos.php"
)

// ErrNotFound is the error returned when a ROM isn't found.
var ErrNotFound = errors.New("not found")

// DevInfo is the information about the developer and used across APIs.
type DevInfo struct {
	ID       string
	Password string
	Name     string
}

// UserInfo is information about the user making the call.
type UserInfo struct {
	ID       string
	Password string
}

type UserInfoResp struct {
	ID              string `xml:"ssuser>id"`
	Level           int    `xml:"ssuser>niveau"`
	Contribution    int    `xml:"ssuser>contribution"`
	UploadedSystems int    `xml:"ssuser>uploadsysteme"`
	UploadedInfo    int    `xml:"ssuser>uploadinfos"`
	ROMsAssociated  int    `xml:"ssuser>romasso"`
	UpdatedMedia    int    `xml:"ssuser>uploadmedia"`
	FavoriteRegion  string `xml:"ssuser>favregion"`
	MaxThreads      int    `xml:"ssuser>maxthreads"`
}

// GameInfoReq is the information we use in the GameInfo command.
type GameInfoReq struct {
	Name    string
	SHA1    string
	RomType string
}

type SafeStringMap struct {
	Map map[string]string
}

func (s *SafeStringMap) UnmarshalJSON(b []byte) error {
	if s.Map == nil {
		s.Map = make(map[string]string)
	}
	x := make(map[string]json.RawMessage)
	if err := json.Unmarshal(b, &x); err != nil {
		log.Print("json: %v", err)
		return nil
	}
	for k, v := range x {
		var y string
		if err := json.Unmarshal(v, &y); err == nil {
			s.Map[k] = y
		}
	}
	return nil
}

type BoxArt struct {
	Box2D SafeStringMap `json:"media_boxs2d"`
	Box3D SafeStringMap `json:"media_boxs3d"`
}

type Media struct {
	Screenshot    string        `json:"media_screenshot"`
	ScreenMarquee string        `json:"media_screenmarquee"`
	Marquee       string        `json:"media_marquee"`
	Video         string        `json:"media_video"`
	Flyers        SafeStringMap `json:"media_flyers"`
	BoxArt        BoxArt        `json:"media_boxs"`
}

func getPrefix(m map[string]string, pre string) (string, bool) {
	for k, v := range m {
		if strings.HasPrefix(k, pre) {
			return v, true
		}
	}
	return "", false
}

func getSuffix(m map[string]string, pre string, suf []string) (string, bool) {
	if m == nil {
		return "", false
	}
	for _, x := range suf {
		if i, ok := m[pre+x]; ok {
			return i, true
		}
		if x == "xx" {
			if i, ok := getPrefix(m, pre); ok {
				return i, true
			}
		}
	}
	return "", false
}

func (m Media) Box2D(r []string) (string, bool) {
	return getSuffix(m.BoxArt.Box2D.Map, pre2D, r)
}

func (m Media) Box3D(r []string) (string, bool) {
	return getSuffix(m.BoxArt.Box3D.Map, pre3D, r)
}

func (m Media) Flyer(r []string) (string, bool) {
	return getSuffix(m.Flyers.Map, preFlyer, r)
}

type ROM struct {
	FileName   string `json:"romfilename"`
	SHA1       string `json:"romsha1"`
	RegionsRaw string `json:"romregions"`
}

func (r ROM) Regions() []string {
	var x []string
	for _, p := range strings.Split(r.RegionsRaw, ",") {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		x = append(x, p)
	}
	return x
}

type Game struct {
	Synopsis  SafeStringMap              `json:"synopsis"`
	ID        string                     `json:"id"`
	Name      string                     `json:"nom"`
	Names     SafeStringMap              `json:"noms"`
	Regions   []string                   `json:"regionshortnames"`
	Publisher string                     `json:"editeur"`
	Developer string                     `json:"developpeur"`
	Players   string                     `json:"joueurs"`
	Rating    string                     `json:"note"`
	Dates     SafeStringMap              `json:"dates"`
	Genres    map[string]json.RawMessage `json:"genres:`
	Media     Media                      `json:"medias"`
	ROMs      []ROM                      `json:"roms"`
	genres    map[string]string
}

func (g Game) Date(r []string) (string, bool) {
	return getSuffix(g.Dates.Map, preDate, r)
}

func (g *Game) decodeGenre() {
	g.genres = make(map[string]string)
	for k, v := range g.Genres {
		if strings.HasSuffix(k, "_medias") || strings.HasSuffix(k, "_id") {
			continue
		}
		s := []string{}
		if err := json.Unmarshal(v, &s); err == nil {
			g.genres[k] = strings.Join(s, " / ")
		}
	}
}

func (g Game) Genre(l []string) (string, bool) {
	if g.genres == nil {
		g.decodeGenre()
	}
	return getSuffix(g.genres, preGenre, l)
}

func (g Game) Desc(l []string) (string, bool) {
	return getSuffix(g.Synopsis.Map, preSynopsis, l)
}

func (g Game) ROM(req GameInfoReq) ROM {
	for _, x := range g.ROMs {
		if strings.ToLower(x.SHA1) == strings.ToLower(req.SHA1) {
			return x
		}
	}
	for _, x := range g.ROMs {
		if x.FileName == req.Name {
			return x
		}
	}
	return ROM{}
}

type Response struct {
	Game Game `json:"jeu"`
}

type GameInfoResp struct {
	Response Response `json:"response"`
}

func SanitizeURL(s string) string {
	u, err := url.Parse(s)
	if err != nil {
		return s
	}
	v := u.Query()
	v.Set("devid", "xxx")
	v.Set("devpassword", "yyy")
	v.Set("softname", "zzz")
	v.Del("ssid")
	v.Del("sspassword")
	u.RawQuery = v.Encode()
	return u.String()
}

func User(dev DevInfo, user UserInfo) (*UserInfoResp, error) {
	u, err := url.Parse(baseURL)
	u.Path = userInfoPath
	q := url.Values{}
	q.Set("output", "xml")
	q.Set("devid", dev.ID)
	q.Set("devpassword", dev.Password)
	if dev.Name != "" {
		q.Set("softname", dev.Name)
	}
	q.Set("ssid", user.ID)
	q.Set("sspassword", user.Password)
	u.RawQuery = q.Encode()
	resp, err := http.Get(u.String())
	if err != nil {
		v := u.Query()
		v.Set("devid", "xxx")
		v.Set("devpassword", "yyy")
		v.Set("softname", "zzz")
		u.RawQuery = v.Encode()
		return nil, fmt.Errorf("getting game url:%s, error:%s", u, err)
	}
	defer resp.Body.Close()
	r := &UserInfoResp{}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if err := xml.Unmarshal(b, r); err != nil {
		return nil, fmt.Errorf("ss: cannot parse response: %q", err)
	}
	return r, nil
}

func Threads(dev DevInfo, user UserInfo) int {
	if user.ID == "" || user.Password == "" {
		return 1
	}
	i, err := User(dev, user)
	if err != nil {
		log.Print("error getting allowed threads defaulting to 1")
		return 1
	}
	if i.MaxThreads < 1 {
		return 1
	}
	return i.MaxThreads
}

// GameInfo is the call to get game info.
func GameInfo(dev DevInfo, user UserInfo, req GameInfoReq) (*GameInfoResp, error) {
	u, err := url.Parse(baseURL)
	u.Path = gameInfoPath
	q := url.Values{}
	q.Set("output", "json")
	q.Set("devid", dev.ID)
	q.Set("devpassword", dev.Password)
	if dev.Name != "" {
		q.Set("softname", dev.Name)
	}
	if user.ID != "" {
		q.Set("ssid", user.ID)
	}
	if user.Password != "" {
		q.Set("sspassword", user.Password)
	}
	if req.SHA1 != "" {
		q.Set("sha1", req.SHA1)
	}
	if req.RomType == "" {
		q.Set("romtype", "rom")
	} else {
		q.Set("romtype", req.RomType)
	}
	if req.Name != "" {
		q.Set("romnom", req.Name)
	}
	u.RawQuery = q.Encode()
	resp, err := http.Get(u.String())
	if err != nil {
		if uerr, ok := err.(*url.Error); ok {
			uerr.URL = SanitizeURL(uerr.URL)
			return nil, uerr
		}
		return nil, err
	}
	defer resp.Body.Close()
	r := &GameInfoResp{}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if bytes.HasPrefix(b, []byte("Erreur : Rom/Iso/Dossier non trouv")) {
		return nil, ErrNotFound
	}
	if err := json.Unmarshal(b, r); err != nil {
		if err.Error() == "invalid character 'm' looking for beginning of value" {
			return nil, fmt.Errorf("ss: cannot parse response(%s): %q", string(b), err)
		}
		if err.Error() == "invalid character 'A' looking for beginning of value" {
			return nil, fmt.Errorf("ss: cannot parse response(%s): %q", string(b), err)
		}
		return nil, fmt.Errorf("ss: cannot parse response: %q", err)
	}
	return r, nil
}
<|endoftext|>


<commit_before>package nra

// rootSegment is used as the parentID for root segments.
const rootSegment int64 = 0

type SegmentStack struct {
	s []int64
}

func NewSegmentStack() *SegmentStack {
	return &SegmentStack{s: []int64{}}
}

// Push pushes a segment id onto the segment stack.
func (s *SegmentStack) Push(id int64) {
	s.s = append(s.s, id)
}

// Pop pops a segment id off of the segment stack. It returns false if the stack is empty.
func (s *SegmentStack) Pop() (int64, bool) {
	if s.Len() == 0 {
		return rootSegment, false
	}
	id := s.s[s.Len()-1]
	s.s = s.s[:s.Len()-1]
	return id, true
}

// Peek returns id from the top of the stack. It returns rootSegment if the stack is empty.
func (s *SegmentStack) Peek() int64 {
	if s.Len() == 0 {
		return rootSegment
	}
	return s.s[s.Len()-1]
}

// Len returns the length of the stack.
func (s *SegmentStack) Len() int {
	return len(s.s)
}
<commit_msg>Add a mutex<commit_after>package nra

import "sync"

// rootSegment is used as the parentID for root segments.
const rootSegment int64 = 0

type SegmentStack struct {
	sync.Mutex
	s []int64
}

func NewSegmentStack() *SegmentStack {
	return &SegmentStack{s: []int64{}}
}

// Push pushes a segment id onto the segment stack.
func (s *SegmentStack) Push(id int64) {
	s.Lock()
	defer s.Unlock()
	s.s = append(s.s, id)
}

// Pop pops a segment id off of the segment stack. It returns false if the stack is empty.
func (s *SegmentStack) Pop() (int64, bool) {
	s.Lock()
	defer s.Unlock()
	if s.Len() == 0 {
		return rootSegment, false
	}
	id := s.s[s.Len()-1]
	s.s = s.s[:s.Len()-1]
	return id, true
}

// Peek returns id from the top of the stack. It returns rootSegment if the stack is empty.
func (s *SegmentStack) Peek() int64 {
	if s.Len() == 0 {
		return rootSegment
	}
	return s.s[s.Len()-1]
}

// Len returns the length of the stack.
func (s *SegmentStack) Len() int {
	return len(s.s)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package aggregator

import (
	"errors"
	"math"
	"strconv"
	"sync"
	"time"

	"github.com/m3db/m3metrics/metadata"
	"github.com/m3db/m3metrics/metric/aggregated"
	"github.com/m3db/m3metrics/metric/unaggregated"
	"github.com/m3db/m3x/clock"

	"github.com/uber-go/tally"
)

var (
	errAggregatorShardClosed       = errors.New("aggregator shard is closed")
	errAggregatorShardNotWriteable = errors.New("aggregator shard is not writeable")
)

type addUntimedFn func(
	metric unaggregated.MetricUnion,
	metadatas metadata.StagedMetadatas,
) error

type addForwardedFn func(
	metric aggregated.ForwardedMetric,
	metadata metadata.ForwardMetadata,
) error

type aggregatorShardMetrics struct {
	notWriteableErrors tally.Counter
}

func newAggregatorShardMetrics(scope tally.Scope) aggregatorShardMetrics {
	return aggregatorShardMetrics{
		notWriteableErrors: scope.Counter("not-writeable-errors"),
	}
}

type aggregatorShard struct {
	sync.RWMutex

	shard                            uint32
	nowFn                            clock.NowFn
	bufferDurationBeforeShardCutover time.Duration
	bufferDurationAfterShardCutoff   time.Duration
	cutoverNanos                     int64
	cutoffNanos                      int64
	earliestWritableNanos            int64
	latestWriteableNanos             int64

	closed         bool
	metricMap      *metricMap
	metrics        aggregatorShardMetrics
	addUntimedFn   addUntimedFn
	addForwardedFn addForwardedFn
}

func newAggregatorShard(shard uint32, opts Options) *aggregatorShard {
	// NB(xichen): instead of sharding a global time lock, each shard has
	// its own time lock to ensure for an aggregation window, all metrics
	// owned by the shard are aggregated before they are flushed.
	opts = opts.SetTimeLock(&sync.RWMutex{})
	scope := opts.InstrumentOptions().MetricsScope().SubScope("shard").Tagged(
		map[string]string{"shard": strconv.Itoa(int(shard))},
	)
	s := &aggregatorShard{
		shard: shard,
		nowFn: opts.ClockOptions().NowFn(),
		bufferDurationBeforeShardCutover: opts.BufferDurationBeforeShardCutover(),
		bufferDurationAfterShardCutoff:   opts.BufferDurationAfterShardCutoff(),
		metricMap:                        newMetricMap(shard, opts),
		metrics:                          newAggregatorShardMetrics(scope),
	}
	s.addUntimedFn = s.metricMap.AddUntimed
	s.addForwardedFn = s.metricMap.AddForwarded
	return s
}

func (s *aggregatorShard) ID() uint32 { return s.shard }

func (s *aggregatorShard) CutoffNanos() int64 {
	s.RLock()
	cutoffNanos := s.cutoffNanos
	s.RUnlock()
	return cutoffNanos
}

func (s *aggregatorShard) IsWritable() bool {
	s.RLock()
	isWritable := s.isWritableWithLock()
	s.RUnlock()
	return isWritable
}

func (s *aggregatorShard) IsCutoff() bool {
	nowNanos := s.nowFn().UnixNano()
	s.RLock()
	isCutoff := nowNanos >= s.cutoffNanos
	s.RUnlock()
	return isCutoff
}

func (s *aggregatorShard) SetWriteableRange(rng timeRange) {
	var (
		cutoverNanos  = rng.cutoverNanos
		cutoffNanos   = rng.cutoffNanos
		earliestNanos = int64(0)
		latestNanos   = int64(math.MaxInt64)
	)
	if cutoverNanos >= int64(s.bufferDurationBeforeShardCutover) {
		earliestNanos = cutoverNanos - int64(s.bufferDurationBeforeShardCutover)
	}
	if cutoffNanos <= math.MaxInt64-int64(s.bufferDurationAfterShardCutoff) {
		latestNanos = cutoffNanos + int64(s.bufferDurationAfterShardCutoff)
	}
	s.Lock()
	s.cutoverNanos = cutoverNanos
	s.cutoffNanos = cutoffNanos
	s.earliestWritableNanos = earliestNanos
	s.latestWriteableNanos = latestNanos
	s.Unlock()
}

func (s *aggregatorShard) AddUntimed(
	metric unaggregated.MetricUnion,
	metadatas metadata.StagedMetadatas,
) error {
	s.RLock()
	if s.closed {
		s.RUnlock()
		return errAggregatorShardClosed
	}
	if !s.isWritableWithLock() {
		s.RUnlock()
		s.metrics.notWriteableErrors.Inc(1)
		return errAggregatorShardNotWriteable
	}
	err := s.addUntimedFn(metric, metadatas)
	s.RUnlock()
	return err
}

func (s *aggregatorShard) AddForwarded(
	metric aggregated.ForwardedMetric,
	metadata metadata.ForwardMetadata,
) error {
	s.RLock()
	if s.closed {
		s.RUnlock()
		return errAggregatorShardClosed
	}
	if !s.isWritableWithLock() {
		s.RUnlock()
		s.metrics.notWriteableErrors.Inc(1)
		return errAggregatorShardNotWriteable
	}
	err := s.addForwardedFn(metric, metadata)
	s.RUnlock()
	return err
}

func (s *aggregatorShard) Tick(target time.Duration) tickResult {
	return s.metricMap.Tick(target)
}

func (s *aggregatorShard) Close() {
	s.Lock()
	defer s.Unlock()

	if s.closed {
		return
	}
	s.closed = true
	s.metricMap.Close()
}

func (s *aggregatorShard) isWritableWithLock() bool {
	nowNanos := s.nowFn().UnixNano()
	return nowNanos >= s.earliestWritableNanos && nowNanos < s.latestWriteableNanos
}

type timeRange struct {
	cutoverNanos int64
	cutoffNanos  int64
}
<commit_msg>Add shard level metric (#155)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package aggregator

import (
	"errors"
	"math"
	"strconv"
	"sync"
	"time"

	"github.com/m3db/m3metrics/metadata"
	"github.com/m3db/m3metrics/metric/aggregated"
	"github.com/m3db/m3metrics/metric/unaggregated"
	"github.com/m3db/m3x/clock"

	"github.com/uber-go/tally"
)

var (
	errAggregatorShardClosed       = errors.New("aggregator shard is closed")
	errAggregatorShardNotWriteable = errors.New("aggregator shard is not writeable")
)

type addUntimedFn func(
	metric unaggregated.MetricUnion,
	metadatas metadata.StagedMetadatas,
) error

type addForwardedFn func(
	metric aggregated.ForwardedMetric,
	metadata metadata.ForwardMetadata,
) error

type aggregatorShardMetrics struct {
	notWriteableErrors tally.Counter
	writeSucccess      tally.Counter
}

func newAggregatorShardMetrics(scope tally.Scope) aggregatorShardMetrics {
	return aggregatorShardMetrics{
		notWriteableErrors: scope.Counter("not-writeable-errors"),
		writeSucccess:      scope.Counter("write-success"),
	}
}

type aggregatorShard struct {
	sync.RWMutex

	shard                            uint32
	nowFn                            clock.NowFn
	bufferDurationBeforeShardCutover time.Duration
	bufferDurationAfterShardCutoff   time.Duration
	cutoverNanos                     int64
	cutoffNanos                      int64
	earliestWritableNanos            int64
	latestWriteableNanos             int64

	closed         bool
	metricMap      *metricMap
	metrics        aggregatorShardMetrics
	addUntimedFn   addUntimedFn
	addForwardedFn addForwardedFn
}

func newAggregatorShard(shard uint32, opts Options) *aggregatorShard {
	// NB(xichen): instead of sharding a global time lock, each shard has
	// its own time lock to ensure for an aggregation window, all metrics
	// owned by the shard are aggregated before they are flushed.
	opts = opts.SetTimeLock(&sync.RWMutex{})
	scope := opts.InstrumentOptions().MetricsScope().SubScope("shard").Tagged(
		map[string]string{"shard": strconv.Itoa(int(shard))},
	)
	s := &aggregatorShard{
		shard: shard,
		nowFn: opts.ClockOptions().NowFn(),
		bufferDurationBeforeShardCutover: opts.BufferDurationBeforeShardCutover(),
		bufferDurationAfterShardCutoff:   opts.BufferDurationAfterShardCutoff(),
		metricMap:                        newMetricMap(shard, opts),
		metrics:                          newAggregatorShardMetrics(scope),
	}
	s.addUntimedFn = s.metricMap.AddUntimed
	s.addForwardedFn = s.metricMap.AddForwarded
	return s
}

func (s *aggregatorShard) ID() uint32 { return s.shard }

func (s *aggregatorShard) CutoffNanos() int64 {
	s.RLock()
	cutoffNanos := s.cutoffNanos
	s.RUnlock()
	return cutoffNanos
}

func (s *aggregatorShard) IsWritable() bool {
	s.RLock()
	isWritable := s.isWritableWithLock()
	s.RUnlock()
	return isWritable
}

func (s *aggregatorShard) IsCutoff() bool {
	nowNanos := s.nowFn().UnixNano()
	s.RLock()
	isCutoff := nowNanos >= s.cutoffNanos
	s.RUnlock()
	return isCutoff
}

func (s *aggregatorShard) SetWriteableRange(rng timeRange) {
	var (
		cutoverNanos  = rng.cutoverNanos
		cutoffNanos   = rng.cutoffNanos
		earliestNanos = int64(0)
		latestNanos   = int64(math.MaxInt64)
	)
	if cutoverNanos >= int64(s.bufferDurationBeforeShardCutover) {
		earliestNanos = cutoverNanos - int64(s.bufferDurationBeforeShardCutover)
	}
	if cutoffNanos <= math.MaxInt64-int64(s.bufferDurationAfterShardCutoff) {
		latestNanos = cutoffNanos + int64(s.bufferDurationAfterShardCutoff)
	}
	s.Lock()
	s.cutoverNanos = cutoverNanos
	s.cutoffNanos = cutoffNanos
	s.earliestWritableNanos = earliestNanos
	s.latestWriteableNanos = latestNanos
	s.Unlock()
}

func (s *aggregatorShard) AddUntimed(
	metric unaggregated.MetricUnion,
	metadatas metadata.StagedMetadatas,
) error {
	s.RLock()
	if s.closed {
		s.RUnlock()
		return errAggregatorShardClosed
	}
	if !s.isWritableWithLock() {
		s.RUnlock()
		s.metrics.notWriteableErrors.Inc(1)
		return errAggregatorShardNotWriteable
	}
	err := s.addUntimedFn(metric, metadatas)
	s.RUnlock()
	if err != nil {
		return err
	}
	s.metrics.writeSucccess.Inc(1)
	return nil
}

func (s *aggregatorShard) AddForwarded(
	metric aggregated.ForwardedMetric,
	metadata metadata.ForwardMetadata,
) error {
	s.RLock()
	if s.closed {
		s.RUnlock()
		return errAggregatorShardClosed
	}
	if !s.isWritableWithLock() {
		s.RUnlock()
		s.metrics.notWriteableErrors.Inc(1)
		return errAggregatorShardNotWriteable
	}
	err := s.addForwardedFn(metric, metadata)
	s.RUnlock()
	if err != nil {
		return err
	}
	s.metrics.writeSucccess.Inc(1)
	return nil
}

func (s *aggregatorShard) Tick(target time.Duration) tickResult {
	return s.metricMap.Tick(target)
}

func (s *aggregatorShard) Close() {
	s.Lock()
	defer s.Unlock()

	if s.closed {
		return
	}
	s.closed = true
	s.metricMap.Close()
}

func (s *aggregatorShard) isWritableWithLock() bool {
	nowNanos := s.nowFn().UnixNano()
	return nowNanos >= s.earliestWritableNanos && nowNanos < s.latestWriteableNanos
}

type timeRange struct {
	cutoverNanos int64
	cutoffNanos  int64
}
<|endoftext|>