<commit_before>package common

import (
	"encoding/json"
	"errors"
	"io"

	"github.com/fractalplatform/fractal/utils/rlp"
)

type AuthorType uint8

const (
	AccountNameType AuthorType = iota
	PubKeyType
	AddressType
)

type (
	Author struct {
		Owner  `json:"owner"`
		Weight uint64 `json:"weight"`
	}
	Owner interface {
		String() string
	}
)

type AccountAuthor struct {
	Account Name
}
type StorageAuthor struct {
	Type    AuthorType
	DataRaw rlp.RawValue
	Weight  uint64
}

type AuthorJSON struct {
	Type     AuthorType
	OwnerStr string
	Weight   uint64
}

func NewAuthor(owner Owner, weight uint64) *Author {
	return &Author{Owner: owner, Weight: weight}
}

func (a *Author) GetWeight() uint64 {
	return a.Weight
}

func (a *Author) EncodeRLP(w io.Writer) error {
	storageAuthor, err := a.encode()
	if err != nil {
		return err
	}
	return rlp.Encode(w, storageAuthor)
}

func (a *Author) encode() (*StorageAuthor, error) {
	switch aTy := a.Owner.(type) {
	case Name:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    AccountNameType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	case PubKey:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    PubKeyType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	case Address:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    AddressType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	}
	return nil, errors.New("Author encode failed")
}

func (a *Author) DecodeRLP(s *rlp.Stream) error {
	storageAuthor := new(StorageAuthor)
	err := s.Decode(storageAuthor)
	if err != nil {
		return err
	}
	return a.decode(storageAuthor)
}

func (a *Author) decode(sa *StorageAuthor) error {
	switch sa.Type {
	case AccountNameType:
		var name Name
		if err := rlp.DecodeBytes(sa.DataRaw, &name); err != nil {
			return err
		}
		a.Owner = name
		a.Weight = sa.Weight
		return nil
	case PubKeyType:
		var pubKey PubKey
		if err := rlp.DecodeBytes(sa.DataRaw, &pubKey); err != nil {
			return err
		}
		a.Owner = pubKey
		a.Weight = sa.Weight
		return nil
	case AddressType:
		var address Address
		if err := rlp.DecodeBytes(sa.DataRaw, &address); err != nil {
			return err
		}
		a.Owner = address
		a.Weight = sa.Weight
		return nil
	}
	return errors.New("Author decode failed")
}

func (a *Author) MarshalJSON() ([]byte, error) {
	switch aTy := a.Owner.(type) {
	case Name:
		return json.Marshal(&AuthorJSON{Type: AccountNameType, OwnerStr: aTy.String(), Weight: a.Weight})
	case PubKey:
		return json.Marshal(&AuthorJSON{Type: PubKeyType, OwnerStr: aTy.String(), Weight: a.Weight})
	case Address:
		return json.Marshal(&AuthorJSON{Type: AddressType, OwnerStr: aTy.String(), Weight: a.Weight})
	}
	return nil, errors.New("Author marshal failed")
}

func (a *Author) UnmarshalJSON(data []byte) error {
	aj := &AuthorJSON{}
	if err := json.Unmarshal(data, aj); err != nil {
		return err
	}
	switch aj.Type {
	case AccountNameType:
		a.Owner = Name(aj.OwnerStr)
		a.Weight = aj.Weight
	case PubKeyType:
		a.Owner = HexToPubKey(aj.OwnerStr)
		a.Weight = aj.Weight
	case AddressType:
		a.Owner = HexToAddress(aj.OwnerStr)
		a.Weight = aj.Weight
	}
	return nil
}
<commit_msg>change author json (#229)<commit_after>package common

import (
	"encoding/json"
	"errors"
	"io"

	"github.com/fractalplatform/fractal/utils/rlp"
)

type AuthorType uint8

const (
	AccountNameType AuthorType = iota
	PubKeyType
	AddressType
)

type (
	Author struct {
		Owner  `json:"owner"`
		Weight uint64 `json:"weight"`
	}
	Owner interface {
		String() string
	}
)

type AccountAuthor struct {
	Account Name
}
type StorageAuthor struct {
	Type    AuthorType
	DataRaw rlp.RawValue
	Weight  uint64
}

type AuthorJSON struct {
	authorType AuthorType
	OwnerStr   string `json:"owner"`
	Weight     uint64 `json:"weight"`
}

func NewAuthor(owner Owner, weight uint64) *Author {
	return &Author{Owner: owner, Weight: weight}
}

func (a *Author) GetWeight() uint64 {
	return a.Weight
}

func (a *Author) EncodeRLP(w io.Writer) error {
	storageAuthor, err := a.encode()
	if err != nil {
		return err
	}
	return rlp.Encode(w, storageAuthor)
}

func (a *Author) encode() (*StorageAuthor, error) {
	switch aTy := a.Owner.(type) {
	case Name:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    AccountNameType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	case PubKey:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    PubKeyType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	case Address:
		value, err := rlp.EncodeToBytes(&aTy)
		if err != nil {
			return nil, err
		}
		return &StorageAuthor{
			Type:    AddressType,
			DataRaw: value,
			Weight:  a.Weight,
		}, nil
	}
	return nil, errors.New("Author encode failed")
}

func (a *Author) DecodeRLP(s *rlp.Stream) error {
	storageAuthor := new(StorageAuthor)
	err := s.Decode(storageAuthor)
	if err != nil {
		return err
	}
	return a.decode(storageAuthor)
}

func (a *Author) decode(sa *StorageAuthor) error {
	switch sa.Type {
	case AccountNameType:
		var name Name
		if err := rlp.DecodeBytes(sa.DataRaw, &name); err != nil {
			return err
		}
		a.Owner = name
		a.Weight = sa.Weight
		return nil
	case PubKeyType:
		var pubKey PubKey
		if err := rlp.DecodeBytes(sa.DataRaw, &pubKey); err != nil {
			return err
		}
		a.Owner = pubKey
		a.Weight = sa.Weight
		return nil
	case AddressType:
		var address Address
		if err := rlp.DecodeBytes(sa.DataRaw, &address); err != nil {
			return err
		}
		a.Owner = address
		a.Weight = sa.Weight
		return nil
	}
	return errors.New("Author decode failed")
}

func (a *Author) MarshalJSON() ([]byte, error) {
	switch aTy := a.Owner.(type) {
	case Name:
		return json.Marshal(&AuthorJSON{authorType: AccountNameType, OwnerStr: aTy.String(), Weight: a.Weight})
	case PubKey:
		return json.Marshal(&AuthorJSON{authorType: PubKeyType, OwnerStr: aTy.String(), Weight: a.Weight})
	case Address:
		return json.Marshal(&AuthorJSON{authorType: AddressType, OwnerStr: aTy.String(), Weight: a.Weight})
	}
	return nil, errors.New("Author marshal failed")
}

func (a *Author) UnmarshalJSON(data []byte) error {
	aj := &AuthorJSON{}
	if err := json.Unmarshal(data, aj); err != nil {
		return err
	}
	switch aj.authorType {
	case AccountNameType:
		a.Owner = Name(aj.OwnerStr)
		a.Weight = aj.Weight
	case PubKeyType:
		a.Owner = HexToPubKey(aj.OwnerStr)
		a.Weight = aj.Weight
	case AddressType:
		a.Owner = HexToAddress(aj.OwnerStr)
		a.Weight = aj.Weight
	}
	return nil
}
<|endoftext|>


<commit_before>package transports_test

import(
	"github.com/matiasinsaurralde/transports/marshalers"
	// "github.com/matiasinsaurralde/transports/marshalers/protos"

	"strings"
	"net/http"
	"net/url"
	"testing"
	"log"
)

const TestRequestUrl string = "http://whatismyip.akamai.com/"

var request http.Request

func init() {

	log.Println("init")

	url, _ := url.Parse( TestRequestUrl)

	request = http.Request{
		Method: "GET",
		URL: url,
		Proto: "HTTP/1.0",
	}
}

func TestHttpRequestMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpRequest")
	}
}

func TestHttpResponseMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpResponse")
	}
}

func TestUnsupportedType( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}

	var v UnknownType
	v = UnknownType{"Value"}

	var i interface{}
	i = v

	err, _ := marshaler.Marshal(&i)

	exists := strings.Index(err.Error(), transports.MarshalerTypeNotSupportedError)

	if exists < 0 {
		t.Fatal("Unsupported type doesn't break the Protobuf marshaler")
	}
}

func TestNilInput( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	err, _ := marshaler.Marshal(nil)

	if strings.Index( err.Error(), transports.MarshalerNilTypeError ) < 0 {
		t.Fatal("Nil type doesn't break the Protobuf marshaler")
	}
}
<commit_msg>Removing Test Request URL constant<commit_after>package transports_test

import(
	"github.com/matiasinsaurralde/transports/marshalers"
	// "github.com/matiasinsaurralde/transports/marshalers/protos"

	"strings"
	"net/http"
	"net/url"
	"testing"
)

var request http.Request

func init() {

	url, _ := url.Parse( "http://whatismyip.akamai.com/")

	request = http.Request{
		Method: "GET",
		URL: url,
		Proto: "HTTP/1.0",
	}
}

func TestHttpRequestMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpRequest")
	}
}

func TestHttpResponseMarshal( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	var i interface{}
	i = request
	_, err := marshaler.Marshal(&i)
	if err != nil {
		t.Fatal("Can't marshal HttpResponse")
	}
}

func TestUnsupportedType( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}

	var v UnknownType
	v = UnknownType{"Value"}

	var i interface{}
	i = v

	err, _ := marshaler.Marshal(&i)

	exists := strings.Index(err.Error(), transports.MarshalerTypeNotSupportedError)

	if exists < 0 {
		t.Fatal("Unsupported type doesn't break the Protobuf marshaler")
	}
}

func TestNilInput( t *testing.T ) {
	var marshaler transports.Marshaler
	marshaler = transports.ProtobufMarshaler{}
	err, _ := marshaler.Marshal(nil)

	if strings.Index( err.Error(), transports.MarshalerNilTypeError ) < 0 {
		t.Fatal("Nil type doesn't break the Protobuf marshaler")
	}
}
<|endoftext|>


<commit_before>package common

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// Twitch twitch chat client
type Twitch struct {
	connLock       sync.Mutex
	sendLock       sync.Mutex
	conn           *websocket.Conn
	ChLock         sync.RWMutex
	channels       []string
	messages       chan *Message
	MessagePattern *regexp.Regexp
	SubPattern     *regexp.Regexp
	quit           chan struct{}
}

