<commit_before>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// notified when handleRelease() finishes its task
	stoppedCh chan struct{}
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
		stoppedCh:  make(chan struct{}),
	}
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	log.Print("worker: starting delayd2 process")

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		go w.handleConsume()
	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		go w.relayWorker()
	}

	go w.handleAdoptOrphans()
	go w.handleKeepAlive()
	go w.handleMarkActive()
	go w.handleRelease()

	select {
	case <-w.shutdownCh:
		log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	}

	select {
	case <-w.stoppedCh:
		log.Print("worker: the process finished.")
	}
	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) handleKeepAlive() {
	for range time.Tick(1 * time.Second) {
		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// handleConsume consumes messages in SQS.
func (w *Worker) handleConsume() {
	for {
		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) handleAdoptOrphans() {
	for range time.Tick(10 * time.Second) {
		begin := time.Now()

		n, err := w.driver.AdoptOrphans()

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// handleMarkActive marks coming messages in queue.
func (w *Worker) handleMarkActive() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down marking handler")
			return
		default:
		}

		begin := time.Now().Truncate(time.Second)

		n, err := w.markActive(begin)

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// handleRelease releases messages that is ready to fire.
func (w *Worker) handleRelease() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			close(w.relayCh)
			close(w.stoppedCh)
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker() {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()
	}
	log.Print("relayworker: closed")
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<commit_msg>worker: Use sync.WaitGroup to wait for all worker goroutine to be finished<commit_after>// delayd2 is an available setTimeout() service for scheduling message sends
package delayd2

import (
	"log"
	"runtime"
	"sync"
	"time"

	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/pmylund/go-cache"
)

// QueueMessage is a message queued in the database.
type QueueMessage struct {
	QueueID   string
	WorkerID  string
	ReleaseAt time.Time
	RelayTo   string
	Payload   string
}

// WorkerConfig is configuration for Worker.
type WorkerConfig struct {
	ID string

	LeaveMessagesOrphanedAtShutdown bool

	NumConsumerFactor int
	NumRelayFactor    int
}

// Worker is the delayd2 worker.
type Worker struct {
	config *WorkerConfig

	driver   Driver
	consumer *Consumer
	relay    *Relay

	succededIDsCache *cache.Cache

	relayCh chan releaseJob

	// signaled from external
	shutdownCh chan struct{}

	// wait for all worker goroutine to be finished
	stoppped sync.WaitGroup
}

// NewWorker creates a new worker.
func NewWorker(c *WorkerConfig, driver Driver, consumer *Consumer, relay *Relay) *Worker {
	return &Worker{
		driver:   driver,
		consumer: consumer,
		relay:    relay,

		// Cache succededIDs to prepare for transient errors in the database (e.g. failover/network problems)
		succededIDsCache: cache.New(15*time.Minute, 1*time.Minute),

		config: c,

		relayCh:    make(chan releaseJob, runtime.NumCPU()*c.NumRelayFactor),
		shutdownCh: make(chan struct{}),
	}
}

func (w *Worker) runWorker(f func()) {
	w.stoppped.Add(1)
	go func() {
		f()
		w.stoppped.Done()
	}()
}

// Run starts the worker process.
func (w *Worker) Run() error {
	log.Print("worker: starting delayd2 worker process")

	log.Print("worker: registering delayd2 worker process")
	if err := w.driver.RegisterSession(); err != nil {
		return err
	}

	defer func() {
		log.Print("worker: resetting aborted messages remaining in active queue")
		if _, err := w.driver.ResetActive(); err != nil {
			log.Print("worker: unable to reset messages in active queue:", err)
		}

		if w.config.LeaveMessagesOrphanedAtShutdown {
			log.Print("worker: leaving messages as orphaned")

			if err := w.driver.MarkOrphaned(); err != nil {
				log.Print("worker: unable to mark messages as orphaned:", err)
			}
		}

		log.Print("worker: deregistering delayd2 worker process")
		if err := w.driver.DeregisterSession(); err != nil {
			log.Print("worker: unable to deregister session:", err)
		}
	}()

	nCPU := runtime.NumCPU()
	for i := 0; i < nCPU*w.config.NumConsumerFactor; i++ {
		log.Print("worker: launching consumer process")
		w.runWorker(func() { w.handleConsume() })

	}

	for i := 0; i < nCPU*w.config.NumRelayFactor; i++ {
		log.Print("worker: launching relay worker process")
		w.runWorker(func() { w.relayWorker() })
	}

	w.runWorker(func() { w.handleAdoptOrphans() })
	w.runWorker(func() { w.handleKeepAlive() })
	w.runWorker(func() { w.handleMarkActive() })
	w.runWorker(func() { w.handleRelease() })

	log.Print("worker: starting delayd2 process")

	select {
	case <-w.shutdownCh:
		log.Print("worker: receiving shutdown signal. waiting for the process finished.")
	}

	w.stoppped.Wait()
	log.Print("worker: the process finished.")
	return nil
}

func (w *Worker) Stop() {
	close(w.shutdownCh)
}

func (w *Worker) consume() (int64, error) {
	return w.consumer.ConsumeMessages()
}

func (w *Worker) handleKeepAlive() {
	for range time.Tick(1 * time.Second) {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down keepalive worker")
			return
		default:
		}

		if err := w.driver.KeepAliveSession(); err != nil {
			log.Printf("worker: unable to keep alived: %s", err)
		}
	}
}

// handleConsume consumes messages in SQS.
func (w *Worker) handleConsume() {
	for {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down consuming worker")
			return
		default:
		}

		begin := time.Now()
		n, err := w.consume()
		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to consume messages: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages consumed in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) handleAdoptOrphans() {
	for range time.Tick(10 * time.Second) {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down adopting worker")
			return
		default:
		}

		begin := time.Now()

		n, err := w.driver.AdoptOrphans()

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to adopt orphans: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d orphans adopted in %s", n, end.Sub(begin))
		}
	}
}

func (w *Worker) markActive(begin time.Time) (int64, error) {
	return w.driver.MarkActive(begin)
}

// handleMarkActive marks coming messages in queue.
func (w *Worker) handleMarkActive() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			log.Print("worker: shutting down marking worker")
			return
		default:
		}

		begin := time.Now().Truncate(time.Second)

		n, err := w.markActive(begin)

		end := time.Now()

		if err != nil {
			log.Printf("worker: unable to mark messages active: %s", err)
			continue
		}

		if n > 0 {
			log.Printf("worker: %d messages marked as active in %s", n, end.Sub(begin))
		}
		time.Sleep(1 * time.Second)
	}
}

// handleRelease releases messages that is ready to fire.
func (w *Worker) handleRelease() {
	for range time.Tick(10 * time.Millisecond) {
		select {
		case <-w.shutdownCh:
			close(w.relayCh)
			return
		default:
		}

		if err := w.release(); err != nil {
			log.Printf("worker: unable to release messages: %s", err)
			continue
		}
	}
}

type releaseJob struct {
	*sync.WaitGroup

	RelayTo  string
	Messages []releaseMessage
}

type releaseMessage struct {
	QueueID string
	Payload string
}

// release releases active messages. It returns the number of released message n.
func (w *Worker) release() error {
	messages, err := w.driver.GetActiveMessages()

	if err != nil {
		return err
	}

	// we need QueueID to delete from queue
	batchMap := BuildBatchMap(messages)

	var wg sync.WaitGroup
	for relayTo, batchMsgs := range batchMap {
		for _, rms := range batchMsgs {
			wg.Add(1)
			w.relayCh <- releaseJob{
				WaitGroup: &wg,

				RelayTo:  relayTo,
				Messages: rms,
			}
		}
	}
	wg.Wait()
	return nil
}

func (w *Worker) relayWorker() {
	for rj := range w.relayCh {
		begin := time.Now()
		n := w.releaseBatch(rj)
		end := time.Now()

		if n > 0 {
			log.Printf("worker: %d messages relayed in %s", n, end.Sub(begin))
		}
		rj.Done()
	}
	log.Print("relayworker: closed")
}

func (w *Worker) releaseBatch(rj releaseJob) int64 {
	payloads := make([]string, 0, len(rj.Messages))
	for _, m := range rj.Messages {
		if _, found := w.succededIDsCache.Get(m.QueueID); found {
			// skip since we already relayed
			log.Printf("worker: seeing queue_id in the cache so continueing...")
			continue
		}
		payloads = append(payloads, m.Payload)
	}

	var n int64
	failedIndex := make(map[int]struct{})
	if len(payloads) > 0 {
		if err := w.relay.Relay(rj.RelayTo, payloads); err != nil {
			// only print log here in batch operation
			// TODO: a dead letter queue support
			berrs, batchOk := queue.IsBatchError(err)
			if !batchOk {
				log.Printf("worker: unable to send message due to non batch error. skipping: %s", err)
				return 0
			}

			for _, berr := range berrs {
				if berr.SenderFault {
					log.Printf("worker: unable to send message due to sender's fault: skipping: %s", berr.Message)
				}
				failedIndex[berr.Index] = struct{}{}
			}
		}
	}

	succededIDs := make([]string, 0, len(rj.Messages))
	for i, m := range rj.Messages {
		if _, failed := failedIndex[i]; failed {
			continue
		}
		succededIDs = append(succededIDs, m.QueueID)

		// remember succededIDs for a while to prevent us from relaying message in transient error
		w.succededIDsCache.Set(m.QueueID, struct{}{}, cache.DefaultExpiration)
	}

	if err := w.driver.RemoveMessages(succededIDs); err != nil {
		log.Printf("worker: unable to remove messages from queue: %s", err)
	} else {
		// reset succededIDs here since we succeeded to remove messages from the database
		// When we succeeded to remove messages, the messages will not appear again so it's safe.
		for _, id := range succededIDs {
			w.succededIDsCache.Delete(id)
		}
		n += int64(len(succededIDs))
	}

	return n
}

func BuildBatchMap(messages []*QueueMessage) map[string][][]releaseMessage {
	batchMap := make(map[string][][]releaseMessage)
	for _, m := range messages {
		rm := releaseMessage{
			QueueID: m.QueueID,
			Payload: m.Payload,
		}

		var i int
		i = len(batchMap[m.RelayTo])
		if i == 0 || len(batchMap[m.RelayTo][i-1]) >= 10 {
			batchMap[m.RelayTo] = append(batchMap[m.RelayTo], []releaseMessage{})
		}

		i = len(batchMap[m.RelayTo])

		batchMap[m.RelayTo][i-1] = append(batchMap[m.RelayTo][i-1], rm)
	}
	return batchMap
}
<|endoftext|>


<commit_before>package wschan

import (
	"fmt"
	"time"
	"errors"
	"net/http"
	ws "github.com/gorilla/websocket"
)

type (
	WSChan struct {
		Input	<-chan []byte
		Output	chan<- []byte
		Status	<-chan Status
		Command	chan<- byte
	}
	Status struct {
		State	byte
		Error	error
	}
)

const (
	// states
	DISCONNECTED byte = iota
	CONNECTING
	CONNECTED
	WAITING
)
const (
	// commands
	QUIT       byte = 16 + iota
	PING
	USE_TEXT
	USE_BINARY
)