// NewTwitch new twitch chat client
func NewTwitch() *Twitch {
	return &Twitch{
		channels: make([]string, 0),
		messages: make(chan *Message, MessageBufferSize),
		// > @badges=global_mod/1,turbo/1;color=#0D4200;display-name=dallas;emotes=25:0-4,12-16/1902:6-10;mod=0;room-id=1337;
		//subscriber=0;turbo=1;user-id=1337;user-type=global_mod :ronni!ronni@ronni.tmi.twitch.tv PRIVMSG #dallas :Kappa Keepo Kappa
		MessagePattern: regexp.MustCompile(`user-type=.+:([a-z0-9_-]+)\!.+\.tmi\.twitch\.tv PRIVMSG #([a-z0-9_-]+) :(.+)`),
		// > @badges=staff/1,broadcaster/1,turbo/1;color=#008000;display-name=ronni;emotes=;mod=0;msg-id=resub;msg-param-months=6;
		// msg-param-sub-plan=Prime;msg-param-sub-plan-name=Prime;room-id=1337;subscriber=1;system-msg=ronni\shas\ssubscribed\sfor\s6\smonths!;
		// login=ronni;turbo=1;user-id=1337;user-type=staff :tmi.twitch.tv USERNOTICE #dallas :Great stream -- keep it up!
		SubPattern: regexp.MustCompile(`msg-id=(sub|resub);.+;system-msg=(.+);tmi-sent-ts.+ \:tmi\.twitch\.tv USERNOTICE #([a-z0-9_-]+)( :.+)?`),
		quit:       make(chan struct{}, 2),
	}
}

func (c *Twitch) connect() {
	conf := GetConfig()
	dialer := websocket.Dialer{HandshakeTimeout: HandshakeTimeout}
	headers := http.Header{"Origin": []string{conf.Twitch.OriginURL}}

	var err error
	c.connLock.Lock()
	c.conn, _, err = dialer.Dial(GetConfig().Twitch.SocketURL, headers)
	c.connLock.Unlock()
	if err != nil {
		log.Printf("error connecting to twitch ws %s", err)
		c.reconnect()
		return
	}

	if conf.Twitch.OAuth == "" || conf.Twitch.Nick == "" {
		log.Println("missing OAuth or Nick, using justinfan659 as login data")
		conf.Twitch.OAuth = "justinfan659"
		conf.Twitch.Nick = "justinfan659"
	}

	c.send("PASS " + conf.Twitch.OAuth)
	c.send("NICK " + conf.Twitch.Nick)
	c.send("CAP REQ :twitch.tv/tags")
	c.send("CAP REQ :twitch.tv/commands")

	for _, ch := range c.channels {
		select {
		case <-c.quit:
			return
		default:
		}
		ch = strings.ToLower(ch)
		log.Printf("joining %s", ch)
		err := c.send("JOIN #" + ch)
		if err != nil {
			log.Println("failed to join", ch, "after freshly re/connecting to the websocket")
		}
	}
}

func (c *Twitch) reconnect() {
	c.connLock.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connLock.Unlock()

	time.Sleep(SocketReconnectDelay)
	c.connect()
}

// Run connect and start message read loop
func (c *Twitch) Run() {
	pinger := time.NewTicker(4 * time.Minute)
	c.connect()
	go c.rejoinHandler()

	pingTicker := time.NewTicker(5 * time.Minute)

	go func() {
		for {
			select {
			case <-c.quit:
				close(c.messages)
				return
			case <-pingTicker.C:
				err := c.send("PING")
				if err != nil {
					c.reconnect()
					continue
				}
			default:
			}
			c.connLock.Lock()
			err := c.conn.SetReadDeadline(time.Now().Add(SocketReadTimeout))
			c.connLock.Unlock()
			if err != nil {
				log.Printf("error setting the ReadDeadline: %v", err)
				c.reconnect()
				continue
			}

			c.connLock.Lock()
			_, msg, err := c.conn.ReadMessage()
			c.connLock.Unlock()
			if err != nil {
				log.Printf("error reading message: %v", err)
				c.reconnect()
				continue
			}

			if strings.Index(string(msg), "PING") == 0 {
				err := c.send("PONG")
				if err != nil {
					log.Printf("error sending PONG: %v", err)
					c.reconnect()
				}
				continue
			}

			s := c.SubPattern.FindAllStringSubmatch(string(msg), -1)
			for _, v := range s {
				data := strings.Replace(v[2], "\\s", " ", -1)
				if v[4] != "" {
					data += " [SubMessage]: " + v[4][2:]
				}
				m := &Message{
					Type:    "MSG",
					Channel: v[3],
					Nick:    "twitchnotify",
					Data:    data,
					Time:    time.Now().UTC(),
				}

				select {
				case <-pinger.C:
					c.send("PING :tmi.twitch.tv")
				case c.messages <- m:
				default:
					log.Println("error messages channel full :(")
				}
			}

			l := c.MessagePattern.FindAllStringSubmatch(string(msg), -1)
			for _, v := range l {
				data := strings.TrimSpace(v[3])
				data = strings.Replace(data, "ACTION", "/me", -1)
				data = strings.Replace(data, "", "", -1)
				m := &Message{
					Type:    "MSG",
					Channel: v[2],
					Nick:    v[1],
					Data:    data,
					Time:    time.Now().UTC(),
				}

				select {
				case <-pinger.C:
					c.send("PING :tmi.twitch.tv")
				case c.messages <- m:
				default:
					log.Println("error messages channel full :(")
				}
			}
		}
	}()
}

// Channels ...
func (c *Twitch) Channels() []string {
	return c.channels
}

// Messages channel accessor
func (c *Twitch) Messages() <-chan *Message {
	return c.messages
}

// Message send a message to a channel
func (c *Twitch) Message(ch, payload string) error {
	return c.send(fmt.Sprintf("PRIVMSG #%s :%s", ch, payload))
}

func (c *Twitch) send(m string) error {
	c.sendLock.Lock()
	err := c.conn.SetWriteDeadline(time.Now().Add(SocketWriteTimeout))
	c.sendLock.Unlock()
	if err != nil {
		return fmt.Errorf("error setting SetWriteDeadline %s", err)
	}
	c.sendLock.Lock()
	err = c.conn.WriteMessage(websocket.TextMessage, []byte(m+"\r\n"))
	c.sendLock.Unlock()
	if err != nil {
		return fmt.Errorf("error sending message %s", err)
	}
	time.Sleep(SocketWriteDebounce)
	return nil
}

// Join channel
func (c *Twitch) Join(ch string) error {
	ch = strings.ToLower(ch)
	err := c.send("JOIN #" + ch)
	if err != nil {
		c.reconnect()
		return err
	}
	c.ChLock.Lock()
	defer c.ChLock.Unlock()
	if inSlice(c.channels, ch) {
		return errors.New("already in channel")
	}
	c.channels = append(c.channels, ch)
	return nil
}

// Leave channel
func (c *Twitch) Leave(ch string) error {
	ch = strings.ToLower(ch)
	err := c.send("PART #" + ch)
	if err != nil {
		log.Printf("error leaving channel: %s", err)
		c.reconnect()
	}
	return c.removeChannel(ch)
}

func (c *Twitch) removeChannel(ch string) error {
	c.ChLock.Lock()
	defer c.ChLock.Unlock()
	sort.Strings(c.channels)
	i := sort.SearchStrings(c.channels, ch)
	if i < len(c.channels) && c.channels[i] == ch {
		c.channels = append(c.channels[:i], c.channels[i+1:]...)
		return nil
	}
	return errors.New("not in channel")
}

func (c *Twitch) rejoinHandler() {
	const interval = 2 * time.Hour
	ticker := time.NewTicker(interval)
	for {
		select {
		case <-c.quit:
			ticker.Stop()
			return
		case <-ticker.C:
			c.ChLock.RLock()
			for _, ch := range c.channels {
				ch = strings.ToLower(ch)
				if err := c.send("JOIN #" + ch); err != nil {
					log.Println(err)
					continue
				}
			}
			c.ChLock.RUnlock()
		}
	}
}

// Stop stops the chats
func (c *Twitch) Stop(wg *sync.WaitGroup) {
	close(c.quit)
	c.connLock.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connLock.Unlock()
	wg.Done()
}

func inSlice(s []string, v string) bool {
	for _, sv := range s {
		if strings.EqualFold(sv, v) {
			return true
		}
	}
	return false
}
<commit_msg>add subgift<commit_after>package common

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// Twitch twitch chat client
type Twitch struct {
	connLock       sync.Mutex
	sendLock       sync.Mutex
	conn           *websocket.Conn
	ChLock         sync.RWMutex
	channels       []string
	messages       chan *Message
	MessagePattern *regexp.Regexp
	SubPattern     *regexp.Regexp
	quit           chan struct{}
}

// NewTwitch new twitch chat client
func NewTwitch() *Twitch {
	return &Twitch{
		channels: make([]string, 0),
		messages: make(chan *Message, MessageBufferSize),
		// > @badges=global_mod/1,turbo/1;color=#0D4200;display-name=dallas;emotes=25:0-4,12-16/1902:6-10;mod=0;room-id=1337;
		//subscriber=0;turbo=1;user-id=1337;user-type=global_mod :ronni!ronni@ronni.tmi.twitch.tv PRIVMSG #dallas :Kappa Keepo Kappa
		MessagePattern: regexp.MustCompile(`user-type=.+:([a-z0-9_-]+)\!.+\.tmi\.twitch\.tv PRIVMSG #([a-z0-9_-]+) :(.+)`),
		// > @badges=staff/1,broadcaster/1,turbo/1;color=#008000;display-name=ronni;emotes=;mod=0;msg-id=resub;msg-param-months=6;
		// msg-param-sub-plan=Prime;msg-param-sub-plan-name=Prime;room-id=1337;subscriber=1;system-msg=ronni\shas\ssubscribed\sfor\s6\smonths!;
		// login=ronni;turbo=1;user-id=1337;user-type=staff :tmi.twitch.tv USERNOTICE #dallas :Great stream -- keep it up!
		SubPattern: regexp.MustCompile(`msg-id=(sub|resub|subgift);.+;system-msg=(.+);tmi-sent-ts.+ \:tmi\.twitch\.tv USERNOTICE #([a-z0-9_-]+)( :.+)?`),
		quit:       make(chan struct{}, 2),
	}
}

func (c *Twitch) connect() {
	conf := GetConfig()
	dialer := websocket.Dialer{HandshakeTimeout: HandshakeTimeout}
	headers := http.Header{"Origin": []string{conf.Twitch.OriginURL}}

	var err error
	c.connLock.Lock()
	c.conn, _, err = dialer.Dial(GetConfig().Twitch.SocketURL, headers)
	c.connLock.Unlock()
	if err != nil {
		log.Printf("error connecting to twitch ws %s", err)
		c.reconnect()
		return
	}

	if conf.Twitch.OAuth == "" || conf.Twitch.Nick == "" {
		log.Println("missing OAuth or Nick, using justinfan659 as login data")
		conf.Twitch.OAuth = "justinfan659"
		conf.Twitch.Nick = "justinfan659"
	}

	c.send("PASS " + conf.Twitch.OAuth)
	c.send("NICK " + conf.Twitch.Nick)
	c.send("CAP REQ :twitch.tv/tags")
	c.send("CAP REQ :twitch.tv/commands")

	for _, ch := range c.channels {
		select {
		case <-c.quit:
			return
		default:
		}
		ch = strings.ToLower(ch)
		log.Printf("joining %s", ch)
		err := c.send("JOIN #" + ch)
		if err != nil {
			log.Println("failed to join", ch, "after freshly re/connecting to the websocket")
		}
	}
}

func (c *Twitch) reconnect() {
	c.connLock.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connLock.Unlock()

	time.Sleep(SocketReconnectDelay)
	c.connect()
}