func New(url string, headers http.Header) *WSChan {
	inp_ch := make(chan []byte, 8)
	out_ch := make(chan []byte, 8)
	sts_ch := make(chan Status, 2)
	cmd_ch := make(chan byte,   2)

	con_return_ch := make(chan *ws.Conn, 1)
	con_cancel_ch := make(chan bool,     1)

	r_error_ch    := make(chan error, 1)
	w_error_ch    := make(chan error, 1)
	w_control_ch  := make(chan byte,  1)

	io_event_ch   := make(chan bool, 2)

	connect := func() {
		for {
			sts_ch <- Status{State:CONNECTING}
			dialer := ws.Dialer{HandshakeTimeout: 5*time.Second}
			conn, _, err := dialer.Dial(url, headers)
			if err == nil {
				conn.SetPongHandler(func(string) error {io_event_ch <- true; return nil})
				con_return_ch <- conn
				sts_ch <- Status{State:CONNECTED}
				return
			} else {
				sts_ch <- Status{DISCONNECTED, err}
			}

			sts_ch <- Status{State:WAITING}
			select {
			case <- time.After(34*time.Second):
			case <- con_cancel_ch:
				sts_ch <- Status{DISCONNECTED, errors.New("cancelled")}
				return
			}
		}
	}
	keep_alive := func() {
		dur   := 34 * time.Second
		timer := time.NewTimer(dur)
		timer.Stop()

		loop:
		for {
			select {
			case _, ok := <-io_event_ch:
				if ok {
					timer.Reset(dur)
				} else {
					timer.Stop()
					break loop
				}
			case <-timer.C:
				timer.Reset(dur)
				// non-blocking PING request
				select {
				case w_control_ch <- PING:
				default:
				}
			}
		}
	}
	read := func(conn *ws.Conn) {
		for {
			if _, msg, err := conn.ReadMessage(); err == nil {
				io_event_ch <- true
				inp_ch <- msg
			} else {
				r_error_ch <- err
				break
			}
		}
	}
	write := func(conn *ws.Conn, msg_type int) {
		loop:
		for {
			select {
			case msg, ok := <-out_ch:
				if ok {
					io_event_ch <- true
					if err := conn.SetWriteDeadline(time.Now().Add(3*time.Second)); err != nil {
						w_error_ch <- err
						break loop
					}
					if err := conn.WriteMessage(msg_type, msg); err != nil {
						w_error_ch <- err
						break loop
					}
					conn.SetWriteDeadline(time.Time{})  // reset write deadline
				} else {
					w_error_ch <- errors.New("out_ch closed")
					break loop
				}
			case cmd, ok := <-w_control_ch:
				if !ok {
					w_error_ch <- errors.New("w_control_ch closed")
					break loop
				} else {
					switch cmd {
					case QUIT:
						w_error_ch <- errors.New("cancelled")
						break loop
					case PING:
						if err := conn.WriteControl(ws.PingMessage, []byte{}, time.Now().Add(3*time.Second)); err != nil {
							w_error_ch <- errors.New("cancelled")
							break loop
						}
					case USE_TEXT:
						msg_type = ws.TextMessage
					case USE_BINARY:
						msg_type = ws.BinaryMessage
					}
				}
			}
		}
	}

	go func() {
		// local state
		var conn *ws.Conn
		reading	 := false
		writing  := false
		msg_type := ws.BinaryMessage  // use Binary messages by default

		defer func() {
			if conn != nil {conn.Close()}  // this also makes reader to exit

			// close local output channels
			close(con_cancel_ch)  // this makes connect    to exit
			close(w_control_ch)   // this makes write      to exit
			close(io_event_ch)    // this makes keep_alive to exit

			// drain input channels
			<-time.After(50*time.Millisecond)  // small pause to let things react

			drain_loop:
			for {
				select {
				case _, ok := <-out_ch:
					if !ok {out_ch = nil}
				case _, ok := <-cmd_ch:
					if !ok {inp_ch = nil}
				case conn, ok := <-con_return_ch:
					if conn != nil {conn.Close()}
					if !ok {con_return_ch = nil}
				case _, ok := <-r_error_ch:
					if !ok {r_error_ch = nil}
				case _, ok := <-w_error_ch:
					if !ok {w_error_ch = nil}
				default:
					break drain_loop
				}
			}

			// close output channels
			close(inp_ch)
			close(sts_ch)
		}()

		go connect()
		go keep_alive()

		main_loop:
		for {
			select {
			case conn = <-con_return_ch:
				if conn == nil {
					break main_loop
				}
				reading = true; go read(conn)
				writing = true; go write(conn, msg_type)

			case err := <-r_error_ch:
				reading = false
				if writing {
					// write goroutine is still active
					sts_ch <- Status{DISCONNECTED, err}
					w_control_ch <- QUIT  // ask write to exit
				} else {
					// both read and write goroutines have exited
					if conn != nil {
						conn.Close()
						conn = nil
					}
					go connect()
				}
			case err := <-w_error_ch:
				// write goroutine has exited
				writing = false
				if reading {
					// read goroutine is still active
					sts_ch <- Status{DISCONNECTED, err}
					if conn != nil {
						conn.Close()  // this also makes read to exit
						conn = nil
					}
				} else {
					// both read and write goroutines have exited
					go connect()
				}
			case cmd, ok := <-cmd_ch:
				switch {
				case !ok || cmd == QUIT:
					if reading || writing || conn != nil {sts_ch <- Status{DISCONNECTED, nil}}
					break main_loop  // defer should clean everything up
				case cmd == PING:
					if conn != nil && writing {w_control_ch <- cmd}
				case cmd == USE_TEXT:
					msg_type = ws.TextMessage
					if writing {w_control_ch <- cmd}
				case cmd == USE_BINARY:
					msg_type = ws.BinaryMessage
					if writing {w_control_ch <- cmd}
				default:
					panic(fmt.Sprintf("unsupported command: %v", cmd))
				}
			}
		}
	}()

	return &WSChan{inp_ch, out_ch, sts_ch, cmd_ch}
}
<commit_msg>added State and Command types<commit_after>package wschan

import (
	"fmt"
	"time"
	"errors"
	"net/http"
	ws "github.com/gorilla/websocket"
	"github.com/aglyzov/log15"
)

var Log = log15.New("pkg", "wschan")
type State   byte
type Command byte

type (
	WSChan struct {
		Input	<-chan []byte
		Output	chan<- []byte
		Status	<-chan Status
		Command	chan<- Command
	}
	Status struct {
		State	State
		Error	error
	}
)

const (
	// states
	DISCONNECTED State = iota
	CONNECTING
	CONNECTED
	WAITING
)
const (
	// commands
	QUIT       Command = 16 + iota
	PING
	USE_TEXT
	USE_BINARY
)

func (s State) String() string {
	switch s {
	case DISCONNECTED:  return "DISCONNECTED"
	case CONNECTING:    return "CONNECTING"
	case CONNECTED:     return "CONNECTED"
	case WAITING:       return "WAITING"
	}
	return fmt.Sprintf("UNKNOWN STATUS %v", s)
}

func (c Command) String() string {
	switch c {
	case QUIT:        return "QUIT"
	case PING:        return "PING"
	case USE_TEXT:    return "USE_TEXT"
	case USE_BINARY:  return "USE_BINARY"
	}
	return fmt.Sprintf("UNKNOWN COMMAND %v", c)
}

func New(url string, headers http.Header) *WSChan {
	inp_ch := make(chan []byte,  8)
	out_ch := make(chan []byte,  8)
	sts_ch := make(chan Status,  2)
	cmd_ch := make(chan Command, 2)

	con_return_ch := make(chan *ws.Conn, 1)
	con_cancel_ch := make(chan bool,     1)

	r_error_ch    := make(chan error,    1)
	w_error_ch    := make(chan error,    1)
	w_control_ch  := make(chan Command,  1)

	io_event_ch   := make(chan bool,     2)

	connect := func() {
		for {
			sts_ch <- Status{State:CONNECTING}
			dialer := ws.Dialer{HandshakeTimeout: 5*time.Second}
			conn, _, err := dialer.Dial(url, headers)
			if err == nil {
				conn.SetPongHandler(func(string) error {io_event_ch <- true; return nil})
				con_return_ch <- conn
				sts_ch <- Status{State:CONNECTED}
				return
			} else {
				sts_ch <- Status{DISCONNECTED, err}
			}

			sts_ch <- Status{State:WAITING}
			select {
			case <- time.After(34*time.Second):
			case <- con_cancel_ch:
				sts_ch <- Status{DISCONNECTED, errors.New("cancelled")}
				return
			}
		}
	}
	keep_alive := func() {
		dur   := 34 * time.Second
		timer := time.NewTimer(dur)
		timer.Stop()

		loop:
		for {
			select {
			case _, ok := <-io_event_ch:
				if ok {
					timer.Reset(dur)
				} else {
					timer.Stop()
					break loop
				}
			case <-timer.C:
				timer.Reset(dur)
				// non-blocking PING request
				select {
				case w_control_ch <- PING:
				default:
				}
			}
		}
	}
	read := func(conn *ws.Conn) {
		for {
			if _, msg, err := conn.ReadMessage(); err == nil {
				io_event_ch <- true
				inp_ch <- msg
			} else {
				r_error_ch <- err
				break
			}
		}
	}
	write := func(conn *ws.Conn, msg_type int) {
		loop:
		for {
			select {
			case msg, ok := <-out_ch:
				if ok {
					io_event_ch <- true
					if err := conn.SetWriteDeadline(time.Now().Add(3*time.Second)); err != nil {
						w_error_ch <- err
						break loop
					}
					if err := conn.WriteMessage(msg_type, msg); err != nil {
						w_error_ch <- err
						break loop
					}
					conn.SetWriteDeadline(time.Time{})  // reset write deadline
				} else {
					w_error_ch <- errors.New("out_ch closed")
					break loop
				}
			case cmd, ok := <-w_control_ch:
				if !ok {
					w_error_ch <- errors.New("w_control_ch closed")
					break loop
				} else {
					switch cmd {
					case QUIT:
						w_error_ch <- errors.New("cancelled")
						break loop
					case PING:
						if err := conn.WriteControl(ws.PingMessage, []byte{}, time.Now().Add(3*time.Second)); err != nil {
							w_error_ch <- errors.New("cancelled")
							break loop
						}
					case USE_TEXT:
						msg_type = ws.TextMessage
					case USE_BINARY:
						msg_type = ws.BinaryMessage
					}
				}
			}
		}
	}

	go func() {
		// local state
		var conn *ws.Conn
		reading	 := false
		writing  := false
		msg_type := ws.BinaryMessage  // use Binary messages by default

		defer func() {
			if conn != nil {conn.Close()}  // this also makes reader to exit

			// close local output channels
			close(con_cancel_ch)  // this makes connect    to exit
			close(w_control_ch)   // this makes write      to exit
			close(io_event_ch)    // this makes keep_alive to exit

			// drain input channels
			<-time.After(50*time.Millisecond)  // small pause to let things react

			drain_loop:
			for {
				select {
				case _, ok := <-out_ch:
					if !ok {out_ch = nil}
				case _, ok := <-cmd_ch:
					if !ok {inp_ch = nil}
				case conn, ok := <-con_return_ch:
					if conn != nil {conn.Close()}
					if !ok {con_return_ch = nil}
				case _, ok := <-r_error_ch:
					if !ok {r_error_ch = nil}
				case _, ok := <-w_error_ch:
					if !ok {w_error_ch = nil}
				default:
					break drain_loop
				}
			}

			// close output channels
			close(inp_ch)
			close(sts_ch)
		}()

		go connect()
		go keep_alive()

		main_loop:
		for {
			select {
			case conn = <-con_return_ch:
				if conn == nil {
					break main_loop
				}
				reading = true; go read(conn)
				writing = true; go write(conn, msg_type)

			case err := <-r_error_ch:
				reading = false
				if writing {
					// write goroutine is still active
					sts_ch <- Status{DISCONNECTED, err}
					w_control_ch <- QUIT  // ask write to exit
				} else {
					// both read and write goroutines have exited
					if conn != nil {
						conn.Close()
						conn = nil
					}
					go connect()
				}
			case err := <-w_error_ch:
				// write goroutine has exited
				writing = false
				if reading {
					// read goroutine is still active
					sts_ch <- Status{DISCONNECTED, err}
					if conn != nil {
						conn.Close()  // this also makes read to exit
						conn = nil
					}
				} else {
					// both read and write goroutines have exited
					go connect()
				}
			case cmd, ok := <-cmd_ch:
				switch {
				case !ok || cmd == QUIT:
					if reading || writing || conn != nil {sts_ch <- Status{DISCONNECTED, nil}}
					break main_loop  // defer should clean everything up
				case cmd == PING:
					if conn != nil && writing {w_control_ch <- cmd}
				case cmd == USE_TEXT:
					msg_type = ws.TextMessage
					if writing {w_control_ch <- cmd}
				case cmd == USE_BINARY:
					msg_type = ws.BinaryMessage
					if writing {w_control_ch <- cmd}
				default:
					panic(fmt.Sprintf("unsupported command: %v", cmd))
				}
			}
		}
	}()

	return &WSChan{inp_ch, out_ch, sts_ch, cmd_ch}
}
<|endoftext|>


<commit_before>// Package global provides information on global variables and status
package global

import (
	"database/sql"
	"strings"

	"github.com/sjmudd/ps-top/logger"
)

const (
	showCompatibility56Error    = "Error 3167: The 'INFORMATION_SCHEMA.GLOBAL_VARIABLES' feature is disabled; see the documentation for 'show_compatibility_56'"
	globalVariablesNotInISError = "Error 1109: Unknown table 'GLOBAL_VARIABLES' in information_schema"
)

// We expect to use I_S to query Global Variables. 5.7 now wants us to use P_S,
// so this variable will be changed if we see the show_compatibility_56 error message
var seenCompatibiltyError = false

func selectVariablesFrom(seenError bool) string {
	if !seenError {
		return "INFORMATION_SCHEMA.GLOBAL_VARIABLES"
	}
	return "performance_schema.global_variables"
}

// Variables holds the handle and variables collected from the database
type Variables struct {
	dbh       *sql.DB
	variables map[string]string
}

// NewVariables returns a pointer to an initialised Variables structure
func NewVariables(dbh *sql.DB) *Variables {
	if dbh == nil {
		logger.Fatal("NewVariables(): dbh == nil")
	}
	v := &Variables{dbh: dbh}
	v.selectAll()

	return v
}

// Get returns the value of the given variable
func (v Variables) Get(key string) string {
	var result string
	var ok bool

	if result, ok = v.variables[key]; !ok {
		result = ""
	}

	return result
}

// selectAll() collects all variables from the database and stores for later use.
// - all returned keys are lower-cased.
func (v *Variables) selectAll() {
	hashref := make(map[string]string)

	query := "SELECT VARIABLE_NAME, VARIABLE_VALUE FROM " + selectVariablesFrom(seenCompatibiltyError)
	logger.Println("query:", query)

	rows, err := v.dbh.Query(query)
	if err != nil {
		if !seenCompatibiltyError && err.Error() == showCompatibility56Error {
			logger.Println("selectAll() I_S query failed, trying with P_S")
			seenCompatibiltyError = true
			query = "SELECT VARIABLE_NAME, VARIABLE_VALUE FROM " + selectVariablesFrom(seenCompatibiltyError)
			logger.Println("query:", query)

			rows, err = v.dbh.Query(query)
		}
		if err != nil {
			logger.Fatal("selectAll() query failed with:", err)
		}
	}
	logger.Println("selectAll() query succeeded")
	defer rows.Close()

	for rows.Next() {
		var variable, value string
		if err := rows.Scan(&variable, &value); err != nil {
			logger.Fatal(err)
		}
		hashref[strings.ToLower(variable)] = value
	}
	if err := rows.Err(); err != nil {
		logger.Fatal(err)
	}
	logger.Println("selectAll() result has", len(hashref), "rows")

	v.variables = hashref
}
<commit_msg>Catch GLOBAL_VARIABLES table not in information_schema error<commit_after>// Package global provides information on global variables and status
package global

import (
	"database/sql"
	"strings"

	"github.com/sjmudd/ps-top/logger"
)

const (
	showCompatibility56Error    = "Error 3167: The 'INFORMATION_SCHEMA.GLOBAL_VARIABLES' feature is disabled; see the documentation for 'show_compatibility_56'"
	globalVariablesNotInISError = "Error 1109: Unknown table 'GLOBAL_VARIABLES' in information_schema"
)

// We expect to use I_S to query Global Variables. 5.7 now wants us to use P_S,
// so this variable will be changed if we see the show_compatibility_56 error message
var seenCompatibiltyError = false

func selectVariablesFrom(seenError bool) string {
	if !seenError {
		return "INFORMATION_SCHEMA.GLOBAL_VARIABLES"
	}
	return "performance_schema.global_variables"
}

// Variables holds the handle and variables collected from the database
type Variables struct {
	dbh       *sql.DB
	variables map[string]string
}

// NewVariables returns a pointer to an initialised Variables structure
func NewVariables(dbh *sql.DB) *Variables {
	if dbh == nil {
		logger.Fatal("NewVariables(): dbh == nil")
	}
	v := &Variables{dbh: dbh}
	v.selectAll()

	return v
}

// Get returns the value of the given variable
func (v Variables) Get(key string) string {
	var result string
	var ok bool

	if result, ok = v.variables[key]; !ok {
		result = ""
	}

	return result
}

// selectAll() collects all variables from the database and stores for later use.
// - all returned keys are lower-cased.
func (v *Variables) selectAll() {
	hashref := make(map[string]string)

	query := "SELECT VARIABLE_NAME, VARIABLE_VALUE FROM " + selectVariablesFrom(seenCompatibiltyError)
	logger.Println("query:", query)

	rows, err := v.dbh.Query(query)
	if err != nil {
		if !seenCompatibiltyError && (err.Error() == showCompatibility56Error || err.Error() == globalVariablesNotInISError) {
			logger.Println("selectAll() I_S query failed, trying with P_S")
			seenCompatibiltyError = true
			query = "SELECT VARIABLE_NAME, VARIABLE_VALUE FROM " + selectVariablesFrom(seenCompatibiltyError)
			logger.Println("query:", query)

			rows, err = v.dbh.Query(query)
		}
		if err != nil {
			logger.Fatal("selectAll() query failed with:", err)
		}
	}
	logger.Println("selectAll() query succeeded")
	defer rows.Close()

	for rows.Next() {
		var variable, value string
		if err := rows.Scan(&variable, &value); err != nil {
			logger.Fatal(err)
		}
		hashref[strings.ToLower(variable)] = value
	}
	if err := rows.Err(); err != nil {
		logger.Fatal(err)
	}
	logger.Println("selectAll() result has", len(hashref), "rows")

	v.variables = hashref
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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 gnmit is a single-target gNMI collector implementation that can be
// used as an on-device/fake device implementation. It supports the Subscribe RPC
// using the libraries from openconfig/gnmi.
package gnmit

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"reflect"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/openconfig/gnmi/cache"
	"github.com/openconfig/lemming/gnmi/subscribe"
	"github.com/openconfig/ygot/ygot"
	"github.com/openconfig/ygot/ytypes"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/prototext"
	"google.golang.org/protobuf/proto"

	gpb "github.com/openconfig/gnmi/proto/gnmi"
)

var (
	// metadataUpdatePeriod is the period of time after which the metadata for the collector
	// is updated to the client.
	metadataUpdatePeriod = 30 * time.Second
	// sizeUpdatePeriod is the period of time after which the storage size information for
	// the collector is updated to the client.
	sizeUpdatePeriod = 30 * time.Second
)

// periodic runs the function fn every period.
func periodic(period time.Duration, fn func()) {
	if period == 0 {
		return
	}
	t := time.NewTicker(period)
	defer t.Stop()
	for range t.C {
		fn()
	}
}

// Queue is an interface that represents a possibly coalescing queue of updates.
type Queue interface {
	Next(ctx context.Context) (interface{}, uint32, error)
	Len() int
	Close()
}

// UpdateFn is a function that takes in a gNMI Notification object and updates
// a gNMI datastore with it.
type UpdateFn func(*gpb.Notification) error

// TaskRoutine is a reactor function that listens for updates from a queue,
// emits updates via an update function. It does this on a target (string
// parameter), and also has a final clean-up function to call when it finishes
// processing.
type TaskRoutine func(Queue, UpdateFn, string, func()) error

// Task defines a particular task that runs on the gNMI datastore.
type Task struct {
	Run    TaskRoutine
	Paths  []ygot.PathStruct
	Prefix *gpb.Path
}

// GNMIServer implements the gNMI server interface.
type GNMIServer struct {
	// The subscribe Server implements only Subscribe for gNMI.
	*subscribe.Server
	c *Collector
}

// RegisterTask starts up a task on the gNMI datastore.
func (s *GNMIServer) RegisterTask(task Task) error {
	var paths []*gpb.Path
	for _, p := range task.Paths {
		path, _, err := ygot.ResolvePath(p)
		if err != nil {
			return fmt.Errorf("gnmit: cannot register task: %v", err)
		}
		paths = append(paths, path)
	}
	queue, remove, err := s.Server.SubscribeLocal(s.c.name, paths, task.Prefix)
	if err != nil {
		return err
	}
	return task.Run(queue, s.c.cache.GnmiUpdate, s.c.name, remove)
}