// Run connect and start message read loop
func (c *Twitch) Run() {
	pinger := time.NewTicker(4 * time.Minute)
	c.connect()
	go c.rejoinHandler()

	pingTicker := time.NewTicker(5 * time.Minute)

	go func() {
		for {
			select {
			case <-c.quit:
				close(c.messages)
				return
			case <-pingTicker.C:
				err := c.send("PING")
				if err != nil {
					c.reconnect()
					continue
				}
			default:
			}
			c.connLock.Lock()
			err := c.conn.SetReadDeadline(time.Now().Add(SocketReadTimeout))
			c.connLock.Unlock()
			if err != nil {
				log.Printf("error setting the ReadDeadline: %v", err)
				c.reconnect()
				continue
			}

			c.connLock.Lock()
			_, msg, err := c.conn.ReadMessage()
			c.connLock.Unlock()
			if err != nil {
				log.Printf("error reading message: %v", err)
				c.reconnect()
				continue
			}

			if strings.Index(string(msg), "PING") == 0 {
				err := c.send("PONG")
				if err != nil {
					log.Printf("error sending PONG: %v", err)
					c.reconnect()
				}
				continue
			}

			s := c.SubPattern.FindAllStringSubmatch(string(msg), -1)
			for _, v := range s {
				data := strings.Replace(v[2], "\\s", " ", -1)
				if v[4] != "" {
					data += " [SubMessage]: " + v[4][2:]
				}
				m := &Message{
					Type:    "MSG",
					Channel: v[3],
					Nick:    "twitchnotify",
					Data:    data,
					Time:    time.Now().UTC(),
				}

				select {
				case <-pinger.C:
					c.send("PING :tmi.twitch.tv")
				case c.messages <- m:
				default:
					log.Println("error messages channel full :(")
				}
			}

			l := c.MessagePattern.FindAllStringSubmatch(string(msg), -1)
			for _, v := range l {
				data := strings.TrimSpace(v[3])
				data = strings.Replace(data, "ACTION", "/me", -1)
				data = strings.Replace(data, "", "", -1)
				m := &Message{
					Type:    "MSG",
					Channel: v[2],
					Nick:    v[1],
					Data:    data,
					Time:    time.Now().UTC(),
				}

				select {
				case <-pinger.C:
					c.send("PING :tmi.twitch.tv")
				case c.messages <- m:
				default:
					log.Println("error messages channel full :(")
				}
			}
		}
	}()
}

// Channels ...
func (c *Twitch) Channels() []string {
	return c.channels
}

// Messages channel accessor
func (c *Twitch) Messages() <-chan *Message {
	return c.messages
}

// Message send a message to a channel
func (c *Twitch) Message(ch, payload string) error {
	return c.send(fmt.Sprintf("PRIVMSG #%s :%s", ch, payload))
}

func (c *Twitch) send(m string) error {
	c.sendLock.Lock()
	err := c.conn.SetWriteDeadline(time.Now().Add(SocketWriteTimeout))
	c.sendLock.Unlock()
	if err != nil {
		return fmt.Errorf("error setting SetWriteDeadline %s", err)
	}
	c.sendLock.Lock()
	err = c.conn.WriteMessage(websocket.TextMessage, []byte(m+"\r\n"))
	c.sendLock.Unlock()
	if err != nil {
		return fmt.Errorf("error sending message %s", err)
	}
	time.Sleep(SocketWriteDebounce)
	return nil
}

// Join channel
func (c *Twitch) Join(ch string) error {
	ch = strings.ToLower(ch)
	err := c.send("JOIN #" + ch)
	if err != nil {
		c.reconnect()
		return err
	}
	c.ChLock.Lock()
	defer c.ChLock.Unlock()
	if inSlice(c.channels, ch) {
		return errors.New("already in channel")
	}
	c.channels = append(c.channels, ch)
	return nil
}

// Leave channel
func (c *Twitch) Leave(ch string) error {
	ch = strings.ToLower(ch)
	err := c.send("PART #" + ch)
	if err != nil {
		log.Printf("error leaving channel: %s", err)
		c.reconnect()
	}
	return c.removeChannel(ch)
}

func (c *Twitch) removeChannel(ch string) error {
	c.ChLock.Lock()
	defer c.ChLock.Unlock()
	sort.Strings(c.channels)
	i := sort.SearchStrings(c.channels, ch)
	if i < len(c.channels) && c.channels[i] == ch {
		c.channels = append(c.channels[:i], c.channels[i+1:]...)
		return nil
	}
	return errors.New("not in channel")
}

func (c *Twitch) rejoinHandler() {
	const interval = 2 * time.Hour
	ticker := time.NewTicker(interval)
	for {
		select {
		case <-c.quit:
			ticker.Stop()
			return
		case <-ticker.C:
			c.ChLock.RLock()
			for _, ch := range c.channels {
				ch = strings.ToLower(ch)
				if err := c.send("JOIN #" + ch); err != nil {
					log.Println(err)
					continue
				}
			}
			c.ChLock.RUnlock()
		}
	}
}

// Stop stops the chats
func (c *Twitch) Stop(wg *sync.WaitGroup) {
	close(c.quit)
	c.connLock.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connLock.Unlock()
	wg.Done()
}

func inSlice(s []string, v string) bool {
	for _, sv := range s {
		if strings.EqualFold(sv, v) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// Tideland Go Cells - Behaviors - Unit Tests - Event Rate
//
// Copyright (C) 2010-2016 Frank Mueller / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package behaviors_test

//--------------------
// IMPORTS
//--------------------

import (
	"math/rand"
	"testing"
	"time"

	"github.com/tideland/golib/audit"

	"github.com/tideland/gocells/behaviors"
	"github.com/tideland/gocells/cells"
)

//--------------------
// TESTS
//--------------------

// TestRateBehavior tests the event rate behavior.
func TestRateBehavior(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	env := cells.NewEnvironment("rate-behavior")
	defer env.Stop()

	matches := func(event cells.Event) bool {
		return event.Topic() == "now"
	}
	topics := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "now"}

	env.StartCell("rater", behaviors.NewRateBehavior(matches, 5))
	env.StartCell("collector", behaviors.NewCollectorBehavior(1000))
	env.Subscribe("rater", "collector")

	for i := 0; i < 1000; i++ {
		topic := topics[rand.Intn(len(topics))]
		env.EmitNew("rater", topic, nil)
		time.Sleep(time.Duration(rand.Intn(3)) * time.Millisecond)
	}

	collected, err := env.Request("collector", cells.CollectedTopic, nil, cells.DefaultTimeout)
	assert.Nil(err)
	events := collected.([]behaviors.EventData)
	assert.True(len(events) <= 1000)
	for _, event := range events {
		assert.Equal(event.Topic, "event-rate!")
	}
}

// EOF
<commit_msg>More rate behavior testing<commit_after>// Tideland Go Cells - Behaviors - Unit Tests - Event Rate
//
// Copyright (C) 2010-2016 Frank Mueller / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package behaviors_test

//--------------------
// IMPORTS
//--------------------

import (
	"math/rand"
	"testing"
	"time"

	"github.com/tideland/golib/audit"

	"github.com/tideland/gocells/behaviors"
	"github.com/tideland/gocells/cells"
)

//--------------------
// TESTS
//--------------------

// TestRateBehavior tests the event rate behavior.
func TestRateBehavior(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	env := cells.NewEnvironment("rate-behavior")
	defer env.Stop()

	matches := func(event cells.Event) bool {
		return event.Topic() == "now"
	}
	topics := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "now"}

	env.StartCell("rater", behaviors.NewRateBehavior(matches, 5))
	env.StartCell("collector", behaviors.NewCollectorBehavior(1000))
	env.Subscribe("rater", "collector")

	for i := 0; i < 1000; i++ {
		topic := topics[rand.Intn(len(topics))]
		env.EmitNew("rater", topic, nil)
		time.Sleep(time.Duration(rand.Intn(3)) * time.Millisecond)
	}

	collected, err := env.Request("collector", cells.CollectedTopic, nil, cells.DefaultTimeout)
	assert.Nil(err)
	events := collected.([]behaviors.EventData)
	assert.True(len(events) <= 1000)
	for _, event := range events {
		assert.Equal(event.Topic, "event-rate!")
		_, ok := event.Payload.GetDuration(behaviors.EventRateAveragePayload)
		assert.True(ok)
		_, ok = event.Payload.GetDuration(behaviors.EventRateAveragePayload)
		assert.True(ok)
		_, ok = event.Payload.GetDuration(behaviors.EventRateLowPayload)
		assert.True(ok)
	}
}

// EOF
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/config"
)

type graphNodeExpandedResource struct {
	Index    int
	Resource *config.Resource
	Path     []string
}

func (n *graphNodeExpandedResource) Name() string {
	if n.Index == -1 {
		return n.Resource.Id()
	}

	return fmt.Sprintf("%s #%d", n.Resource.Id(), n.Index)
}

// GraphNodeAddressable impl.
func (n *graphNodeExpandedResource) ResourceAddress() *ResourceAddress {
	// We want this to report the logical index properly, so we must undo the
	// special case from the expand
	index := n.Index
	if index == -1 {
		index = 0
	}
	return &ResourceAddress{
		Path:         n.Path[1:],
		Index:        index,
		InstanceType: TypePrimary,
		Name:         n.Resource.Name,
		Type:         n.Resource.Type,
		Mode:         n.Resource.Mode,
	}
}

// GraphNodeDependable impl.
func (n *graphNodeExpandedResource) DependableName() []string {
	return []string{
		n.Resource.Id(),
		n.stateId(),
	}
}

// GraphNodeDependent impl.
func (n *graphNodeExpandedResource) DependentOn() []string {
	configNode := &GraphNodeConfigResource{Resource: n.Resource}
	result := configNode.DependentOn()

	// Walk the variables to find any count-specific variables we depend on.
	configNode.VarWalk(func(v config.InterpolatedVariable) {
		rv, ok := v.(*config.ResourceVariable)
		if !ok {
			return
		}

		// We only want ourselves
		if rv.ResourceId() != n.Resource.Id() {
			return
		}

		// If this isn't a multi-access (which shouldn't be allowed but
		// is verified elsewhere), then we depend on the specific count
		// of this resource, ignoring ourself (which again should be
		// validated elsewhere).
		if rv.Index > -1 {
			id := fmt.Sprintf("%s.%d", rv.ResourceId(), rv.Index)
			if id != n.stateId() && id != n.stateId()+".0" {
				result = append(result, id)
			}
		}
	})

	return result
}

// GraphNodeProviderConsumer
func (n *graphNodeExpandedResource) ProvidedBy() []string {
	return []string{resourceProvider(n.Resource.Type, n.Resource.Provider)}
}

func (n *graphNodeExpandedResource) StateDependencies() []string {
	depsRaw := n.DependentOn()
	deps := make([]string, 0, len(depsRaw))
	for _, d := range depsRaw {
		// Ignore any variable dependencies
		if strings.HasPrefix(d, "var.") {
			continue
		}

		// This is sad. The dependencies are currently in the format of
		// "module.foo.bar" (the full field). This strips the field off.
		if strings.HasPrefix(d, "module.") {
			parts := strings.SplitN(d, ".", 3)
			d = strings.Join(parts[0:2], ".")
		}
		deps = append(deps, d)
	}

	return deps
}