// New returns a new collector server implementation that can be registered on
// an existing gRPC server. It takes a string indicating the hostname of the
// target, a boolean indicating whether metadata should be sent, and a slice of
// tasks that are to be launched to run on the server.
func NewServer(ctx context.Context, hostname string, sendMeta bool, tasks []Task) (*Collector, *GNMIServer, error) {
	c := &Collector{
		inCh: make(chan *gpb.SubscribeResponse),
		name: hostname,
	}

	c.cache = cache.New([]string{hostname})
	t := c.cache.GetTarget(hostname)

	if sendMeta {
		go periodic(metadataUpdatePeriod, c.cache.UpdateMetadata)
		go periodic(sizeUpdatePeriod, c.cache.UpdateSize)
	}
	t.Connect()

	// start our single collector from the input channel.
	go func() {
		for {
			select {
			case msg := <-c.inCh:
				if err := c.handleUpdate(msg); err != nil {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	subscribeSrv, err := subscribe.NewServer(c.cache)
	if err != nil {
		return nil, nil, fmt.Errorf("could not instantiate gNMI server: %v", err)
	}

	gnmiserver := &GNMIServer{
		Server: subscribeSrv, // use the 'subscribe' implementation.
		c:      c,
	}

	for _, t := range tasks {
		if err := gnmiserver.RegisterTask(t); err != nil {
			return nil, nil, err
		}
	}
	c.cache.SetClient(subscribeSrv.Update)

	return c, gnmiserver, nil
}

// New returns a new collector that listens on the specified addr (in the form host:port),
// supporting a single downstream target named hostname. sendMeta controls whether the
// metadata *other* than meta/sync and meta/connected is sent by the collector.
//
// New returns the new collector, the address it is listening on in the form hostname:port
// or any errors encounted whilst setting it up.
func New(ctx context.Context, addr, hostname string, sendMeta bool, tasks []Task, opts ...grpc.ServerOption) (*Collector, string, error) {
	c, gnmiserver, err := NewServer(ctx, hostname, sendMeta, tasks)
	if err != nil {
		return nil, "", err
	}
	srv := grpc.NewServer(opts...)
	gpb.RegisterGNMIServer(srv, gnmiserver)
	// Forward streaming updates to clients.
	// Register listening port and start serving.
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, "", fmt.Errorf("failed to listen: %v", err)
	}

	go func() {
		if err := srv.Serve(lis); err != nil {
			log.Errorf("Error while serving gnmi target: %v", err)
		}
	}()
	c.stopFn = srv.GracefulStop
	return c, lis.Addr().String(), nil
}

type populateDefaultser interface {
	PopulateDefaults()
}

// New returns a new collector that listens on the specified addr (in the form host:port),
// supporting a single downstream target named hostname. sendMeta controls whether the
// metadata *other* than meta/sync and meta/connected is sent by the collector.
//
// New returns the new collector, the address it is listening on in the form hostname:port
// or any errors encounted whilst setting it up.
//
// NewSettable is different from New in that the returned collector is
// schema-aware and supports gNMI Set. Currently it is not possible to change
// the schema of a Collector after it is created.
func NewSettable(ctx context.Context, addr string, hostname string, sendMeta bool, schema *ytypes.Schema, tasks []Task, opts ...grpc.ServerOption) (*Collector, string, error) {
	if !schema.IsValid() {
		return nil, "", fmt.Errorf("cannot obtain valid schema for GoStructs: %v", schema)
	}
	// Initialize the root with default values.
	schema.Root.(populateDefaultser).PopulateDefaults()
	if err := schema.Root.Validate(); err != nil {
		return nil, "", fmt.Errorf("default root of input schema fails validation: %v", err)
	}

	// FIXME(wenbli): initialize the collector with default values.
	collector, addr, err := New(ctx, addr, hostname, sendMeta, tasks, opts...)
	if err != nil {
		return nil, "", err
	}
	collector.schema = schema
	return collector, addr, nil
}

// Stop halts the running collector.
func (c *Collector) Stop() {
	c.stopFn()
}

// handleUpdate handles an input gNMI SubscribeResponse that is received by
// the target.
func (c *Collector) handleUpdate(resp *gpb.SubscribeResponse) error {
	t := c.cache.GetTarget(c.name)
	switch v := resp.Response.(type) {
	case *gpb.SubscribeResponse_Update:
		return t.GnmiUpdate(v.Update)
	case *gpb.SubscribeResponse_SyncResponse:
		t.Sync()
	case *gpb.SubscribeResponse_Error:
		return fmt.Errorf("error in response: %s", v)
	default:
		return fmt.Errorf("unknown response %T: %s", v, v)
	}
	return nil
}

// Collector is a basic gNMI target that supports only the Subscribe
// RPC, and acts as a cache for exactly one target.
type Collector struct {
	cache *cache.Cache

	smu    sync.Mutex
	schema *ytypes.Schema

	// name is the hostname of the client.
	name string
	// inCh is a channel use to write new SubscribeResponses to the client.
	inCh chan *gpb.SubscribeResponse
	// stopFn is the function used to stop the server.
	stopFn func()
}

func (c *Collector) Schema() *ytypes.Schema {
	c.smu.Lock()
	defer c.smu.Unlock()
	return c.schema
}

// TargetUpdate provides an input gNMI SubscribeResponse to update the
// cache and clients with.
func (c *Collector) TargetUpdate(m *gpb.SubscribeResponse) {
	c.inCh <- m
}

// setNode is a function that is able to unmarshal either a JSON-encoded value or a gNMI-encoded value.
// TODO(wenbli): This functionality should be upstreamed to ygot somehow.
func setNode(schema *ytypes.Schema, goStruct ygot.GoStruct, update *gpb.Update) error {
	nodeName := reflect.TypeOf(goStruct).Elem().Name()
	nodeI, targetSchema, err := ytypes.GetOrCreateNode(schema.SchemaTree[nodeName], goStruct, update.Path)
	if err != nil {
		return fmt.Errorf("gnmit: failed to GetOrCreate a node: %v", err)
	}

	// TODO(wenbli): Populate default values using PopulateDefaults.
	jsonUpdate := update.Val.GetJsonIetfVal()
	if jsonUpdate == nil {
		if err := ytypes.SetNode(schema.SchemaTree[nodeName], goStruct, update.Path, update.Val); err != nil {
			return fmt.Errorf("gnmit: SetNode failed on leaf node: %v", err)
		}
		return nil
	}

	node, ok := nodeI.(ygot.GoStruct)
	path := proto.Clone(update.Path).(*gpb.Path)
	for i := len(path.Elem) - 1; i >= 0 && !ok; i-- {
		path.Elem = path.Elem[:i]
		nodeI, _, err := ytypes.GetOrCreateNode(schema.SchemaTree[nodeName], goStruct, path)
		if err != nil {
			continue
		}
		node, ok = nodeI.(ygot.GoStruct)
	}
	if ok {
		var jsonTree interface{}
		if err := json.Unmarshal(jsonUpdate, &jsonTree); err != nil {
			return err
		}
		return ytypes.Unmarshal(targetSchema, node, jsonTree)
	}

	return fmt.Errorf("gnmit: cannot find GoStruct parent into which to ummarshal update message: %s", prototext.Format(update))
}

// Set is a prototype for a gNMI Set operation.
// TODO(wenbli): This function is too complex to stand on its own. We should
// split this into stages, each encapsulated within its own function.
func (s *GNMIServer) Set(ctx context.Context, req *gpb.SetRequest) (*gpb.SetResponse, error) {
	if s.c.schema == nil {
		return s.UnimplementedGNMIServer.Set(ctx, req)
	}
	// Create a copy so that we can rollback the transaction when validation fails.
	dirtyRootG, err := ygot.DeepCopy(s.c.Schema().Root)
	if err != nil {
		return nil, fmt.Errorf("gnmit: failed to ygot.DeepCopy the cached root object: %v", err)
	}
	dirtyRoot, ok := dirtyRootG.(ygot.ValidatedGoStruct)
	if !ok {
		return nil, fmt.Errorf("gnmit: cannot convert root object to ValidatedGoStruct")
	}
	// Operate at the prefix level.
	nodeI, _, err := ytypes.GetOrCreateNode(s.c.Schema().RootSchema(), dirtyRoot, req.Prefix)
	if err != nil {
		return nil, fmt.Errorf("gnmit: failed to GetOrCreate the prefix node: %v", err)
	}
	node, ok := nodeI.(ygot.GoStruct)
	if !ok {
		return nil, fmt.Errorf("gnmit: prefix path points to a non-GoStruct, this is not allowed: %T, %v", nodeI, nodeI)
	}
	nodeName := reflect.TypeOf(nodeI).Elem().Name()

	// TODO(wenbli): Reject paths that try to modify read-only values.
	// TODO(wenbli): Question: what to do if there are operational-state values in a container that is specified to be replaced or deleted?

	// Process deletes, then replace, then updates.
	for _, path := range req.Delete {
		if err := ytypes.DeleteNode(s.c.Schema().SchemaTree[nodeName], node, path); err != nil {
			return nil, fmt.Errorf("gnmit: DeleteNode error: %v", err)
		}
	}
	for _, update := range req.Replace {
		if err := ytypes.DeleteNode(s.c.Schema().SchemaTree[nodeName], node, update.Path); err != nil {
			return nil, fmt.Errorf("gnmit: DeleteNode error: %v", err)
		}
		if err := setNode(s.c.schema, node, update); err != nil {
			return nil, err
		}
	}
	for _, update := range req.Update {
		if err := setNode(s.c.schema, node, update); err != nil {
			return nil, err
		}
	}

	if err := dirtyRoot.Validate(); err != nil {
		return nil, fmt.Errorf("gnmit: invalid SetRequest: %v", err)
	}

	n, err := ygot.Diff(s.c.Schema().Root, dirtyRoot)
	if err != nil {
		return nil, fmt.Errorf("gnmit: error while creating update notification for Set: %v", err)
	}
	n.Timestamp = time.Now().UnixNano()
	n.Prefix = &gpb.Path{Origin: req.Prefix.Origin, Target: s.c.name}
	if n.Prefix.Origin == "" {
		n.Prefix.Origin = "openconfig"
	}

	// Update cache
	t := s.c.cache.GetTarget(s.c.name)
	if err := t.GnmiUpdate(n); err != nil {
		return nil, err
	}
	// TODO(wenbli): Should handle updates one at a time to avoid partial updates being reflected in the cache when an error occurs. AKA we should support transactional semantics.
	s.c.Schema().Root = dirtyRoot
	// TODO(wenbli): Currently the SetResponse is not filled.
	return &gpb.SetResponse{}, nil
}
<commit_msg>Update to new definition of ygot.GoStruct.<commit_after>// Copyright 2021 Google LLC
//
// 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 gnmit is a single-target gNMI collector implementation that can be
// used as an on-device/fake device implementation. It supports the Subscribe RPC
// using the libraries from openconfig/gnmi.
package gnmit

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"reflect"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/openconfig/gnmi/cache"
	"github.com/openconfig/lemming/gnmi/subscribe"
	"github.com/openconfig/ygot/ygot"
	"github.com/openconfig/ygot/ytypes"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/prototext"
	"google.golang.org/protobuf/proto"

	gpb "github.com/openconfig/gnmi/proto/gnmi"
)

var (
	// metadataUpdatePeriod is the period of time after which the metadata for the collector
	// is updated to the client.
	metadataUpdatePeriod = 30 * time.Second
	// sizeUpdatePeriod is the period of time after which the storage size information for
	// the collector is updated to the client.
	sizeUpdatePeriod = 30 * time.Second
)

// periodic runs the function fn every period.
func periodic(period time.Duration, fn func()) {
	if period == 0 {
		return
	}
	t := time.NewTicker(period)
	defer t.Stop()
	for range t.C {
		fn()
	}
}

// Queue is an interface that represents a possibly coalescing queue of updates.
type Queue interface {
	Next(ctx context.Context) (interface{}, uint32, error)
	Len() int
	Close()
}

// UpdateFn is a function that takes in a gNMI Notification object and updates
// a gNMI datastore with it.
type UpdateFn func(*gpb.Notification) error

// TaskRoutine is a reactor function that listens for updates from a queue,
// emits updates via an update function. It does this on a target (string
// parameter), and also has a final clean-up function to call when it finishes
// processing.
type TaskRoutine func(Queue, UpdateFn, string, func()) error

// Task defines a particular task that runs on the gNMI datastore.
type Task struct {
	Run    TaskRoutine
	Paths  []ygot.PathStruct
	Prefix *gpb.Path
}

// GNMIServer implements the gNMI server interface.
type GNMIServer struct {
	// The subscribe Server implements only Subscribe for gNMI.
	*subscribe.Server
	c *Collector
}

// RegisterTask starts up a task on the gNMI datastore.
func (s *GNMIServer) RegisterTask(task Task) error {
	var paths []*gpb.Path
	for _, p := range task.Paths {
		path, _, err := ygot.ResolvePath(p)
		if err != nil {
			return fmt.Errorf("gnmit: cannot register task: %v", err)
		}
		paths = append(paths, path)
	}
	queue, remove, err := s.Server.SubscribeLocal(s.c.name, paths, task.Prefix)
	if err != nil {
		return err
	}
	return task.Run(queue, s.c.cache.GnmiUpdate, s.c.name, remove)
}

// New returns a new collector server implementation that can be registered on
// an existing gRPC server. It takes a string indicating the hostname of the
// target, a boolean indicating whether metadata should be sent, and a slice of
// tasks that are to be launched to run on the server.
func NewServer(ctx context.Context, hostname string, sendMeta bool, tasks []Task) (*Collector, *GNMIServer, error) {
	c := &Collector{
		inCh: make(chan *gpb.SubscribeResponse),
		name: hostname,
	}

	c.cache = cache.New([]string{hostname})
	t := c.cache.GetTarget(hostname)

	if sendMeta {
		go periodic(metadataUpdatePeriod, c.cache.UpdateMetadata)
		go periodic(sizeUpdatePeriod, c.cache.UpdateSize)
	}
	t.Connect()

	// start our single collector from the input channel.
	go func() {
		for {
			select {
			case msg := <-c.inCh:
				if err := c.handleUpdate(msg); err != nil {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	subscribeSrv, err := subscribe.NewServer(c.cache)
	if err != nil {
		return nil, nil, fmt.Errorf("could not instantiate gNMI server: %v", err)
	}

	gnmiserver := &GNMIServer{
		Server: subscribeSrv, // use the 'subscribe' implementation.
		c:      c,
	}

	for _, t := range tasks {
		if err := gnmiserver.RegisterTask(t); err != nil {
			return nil, nil, err
		}
	}
	c.cache.SetClient(subscribeSrv.Update)

	return c, gnmiserver, nil
}

// New returns a new collector that listens on the specified addr (in the form host:port),
// supporting a single downstream target named hostname. sendMeta controls whether the
// metadata *other* than meta/sync and meta/connected is sent by the collector.
//
// New returns the new collector, the address it is listening on in the form hostname:port
// or any errors encounted whilst setting it up.
func New(ctx context.Context, addr, hostname string, sendMeta bool, tasks []Task, opts ...grpc.ServerOption) (*Collector, string, error) {
	c, gnmiserver, err := NewServer(ctx, hostname, sendMeta, tasks)
	if err != nil {
		return nil, "", err
	}
	srv := grpc.NewServer(opts...)
	gpb.RegisterGNMIServer(srv, gnmiserver)
	// Forward streaming updates to clients.
	// Register listening port and start serving.
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, "", fmt.Errorf("failed to listen: %v", err)
	}

	go func() {
		if err := srv.Serve(lis); err != nil {
			log.Errorf("Error while serving gnmi target: %v", err)
		}
	}()
	c.stopFn = srv.GracefulStop
	return c, lis.Addr().String(), nil
}

type populateDefaultser interface {
	PopulateDefaults()
}

// New returns a new collector that listens on the specified addr (in the form host:port),
// supporting a single downstream target named hostname. sendMeta controls whether the
// metadata *other* than meta/sync and meta/connected is sent by the collector.
//
// New returns the new collector, the address it is listening on in the form hostname:port
// or any errors encounted whilst setting it up.
//
// NewSettable is different from New in that the returned collector is
// schema-aware and supports gNMI Set. Currently it is not possible to change
// the schema of a Collector after it is created.
func NewSettable(ctx context.Context, addr string, hostname string, sendMeta bool, schema *ytypes.Schema, tasks []Task, opts ...grpc.ServerOption) (*Collector, string, error) {
	if !schema.IsValid() {
		return nil, "", fmt.Errorf("cannot obtain valid schema for GoStructs: %v", schema)
	}

	vr, ok := schema.Root.(ygot.ValidatedGoStruct)
	if err != nil {
		return nil, "", fmt.Errorf("invalid schema root, %v", ok)
	}
	// Initialize the root with default values.
	schema.Root.(populateDefaultser).PopulateDefaults()
	if err := vr.Validate(); err != nil {
		return nil, "", fmt.Errorf("default root of input schema fails validation: %v", err)
	}

	// FIXME(wenbli): initialize the collector with default values.
	collector, addr, err := New(ctx, addr, hostname, sendMeta, tasks, opts...)
	if err != nil {
		return nil, "", err
	}
	collector.schema = schema
	return collector, addr, nil
}

// Stop halts the running collector.
func (c *Collector) Stop() {
	c.stopFn()
}

// handleUpdate handles an input gNMI SubscribeResponse that is received by
// the target.
func (c *Collector) handleUpdate(resp *gpb.SubscribeResponse) error {
	t := c.cache.GetTarget(c.name)
	switch v := resp.Response.(type) {
	case *gpb.SubscribeResponse_Update:
		return t.GnmiUpdate(v.Update)
	case *gpb.SubscribeResponse_SyncResponse:
		t.Sync()
	case *gpb.SubscribeResponse_Error:
		return fmt.Errorf("error in response: %s", v)
	default:
		return fmt.Errorf("unknown response %T: %s", v, v)
	}
	return nil
}

// Collector is a basic gNMI target that supports only the Subscribe
// RPC, and acts as a cache for exactly one target.
type Collector struct {
	cache *cache.Cache

	smu    sync.Mutex
	schema *ytypes.Schema

	// name is the hostname of the client.
	name string
	// inCh is a channel use to write new SubscribeResponses to the client.
	inCh chan *gpb.SubscribeResponse
	// stopFn is the function used to stop the server.
	stopFn func()
}

func (c *Collector) Schema() *ytypes.Schema {
	c.smu.Lock()
	defer c.smu.Unlock()
	return c.schema
}

// TargetUpdate provides an input gNMI SubscribeResponse to update the
// cache and clients with.
func (c *Collector) TargetUpdate(m *gpb.SubscribeResponse) {
	c.inCh <- m
}

// setNode is a function that is able to unmarshal either a JSON-encoded value or a gNMI-encoded value.
// TODO(wenbli): This functionality should be upstreamed to ygot somehow.
func setNode(schema *ytypes.Schema, goStruct ygot.GoStruct, update *gpb.Update) error {
	nodeName := reflect.TypeOf(goStruct).Elem().Name()
	nodeI, targetSchema, err := ytypes.GetOrCreateNode(schema.SchemaTree[nodeName], goStruct, update.Path)
	if err != nil {
		return fmt.Errorf("gnmit: failed to GetOrCreate a node: %v", err)
	}

	// TODO(wenbli): Populate default values using PopulateDefaults.
	jsonUpdate := update.Val.GetJsonIetfVal()
	if jsonUpdate == nil {
		if err := ytypes.SetNode(schema.SchemaTree[nodeName], goStruct, update.Path, update.Val); err != nil {
			return fmt.Errorf("gnmit: SetNode failed on leaf node: %v", err)
		}
		return nil
	}

	node, ok := nodeI.(ygot.GoStruct)
	path := proto.Clone(update.Path).(*gpb.Path)
	for i := len(path.Elem) - 1; i >= 0 && !ok; i-- {
		path.Elem = path.Elem[:i]
		nodeI, _, err := ytypes.GetOrCreateNode(schema.SchemaTree[nodeName], goStruct, path)
		if err != nil {
			continue
		}
		node, ok = nodeI.(ygot.GoStruct)
	}
	if ok {
		var jsonTree interface{}
		if err := json.Unmarshal(jsonUpdate, &jsonTree); err != nil {
			return err
		}
		return ytypes.Unmarshal(targetSchema, node, jsonTree)
	}

	return fmt.Errorf("gnmit: cannot find GoStruct parent into which to ummarshal update message: %s", prototext.Format(update))
}

// Set is a prototype for a gNMI Set operation.
// TODO(wenbli): This function is too complex to stand on its own. We should
// split this into stages, each encapsulated within its own function.
func (s *GNMIServer) Set(ctx context.Context, req *gpb.SetRequest) (*gpb.SetResponse, error) {
	if s.c.schema == nil {
		return s.UnimplementedGNMIServer.Set(ctx, req)
	}
	// Create a copy so that we can rollback the transaction when validation fails.
	dirtyRootG, err := ygot.DeepCopy(s.c.Schema().Root)
	if err != nil {
		return nil, fmt.Errorf("gnmit: failed to ygot.DeepCopy the cached root object: %v", err)
	}
	dirtyRoot, ok := dirtyRootG.(ygot.ValidatedGoStruct)
	if !ok {
		return nil, fmt.Errorf("gnmit: cannot convert root object to ValidatedGoStruct")
	}
	// Operate at the prefix level.
	nodeI, _, err := ytypes.GetOrCreateNode(s.c.Schema().RootSchema(), dirtyRoot, req.Prefix)
	if err != nil {
		return nil, fmt.Errorf("gnmit: failed to GetOrCreate the prefix node: %v", err)
	}
	node, ok := nodeI.(ygot.GoStruct)
	if !ok {
		return nil, fmt.Errorf("gnmit: prefix path points to a non-GoStruct, this is not allowed: %T, %v", nodeI, nodeI)
	}
	nodeName := reflect.TypeOf(nodeI).Elem().Name()

	// TODO(wenbli): Reject paths that try to modify read-only values.
	// TODO(wenbli): Question: what to do if there are operational-state values in a container that is specified to be replaced or deleted?

	// Process deletes, then replace, then updates.
	for _, path := range req.Delete {
		if err := ytypes.DeleteNode(s.c.Schema().SchemaTree[nodeName], node, path); err != nil {
			return nil, fmt.Errorf("gnmit: DeleteNode error: %v", err)
		}
	}
	for _, update := range req.Replace {
		if err := ytypes.DeleteNode(s.c.Schema().SchemaTree[nodeName], node, update.Path); err != nil {
			return nil, fmt.Errorf("gnmit: DeleteNode error: %v", err)
		}
		if err := setNode(s.c.schema, node, update); err != nil {
			return nil, err
		}
	}
	for _, update := range req.Update {
		if err := setNode(s.c.schema, node, update); err != nil {
			return nil, err
		}
	}

	if err := dirtyRoot.Validate(); err != nil {
		return nil, fmt.Errorf("gnmit: invalid SetRequest: %v", err)
	}

	n, err := ygot.Diff(s.c.Schema().Root, dirtyRoot)
	if err != nil {
		return nil, fmt.Errorf("gnmit: error while creating update notification for Set: %v", err)
	}
	n.Timestamp = time.Now().UnixNano()
	n.Prefix = &gpb.Path{Origin: req.Prefix.Origin, Target: s.c.name}
	if n.Prefix.Origin == "" {
		n.Prefix.Origin = "openconfig"
	}

	// Update cache
	t := s.c.cache.GetTarget(s.c.name)
	if err := t.GnmiUpdate(n); err != nil {
		return nil, err
	}
	// TODO(wenbli): Should handle updates one at a time to avoid partial updates being reflected in the cache when an error occurs. AKA we should support transactional semantics.
	s.c.Schema().Root = dirtyRoot
	// TODO(wenbli): Currently the SetResponse is not filled.
	return &gpb.SetResponse{}, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"encoding/base64"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strings"

	"github.com/buildkite/agent/logger"
)

type ArtifactoryDownloaderConfig struct {
	// The Artifactory repository name and the path, e.g rt://my-repo-name/foo/bar
	Repository string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder,
	// also it's location in the repo
	Path string

	// How many times should it retry the download before giving up
	Retries int

	// If failed responses should be dumped to the log
	DebugHTTP bool
}

type ArtifactoryDownloader struct {
	// The download config
	conf ArtifactoryDownloaderConfig

	// The logger instance to use
	logger logger.Logger
}

func NewArtifactoryDownloader(l logger.Logger, c ArtifactoryDownloaderConfig) *ArtifactoryDownloader {
	return &ArtifactoryDownloader{
		conf:   c,
		logger: l,
	}
}

func (d ArtifactoryDownloader) Start() error {
	// Pull environment variables
	stringURL := os.Getenv("BUILDKITE_ARTIFACTORY_URL")
	username := os.Getenv("BUILDKITE_ARTIFACTORY_USER")
	password := os.Getenv("BUILDKITE_ARTIFACTORY_PASSWORD")
	if stringURL == "" || username == "" || password == "" {
		return errors.New("Must set BUILDKITE_ARTIFACTORY_URL, BUILDKITE_ARTIFACTORY_USER, BUILDKITE_ARTIFACTORY_PASSWORD when using rt:// path")
	}

	// create full URL
	fullURL := fmt.Sprintf("%s/%s/%s",
		strings.TrimSuffix(stringURL, "/"),
		d.RepositoryName(),
		d.RepositoryFileLocation(),
	)

	// create headers map
	headers := map[string]string{
		"Authorization": fmt.Sprintf("Basic %s", getBasicAuthHeader(username, password)),
	}

	// We can now cheat and pass the URL onto our regular downloader
	return NewDownload(d.logger, http.DefaultClient, DownloadConfig{
		URL:         fullURL,
		Path:        d.conf.Path,
		Destination: d.conf.Destination,
		Retries:     d.conf.Retries,
		Headers:     headers,
		DebugHTTP:   d.conf.DebugHTTP,
	}).Start()
}

func (d ArtifactoryDownloader) RepositoryFileLocation() string {
	if d.RepositoryPath() != "" {
		return strings.TrimSuffix(d.RepositoryPath(), "/") + "/" + strings.TrimPrefix(d.conf.Path, "/")
	} else {
		return d.conf.Path
	}
}

func (d ArtifactoryDownloader) RepositoryPath() string {
	return strings.Join(d.destinationParts()[1:len(d.destinationParts())], "/")
}

func (d ArtifactoryDownloader) RepositoryName() string {
	return d.destinationParts()[0]
}

func (d ArtifactoryDownloader) destinationParts() []string {
	trimmed := strings.TrimPrefix(d.conf.Repository, "rt://")

	return strings.Split(trimmed, "/")
}

func getBasicAuthHeader(username, password string) string {
	auth := username + ":" + password
	return base64.StdEncoding.EncodeToString([]byte(auth))
}
<commit_msg>update path parsing to support windows artifact download  (#1014)<commit_after>package agent

import (
	"encoding/base64"
	"errors"
	"fmt"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/buildkite/agent/logger"
)

type ArtifactoryDownloaderConfig struct {
	// The Artifactory repository name and the path, e.g rt://my-repo-name/foo/bar
	Repository string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder,
	// also it's location in the repo
	Path string

	// How many times should it retry the download before giving up
	Retries int

	// If failed responses should be dumped to the log
	DebugHTTP bool
}

type ArtifactoryDownloader struct {
	// The download config
	conf ArtifactoryDownloaderConfig

	// The logger instance to use
	logger logger.Logger
}

func NewArtifactoryDownloader(l logger.Logger, c ArtifactoryDownloaderConfig) *ArtifactoryDownloader {
	return &ArtifactoryDownloader{
		conf:   c,
		logger: l,
	}
}

func (d ArtifactoryDownloader) Start() error {
	// Pull environment variables
	stringURL := os.Getenv("BUILDKITE_ARTIFACTORY_URL")
	username := os.Getenv("BUILDKITE_ARTIFACTORY_USER")
	password := os.Getenv("BUILDKITE_ARTIFACTORY_PASSWORD")
	if stringURL == "" || username == "" || password == "" {
		return errors.New("Must set BUILDKITE_ARTIFACTORY_URL, BUILDKITE_ARTIFACTORY_USER, BUILDKITE_ARTIFACTORY_PASSWORD when using rt:// path")
	}

	// create full URL
	fullURL := fmt.Sprintf("%s/%s/%s",
		strings.TrimSuffix(stringURL, "/"),
		d.RepositoryName(),
		d.RepositoryFileLocation(),
	)

	// create headers map
	headers := map[string]string{
		"Authorization": fmt.Sprintf("Basic %s", getBasicAuthHeader(username, password)),
	}

	// We can now cheat and pass the URL onto our regular downloader
	return NewDownload(d.logger, http.DefaultClient, DownloadConfig{
		URL:         fullURL,
		Path:        d.conf.Path,
		Destination: d.conf.Destination,
		Retries:     d.conf.Retries,
		Headers:     headers,
		DebugHTTP:   d.conf.DebugHTTP,
	}).Start()
}

func (d ArtifactoryDownloader) RepositoryFileLocation() string {
	if d.RepositoryPath() != "" {
		return path.Join(strings.TrimSuffix(d.RepositoryPath(), "/"),"/" , strings.TrimPrefix(filepath.ToSlash(d.conf.Path), "/"))
	} else {
		return d.conf.Path
	}
}

func (d ArtifactoryDownloader) RepositoryPath() string {
	return strings.Join(d.destinationParts()[1:len(d.destinationParts())], "/")
}

func (d ArtifactoryDownloader) RepositoryName() string {
	return d.destinationParts()[0]
}

func (d ArtifactoryDownloader) destinationParts() []string {
	trimmed := strings.TrimPrefix(d.conf.Repository, "rt://")

	return strings.Split(trimmed, "/")
}

func getBasicAuthHeader(username, password string) string {
	auth := username + ":" + password
	return base64.StdEncoding.EncodeToString([]byte(auth))
}
<|endoftext|>


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

// Package main implements an injection function for resource reservations and
// is run with `kustomize config run -- DIR/`.
package main

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"

	multierr "github.com/hashicorp/go-multierror"
	"sigs.k8s.io/kustomize/kyaml/kio"
	"sigs.k8s.io/kustomize/kyaml/kio/kioutil"
	"sigs.k8s.io/kustomize/kyaml/yaml"
)

const (
	loggingServiceValue    = "logging.googleapis.com/kubernetes"
	monitoringServiceValue = "monitoring.googleapis.com/kubernetes"
	containerClusterKind   = "ContainerCluster"
	containerNodePoolKind  = "ContainerNodePool"
	apiGroup               = "container.cnrm.cloud.google.com"
)

var supportedReleaseChannels = []string{"REGULAR", "RAPID", "STABLE"}
var nodeVersionRegex = regexp.MustCompile(`(\d+).(\d+).(\d+)-gke.(\d+)`)
var n1CustomMachineTypeRegex = regexp.MustCompile(`(\w+)-(\d+)-(\d+)`)
var customMachineTypeRegex = regexp.MustCompile(`(\w+)-(\w+)-(\d+)-(\d+)`)
var machineTypeRegex = regexp.MustCompile(`(\w+)-(\w+)-(\d+)`)

func main() {
	rw := &kio.ByteReadWriter{Reader: os.Stdin, Writer: os.Stdout, KeepReaderAnnotations: true}
	p := kio.Pipeline{
		Inputs:  []kio.Reader{rw},       // read the inputs into a slice
		Filters: []kio.Filter{filter{}}, // run the enableASM into the inputs
		Outputs: []kio.Writer{rw}}       // copy the inputs to the output
	if err := p.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

// filter implements kio.Filter
type filter struct{}

// Filter injects new filters into container cluster and nodepool.
func (filter) Filter(in []*yaml.RNode) ([]*yaml.RNode, error) {
	var errList []error

	for _, r := range in {
		if errs := validate(r); len(errs) != 0 {
			errList = append(errList, errs...)
		}
	}
	if errs := multierr.Append(nil, errList...); errs != nil {
		return nil, errs
	}
	return in, nil
}

// https://cloud.google.com/service-mesh/docs/gke-install-existing-cluster#setting_up_your_project
func validate(r *yaml.RNode) []error {

	var errList []error

	meta, err := r.GetMeta()
	if err != nil {
		errList = append(errList, err)
		return errList
	}

	if strings.HasPrefix(meta.ApiVersion, apiGroup) && meta.Kind == containerClusterKind {

		// validate if Cloud Monitoring and Logging are enabled
		if err := validateNodeValue(r, meta, loggingServiceValue, "spec", "loggingService"); err != nil {
			errList = append(errList, err)
		}
		if err := validateNodeValue(r, meta, monitoringServiceValue, "spec", "monitoringService"); err != nil {
			errList = append(errList, err)
		}

		// validate if Workload Identity is enabled
		if _, err := validateNodeExists(r, meta, "spec", "workloadIdentity", "identityNamespace"); err != nil {
			errList = append(errList, err)
		}

		// validate if mesh_id lable is set
		if _, err := validateNodeExists(r, meta, "spec", "labels", "mesh_id"); err != nil {
			errList = append(errList, err)
		}

		// validate release channel
		if err := validateNodeValueIn(r, meta, supportedReleaseChannels, "spec", "releaseChannel", "channel"); err != nil {
			errList = append(errList, err)
		}

		// validate master node version
		if err := validateMasterNodeVersion(r, meta); err != nil {
			errList = append(errList, err)
		}

		// validate machine type
		if err := validateMachineType(r, meta); err != nil {
			errList = append(errList, err)
		}
	}

	if strings.HasPrefix(meta.ApiVersion, apiGroup) && meta.Kind == containerNodePoolKind {
		if err := validateNodeValueGreaterThan(r, meta, 3, "ASM requires at least four nodes. If you need to add nodes, see https://bit.ly/2RnVL2T", "spec", "nodeCount"); err != nil {
			errList = append(errList, err)
		}

		// validate machine type
		if err := validateMachineType(r, meta); err != nil {
			errList = append(errList, err)
		}
	}
	return errList
}

func validateNodeValue(r *yaml.RNode, meta yaml.ResourceMeta, expected string, path ...string) error {
	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	if value != expected {
		return fmt.Errorf(
			"unsupported %s value in %s %s (%s [%s]), expected: %s, actual: %s",
			strings.Join(path, "."),
			meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation],
			expected,
			value)
	}
	return nil
}

func validateNodeExists(r *yaml.RNode, meta yaml.ResourceMeta, path ...string) (*yaml.RNode, error) {
	node, err := r.Pipe(yaml.Lookup(path...))
	if err != nil {
		s, _ := r.String()
		return nil, fmt.Errorf("%v: %s", err, s)
	}
	if node == nil {
		return node, fmt.Errorf(
			"%s missing in %s %s (%s [%s])",
			strings.Join(path, "."),
			meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}
	return node, nil
}

func validateNodeValueIn(r *yaml.RNode, meta yaml.ResourceMeta, expected []string, path ...string) error {
	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	if contains(expected, value) {
		return nil
	}
	return fmt.Errorf(
		"unsupported %s value in %s %s (%s [%s]), expected: %s, actual: %s",
		strings.Join(path, "."),
		meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation],
		strings.Join(expected, ","),
		value)
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

func validateNodeValueGreaterThan(r *yaml.RNode, meta yaml.ResourceMeta, min int, errorInfo string, path ...string) error {
	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	count, err := strconv.Atoi(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	if count <= min {
		return fmt.Errorf(
			"%s is %d in %s %s (%s [%s]). %s",
			strings.Join(path, "."),
			count,
			meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation],
			errorInfo)
	}
	return nil
}

func validateMasterNodeVersion(r *yaml.RNode, meta yaml.ResourceMeta) error {
	node, err := validateNodeExists(r, meta, "status", "masterVersion")
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	version := nodeVersionRegex.FindStringSubmatch(value)
	if len(version) < 5 {
		return fmt.Errorf("unknown masterVersion format: %s in %s %s (%s [%s])",
			value, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}

	g2, err := strconv.Atoi(version[2])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	g3, err := strconv.Atoi(version[3])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	g4, err := strconv.Atoi(version[4])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	switch g2 {
	case 13:
		if g3 < 11 || (g3 == 11 && g4 < 14) {
			return unsupportedGKEVersionError(value, meta)
		}
	case 14:
		if g3 < 8 || (g3 == 8 && g4 < 18) {
			return unsupportedGKEVersionError(value, meta)
		}
	case 15:
		if g3 < 4 || (g3 == 4 && g4 < 15) {
			return unsupportedGKEVersionError(value, meta)
		}
	default:
		return unsupportedGKEVersionError(value, meta)
	}
	return nil

}

func validateMachineType(r *yaml.RNode, meta yaml.ResourceMeta) error {
	node, err := validateNodeExists(r, meta, "spec", "nodeConfig", "machineType")
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	// n1 custom
	if strings.HasPrefix(value, "custom") {
		if err := validateCustomMachineType(r, meta, n1CustomMachineTypeRegex, value, 4, 2); err != nil {
			return err
		}
	} else if strings.Contains(value, "custom") { // n2, e2 custom
		if err := validateCustomMachineType(r, meta, customMachineTypeRegex, value, 5, 3); err != nil {
			return err
		}
	} else if strings.Contains(value, "micro") || strings.Contains(value, "small") || strings.Contains(value, "medium") {
		return unsupportedMachineTypeError(value, meta)
	} else {
		if err := validateCustomMachineType(r, meta, machineTypeRegex, value, 4, 3); err != nil {
			return err
		}
	}

	return nil
}

func validateCustomMachineType(r *yaml.RNode, meta yaml.ResourceMeta, reg *regexp.Regexp, value string, groupLen, vcpuIndex int) error {
	machineType := reg.FindStringSubmatch(value)
	if len(machineType) < groupLen {
		return fmt.Errorf("invalid machineType format: %s in %s %s (%s [%s])", value, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}

	vcpu, err := strconv.Atoi(machineType[vcpuIndex])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	if vcpu < 4 {
		return unsupportedMachineTypeError(value, meta)
	}
	return nil
}

func unsupportedMachineTypeError(value string, meta yaml.ResourceMeta) error {
	return fmt.Errorf("unsupported machine type: %s in %s %s (%s [%s]). The minimum machine type is n1-standard-4, "+
		"which has four vCPUs. If the machine type for your cluster doesn't have at least four vCPUs, "+
		"change the machine type as described here https://bit.ly/2V0KPdu", value, meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation])
}

func unsupportedGKEVersionError(version string, meta yaml.ResourceMeta) error {
	return fmt.Errorf("unsupported GKE version %s in %s %s (%s [%s]). If you need to upgrade "+
		"your cluster to a supported version (https://bit.ly/2XsilLq), see https://bit.ly/34vygKy",
		version, meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation])
}
<commit_msg>Extract value but ignore comments<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package main implements an injection function for resource reservations and
// is run with `kustomize config run -- DIR/`.
package main

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"

	multierr "github.com/hashicorp/go-multierror"
	"sigs.k8s.io/kustomize/kyaml/kio"
	"sigs.k8s.io/kustomize/kyaml/kio/kioutil"
	"sigs.k8s.io/kustomize/kyaml/yaml"
)

const (
	loggingServiceValue    = "logging.googleapis.com/kubernetes"
	monitoringServiceValue = "monitoring.googleapis.com/kubernetes"
	containerClusterKind   = "ContainerCluster"
	containerNodePoolKind  = "ContainerNodePool"
	apiGroup               = "container.cnrm.cloud.google.com"
)

var supportedReleaseChannels = []string{"REGULAR", "RAPID", "STABLE"}
var nodeVersionRegex = regexp.MustCompile(`(\d+).(\d+).(\d+)-gke.(\d+)`)
var n1CustomMachineTypeRegex = regexp.MustCompile(`(\w+)-(\d+)-(\d+)`)
var customMachineTypeRegex = regexp.MustCompile(`(\w+)-(\w+)-(\d+)-(\d+)`)
var machineTypeRegex = regexp.MustCompile(`(\w+)-(\w+)-(\d+)`)
var valueWithComment = regexp.MustCompile(`\s*(\S+)(\s+#.*)?`)

func main() {
	rw := &kio.ByteReadWriter{Reader: os.Stdin, Writer: os.Stdout, KeepReaderAnnotations: true}
	p := kio.Pipeline{
		Inputs:  []kio.Reader{rw},       // read the inputs into a slice
		Filters: []kio.Filter{filter{}}, // run the enableASM into the inputs
		Outputs: []kio.Writer{rw}}       // copy the inputs to the output
	if err := p.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}

// filter implements kio.Filter
type filter struct{}

// Filter injects new filters into container cluster and nodepool.
func (filter) Filter(in []*yaml.RNode) ([]*yaml.RNode, error) {
	var errList []error

	for _, r := range in {
		if errs := validate(r); len(errs) != 0 {
			errList = append(errList, errs...)
		}
	}
	if errs := multierr.Append(nil, errList...); errs != nil {
		return nil, errs
	}
	return in, nil
}

// https://cloud.google.com/service-mesh/docs/gke-install-existing-cluster#setting_up_your_project
func validate(r *yaml.RNode) []error {

	var errList []error

	meta, err := r.GetMeta()
	if err != nil {
		errList = append(errList, err)
		return errList
	}

	if strings.HasPrefix(meta.ApiVersion, apiGroup) && meta.Kind == containerClusterKind {

		// validate if Cloud Monitoring and Logging are enabled
		if err := validateNodeValue(r, meta, loggingServiceValue, "spec", "loggingService"); err != nil {
			errList = append(errList, err)
		}
		if err := validateNodeValue(r, meta, monitoringServiceValue, "spec", "monitoringService"); err != nil {
			errList = append(errList, err)
		}

		// validate if Workload Identity is enabled
		if _, err := validateNodeExists(r, meta, "spec", "workloadIdentity", "identityNamespace"); err != nil {
			errList = append(errList, err)
		}

		// validate if mesh_id lable is set
		if _, err := validateNodeExists(r, meta, "spec", "labels", "mesh_id"); err != nil {
			errList = append(errList, err)
		}

		// validate release channel
		if err := validateReleaseChannel(r, meta, supportedReleaseChannels, "spec", "releaseChannel", "channel"); err != nil {
			errList = append(errList, err)
		}

		// validate master node version
		if err := validateMasterNodeVersion(r, meta); err != nil {
			errList = append(errList, err)
		}

		// validate machine type
		if err := validateMachineType(r, meta); err != nil {
			errList = append(errList, err)
		}
	}

	if strings.HasPrefix(meta.ApiVersion, apiGroup) && meta.Kind == containerNodePoolKind {
		if err := validateNodeValueGreaterThan(r, meta, 3, "ASM requires at least four nodes. If you need to add nodes, see https://bit.ly/2RnVL2T", "spec", "nodeCount"); err != nil {
			errList = append(errList, err)
		}

		// validate machine type
		if err := validateMachineType(r, meta); err != nil {
			errList = append(errList, err)
		}
	}
	return errList
}

func validateNodeValue(r *yaml.RNode, meta yaml.ResourceMeta, expected string, path ...string) error {
	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	pathString := strings.Join(path, ".")
	value, err := stripValueComment(node, r, meta, pathString)
	if err != nil {
		return err
	}

	if value != expected {
		return fmt.Errorf(
			"unsupported %s value in %s %s (%s [%s]), expected: %s, actual: %s",
			pathString, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation],
			expected,
			value)
	}
	return nil
}

func validateNodeExists(r *yaml.RNode, meta yaml.ResourceMeta, path ...string) (*yaml.RNode, error) {
	node, err := r.Pipe(yaml.Lookup(path...))
	if err != nil {
		s, _ := r.String()
		return nil, fmt.Errorf("%v: %s", err, s)
	}
	if node == nil {
		return node, fmt.Errorf(
			"%s missing in %s %s (%s [%s])",
			strings.Join(path, "."),
			meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}
	return node, nil
}

func stripValueComment(node *yaml.RNode, r *yaml.RNode, meta yaml.ResourceMeta, path string) (string, error) {
	value, err := node.String()
	if err != nil {
		s, _ := r.String()
		return "", fmt.Errorf("%v: %s", err, s)
	}
	group := valueWithComment.FindStringSubmatch(value)
	if len(group) < 3 {
		return "", fmt.Errorf("unknown format of %s: %s in %s %s (%s [%s])",
			path, value, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}
	return group[1], nil
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

func validateNodeValueGreaterThan(r *yaml.RNode, meta yaml.ResourceMeta, min int, errorInfo string, path ...string) error {
	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	count, err := strconv.Atoi(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	if count <= min {
		return fmt.Errorf(
			"%s is %d in %s %s (%s [%s]). %s",
			strings.Join(path, "."),
			count,
			meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation],
			errorInfo)
	}
	return nil
}

func validateReleaseChannel(r *yaml.RNode, meta yaml.ResourceMeta, expected []string, path ...string) error {

	node, err := validateNodeExists(r, meta, path...)
	if err != nil {
		return err
	}
	pathString := strings.Join(path, ".")
	value, err := stripValueComment(node, r, meta, pathString)
	if err != nil {
		return err
	}

	if contains(expected, value) {
		return nil
	}
	return fmt.Errorf(
		"unsupported %s value in %s %s (%s [%s]), expected: %s, actual: %s",
		strings.Join(path, "."),
		meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation],
		strings.Join(expected, ","),
		value)

}

func validateMasterNodeVersion(r *yaml.RNode, meta yaml.ResourceMeta) error {
	node, err := validateNodeExists(r, meta, "status", "masterVersion")
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	version := nodeVersionRegex.FindStringSubmatch(value)
	if len(version) < 5 {
		return fmt.Errorf("unknown masterVersion format: %s in %s %s (%s [%s])",
			value, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}

	g2, err := strconv.Atoi(version[2])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	g3, err := strconv.Atoi(version[3])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	g4, err := strconv.Atoi(version[4])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	switch g2 {
	case 13:
		if g3 < 11 || (g3 == 11 && g4 < 14) {
			return unsupportedGKEVersionError(value, meta)
		}
	case 14:
		if g3 < 8 || (g3 == 8 && g4 < 18) {
			return unsupportedGKEVersionError(value, meta)
		}
	case 15:
		if g3 < 4 || (g3 == 4 && g4 < 15) {
			return unsupportedGKEVersionError(value, meta)
		}
	default:
		return unsupportedGKEVersionError(value, meta)
	}
	return nil

}

func validateMachineType(r *yaml.RNode, meta yaml.ResourceMeta) error {
	node, err := validateNodeExists(r, meta, "spec", "nodeConfig", "machineType")
	if err != nil {
		return err
	}
	value, err := node.String()
	value = strings.TrimSpace(value)
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}

	// n1 custom
	if strings.HasPrefix(value, "custom") {
		if err := validateCustomMachineType(r, meta, n1CustomMachineTypeRegex, value, 4, 2); err != nil {
			return err
		}
	} else if strings.Contains(value, "custom") { // n2, e2 custom
		if err := validateCustomMachineType(r, meta, customMachineTypeRegex, value, 5, 3); err != nil {
			return err
		}
	} else if strings.Contains(value, "micro") || strings.Contains(value, "small") || strings.Contains(value, "medium") {
		return unsupportedMachineTypeError(value, meta)
	} else {
		if err := validateCustomMachineType(r, meta, machineTypeRegex, value, 4, 3); err != nil {
			return err
		}
	}

	return nil
}

func validateCustomMachineType(r *yaml.RNode, meta yaml.ResourceMeta, reg *regexp.Regexp, value string, groupLen, vcpuIndex int) error {
	machineType := reg.FindStringSubmatch(value)
	if len(machineType) < groupLen {
		return fmt.Errorf("invalid machineType format: %s in %s %s (%s [%s])", value, meta.Kind, meta.Name,
			meta.Annotations[kioutil.PathAnnotation],
			meta.Annotations[kioutil.IndexAnnotation])
	}

	vcpu, err := strconv.Atoi(machineType[vcpuIndex])
	if err != nil {
		s, _ := r.String()
		return fmt.Errorf("%v: %s", err, s)
	}
	if vcpu < 4 {
		return unsupportedMachineTypeError(value, meta)
	}
	return nil
}

func unsupportedMachineTypeError(value string, meta yaml.ResourceMeta) error {
	return fmt.Errorf("unsupported machine type: %s in %s %s (%s [%s]). The minimum machine type is n1-standard-4, "+
		"which has four vCPUs. If the machine type for your cluster doesn't have at least four vCPUs, "+
		"change the machine type as described here https://bit.ly/2V0KPdu", value, meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation])
}

func unsupportedGKEVersionError(version string, meta yaml.ResourceMeta) error {
	return fmt.Errorf("unsupported GKE version %s in %s %s (%s [%s]). If you need to upgrade "+
		"your cluster to a supported version (https://bit.ly/2XsilLq), see https://bit.ly/34vygKy",
		version, meta.Kind, meta.Name,
		meta.Annotations[kioutil.PathAnnotation],
		meta.Annotations[kioutil.IndexAnnotation])
}
<|endoftext|>


<commit_before>package decktet

import (
	"fmt"
	"strconv"

	"github.com/frenata/deck"
)

/*
type Card interface {
	String() string
}

type Player interface {
	AddCard(c Card)
	String() string
}
*/

type DecktetCard struct {
	name  string
	rank  rank
	suits []suit
	cats  []category
}

func (d *DecktetCard) String() string {
	s := d.rank.String() + " of "

	switch len(d.suits) {
	case 1:
		return s + fmt.Sprint(d.suits[0])
	case 2:
		return s + fmt.Sprintf("%v, %v", d.suits[0], d.suits[1])
	case 3:
		return s + fmt.Sprintf("%v, %v, %v", d.suits[0], d.suits[1], d.suits[2])
	default:
		return "E"
	}
}

func ShortPrint(cards []*DecktetCard) string {
	var s string
	for _, c := range cards {
		s += ShortPrintCard(c) + " "
	}
	return s
}

func ShortPrintCard(card *DecktetCard) string {
	s := shortRank(card.rank)
	for _, suit := range card.suits {
		s += shortSuit(suit)
	}
	if len(card.suits) == 1 {
		s += " "
	}
	return s
}

func SuitMatch(one, two *DecktetCard) bool {
	for _, s1 := range one.suits {
		for _, s2 := range two.suits {
			if s1 == s2 {
				return true
			}
		}
	}
	return false
}

func PopDecktetCard(c *DecktetCard, s []*DecktetCard) bool {
	for i, v := range s {
		if c == v {
			//fmt.Println(s)
			s = append(s[:i], s[i+1:]...)
			return true
		}
	}
	return false
}

func HasSuit(c *DecktetCard, s suit) bool {
	for _, x := range c.Suits() {
		if x == s {
			return true
		}
	}
	return false
}

func (d *DecktetCard) Cats() []category {
	return d.cats
}
func (d *DecktetCard) Suits() []suit {
	return d.suits
}
func (d *DecktetCard) Rank() rank {
	return d.rank
}
func (d *DecktetCard) Name() string {
	return d.name
}

type rank int
type suit string
type category string

func shortRank(r rank) string {
	switch r {
	case Ace:
		return "A"
	case Pawn:
		return "P"
	case Court:
		return "C"
	case Crown:
		return "X"
	default:
		return strconv.Itoa(int(r))
	}
}

func shortSuit(s suit) string {
	switch s {
	case Wyrms:
		return "Y"
	default:
		return string(s[:1])
	}
}

func (r rank) String() string {
	switch r {
	case Ace:
		return "Ace"
	case Two:
		return "Two"
	case Three:
		return "Three"
	case Four:
		return "Four"
	case Five:
		return "Five"
	case Six:
		return "Six"
	case Seven:
		return "Seven"
	case Eight:
		return "Eight"
	case Nine:
		return "Nine"
	case Pawn:
		return "Pawn"
	case Court:
		return "Court"
	case Crown:
		return "Crown"
	default:
		return ""
	}
}

const (
	_        = iota
	Ace rank = iota
	Two
	Three
	Four
	Five
	Six
	Seven
	Eight
	Nine
	Pawn
	Court
	Crown

	Suns   suit = "Suns"
	Moons  suit = "Moons"
	Waves  suit = "Waves"
	Leaves suit = "Leaves"
	Knots  suit = "Knots"
	Wyrms  suit = "Wyrms"

	Person category = "Person"
	Place  category = "Place"
	Event  category = "Event"
)

var BasicRanks = [10]rank{Ace, Two, Three, Four, Five, Six, Seven, Eight, Nine, Crown}

func BasicDeck() *deck.Deck {
	bd := *basicDeck
	return &bd
}

var basicDeck = newDecktet([]*DecktetCard{
	{"", Ace, []suit{Suns}, nil},
	{"", Ace, []suit{Moons}, nil},
	{"", Ace, []suit{Waves}, nil},
	{"", Ace, []suit{Leaves}, nil},
	{"", Ace, []suit{Knots}, nil},
	{"", Ace, []suit{Wyrms}, nil},
	{"Author", Two, []suit{Moons, Knots}, []category{Person}},
	{"Origin", Two, []suit{Waves, Leaves}, []category{Event, Place}},
	{"Desert", Two, []suit{Suns, Wyrms}, []category{Place}},
	{"Journey", Three, []suit{Moons, Waves}, []category{Event}},
	{"Savage", Three, []suit{Leaves, Wyrms}, []category{Person}},
	{"Painter", Three, []suit{Suns, Knots}, []category{Person}},
	{"Mountain", Four, []suit{Moons, Suns}, []category{Place}},
	{"Battle", Four, []suit{Wyrms, Knots}, []category{Event}},
	{"Sailor", Four, []suit{Waves, Leaves}, []category{Person}},
	{"Discovery", Five, []suit{Suns, Waves}, []category{Event}},
	{"Soldier", Five, []suit{Wyrms, Knots}, []category{Person}},
	{"Forest", Five, []suit{Moons, Leaves}, []category{Place}},
	{"Penitent", Six, []suit{Suns, Wyrms}, []category{Person}},
	{"Lunatic", Six, []suit{Moons, Waves}, []category{Person}},
	{"Market", Six, []suit{Leaves, Knots}, []category{Event, Place}},
	{"Cave", Seven, []suit{Waves, Wyrms}, []category{Place}},
	{"Castle", Seven, []suit{Suns, Knots}, []category{Place}},
	{"Chance Meeting", Seven, []suit{Moons, Leaves}, []category{Event}},
	{"Betrayal", Eight, []suit{Wyrms, Knots}, []category{Event}},
	{"Mill", Eight, []suit{Waves, Leaves}, []category{Place}},
	{"Diplomat", Eight, []suit{Moons, Suns}, []category{Person}},
	{"Merchant", Nine, []suit{Leaves, Knots}, []category{Person}},
	{"Darkness", Nine, []suit{Waves, Wyrms}, []category{Place}},
	{"Pact", Nine, []suit{Moons, Suns}, []category{Event}},
	{"Calamity", Crown, []suit{Wyrms}, []category{Event}},
	{"Huntress", Crown, []suit{Moons}, []category{Person}},
	{"Bard", Crown, []suit{Suns}, []category{Person}},
	{"Sea", Crown, []suit{Waves}, []category{Place}},
	{"Windfall", Crown, []suit{Knots}, []category{Event}},
	{"End", Crown, []suit{Leaves}, []category{Event, Place}},
})

func newDecktet(dc []*DecktetCard) *deck.Deck {
	c := make([]deck.Card, len(dc))

	for i := range dc {
		c[i] = dc[i]
	}

	return deck.New(c)
}
<commit_msg>adjusted API to allow for custom decktet decks<commit_after>package decktet

import (
	"fmt"
	"strconv"

	"github.com/frenata/deck"
)

/*
type Card interface {
	String() string
}

type Player interface {
	AddCard(c Card)
	String() string
}
*/

type DecktetCard struct {
	name  string
	rank  rank
	suits []suit
	cats  []category
}

func (d *DecktetCard) String() string {
	s := d.rank.String() + " of "

	switch len(d.suits) {
	case 1:
		return s + fmt.Sprint(d.suits[0])
	case 2:
		return s + fmt.Sprintf("%v, %v", d.suits[0], d.suits[1])
	case 3:
		return s + fmt.Sprintf("%v, %v, %v", d.suits[0], d.suits[1], d.suits[2])
	default:
		return "E"
	}
}

func ShortPrint(cards []*DecktetCard) string {
	var s string
	for _, c := range cards {
		s += ShortPrintCard(c) + " "
	}
	return s
}

func ShortPrintCard(card *DecktetCard) string {
	s := shortRank(card.rank)
	for _, suit := range card.suits {
		s += shortSuit(suit)
	}
	if len(card.suits) == 1 {
		s += " "
	}
	return s
}

func SuitMatch(one, two *DecktetCard) bool {
	for _, s1 := range one.suits {
		for _, s2 := range two.suits {
			if s1 == s2 {
				return true
			}
		}
	}
	return false
}

func PopDecktetCard(c *DecktetCard, s []*DecktetCard) bool {
	for i, v := range s {
		if c == v {
			//fmt.Println(s)
			s = append(s[:i], s[i+1:]...)
			return true
		}
	}
	return false
}

func HasSuit(c *DecktetCard, s suit) bool {
	for _, x := range c.Suits() {
		if x == s {
			return true
		}
	}
	return false
}

func (d *DecktetCard) Cats() []category {
	return d.cats
}
func (d *DecktetCard) Suits() []suit {
	return d.suits
}
func (d *DecktetCard) Rank() rank {
	return d.rank
}
func (d *DecktetCard) Name() string {
	return d.name
}

type rank int
type suit string
type category string

func shortRank(r rank) string {
	switch r {
	case Ace:
		return "A"
	case Pawn:
		return "P"
	case Court:
		return "C"
	case Crown:
		return "X"
	default:
		return strconv.Itoa(int(r))
	}
}

func shortSuit(s suit) string {
	switch s {
	case Wyrms:
		return "Y"
	default:
		return string(s[:1])
	}
}

func (r rank) String() string {
	switch r {
	case Ace:
		return "Ace"
	case Two:
		return "Two"
	case Three:
		return "Three"
	case Four:
		return "Four"
	case Five:
		return "Five"
	case Six:
		return "Six"
	case Seven:
		return "Seven"
	case Eight:
		return "Eight"
	case Nine:
		return "Nine"
	case Pawn:
		return "Pawn"
	case Court:
		return "Court"
	case Crown:
		return "Crown"
	default:
		return ""
	}
}

const (
	_        = iota
	Ace rank = iota
	Two
	Three
	Four
	Five
	Six
	Seven
	Eight
	Nine
	Pawn
	Court
	Crown

	Suns   suit = "Suns"
	Moons  suit = "Moons"
	Waves  suit = "Waves"
	Leaves suit = "Leaves"
	Knots  suit = "Knots"
	Wyrms  suit = "Wyrms"

	Person category = "Person"
	Place  category = "Place"
	Event  category = "Event"
)

var BasicRanks = [10]rank{Ace, Two, Three, Four, Five, Six, Seven, Eight, Nine, Crown}

func BasicDeck() *deck.Deck {
	bd := *basicDeck
	return &bd
}

var basicDeck = NewDecktet([]*DecktetCard{
	{"", Ace, []suit{Suns}, nil},
	{"", Ace, []suit{Moons}, nil},
	{"", Ace, []suit{Waves}, nil},
	{"", Ace, []suit{Leaves}, nil},
	{"", Ace, []suit{Knots}, nil},
	{"", Ace, []suit{Wyrms}, nil},
	{"Author", Two, []suit{Moons, Knots}, []category{Person}},
	{"Origin", Two, []suit{Waves, Leaves}, []category{Event, Place}},
	{"Desert", Two, []suit{Suns, Wyrms}, []category{Place}},
	{"Journey", Three, []suit{Moons, Waves}, []category{Event}},
	{"Savage", Three, []suit{Leaves, Wyrms}, []category{Person}},
	{"Painter", Three, []suit{Suns, Knots}, []category{Person}},
	{"Mountain", Four, []suit{Moons, Suns}, []category{Place}},
	{"Battle", Four, []suit{Wyrms, Knots}, []category{Event}},
	{"Sailor", Four, []suit{Waves, Leaves}, []category{Person}},
	{"Discovery", Five, []suit{Suns, Waves}, []category{Event}},
	{"Soldier", Five, []suit{Wyrms, Knots}, []category{Person}},
	{"Forest", Five, []suit{Moons, Leaves}, []category{Place}},
	{"Penitent", Six, []suit{Suns, Wyrms}, []category{Person}},
	{"Lunatic", Six, []suit{Moons, Waves}, []category{Person}},
	{"Market", Six, []suit{Leaves, Knots}, []category{Event, Place}},
	{"Cave", Seven, []suit{Waves, Wyrms}, []category{Place}},
	{"Castle", Seven, []suit{Suns, Knots}, []category{Place}},
	{"Chance Meeting", Seven, []suit{Moons, Leaves}, []category{Event}},
	{"Betrayal", Eight, []suit{Wyrms, Knots}, []category{Event}},
	{"Mill", Eight, []suit{Waves, Leaves}, []category{Place}},
	{"Diplomat", Eight, []suit{Moons, Suns}, []category{Person}},
	{"Merchant", Nine, []suit{Leaves, Knots}, []category{Person}},
	{"Darkness", Nine, []suit{Waves, Wyrms}, []category{Place}},
	{"Pact", Nine, []suit{Moons, Suns}, []category{Event}},
	{"Calamity", Crown, []suit{Wyrms}, []category{Event}},
	{"Huntress", Crown, []suit{Moons}, []category{Person}},
	{"Bard", Crown, []suit{Suns}, []category{Person}},
	{"Sea", Crown, []suit{Waves}, []category{Place}},
	{"Windfall", Crown, []suit{Knots}, []category{Event}},
	{"End", Crown, []suit{Leaves}, []category{Event, Place}},
})

func NewDecktet(dc []*DecktetCard) *deck.Deck {
	c := make([]deck.Card, len(dc))

	for i := range dc {
		c[i] = dc[i]
	}

	return deck.New(c)
}
<|endoftext|>