// instanceInfo is used for EvalTree.
func (n *graphNodeExpandedResource) instanceInfo() *InstanceInfo {
	return &InstanceInfo{Id: n.stateId(), Type: n.Resource.Type}
}

// stateId is the name used for the state key
func (n *graphNodeExpandedResource) stateId() string {
	if n.Index == -1 {
		return n.Resource.Id()
	}

	return fmt.Sprintf("%s.%d", n.Resource.Id(), n.Index)
}

// GraphNodeStateRepresentative impl.
func (n *graphNodeExpandedResource) StateId() []string {
	return []string{n.stateId()}
}

// graphNodeExpandedResourceDestroy represents an expanded resource that
// is to be destroyed.
type graphNodeExpandedResourceDestroy struct {
	*graphNodeExpandedResource
}

func (n *graphNodeExpandedResourceDestroy) Name() string {
	return fmt.Sprintf("%s (destroy)", n.graphNodeExpandedResource.Name())
}

// GraphNodeEvalable impl.
func (n *graphNodeExpandedResourceDestroy) EvalTree() EvalNode {
	info := n.instanceInfo()
	info.uniqueExtra = "destroy"

	var diffApply *InstanceDiff
	var provider ResourceProvider
	var state *InstanceState
	var err error
	return &EvalOpFilter{
		Ops: []walkOperation{walkApply, walkDestroy},
		Node: &EvalSequence{
			Nodes: []EvalNode{
				// Get the saved diff for apply
				&EvalReadDiff{
					Name: n.stateId(),
					Diff: &diffApply,
				},

				// Filter the diff so we only get the destroy
				&EvalFilterDiff{
					Diff:    &diffApply,
					Output:  &diffApply,
					Destroy: true,
				},

				// If we're not destroying, then compare diffs
				&EvalIf{
					If: func(ctx EvalContext) (bool, error) {
						if diffApply != nil && diffApply.GetDestroy() {
							return true, nil
						}

						return true, EvalEarlyExitError{}
					},
					Then: EvalNoop{},
				},

				// Load the instance info so we have the module path set
				&EvalInstanceInfo{Info: info},

				&EvalGetProvider{
					Name:   n.ProvidedBy()[0],
					Output: &provider,
				},
				&EvalReadState{
					Name:   n.stateId(),
					Output: &state,
				},
				&EvalRequireState{
					State: &state,
				},
				// Make sure we handle data sources properly.
				&EvalIf{
					If: func(ctx EvalContext) (bool, error) {
						if n.Resource.Mode == config.DataResourceMode {
							return true, nil
						}

						return false, nil
					},

					Then: &EvalReadDataApply{
						Info:     info,
						Diff:     &diffApply,
						Provider: &provider,
						Output:   &state,
					},
					Else: &EvalApply{
						Info:     info,
						State:    &state,
						Diff:     &diffApply,
						Provider: &provider,
						Output:   &state,
						Error:    &err,
					},
				},
				&EvalWriteState{
					Name:         n.stateId(),
					ResourceType: n.Resource.Type,
					Provider:     n.Resource.Provider,
					Dependencies: n.StateDependencies(),
					State:        &state,
				},
				&EvalApplyPost{
					Info:  info,
					State: &state,
					Error: &err,
				},
			},
		},
	}
}
<commit_msg>terraform: keep pruning out lines<commit_after>package terraform

import (
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/config"
)

type graphNodeExpandedResource struct {
	Index    int
	Resource *config.Resource
	Path     []string
}

func (n *graphNodeExpandedResource) Name() string {
	if n.Index == -1 {
		return n.Resource.Id()
	}

	return fmt.Sprintf("%s #%d", n.Resource.Id(), n.Index)
}

// GraphNodeDependent impl.
func (n *graphNodeExpandedResource) DependentOn() []string {
	configNode := &GraphNodeConfigResource{Resource: n.Resource}
	result := configNode.DependentOn()

	// Walk the variables to find any count-specific variables we depend on.
	configNode.VarWalk(func(v config.InterpolatedVariable) {
		rv, ok := v.(*config.ResourceVariable)
		if !ok {
			return
		}

		// We only want ourselves
		if rv.ResourceId() != n.Resource.Id() {
			return
		}

		// If this isn't a multi-access (which shouldn't be allowed but
		// is verified elsewhere), then we depend on the specific count
		// of this resource, ignoring ourself (which again should be
		// validated elsewhere).
		if rv.Index > -1 {
			id := fmt.Sprintf("%s.%d", rv.ResourceId(), rv.Index)
			if id != n.stateId() && id != n.stateId()+".0" {
				result = append(result, id)
			}
		}
	})

	return result
}

func (n *graphNodeExpandedResource) StateDependencies() []string {
	depsRaw := n.DependentOn()
	deps := make([]string, 0, len(depsRaw))
	for _, d := range depsRaw {
		// Ignore any variable dependencies
		if strings.HasPrefix(d, "var.") {
			continue
		}

		// This is sad. The dependencies are currently in the format of
		// "module.foo.bar" (the full field). This strips the field off.
		if strings.HasPrefix(d, "module.") {
			parts := strings.SplitN(d, ".", 3)
			d = strings.Join(parts[0:2], ".")
		}
		deps = append(deps, d)
	}

	return deps
}

// stateId is the name used for the state key
func (n *graphNodeExpandedResource) stateId() string {
	if n.Index == -1 {
		return n.Resource.Id()
	}

	return fmt.Sprintf("%s.%d", n.Resource.Id(), n.Index)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 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 tchannel_test

import (
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/stretchr/testify/assert"
	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"
	"golang.org/x/net/context"
)

func TestRequestStateRetry(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	server := testutils.NewServer(t, nil)
	defer server.Close()
	server.Register(raw.Wrap(newTestHandler(t)), "echo")

	client := testutils.NewClient(t, nil)
	defer client.Close()

	counter := 0
	sc := client.GetSubChannel(server.PeerInfo().ServiceName)
	err := client.RunWithRetry(ctx, func(ctx context.Context, rs *RequestState) error {
		defer func() { counter++ }()

		assert.Equal(t, counter, len(rs.SelectedPeers), "SelectedPeers should not be reused")

		if counter < 4 {
			client.Peers().Add(testutils.GetClosedHostPort(t))
		} else {
			client.Peers().Add(server.PeerInfo().HostPort)
		}

		_, err := raw.CallV2(ctx, sc, raw.CArgs{
			Method:   "echo",
			CallOptions: &CallOptions{RequestState: rs},
		})
		return err
	})
	assert.NoError(t, err, "RunWithRetry should succeed")
	assert.Equal(t, 5, counter, "RunWithRetry should retry 5 times")
}
<commit_msg>Update TestRequestStateRetry to account for host<commit_after>// Copyright (c) 2015 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 tchannel_test

import (
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/stretchr/testify/assert"
	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"
	"golang.org/x/net/context"
)

func TestRequestStateRetry(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	server := testutils.NewServer(t, nil)
	defer server.Close()
	server.Register(raw.Wrap(newTestHandler(t)), "echo")

	client := testutils.NewClient(t, nil)
	defer client.Close()

	counter := 0
	sc := client.GetSubChannel(server.PeerInfo().ServiceName)
	err := client.RunWithRetry(ctx, func(ctx context.Context, rs *RequestState) error {
		defer func() { counter++ }()

		expectedPeers := counter
		if expectedPeers > 0 {
			// An entry is also added for each host.
			expectedPeers++
		}

		assert.Equal(t, expectedPeers, len(rs.SelectedPeers), "SelectedPeers should not be reused")

		if counter < 4 {
			client.Peers().Add(testutils.GetClosedHostPort(t))
		} else {
			client.Peers().Add(server.PeerInfo().HostPort)
		}

		_, err := raw.CallV2(ctx, sc, raw.CArgs{
			Method:      "echo",
			CallOptions: &CallOptions{RequestState: rs},
		})
		return err
	})
	assert.NoError(t, err, "RunWithRetry should succeed")
	assert.Equal(t, 5, counter, "RunWithRetry should retry 5 times")
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package compress

import (
	"bytes"
	"io/ioutil"
	"log"
	/*"fmt"*/
	"strconv"
	"testing"
)

var TESTS = [...]string{
	"SIX.MIXED.PIXIES.SIFT.SIXTY.PIXIE.DUST.BOXES",
	"output[j], j, k, minor[k] = s[k], j - 1, major[s[k]] + minor[k], -1",
	"now is the time for the truly nice people to come to the party",
	"Wild Beasts and Their Ways, Reminiscences of Europe, Asia, Africa and America — Volume 1",
	"EEEIT..SXXIT.ESDIXOS..IISXMBSIYSIDXTXIUPF.P.",
	"MXIOTXD.SI.SSTEDXIUIX.X.I.XSPISSFYBPEETEI.I.",
}

func TestSuffixTree(t *testing.T) {
	test := func(input string) {
		tree := BuildSuffixTree([]uint8(input))
		edges, nodes := tree.edges, tree.nodes

		for _, edge := range edges {
			if edge.first_index > edge.last_index {
				t.Errorf("first_index is greater than last_index")
			}
			end := nodes[edge.end_node]
			if end != -1 {
				edge.last_index++
			}
			/*fmt.Printf("%v %v %v '%v'\n", edge.start_node, edge.end_node, end, input[edge.first_index:edge.last_index])*/
		}

		if index := tree.Index(input); index != 0 {
			t.Errorf("index of %v is %v; should be 0", input, index)
		}
	}
	test("banana")
	test("the frightened Mouse splashed his way through the")
}

func TestBurrowsWheeler(t *testing.T) {
	test := func(input string) {
		buffer := make([]byte, len(input))
		copy(buffer, input)
		tree := BuildSuffixTree(buffer)
		bw, sentinel := tree.BurrowsWheelerCoder()
		index, out_buffer := 0, make([]byte, len(buffer)+1)
		for b := range bw {
			out_buffer[index] = b
			index++
		}
		s := <-sentinel
		for b, c := out_buffer[s], s+1; c < len(out_buffer); c++ {
			out_buffer[c], b = b, out_buffer[c]
		}

		/*fmt.Println(strconv.QuoteToASCII(string(out_buffer)))*/
		original := burrowsWheelerDecoder(out_buffer, s)
		if bytes.Compare(buffer, original) != 0 {
			t.Errorf("should be '%v'; got '%v'", input, strconv.QuoteToASCII(string(original)))
		}
	}
	for _, v := range TESTS {
		test(v)
	}
}

const repeated = 10000

func TestBijectiveBurrowsWheeler(t *testing.T) {
	input, output := make(chan []byte), make(chan []byte, 2)
	coder, decoder := BijectiveBurrowsWheelerCoder(input), BijectiveBurrowsWheelerDecoder(output)
	test := func(buffer []byte) {
		for c := 0; c < repeated; c++ {
			input <- buffer
			<-coder.Input
		}

		in := make([]byte, len(buffer))
		for c := 0; c < repeated; c++ {
			output <- in
			output <- nil
			for _, i := range buffer {
				decoder.Output(i)
			}
			copy(buffer, in)
		}
	}
	for _, v := range TESTS {
		buffer := make([]byte, len(v))
		copy(buffer, []byte(v))
		test(buffer)
		if string(buffer) != v {
			t.Errorf("should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(buffer)))
		}
	}
	close(input)
	close(output)
}

func TestMoveToFront(t *testing.T) {
	test := func(buffer []byte) {
		input, output := make(chan []byte), make(chan []byte, 1)
		coder := BijectiveBurrowsWheelerCoder(input).MoveToFrontCoder()
		decoder := BijectiveBurrowsWheelerDecoder(output).MoveToFrontDecoder()

		input <- buffer
		close(input)
		output <- buffer
		close(output)
		for out := range coder.Input {
			for _, symbol := range out {
				decoder.Output(symbol)
			}
		}
	}
	for _, v := range TESTS {
		buffer := make([]byte, len(v))
		copy(buffer, []byte(v))
		test(buffer)
		if string(buffer) != v {
			t.Errorf("inverse should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(buffer)))
		}
	}
}

func TestCode16(t *testing.T) {
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol{'B': {11, 0, 1},
		'I':    {11, 1, 2},
		'L':    {11, 2, 4},
		' ':    {11, 4, 5},
		'G':    {11, 5, 6},
		'A':    {11, 6, 7},
		'T':    {11, 7, 8},
		'E':    {11, 8, 9},
		'S':    {11, 9, 10},
		'\x00': {11, 10, 11}}
	in, buffer := make(chan []Symbol), &bytes.Buffer{}
	go func() {
		input := make([]Symbol, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model{Input: in}.Code(buffer)
	if compressed := [...]byte{120, 253, 188, 155, 248}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint16) Symbol {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol{}
				} else {
					return symbol
				}
			}
		}
		return Symbol{}
	}
	Model{Scale: 11, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestCode32(t *testing.T) {
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol32{'B': {11, 0, 1},
		'I':    {11, 1, 2},
		'L':    {11, 2, 4},
		' ':    {11, 4, 5},
		'G':    {11, 5, 6},
		'A':    {11, 6, 7},
		'T':    {11, 7, 8},
		'E':    {11, 8, 9},
		'S':    {11, 9, 10},
		'\x00': {11, 10, 11}}

	in, buffer := make(chan []Symbol32), &bytes.Buffer{}
	go func() {
		input := make([]Symbol32, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model32{Input: in}.Code(buffer)
	if compressed := [...]byte{120, 254, 27, 129, 174}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint32) Symbol32 {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol32{}
				} else {
					return symbol
				}
			}
		}
		return Symbol32{}
	}
	Model32{Scale: 11, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestFiltered(t *testing.T) {
	testFiltered := func(test string, depth int) {
		t.Log(test, len(test))
		input := make([]uint16, len(test))
		testBytes := []byte(test)
		for i := range testBytes {
			input[i] = uint16(testBytes[i])
		}
		symbols, buffer := make(chan []uint16, 1), &bytes.Buffer{}
		symbols <- input
		close(symbols)
		Coder16{Alphabit: 256, Input: symbols}.FilteredAdaptiveCoder(NewCDF16(depth, true)).Code(buffer)
		t.Log(buffer.Len())

		out, i := make([]byte, len(test)), 0
		output := func(symbol uint16) bool {
			out[i] = byte(symbol)
			i++
			return i >= len(test)
		}
		Coder16{Alphabit: 256, Output: output}.FilteredAdaptiveDecoder(NewCDF16(depth, true)).Decode(buffer)
		t.Log(string(out))
		if string(out) != test {
			t.Errorf("%v != %v", string(out), test)
		}
	}

	d, err := ioutil.ReadFile("bench/alice30.txt")
	if err != nil {
		log.Fatal(err)
	}
	tests := append(TESTS[:], string(d))

	for _, test := range tests {
		for i := 0; i < 3; i++ {
			testFiltered(test, i)
		}
	}
}

func TestMark1(t *testing.T) {
	for _, v := range TESTS {
		output, buffer := make([]byte, len(v)), &bytes.Buffer{}
		Mark1Compress16([]byte(v), buffer)
		Mark1Decompress16(buffer, output)
		if string(output) != v {
			t.Errorf("should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(output)))
		}
	}
}
<commit_msg>Added tests for fixed scale coders<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package compress

import (
	"bytes"
	"io/ioutil"
	"log"
	/*"fmt"*/
	"strconv"
	"testing"
)

var TESTS = [...]string{
	"SIX.MIXED.PIXIES.SIFT.SIXTY.PIXIE.DUST.BOXES",
	"output[j], j, k, minor[k] = s[k], j - 1, major[s[k]] + minor[k], -1",
	"now is the time for the truly nice people to come to the party",
	"Wild Beasts and Their Ways, Reminiscences of Europe, Asia, Africa and America — Volume 1",
	"EEEIT..SXXIT.ESDIXOS..IISXMBSIYSIDXTXIUPF.P.",
	"MXIOTXD.SI.SSTEDXIUIX.X.I.XSPISSFYBPEETEI.I.",
}

func TestSuffixTree(t *testing.T) {
	test := func(input string) {
		tree := BuildSuffixTree([]uint8(input))
		edges, nodes := tree.edges, tree.nodes

		for _, edge := range edges {
			if edge.first_index > edge.last_index {
				t.Errorf("first_index is greater than last_index")
			}
			end := nodes[edge.end_node]
			if end != -1 {
				edge.last_index++
			}
			/*fmt.Printf("%v %v %v '%v'\n", edge.start_node, edge.end_node, end, input[edge.first_index:edge.last_index])*/
		}

		if index := tree.Index(input); index != 0 {
			t.Errorf("index of %v is %v; should be 0", input, index)
		}
	}
	test("banana")
	test("the frightened Mouse splashed his way through the")
}

func TestBurrowsWheeler(t *testing.T) {
	test := func(input string) {
		buffer := make([]byte, len(input))
		copy(buffer, input)
		tree := BuildSuffixTree(buffer)
		bw, sentinel := tree.BurrowsWheelerCoder()
		index, out_buffer := 0, make([]byte, len(buffer)+1)
		for b := range bw {
			out_buffer[index] = b
			index++
		}
		s := <-sentinel
		for b, c := out_buffer[s], s+1; c < len(out_buffer); c++ {
			out_buffer[c], b = b, out_buffer[c]
		}

		/*fmt.Println(strconv.QuoteToASCII(string(out_buffer)))*/
		original := burrowsWheelerDecoder(out_buffer, s)
		if bytes.Compare(buffer, original) != 0 {
			t.Errorf("should be '%v'; got '%v'", input, strconv.QuoteToASCII(string(original)))
		}
	}
	for _, v := range TESTS {
		test(v)
	}
}

const repeated = 10000

func TestBijectiveBurrowsWheeler(t *testing.T) {
	input, output := make(chan []byte), make(chan []byte, 2)
	coder, decoder := BijectiveBurrowsWheelerCoder(input), BijectiveBurrowsWheelerDecoder(output)
	test := func(buffer []byte) {
		for c := 0; c < repeated; c++ {
			input <- buffer
			<-coder.Input
		}

		in := make([]byte, len(buffer))
		for c := 0; c < repeated; c++ {
			output <- in
			output <- nil
			for _, i := range buffer {
				decoder.Output(i)
			}
			copy(buffer, in)
		}
	}
	for _, v := range TESTS {
		buffer := make([]byte, len(v))
		copy(buffer, []byte(v))
		test(buffer)
		if string(buffer) != v {
			t.Errorf("should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(buffer)))
		}
	}
	close(input)
	close(output)
}

func TestMoveToFront(t *testing.T) {
	test := func(buffer []byte) {
		input, output := make(chan []byte), make(chan []byte, 1)
		coder := BijectiveBurrowsWheelerCoder(input).MoveToFrontCoder()
		decoder := BijectiveBurrowsWheelerDecoder(output).MoveToFrontDecoder()

		input <- buffer
		close(input)
		output <- buffer
		close(output)
		for out := range coder.Input {
			for _, symbol := range out {
				decoder.Output(symbol)
			}
		}
	}
	for _, v := range TESTS {
		buffer := make([]byte, len(v))
		copy(buffer, []byte(v))
		test(buffer)
		if string(buffer) != v {
			t.Errorf("inverse should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(buffer)))
		}
	}
}

func TestCode16(t *testing.T) {
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol{'B': {11, 0, 1},
		'I':    {11, 1, 2},
		'L':    {11, 2, 4},
		' ':    {11, 4, 5},
		'G':    {11, 5, 6},
		'A':    {11, 6, 7},
		'T':    {11, 7, 8},
		'E':    {11, 8, 9},
		'S':    {11, 9, 10},
		'\x00': {11, 10, 11}}
	in, buffer := make(chan []Symbol), &bytes.Buffer{}
	go func() {
		input := make([]Symbol, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model{Input: in}.Code(buffer)
	if compressed := [...]byte{120, 253, 188, 155, 248}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint16) Symbol {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol{}
				} else {
					return symbol
				}
			}
		}
		return Symbol{}
	}
	Model{Scale: 11, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestCode16Fixed(t *testing.T) {
	const fixed = 0x10
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol{'B': {0, 0, 1},
		'I':    {0, 1, 2},
		'L':    {0, 2, 4},
		' ':    {0, 4, 5},
		'G':    {0, 5, 6},
		'A':    {0, 6, 7},
		'T':    {0, 7, 8},
		'E':    {0, 8, 9},
		'S':    {0, 9, 10},
		'\x00': {0, 10, 11}}
	in, buffer := make(chan []Symbol), &bytes.Buffer{}
	go func() {
		input := make([]Symbol, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model{Input: in, Fixed: fixed}.Code(buffer)
	if compressed := [...]byte{0, 5, 0, 2, 0, 2, 0, 0, 0, 8, 0, 0, 0, 12, 0, 14, 0, 16, 0, 18, 0, 20, 128}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint16) Symbol {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol{}
				} else {
					return symbol
				}
			}
		}
		return Symbol{}
	}
	Model{Fixed: fixed, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestCode32(t *testing.T) {
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol32{'B': {11, 0, 1},
		'I':    {11, 1, 2},
		'L':    {11, 2, 4},
		' ':    {11, 4, 5},
		'G':    {11, 5, 6},
		'A':    {11, 6, 7},
		'T':    {11, 7, 8},
		'E':    {11, 8, 9},
		'S':    {11, 9, 10},
		'\x00': {11, 10, 11}}

	in, buffer := make(chan []Symbol32), &bytes.Buffer{}
	go func() {
		input := make([]Symbol32, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model32{Input: in}.Code(buffer)
	if compressed := [...]byte{120, 254, 27, 129, 174}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint32) Symbol32 {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol32{}
				} else {
					return symbol
				}
			}
		}
		return Symbol32{}
	}
	Model32{Scale: 11, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestCode32Fixed(t *testing.T) {
	const fixed = 0x10
	test := []byte("GLIB BATES\x00")
	var table = [256]Symbol32{'B': {0, 0, 1},
		'I':    {0, 1, 2},
		'L':    {0, 2, 4},
		' ':    {0, 4, 5},
		'G':    {0, 5, 6},
		'A':    {0, 6, 7},
		'T':    {0, 7, 8},
		'E':    {0, 8, 9},
		'S':    {0, 9, 10},
		'\x00': {0, 10, 11}}

	in, buffer := make(chan []Symbol32), &bytes.Buffer{}
	go func() {
		input := make([]Symbol32, len(test))
		for i, s := range test {
			input[i] = table[s]
		}
		in <- input
		close(in)
	}()
	Model32{Input: in, Fixed: fixed}.Code(buffer)
	if compressed := [...]byte{0, 5, 0, 2, 0, 2, 0, 0, 0, 8, 0, 0, 0, 12, 0, 14, 0, 16, 0, 18, 0, 20, 128}; bytes.Compare(compressed[:], buffer.Bytes()) != 0 {
		t.Errorf("arithmetic coding failed")
	}

	uncompressed, j := make([]byte, len(test)), 0
	lookup := func(code uint32) Symbol32 {
		for i, symbol := range table {
			if code >= symbol.Low && code < symbol.High {
				uncompressed[j], j = byte(i), j+1
				if i == 0 {
					return Symbol32{}
				} else {
					return symbol
				}
			}
		}
		return Symbol32{}
	}
	Model32{Fixed: fixed, Output: lookup}.Decode(buffer)
	if bytes.Compare(test, uncompressed) != 0 {
		t.Errorf("arithmetic decoding failed")
	}
}

func TestFiltered(t *testing.T) {
	testFiltered := func(test string, depth int) {
		t.Log(test, len(test))
		input := make([]uint16, len(test))
		testBytes := []byte(test)
		for i := range testBytes {
			input[i] = uint16(testBytes[i])
		}
		symbols, buffer := make(chan []uint16, 1), &bytes.Buffer{}
		symbols <- input
		close(symbols)
		Coder16{Alphabit: 256, Input: symbols}.FilteredAdaptiveCoder(NewCDF16(depth, true)).Code(buffer)
		t.Log(buffer.Len())

		out, i := make([]byte, len(test)), 0
		output := func(symbol uint16) bool {
			out[i] = byte(symbol)
			i++
			return i >= len(test)
		}
		Coder16{Alphabit: 256, Output: output}.FilteredAdaptiveDecoder(NewCDF16(depth, true)).Decode(buffer)
		t.Log(string(out))
		if string(out) != test {
			t.Errorf("%v != %v", string(out), test)
		}
	}

	d, err := ioutil.ReadFile("bench/alice30.txt")
	if err != nil {
		log.Fatal(err)
	}
	tests := append(TESTS[:], string(d))

	for _, test := range tests {
		for i := 0; i < 3; i++ {
			testFiltered(test, i)
		}
	}
}

func TestMark1(t *testing.T) {
	for _, v := range TESTS {
		output, buffer := make([]byte, len(v)), &bytes.Buffer{}
		Mark1Compress16([]byte(v), buffer)
		Mark1Decompress16(buffer, output)
		if string(output) != v {
			t.Errorf("should be '%v'; got '%v'", v, strconv.QuoteToASCII(string(output)))
		}
	}
}
<|endoftext|>


<commit_before>package deploystack

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"google.golang.org/api/compute/v1"
)

var computeService *compute.Service

// DiskProjects are the list of projects for disk images for Compute Engine
var DiskProjects = LabeledValues{
	LabeledValue{Label: "CentOS", Value: "centos-cloud"},
	LabeledValue{Label: "Container-Optimized OS (COS)", Value: "cos-cloud"},
	LabeledValue{Label: "Debian", Value: "debian-cloud"},
	LabeledValue{Label: "Fedora CoreOS", Value: "fedora-coreos-cloud"},
	LabeledValue{Label: "Red Hat Enterprise Linux (RHEL)", Value: "rhel-cloud"},
	LabeledValue{Label: "Red Hat Enterprise Linux (RHEL) for SAP", Value: "rhel-sap-cloud"},
	LabeledValue{Label: "Rocky Linux", Value: "rocky-linux-cloud"},
	LabeledValue{Label: "SQL Server", Value: "windows-sql-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES)", Value: "suse-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES) for SAP", Value: "suse-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES) BYOS", Value: "suse-byos-cloud"},
	LabeledValue{Label: "Ubuntu LTS", Value: "ubuntu-os-cloud"},
	LabeledValue{Label: "Ubuntu Pro", Value: "ubuntu-os-pro-cloud"},
	LabeledValue{Label: "Windows Server", Value: "windows-cloud"},
}

// regionsCompute will return a list of regions for Compute Engine
func regionsCompute(project string) ([]string, error) {
	resp := []string{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	results, err := svc.Regions.List(project).Do()
	if err != nil {
		return resp, err
	}

	for _, v := range results.Items {
		resp = append(resp, v.Name)
	}

	sort.Strings(resp)

	return resp, nil
}

func getComputeService(project string) (*compute.Service, error) {
	if computeService != nil {
		return computeService, nil
	}

	if err := ServiceEnable(project, "compute.googleapis.com"); err != nil {
		return nil, fmt.Errorf("error activating service for polling: %s", err)
	}

	ctx := context.Background()
	svc, err := compute.NewService(ctx, opts)
	if err != nil {
		return nil, err
	}

	computeService = svc

	return svc, nil
}

// zones will return a list of zones in a given region
func zones(project, region string) ([]string, error) {
	resp := []string{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	filter := fmt.Sprintf("name=%s*", region)

	results, err := svc.Zones.List(project).Filter(filter).Do()
	if err != nil {
		return resp, err
	}

	for _, v := range results.Items {
		resp = append(resp, v.Name)
	}

	sort.Strings(resp)

	return resp, nil
}

func machineTypes(project, zone string) (*compute.MachineTypeList, error) {
	resp := &compute.MachineTypeList{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	results, err := svc.MachineTypes.List(project, zone).Do()
	if err != nil {
		return resp, err
	}

	return results, nil
}

func formatMBToGB(i int64) string {
	return fmt.Sprintf("%d GB", i/1024)
}

// TODO: Write tests for this function
func images(project string) (*compute.ImageList, error) {
	resp := &compute.ImageList{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	results, err := svc.Images.List(project).Do()
	if err != nil {
		return resp, err
	}

	return results, nil
}

func GetListOfMachineTypeFamily(imgs *compute.MachineTypeList) LabeledValues {
	fam := make(map[string]string)
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		parts := strings.Split(v.Name, "-")

		key := fmt.Sprintf("%s %s", parts[0], parts[1])
		fam[key] = fmt.Sprintf("%s-%s", parts[0], parts[1])
	}

	for i, v := range fam {
		if i == "" {
			continue
		}
		lb = append(lb, LabeledValue{v, i})
	}
	lb.sort()
	return lb
}

func GetListOfMachineTypeByFamily(imgs *compute.MachineTypeList, family string) LabeledValues {
	lb := LabeledValues{}

	tempTypes := []compute.MachineType{}

	for _, v := range imgs.Items {
		if strings.Contains(v.Name, family) {
			tempTypes = append(tempTypes, *v)
		}
	}

	sort.Slice(tempTypes, func(i, j int) bool {
		return tempTypes[i].GuestCpus < tempTypes[j].GuestCpus
	})

	for _, v := range tempTypes {
		if strings.Contains(v.Name, family) {
			value := v.Name
			label := fmt.Sprintf("%s %s", v.Name, v.Description)
			lb = append(lb, LabeledValue{value, label})
		}
	}
	return lb
}

func getListOfImageFamilies(imgs *compute.ImageList) LabeledValues {
	fam := make(map[string]bool)
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		fam[v.Family] = false
	}

	for i := range fam {
		if i == "" {
			continue
		}
		lb = append(lb, LabeledValue{i, i})
	}
	lb.sort()
	return lb
}

func getListOfImageTypesByFamily(imgs *compute.ImageList, project, family string) LabeledValues {
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		if v.Family == family {
			value := fmt.Sprintf("%s/%s", project, v.Name)
			lb = append(lb, LabeledValue{value, v.Name})
		}
	}

	last := lb[len(lb)-1]
	last.Label = fmt.Sprintf("%s (Latest)", last.Label)
	lb[len(lb)-1] = last
	lb.sort()

	return lb
}
<commit_msg>Added the ability to find the latest image in the default family<commit_after>package deploystack

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"google.golang.org/api/compute/v1"
)

var computeService *compute.Service

// DiskProjects are the list of projects for disk images for Compute Engine
var DiskProjects = LabeledValues{
	LabeledValue{Label: "CentOS", Value: "centos-cloud"},
	LabeledValue{Label: "Container-Optimized OS (COS)", Value: "cos-cloud"},
	LabeledValue{Label: "Debian", Value: "debian-cloud"},
	LabeledValue{Label: "Fedora CoreOS", Value: "fedora-coreos-cloud"},
	LabeledValue{Label: "Red Hat Enterprise Linux (RHEL)", Value: "rhel-cloud"},
	LabeledValue{Label: "Red Hat Enterprise Linux (RHEL) for SAP", Value: "rhel-sap-cloud"},
	LabeledValue{Label: "Rocky Linux", Value: "rocky-linux-cloud"},
	LabeledValue{Label: "SQL Server", Value: "windows-sql-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES)", Value: "suse-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES) for SAP", Value: "suse-cloud"},
	LabeledValue{Label: "SUSE Linux Enterprise Server (SLES) BYOS", Value: "suse-byos-cloud"},
	LabeledValue{Label: "Ubuntu LTS", Value: "ubuntu-os-cloud"},
	LabeledValue{Label: "Ubuntu Pro", Value: "ubuntu-os-pro-cloud"},
	LabeledValue{Label: "Windows Server", Value: "windows-cloud"},
}

// regionsCompute will return a list of regions for Compute Engine
func regionsCompute(project string) ([]string, error) {
	resp := []string{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	results, err := svc.Regions.List(project).Do()
	if err != nil {
		return resp, err
	}

	for _, v := range results.Items {
		resp = append(resp, v.Name)
	}

	sort.Strings(resp)

	return resp, nil
}

func getComputeService(project string) (*compute.Service, error) {
	if computeService != nil {
		return computeService, nil
	}

	if err := ServiceEnable(project, "compute.googleapis.com"); err != nil {
		return nil, fmt.Errorf("error activating service for polling: %s", err)
	}

	ctx := context.Background()
	svc, err := compute.NewService(ctx, opts)
	if err != nil {
		return nil, err
	}

	computeService = svc

	return svc, nil
}

// zones will return a list of zones in a given region
func zones(project, region string) ([]string, error) {
	resp := []string{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	filter := fmt.Sprintf("name=%s*", region)

	results, err := svc.Zones.List(project).Filter(filter).Do()
	if err != nil {
		return resp, err
	}

	for _, v := range results.Items {
		resp = append(resp, v.Name)
	}

	sort.Strings(resp)

	return resp, nil
}

func machineTypes(project, zone string) (*compute.MachineTypeList, error) {
	resp := &compute.MachineTypeList{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	results, err := svc.MachineTypes.List(project, zone).Do()
	if err != nil {
		return resp, err
	}

	return results, nil
}

func formatMBToGB(i int64) string {
	return fmt.Sprintf("%d GB", i/1024)
}

// TODO: Write tests for this function
func images(project, imageproject string) (*compute.ImageList, error) {
	resp := &compute.ImageList{}

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}
	results, err := svc.Images.List(imageproject).Do()
	if err != nil {
		return resp, err
	}

	return results, nil
}

func GetLatestImage(project, imageproject string) (string, error) {
	resp := ""

	svc, err := getComputeService(project)
	if err != nil {
		return resp, err
	}

	filter := fmt.Sprintf("(family=\"%s\")", DefaultImageFamily)
	results, err := svc.Images.List(imageproject).Filter(filter).Do()
	if err != nil {
		return resp, err
	}

	sort.Slice(results.Items, func(i, j int) bool {
		return results.Items[i].CreationTimestamp > results.Items[j].CreationTimestamp
	})

	for _, v := range results.Items {
		if v.Deprecated == nil {
			return v.Name, nil
		}
	}

	return "", fmt.Errorf("error: could not find ")
}

func GetListOfMachineTypeFamily(imgs *compute.MachineTypeList) LabeledValues {
	fam := make(map[string]string)
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		parts := strings.Split(v.Name, "-")

		key := fmt.Sprintf("%s %s", parts[0], parts[1])
		fam[key] = fmt.Sprintf("%s-%s", parts[0], parts[1])
	}

	for i, v := range fam {
		if i == "" {
			continue
		}
		lb = append(lb, LabeledValue{v, i})
	}
	lb.sort()
	return lb
}

func GetListOfMachineTypeByFamily(imgs *compute.MachineTypeList, family string) LabeledValues {
	lb := LabeledValues{}

	tempTypes := []compute.MachineType{}

	for _, v := range imgs.Items {
		if strings.Contains(v.Name, family) {
			tempTypes = append(tempTypes, *v)
		}
	}

	sort.Slice(tempTypes, func(i, j int) bool {
		return tempTypes[i].GuestCpus < tempTypes[j].GuestCpus
	})

	for _, v := range tempTypes {
		if strings.Contains(v.Name, family) {
			value := v.Name
			label := fmt.Sprintf("%s %s", v.Name, v.Description)
			lb = append(lb, LabeledValue{value, label})
		}
	}
	return lb
}

func getListOfImageFamilies(imgs *compute.ImageList) LabeledValues {
	fam := make(map[string]bool)
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		fam[v.Family] = false
	}

	for i := range fam {
		if i == "" {
			continue
		}
		lb = append(lb, LabeledValue{i, i})
	}
	lb.sort()
	return lb
}

func getListOfImageTypesByFamily(imgs *compute.ImageList, project, family string) LabeledValues {
	lb := LabeledValues{}

	for _, v := range imgs.Items {
		if v.Family == family {
			value := fmt.Sprintf("%s/%s", project, v.Name)
			lb = append(lb, LabeledValue{value, v.Name})
		}
	}

	last := lb[len(lb)-1]
	last.Label = fmt.Sprintf("%s (Latest)", last.Label)
	lb[len(lb)-1] = last
	lb.sort()

	return lb
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 app

import (
	"net/http"
	"testing"
	"time"

	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

var errParseInt = `strconv.ParseInt: parsing "string": invalid syntax`

func TestParseTraceQuery(t *testing.T) {
	timeNow := time.Now()
	const noErr = ""
	tests := []struct {
		urlStr        string
		errMsg        string
		expectedQuery *traceQueryParameters
	}{
		{"", "parameter 'service' is required", nil},
		{"x?service=service&start=string", errParseInt, nil},
		{"x?service=service&end=string", errParseInt, nil},
		{"x?service=service&limit=string", errParseInt, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=20", "cannot not parse minDuration: time: missing unit in duration 20", nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=20s&maxDuration=30", "cannot not parse maxDuration: time: missing unit in duration 30", nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tag=x:y&tag=k&log=k:v&log=k", `malformed 'tag' parameter, expecting key:value, received: k`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=25s&maxDuration=1s", `'maxDuration' should be greater than 'minDuration'`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tag=x:y", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		// tags=JSON with a non-string value 123
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags={"x":123}`, "malformed 'tags' parameter, cannot unmarshal JSON: json: cannot unmarshal number into Go value of type string", nil},
		// tags=JSON
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags={"x":"y"}`, noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		// tags=url_encode(JSON)
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags=%7B%22x%22%3A%22y%22%7D`, noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=10s&maxDuration=20s", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					DurationMin:   10 * time.Second,
					DurationMax:   20 * time.Second,
					Tags:          make(map[string]string),
				},
			},
		},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=10s", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					DurationMin:   10 * time.Second,
					Tags:          make(map[string]string),
				},
			},
		},
		// trace ID in upper/lower case
		{"x?traceID=1f00&traceID=1E00", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					NumTraces:    100,
					StartTimeMin: timeNow,
					StartTimeMax: timeNow,
					Tags:         make(map[string]string),
				},
				traceIDs: []model.TraceID{
					model.NewTraceID(0, 0x1f00),
					model.NewTraceID(0, 0x1e00),
				},
			},
		},
		{"x?traceID=100&traceID=x200", `cannot parse traceID param: strconv.ParseUint: parsing "x200": invalid syntax`,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					StartTimeMin: time.Unix(0, 0),
					StartTimeMax: time.Unix(0, 0),
					NumTraces:    100,
					Tags:         make(map[string]string),
				},
				traceIDs: []model.TraceID{
					model.NewTraceID(0, 0x100),
					model.NewTraceID(0, 0x200),
				},
			},
		},
	}
	for _, tc := range tests {
		test := tc // capture loop var
		t.Run(test.urlStr, func(t *testing.T) {
			request, err := http.NewRequest(http.MethodGet, test.urlStr, nil)
			assert.NoError(t, err)
			parser := &queryParser{
				timeNow: func() time.Time {
					return timeNow
				},
			}
			actualQuery, err := parser.parse(request)
			if test.errMsg == "" {
				assert.NoError(t, err)
				if !assert.Equal(t, test.expectedQuery, actualQuery) {
					for _, s := range pretty.Diff(test.expectedQuery, actualQuery) {
						t.Log(s)
					}
				}
			} else {
				assert.EqualError(t, err, test.errMsg)
			}
		})
	}
}
<commit_msg>Fix error equals (#2429)<commit_after>// Copyright (c) 2019 The Jaeger Authors.
// Copyright (c) 2017 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 app

import (
	"fmt"
	"net/http"
	"regexp"
	"testing"
	"time"

	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/storage/spanstore"
)

var errParseInt = `strconv.ParseInt: parsing "string": invalid syntax`

func TestParseTraceQuery(t *testing.T) {
	timeNow := time.Now()
	const noErr = ""
	tests := []struct {
		urlStr        string
		errMsg        string
		expectedQuery *traceQueryParameters
	}{
		{"", "parameter 'service' is required", nil},
		{"x?service=service&start=string", errParseInt, nil},
		{"x?service=service&end=string", errParseInt, nil},
		{"x?service=service&limit=string", errParseInt, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=20", `cannot not parse minDuration: time: missing unit in duration "?20"?$`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=20s&maxDuration=30", `cannot not parse maxDuration: time: missing unit in duration "?30"?$`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tag=x:y&tag=k&log=k:v&log=k", `malformed 'tag' parameter, expecting key:value, received: k`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=25s&maxDuration=1s", `'maxDuration' should be greater than 'minDuration'`, nil},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tag=x:y", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		// tags=JSON with a non-string value 123
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags={"x":123}`, "malformed 'tags' parameter, cannot unmarshal JSON: json: cannot unmarshal number into Go value of type string", nil},
		// tags=JSON
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags={"x":"y"}`, noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		// tags=url_encode(JSON)
		{`x?service=service&start=0&end=0&operation=operation&limit=200&tag=k:v&tags=%7B%22x%22%3A%22y%22%7D`, noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					Tags:          map[string]string{"k": "v", "x": "y"},
				},
			},
		},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=10s&maxDuration=20s", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					DurationMin:   10 * time.Second,
					DurationMax:   20 * time.Second,
					Tags:          make(map[string]string),
				},
			},
		},
		{"x?service=service&start=0&end=0&operation=operation&limit=200&minDuration=10s", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					ServiceName:   "service",
					OperationName: "operation",
					StartTimeMin:  time.Unix(0, 0),
					StartTimeMax:  time.Unix(0, 0),
					NumTraces:     200,
					DurationMin:   10 * time.Second,
					Tags:          make(map[string]string),
				},
			},
		},
		// trace ID in upper/lower case
		{"x?traceID=1f00&traceID=1E00", noErr,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					NumTraces:    100,
					StartTimeMin: timeNow,
					StartTimeMax: timeNow,
					Tags:         make(map[string]string),
				},
				traceIDs: []model.TraceID{
					model.NewTraceID(0, 0x1f00),
					model.NewTraceID(0, 0x1e00),
				},
			},
		},
		{"x?traceID=100&traceID=x200", `cannot parse traceID param: strconv.ParseUint: parsing "x200": invalid syntax`,
			&traceQueryParameters{
				TraceQueryParameters: spanstore.TraceQueryParameters{
					StartTimeMin: time.Unix(0, 0),
					StartTimeMax: time.Unix(0, 0),
					NumTraces:    100,
					Tags:         make(map[string]string),
				},
				traceIDs: []model.TraceID{
					model.NewTraceID(0, 0x100),
					model.NewTraceID(0, 0x200),
				},
			},
		},
	}
	for _, tc := range tests {
		test := tc // capture loop var
		t.Run(test.urlStr, func(t *testing.T) {
			request, err := http.NewRequest(http.MethodGet, test.urlStr, nil)
			assert.NoError(t, err)
			parser := &queryParser{
				timeNow: func() time.Time {
					return timeNow
				},
			}
			actualQuery, err := parser.parse(request)
			if test.errMsg == "" {
				assert.NoError(t, err)
				if !assert.Equal(t, test.expectedQuery, actualQuery) {
					for _, s := range pretty.Diff(test.expectedQuery, actualQuery) {
						t.Log(s)
					}
				}
			} else {
				matched, matcherr := regexp.MatchString(test.errMsg, err.Error())
				require.NoError(t, matcherr)
				assert.True(t, matched, fmt.Sprintf("Error \"%s\" should match \"%s\"", err.Error(), test.errMsg))
			}
		})
	}
}
<|endoftext|>


<commit_before>package manager

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dhttp"
	"github.com/datawire/dlib/dlog"
	rpc "github.com/telepresenceio/telepresence/rpc/v2/manager"
	"github.com/telepresenceio/telepresence/rpc/v2/systema"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/internal/mutator"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/internal/watchable"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/telepresenceio/telepresence/v2/pkg/k8sapi"
	"github.com/telepresenceio/telepresence/v2/pkg/version"
)

// Main starts up the traffic manager and blocks until it ends
func Main(ctx context.Context, args ...string) error {
	dlog.Infof(ctx, "Traffic Manager %s [pid:%d]", version.Version, os.Getpid())

	ctx, err := managerutil.LoadEnv(ctx)
	if err != nil {
		return fmt.Errorf("failed to LoadEnv: %w", err)
	}

	cfg, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("unable to get the Kubernetes InClusterConfig: %w", err)
	}
	ki, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return fmt.Errorf("unable to create the Kubernetes Interface from InClusterConfig: %w", err)
	}
	ctx = k8sapi.WithK8sInterface(ctx, ki)

	g := dgroup.NewGroup(ctx, dgroup.GroupConfig{
		EnableSignalHandling: true,
	})
	mgr := NewManager(ctx)

	// Serve HTTP (including gRPC)
	g.Go("httpd", mgr.serveHTTP)

	g.Go("agent-injector", mutator.ServeMutator)

	g.Go("intercept-gc", mgr.runInterceptGCLoop)

	// This goroutine is responsible for informing System A of intercepts (and
	// relevant metadata like domains) that have been garbage collected. This
	// ensures System A doesn't list preview URLs + intercepts that no longer
	// exist.
	g.Go("systema-gc", mgr.runSystemAGCLoop)

	// Wait for exit
	return g.Wait()
}

func (m *Manager) serveHTTP(ctx context.Context) error {
	env := managerutil.GetEnv(ctx)
	host := env.ServerHost
	port := env.ServerPort
	opts := []grpc.ServerOption{}
	if mz, ok := env.MaxReceiveSize.AsInt64(); ok {
		opts = append(opts, grpc.MaxRecvMsgSize(int(mz)))
	}

	grpcHandler := grpc.NewServer(opts...)
	httpHandler := http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello World from: %s\n", r.URL.Path)
	}))
	sc := &dhttp.ServerConfig{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
				grpcHandler.ServeHTTP(w, r)
			} else {
				httpHandler.ServeHTTP(w, r)
			}
		}),
	}

	rpc.RegisterManagerServer(grpcHandler, m)
	grpc_health_v1.RegisterHealthServer(grpcHandler, &HealthChecker{})

	return sc.ListenAndServe(ctx, host+":"+port)
}

func (m *Manager) runInterceptGCLoop(ctx context.Context) error {
	// Loop calling Expire
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			m.expire(ctx)
		case <-ctx.Done():
			return nil
		}
	}
}

func (m *Manager) runSystemAGCLoop(ctx context.Context) error {
	for snapshot := range m.state.WatchIntercepts(ctx, nil) {
		for _, update := range snapshot.Updates {
			// Since all intercepts with a domain require a login, we can use
			// presence of the ApiKey in the interceptInfo to determine all
			// intercepts that we need to inform System A of their deletion
			if update.Delete && update.Value.ApiKey != "" {
				if sa, err := m.systema.Get(); err != nil {
					dlog.Errorln(ctx, "systema: acquire connection:", err)
				} else {
					// First we remove the PreviewDomain if it exists
					if update.Value.PreviewDomain != "" {
						err = m.reapDomain(ctx, sa, update)
						if err != nil {
							dlog.Errorln(ctx, "systema: remove domain:", err)
						}
					}
					// Now we inform SystemA of the intercepts removal
					dlog.Debugf(ctx, "systema: remove intercept: %q", update.Value.Id)
					err = m.reapIntercept(ctx, sa, update)
					if err != nil {
						dlog.Errorln(ctx, "systema: remove intercept:", err)
					}

					// Release the connection we got to delete the domain + intercept
					if err := m.systema.Done(); err != nil {
						dlog.Errorln(ctx, "systema: release management connection:", err)
					}
				}
				// Release the refcount on the proxy connection
				if err := m.systema.Done(); err != nil {
					dlog.Errorln(ctx, "systema: release proxy connection:", err)
				}
			}
		}
	}
	return nil
}

// reapDomain informs SystemA that an intercept with a domain has been garbage collected
func (m *Manager) reapDomain(ctx context.Context, sa systema.SystemACRUDClient, interceptUpdate watchable.InterceptMapUpdate) error {
	// we only reapDomains for intercepts that have been deleted
	if !interceptUpdate.Delete {
		return fmt.Errorf("%s is not being deleted, so the domain was not reaped", interceptUpdate.Value.Id)
	}
	dlog.Debugf(ctx, "systema: removing domain: %q", interceptUpdate.Value.PreviewDomain)
	_, err := sa.RemoveDomain(ctx, &systema.RemoveDomainRequest{
		Domain: interceptUpdate.Value.PreviewDomain,
	})

	if err != nil {
		return err
	}
	return nil
}

// reapIntercept informs SystemA that an intercept has been garbage collected
func (m *Manager) reapIntercept(ctx context.Context, sa systema.SystemACRUDClient, interceptUpdate watchable.InterceptMapUpdate) error {
	// we only reapIntercept for intercepts that have been deleted
	if !interceptUpdate.Delete {
		return fmt.Errorf("%s is not being deleted, so the intercept was not reaped", interceptUpdate.Value.Id)
	}
	dlog.Debugf(ctx, "systema: remove intercept: %q", interceptUpdate.Value.Id)
	_, err := sa.RemoveIntercept(ctx, &systema.InterceptRemoval{
		InterceptId: interceptUpdate.Value.Id,
	})

	// We remove the APIKey whether or not the RemoveIntercept call was successful, so
	// let's do that before we check the error.
	if wasRemoved := m.state.RemoveInterceptAPIKey(interceptUpdate.Value.Id); !wasRemoved {
		dlog.Debugf(ctx, "Intercept ID %s had no APIKey", interceptUpdate.Value.Id)
	}

	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Fix unbalanced systemaPool Get/Done calls in runSystemAGCLoop<commit_after>package manager

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/health/grpc_health_v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dhttp"
	"github.com/datawire/dlib/dlog"
	rpc "github.com/telepresenceio/telepresence/rpc/v2/manager"
	"github.com/telepresenceio/telepresence/rpc/v2/systema"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/internal/mutator"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/internal/watchable"
	"github.com/telepresenceio/telepresence/v2/cmd/traffic/cmd/manager/managerutil"
	"github.com/telepresenceio/telepresence/v2/pkg/k8sapi"
	"github.com/telepresenceio/telepresence/v2/pkg/version"
)

// Main starts up the traffic manager and blocks until it ends
func Main(ctx context.Context, args ...string) error {
	dlog.Infof(ctx, "Traffic Manager %s [pid:%d]", version.Version, os.Getpid())

	ctx, err := managerutil.LoadEnv(ctx)
	if err != nil {
		return fmt.Errorf("failed to LoadEnv: %w", err)
	}

	cfg, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("unable to get the Kubernetes InClusterConfig: %w", err)
	}
	ki, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return fmt.Errorf("unable to create the Kubernetes Interface from InClusterConfig: %w", err)
	}
	ctx = k8sapi.WithK8sInterface(ctx, ki)

	g := dgroup.NewGroup(ctx, dgroup.GroupConfig{
		EnableSignalHandling: true,
	})
	mgr := NewManager(ctx)

	// Serve HTTP (including gRPC)
	g.Go("httpd", mgr.serveHTTP)

	g.Go("agent-injector", mutator.ServeMutator)

	g.Go("intercept-gc", mgr.runInterceptGCLoop)

	// This goroutine is responsible for informing System A of intercepts (and
	// relevant metadata like domains) that have been garbage collected. This
	// ensures System A doesn't list preview URLs + intercepts that no longer
	// exist.
	g.Go("systema-gc", mgr.runSystemAGCLoop)

	// Wait for exit
	return g.Wait()
}

func (m *Manager) serveHTTP(ctx context.Context) error {
	env := managerutil.GetEnv(ctx)
	host := env.ServerHost
	port := env.ServerPort
	opts := []grpc.ServerOption{}
	if mz, ok := env.MaxReceiveSize.AsInt64(); ok {
		opts = append(opts, grpc.MaxRecvMsgSize(int(mz)))
	}

	grpcHandler := grpc.NewServer(opts...)
	httpHandler := http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello World from: %s\n", r.URL.Path)
	}))
	sc := &dhttp.ServerConfig{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
				grpcHandler.ServeHTTP(w, r)
			} else {
				httpHandler.ServeHTTP(w, r)
			}
		}),
	}

	rpc.RegisterManagerServer(grpcHandler, m)
	grpc_health_v1.RegisterHealthServer(grpcHandler, &HealthChecker{})

	return sc.ListenAndServe(ctx, host+":"+port)
}

func (m *Manager) runInterceptGCLoop(ctx context.Context) error {
	// Loop calling Expire
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			m.expire(ctx)
		case <-ctx.Done():
			return nil
		}
	}
}

func (m *Manager) runSystemAGCLoop(ctx context.Context) error {
	for snapshot := range m.state.WatchIntercepts(ctx, nil) {
		for _, update := range snapshot.Updates {
			// Since all intercepts with a domain require a login, we can use
			// presence of the ApiKey in the interceptInfo to determine all
			// intercepts that we need to inform System A of their deletion
			if update.Delete && update.Value.ApiKey != "" {
				if sa, err := m.systema.Get(); err != nil {
					dlog.Errorln(ctx, "systema: acquire connection:", err)
				} else {
					// First we remove the PreviewDomain if it exists
					if update.Value.PreviewDomain != "" {
						err = m.reapDomain(ctx, sa, update)
						if err != nil {
							dlog.Errorln(ctx, "systema: remove domain:", err)
						}
					}
					// Now we inform SystemA of the intercepts removal
					dlog.Debugf(ctx, "systema: remove intercept: %q", update.Value.Id)
					err = m.reapIntercept(ctx, sa, update)
					if err != nil {
						dlog.Errorln(ctx, "systema: remove intercept:", err)
					}

					// Release the connection we got to delete the domain + intercept
					if err := m.systema.Done(); err != nil {
						dlog.Errorln(ctx, "systema: release management connection:", err)
					}
				}
			}
		}
	}
	return nil
}

// reapDomain informs SystemA that an intercept with a domain has been garbage collected
func (m *Manager) reapDomain(ctx context.Context, sa systema.SystemACRUDClient, interceptUpdate watchable.InterceptMapUpdate) error {
	// we only reapDomains for intercepts that have been deleted
	if !interceptUpdate.Delete {
		return fmt.Errorf("%s is not being deleted, so the domain was not reaped", interceptUpdate.Value.Id)
	}
	dlog.Debugf(ctx, "systema: removing domain: %q", interceptUpdate.Value.PreviewDomain)
	_, err := sa.RemoveDomain(ctx, &systema.RemoveDomainRequest{
		Domain: interceptUpdate.Value.PreviewDomain,
	})

	if err != nil {
		return err
	}
	return nil
}

// reapIntercept informs SystemA that an intercept has been garbage collected
func (m *Manager) reapIntercept(ctx context.Context, sa systema.SystemACRUDClient, interceptUpdate watchable.InterceptMapUpdate) error {
	// we only reapIntercept for intercepts that have been deleted
	if !interceptUpdate.Delete {
		return fmt.Errorf("%s is not being deleted, so the intercept was not reaped", interceptUpdate.Value.Id)
	}
	dlog.Debugf(ctx, "systema: remove intercept: %q", interceptUpdate.Value.Id)
	_, err := sa.RemoveIntercept(ctx, &systema.InterceptRemoval{
		InterceptId: interceptUpdate.Value.Id,
	})

	// We remove the APIKey whether or not the RemoveIntercept call was successful, so
	// let's do that before we check the error.
	if wasRemoved := m.state.RemoveInterceptAPIKey(interceptUpdate.Value.Id); !wasRemoved {
		dlog.Debugf(ctx, "Intercept ID %s had no APIKey", interceptUpdate.Value.Id)
	}

	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>