<commit_before>package main

import (
	"errors"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"time"
)

var (
	green = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	red   = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	reset = string([]byte{27, 91, 48, 109})
	white = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
)

type Server struct {
	Name      string
	Address   string
	Interval  int
	alerts    []Alert
	log       *log.Logger
	service   *Service
	failCount int
	LastEvent *Event
	wg        sync.WaitGroup
}

func (s *Service) AddServer(name string, address string, interval int, alertNames []string) {

	alerts := []Alert{}
	for _, alertName := range alertNames {
		alerts = append(alerts, s.GetAlert(alertName))
	}

	var wg sync.WaitGroup
	s.servers = append(s.servers, &Server{
		Name:     name,
		Address:  address,
		Interval: interval,
		alerts:   alerts,
		log:      log.New(os.Stdout, name+" ", log.Ldate|log.Ltime),
		service:  s,
		wg:       wg,
	})

}

func (s *Server) Ping() (time.Duration, error) {

	startTime := time.Now()
	s.log.Println("Pinging: ", s.Name)

	req, err := http.NewRequest("GET", s.Address, nil)
	if err != nil {
		return 0, errors.New("redalert ping: failed parsing url in http.NewRequest " + err.Error())
	}

	req.Header.Add("User-Agent", "Redalert/1.0")
	client := &http.Client{}
	resp, err := client.Do(req)

	endTime := time.Now()
	latency := endTime.Sub(startTime)
	s.log.Println(white, "Analytics: ", latency, reset)

	if err != nil {
		return latency, errors.New("redalert ping: failed client.Do " + err.Error())
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return latency, errors.New("redalert ping: non-200 status code. status code was " + strconv.Itoa(resp.StatusCode))
	}
	s.log.Println(green, "OK", reset, s.Name)

	return latency, nil
}

func (s *Server) SchedulePing(stopChan chan bool) {

	go func() {

		var err error
		var event *Event
		var latency time.Duration

		originalDelay := time.Second * time.Duration(s.Interval)
		delay := time.Second * time.Duration(s.Interval)

		for {

			latency, err = s.Ping()

			if err != nil {

				s.log.Println(red, "ERROR: ", err, reset)
				event = NewRedAlert(s, latency)
				s.LastEvent = event

				s.TriggerAlerts(event)

				s.IncrFailCount()
				if s.failCount > 0 {
					delay = time.Second * time.Duration(s.failCount*s.Interval)
				}

			} else {

				event = NewGreenAlert(s, latency)
				isRedalertRecovery := s.LastEvent != nil && s.LastEvent.isRedAlert()
				s.LastEvent = event
				if isRedalertRecovery {
					s.log.Println(green, "RECOVERY: ", reset, s.Name)
					s.TriggerAlerts(event)
				}

				delay = originalDelay
				s.failCount = 0

			}

			select {
			case <-time.After(delay):
			case <-stopChan:
				return
			}
		}
	}()

}

func (s *Server) Monitor() {

	s.service.wg.Add(1)
	s.wg.Add(1)

	stopScheduler := make(chan bool)
	s.SchedulePing(stopScheduler)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		for _ = range sigChan {
			stopScheduler <- true
			s.wg.Done()
		}
	}()

	s.wg.Wait()

	s.service.wg.Done()

}

func (s *Server) TriggerAlerts(event *Event) {

	go func() {

		var err error
		for _, alert := range s.alerts {
			err = alert.Trigger(event)
			if err != nil {
				s.log.Println(red, "CRITICAL: Failure triggering alert ["+alert.Name()+"]: ", err.Error())
			}
		}

	}()
}

func (s *Server) IncrFailCount() {
	s.failCount++
}
<commit_msg>Re-ping before triggering alert.<commit_after>package main

import (
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"time"
)

var (
	green = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	red   = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	reset = string([]byte{27, 91, 48, 109})
	white = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
)

type Server struct {
	Name      string
	Address   string
	Interval  int
	alerts    []Alert
	log       *log.Logger
	service   *Service
	failCount int
	LastEvent *Event
	wg        sync.WaitGroup
}

func (s *Service) AddServer(name string, address string, interval int, alertNames []string) {

	alerts := []Alert{}
	for _, alertName := range alertNames {
		alerts = append(alerts, s.GetAlert(alertName))
	}

	var wg sync.WaitGroup
	s.servers = append(s.servers, &Server{
		Name:     name,
		Address:  address,
		Interval: interval,
		alerts:   alerts,
		log:      log.New(os.Stdout, name+" ", log.Ldate|log.Ltime),
		service:  s,
		wg:       wg,
	})

}

var GlobalClient = http.Client{
	Timeout: time.Duration(10 * time.Second),
}

func (s *Server) Ping() (time.Duration, error) {

	startTime := time.Now()
	s.log.Println("Pinging: ", s.Name)

	req, err := http.NewRequest("GET", s.Address, nil)
	if err != nil {
		return 0, errors.New("redalert ping: failed parsing url in http.NewRequest " + err.Error())
	}

	req.Header.Add("User-Agent", "Redalert/1.0")
	resp, err := GlobalClient.Do(req)

	endTime := time.Now()
	latency := endTime.Sub(startTime)
	s.log.Println(white, "Analytics: ", latency, reset)

	if resp != nil {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}
	if err != nil {
		return latency, errors.New("redalert ping: failed client.Do " + err.Error())
	}

	if resp.StatusCode != http.StatusOK {
		return latency, errors.New("redalert ping: non-200 status code. status code was " + strconv.Itoa(resp.StatusCode))
	}
	s.log.Println(green, "OK", reset, s.Name)

	return latency, nil
}

func (s *Server) SchedulePing(stopChan chan bool) {

	go func() {

		var err error
		var event *Event
		var latency time.Duration

		originalDelay := time.Second * time.Duration(s.Interval)
		delay := time.Second * time.Duration(s.Interval)

		for {

			latency, err = s.Ping()

			if err != nil {

				s.log.Println(red, "ERROR: ", err, reset)
				event = NewRedAlert(s, latency)
				s.LastEvent = event

				// before sending an alert, pause 5 seconds & retry
				// prevent alerts from occaisional errors ('no such host' / 'i/o timeout') on cloud providers
				// todo: adjust sleep to fit with interval
				time.Sleep(5 * time.Second)
				_, rePingErr := s.Ping()
				if rePingErr != nil {

					// re-ping fails (confirms error)

					s.TriggerAlerts(event)

					s.IncrFailCount()
					if s.failCount > 0 {
						delay = time.Second * time.Duration(s.failCount*s.Interval)
					}

				} else {

					// re-ping succeeds (likely false positive, discard last event)

					delay = originalDelay
					s.failCount = 0
					s.LastEvent = nil
				}

			} else {

				event = NewGreenAlert(s, latency)
				isRedalertRecovery := s.LastEvent != nil && s.LastEvent.isRedAlert()
				s.LastEvent = event
				if isRedalertRecovery {
					s.log.Println(green, "RECOVERY: ", reset, s.Name)
					s.TriggerAlerts(event)
				}

				delay = originalDelay
				s.failCount = 0

			}

			select {
			case <-time.After(delay):
			case <-stopChan:
				return
			}
		}
	}()

}

func (s *Server) Monitor() {

	s.service.wg.Add(1)
	s.wg.Add(1)

	stopScheduler := make(chan bool)
	s.SchedulePing(stopScheduler)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		for _ = range sigChan {
			stopScheduler <- true
			s.wg.Done()
		}
	}()

	s.wg.Wait()

	s.service.wg.Done()

}

func (s *Server) TriggerAlerts(event *Event) {

	go func() {

		var err error
		for _, alert := range s.alerts {
			err = alert.Trigger(event)
			if err != nil {
				s.log.Println(red, "CRITICAL: Failure triggering alert ["+alert.Name()+"]: ", err.Error())
			}
		}

	}()
}

func (s *Server) IncrFailCount() {
	s.failCount++
}
<|endoftext|>


<commit_before>package tarantool

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"sync"
)

const greetingSize = 128
const saltSize = 32
const tarantoolVersion = "Tarantool 1.6.8 (Binary)"
const connBufSize = 128 * 1024

type QueryHandler func(query Query) *Result

type IprotoServer struct {
	conn      net.Conn
	reader    *bufio.Reader
	writer    *bufio.Writer
	uuid      string
	salt      []byte // base64-encoded salt
	quit      chan bool
	handler   QueryHandler
	output    chan []byte
	closeOnce sync.Once
}

func NewIprotoServer(uuid string, handler QueryHandler) *IprotoServer {
	return &IprotoServer{
		conn:    nil,
		reader:  nil,
		writer:  nil,
		handler: handler,
		uuid:    uuid,
	}
}

func (s *IprotoServer) Accept(conn net.Conn) {
	s.conn = conn
	s.reader = bufio.NewReaderSize(conn, connBufSize)
	s.writer = bufio.NewWriterSize(conn, connBufSize)
	s.quit = make(chan bool)
	s.output = make(chan []byte, 1024)

	err := s.greet()
	if err != nil {
		conn.Close()
		return
	}

	go s.loop()
}

func (s *IprotoServer) CheckAuth(hash []byte, password string) bool {
	scr, err := scramble(s.salt, password)
	if err != nil {
		return false
	}

	if len(scr) != len(hash) {
		return false
	}

	for i, v := range hash {
		if v != scr[i] {
			return false
		}
	}
	return true
}

func (s *IprotoServer) Close() {
	s.closeOnce.Do(func() {
		close(s.quit)
		s.conn.Close()
	})
}

func (s *IprotoServer) greet() (err error) {
	var line1, line2 string
	var format, greeting string
	var n int

	salt := make([]byte, saltSize)
	_, err = rand.Read(s.salt)
	if err != nil {
		return
	}

	s.salt = []byte(base64.StdEncoding.EncodeToString(salt))

	line1 = fmt.Sprintf("%s %s", tarantoolVersion, s.uuid)
	line2 = fmt.Sprintf("%s", s.salt)

	format = fmt.Sprintf("%%-%ds\n%%-%ds\n", greetingSize/2-1, greetingSize/2-1)
	greeting = fmt.Sprintf(format, line1, line2)

	// send greeting
	n, err = fmt.Fprintf(s.writer, "%s", greeting)
	if err != nil || n != greetingSize {
		return
	}

	return s.writer.Flush()
}

func (s *IprotoServer) loop() {
	go s.read()
	go s.write()
}

func (s *IprotoServer) read() {
	var packet *Packet
	var err error
	var body []byte

	r := s.reader

READER_LOOP:
	for {
		select {
		case <-s.quit:
			break READER_LOOP
		default:
			// read raw bytes
			body, err = readMessage(r)
			if err != nil {
				break READER_LOOP
			}

			packet, err = decodePacket(bytes.NewBuffer(body))
			if err != nil {
				break READER_LOOP
			}

			if packet.request != nil {
				go func(packet *Packet) {
					var res *Result
					var code = byte(packet.code)

					if code == PingRequest {
						s.output <- packIprotoOk(packet.requestID)
					} else {
						res = s.handler(packet.request.(Query))
						body, _ = res.pack(packet.requestID)
						s.output <- body
					}
				}(packet)
			}
		}
	}

	s.Close()
}

func (s *IprotoServer) write() {
	var err error
	var n int

	w := s.writer

WRITER_LOOP:
	for {
		select {
		case messageBody, ok := <-s.output:
			if !ok {
				break WRITER_LOOP
			}
			n, err = w.Write(messageBody)
			if err != nil || n != len(messageBody) {
				break WRITER_LOOP
			}
		case <-s.quit:
			w.Flush()
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case messageBody, ok := <-s.output:
				if !ok {
					break WRITER_LOOP
				}
				n, err = w.Write(messageBody)
				if err != nil || n != len(messageBody) {
					break WRITER_LOOP
				}
			case <-s.quit:
				w.Flush()
				break WRITER_LOOP
			}

		}
	}

	s.Close()
}
<commit_msg>Add OnDisconnect callback<commit_after>package tarantool

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"sync"
)

const greetingSize = 128
const saltSize = 32
const tarantoolVersion = "Tarantool 1.6.8 (Binary)"
const connBufSize = 128 * 1024

type QueryHandler func(query Query) *Result
type OnDisconnectCallback func()

type IprotoServer struct {
	conn          net.Conn
	reader        *bufio.Reader
	writer        *bufio.Writer
	uuid          string
	salt          []byte // base64-encoded salt
	quit          chan bool
	handler       QueryHandler
	on_disconnect OnDisconnectCallback
	output        chan []byte
	closeOnce     sync.Once
}

func NewIprotoServer(uuid string, handler QueryHandler, on_disconnect OnDisconnectCallback) *IprotoServer {
	return &IprotoServer{
		conn:          nil,
		reader:        nil,
		writer:        nil,
		handler:       handler,
		on_disconnect: on_disconnect,
		uuid:          uuid,
	}
}

func (s *IprotoServer) Accept(conn net.Conn) {
	s.conn = conn
	s.reader = bufio.NewReaderSize(conn, connBufSize)
	s.writer = bufio.NewWriterSize(conn, connBufSize)
	s.quit = make(chan bool)
	s.output = make(chan []byte, 1024)

	err := s.greet()
	if err != nil {
		s.Close()
		return
	}

	go s.loop()
}

func (s *IprotoServer) CheckAuth(hash []byte, password string) bool {
	scr, err := scramble(s.salt, password)
	if err != nil {
		return false
	}

	if len(scr) != len(hash) {
		return false
	}

	for i, v := range hash {
		if v != scr[i] {
			return false
		}
	}
	return true
}

func (s *IprotoServer) Close() {
	s.closeOnce.Do(func() {
		if s.on_disconnect != nil {
			s.on_disconnect()
		}
		close(s.quit)
		s.conn.Close()
	})
}

func (s *IprotoServer) greet() (err error) {
	var line1, line2 string
	var format, greeting string
	var n int

	salt := make([]byte, saltSize)
	_, err = rand.Read(s.salt)
	if err != nil {
		return
	}

	s.salt = []byte(base64.StdEncoding.EncodeToString(salt))

	line1 = fmt.Sprintf("%s %s", tarantoolVersion, s.uuid)
	line2 = fmt.Sprintf("%s", s.salt)

	format = fmt.Sprintf("%%-%ds\n%%-%ds\n", greetingSize/2-1, greetingSize/2-1)
	greeting = fmt.Sprintf(format, line1, line2)

	// send greeting
	n, err = fmt.Fprintf(s.writer, "%s", greeting)
	if err != nil || n != greetingSize {
		return
	}

	return s.writer.Flush()
}

func (s *IprotoServer) loop() {
	go s.read()
	go s.write()
}

func (s *IprotoServer) read() {
	var packet *Packet
	var err error
	var body []byte

	r := s.reader

READER_LOOP:
	for {
		select {
		case <-s.quit:
			break READER_LOOP
		default:
			// read raw bytes
			body, err = readMessage(r)
			if err != nil {
				break READER_LOOP
			}

			packet, err = decodePacket(bytes.NewBuffer(body))
			if err != nil {
				break READER_LOOP
			}

			if packet.request != nil {
				go func(packet *Packet) {
					var res *Result
					var code = byte(packet.code)

					if code == PingRequest {
						s.output <- packIprotoOk(packet.requestID)
					} else {
						res = s.handler(packet.request.(Query))
						body, _ = res.pack(packet.requestID)
						s.output <- body
					}
				}(packet)
			}
		}
	}

	s.Close()
}

func (s *IprotoServer) write() {
	var err error
	var n int

	w := s.writer

WRITER_LOOP:
	for {
		select {
		case messageBody, ok := <-s.output:
			if !ok {
				break WRITER_LOOP
			}
			n, err = w.Write(messageBody)
			if err != nil || n != len(messageBody) {
				break WRITER_LOOP
			}
		case <-s.quit:
			w.Flush()
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case messageBody, ok := <-s.output:
				if !ok {
					break WRITER_LOOP
				}
				n, err = w.Write(messageBody)
				if err != nil || n != len(messageBody) {
					break WRITER_LOOP
				}
			case <-s.quit:
				w.Flush()
				break WRITER_LOOP
			}

		}
	}

	s.Close()
}
<|endoftext|>


<commit_before>package webapp

import (
	"log"
	"net/http"
	"path/filepath"
	"text/template"

	"github.com/sdboyer/pipeviz/broker"
	"github.com/sdboyer/pipeviz/represent"
	"github.com/zenazn/goji/web"
	"github.com/zenazn/goji/web/middleware"
)

var (
	assetDir = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "assets")
	jsDir    = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "src")
	tmplDir  = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "tmpl")
)

var (
	// TODO crappily hardcoded, for now
	brokerListen broker.GraphReceiver
	latestGraph  represent.CoreGraph
)

func init() {
	// Subscribe to the master broker and store latest locally as it comes
	brokerListen = broker.Get().Subscribe()
	// FIXME spawning a goroutine in init() used to crappy, is it still?
	go func() {
		for g := range brokerListen {
			latestGraph = g
		}
	}()
}

// Creates a Goji *web.Mux that can act as the http muxer for the frontend app.
func NewMux() *web.Mux {
	m := web.New()

	m.Use(middleware.Logger)
	m.Get("/assets/*", http.StripPrefix("/assets/", http.FileServer(http.Dir(assetDir))))
	m.Get("/js/*", http.StripPrefix("/js/", http.FileServer(http.Dir(jsDir))))
	m.Get("/", RootEntry)

	return m
}

func RootEntry(w http.ResponseWriter, r *http.Request) {
	vars := struct {
		Title string
	}{
		Title: "pipeviz",
	}

	t, err := template.ParseFiles(filepath.Join(tmplDir, "index.html"))
	if err != nil {
		log.Fatal(err)
	}
	t.Execute(w, vars)
}
<commit_msg>Better name for root route<commit_after>package webapp

import (
	"log"
	"net/http"
	"path/filepath"
	"text/template"

	"github.com/sdboyer/pipeviz/broker"
	"github.com/sdboyer/pipeviz/represent"
	"github.com/zenazn/goji/web"
	"github.com/zenazn/goji/web/middleware"
)

var (
	assetDir = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "assets")
	jsDir    = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "src")
	tmplDir  = filepath.Join(defaultBase("github.com/sdboyer/pipeviz/webapp"), "tmpl")
)

var (
	// TODO crappily hardcoded, for now
	brokerListen broker.GraphReceiver
	latestGraph  represent.CoreGraph
)

func init() {
	// Subscribe to the master broker and store latest locally as it comes
	brokerListen = broker.Get().Subscribe()
	// FIXME spawning a goroutine in init() used to crappy, is it still?
	go func() {
		for g := range brokerListen {
			latestGraph = g
		}
	}()
}

// Creates a Goji *web.Mux that can act as the http muxer for the frontend app.
func NewMux() *web.Mux {
	m := web.New()

	m.Use(middleware.Logger)
	m.Get("/assets/*", http.StripPrefix("/assets/", http.FileServer(http.Dir(assetDir))))
	m.Get("/js/*", http.StripPrefix("/js/", http.FileServer(http.Dir(jsDir))))
	m.Get("/", WebRoot)

	return m
}

func WebRoot(w http.ResponseWriter, r *http.Request) {
	vars := struct {
		Title string
	}{
		Title: "pipeviz",
	}

	t, err := template.ParseFiles(filepath.Join(tmplDir, "index.html"))
	if err != nil {
		log.Fatal(err)
	}
	t.Execute(w, vars)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
)

func StartServer() {
	listener, err := net.Listen("tcp", Config.String("net.listen"))
	if err != nil {
		log.Fatal(err)
	}
	defer listener.Close()
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Fatal(err)
		}
		go communicate(conn)
	}
}

func communicate(conn net.Conn) {
	log.Printf("%v connected", conn.RemoteAddr())

	buf := make([]byte, TotalVoxels * 3)
	for {
		read, err := conn.Read(buf[:3])
		if err != nil {
			log.Printf("%v disconnected", conn.RemoteAddr())
			break
		}
		if read != 3 {
			conn.Write([]byte("err"))
			log.Println("Client did not sent 3 command bytes")
			continue
		}
		switch string(buf[:3]) {
		case "frm":
			for completed := 0; completed < TotalVoxels * 3; {
				read, err := conn.Read(buf)
				if err != nil {
					log.Printf("%v disconnected", conn.RemoteAddr())
					break
				}
				for i, b := range buf[:read] {
					DisplayBackBuffer[completed+i] = float32(b) / 256
				}
				completed += read
			}
		case "swp":
			SwapDisplayBuffer()
		default:
			conn.Write([]byte("err"))
		}
	}
}
<commit_msg>Fixed buffer index out of range<commit_after>package main

import (
	"log"
	"net"
)

func StartServer() {
	listener, err := net.Listen("tcp", Config.String("net.listen"))
	if err != nil {
		log.Fatal(err)
	}
	defer listener.Close()
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Fatal(err)
		}
		go communicate(conn)
	}
}

func communicate(conn net.Conn) {
	log.Printf("%v connected", conn.RemoteAddr())

	buf := make([]byte, TotalVoxels * 3)
	for {
		_, err := conn.Read(buf[:3])
		if err != nil {
			log.Printf("%v disconnected", conn.RemoteAddr())
			break
		}
		switch string(buf[:3]) {
		case "frm":
			for completed := 0; completed < TotalVoxels * 3; {
				read, err := conn.Read(buf[:TotalVoxels*3 - completed])
				if err != nil {
					log.Printf("%v disconnected", conn.RemoteAddr())
					break
				}
				if read + completed > TotalVoxels * 3 {
					conn.Close()
					log.Printf("%v disconnected (frm overflow %v)", conn.RemoteAddr(), (completed+read) - TotalVoxels*3)
					return
				}
				for i, b := range buf[:read] {
					DisplayBackBuffer[completed+i] = float32(b) / 256
				}
				completed += read
			}
		case "swp":
			SwapDisplayBuffer()
		default:
			conn.Write([]byte("err\n"))
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"

	"github.com/flynn/go-discover/discover"
	lornec "github.com/flynn/lorne/client"
	"github.com/flynn/lorne/types"
	sampic "github.com/flynn/sampi/client"
	"github.com/flynn/sampi/types"
	strowgerc "github.com/flynn/strowger/client"
	"github.com/flynn/strowger/types"
	"github.com/titanous/go-dockerclient"
	"github.com/titanous/go-tigertonic"
)

func main() {
	var err error
	scheduler, err = sampic.New()
	if err != nil {
		log.Fatal(err)
	}
	disc, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	router, err = strowgerc.New()
	if err != nil {
		log.Fatal(err)
	}

	mux := tigertonic.NewTrieServeMux()
	mux.Handle("PUT", "/apps/{app_id}/domains/{domain}", tigertonic.Marshaled(addDomain))
	mux.Handle("POST", "/apps/{app_id}/formation/{formation_id}", tigertonic.Marshaled(changeFormation))
	mux.Handle("GET", "/apps/{app_id}/jobs", tigertonic.Marshaled(getJobs))
	mux.HandleFunc("GET", "/apps/{app_id}/jobs/{job_id}/logs", getJobLog)
	mux.HandleFunc("POST", "/apps/{app_id}/jobs", runJob)
	http.ListenAndServe("127.0.0.1:1200", tigertonic.Logged(mux, nil))
}

var scheduler *sampic.Client
var disc *discover.Client
var router *strowgerc.Client

type Job struct {
	ID   string `json:"id"`
	Type string `json:"type"`
}

func addDomain(u *url.URL, h http.Header, data *struct{}) (int, http.Header, struct{}, error) {
	q := u.Query()
	if err := router.AddFrontend(&strowger.Config{Service: q.Get("app_id"), HTTPDomain: q.Get("domain")}); err != nil {
		return 500, nil, struct{}{}, err
	}
	return 200, nil, struct{}{}, nil
}

// GET /apps/{app_id}/jobs
func getJobs(u *url.URL, h http.Header) (int, http.Header, []Job, error) {
	state, err := scheduler.State()
	if err != nil {
		return 500, nil, nil, err
	}

	q := u.Query()
	prefix := q.Get("app_id") + "-"
	jobs := make([]Job, 0)
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				typ := strings.Split(job.ID[len(prefix):], ".")[0]
				jobs = append(jobs, Job{ID: job.ID, Type: typ})
			}
		}
	}

	return 200, nil, jobs, nil
}

type Formation struct {
	Quantity int    `json:"quantity"`
	Type     string `json:"type"`
}

func shelfURL() string {
	set, _ := disc.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	return addrs[0]
}

// POST /apps/{app_id}/formation/{formation_id}
func changeFormation(u *url.URL, h http.Header, req *Formation) (int, http.Header, *Formation, error) {
	state, err := scheduler.State()
	if err != nil {
		log.Println("scheduler state error", err)
		return 500, nil, nil, err
	}

	q := u.Query()
	req.Type = q.Get("formation_id")
	prefix := q.Get("app_id") + "-" + req.Type + "."
	var jobs []*sampi.Job
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				if job.Attributes == nil {
					job.Attributes = make(map[string]string)
				}
				job.Attributes["host_id"] = host.ID
				jobs = append(jobs, job)
			}
		}
	}

	if req.Quantity < 0 {
		req.Quantity = 0
	}
	diff := req.Quantity - len(jobs)
	log.Printf("have %d %s, diff %d", len(jobs), req.Type, diff)
	if diff > 0 {
		config := &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          []string{"start", req.Type},
			AttachStdout: true,
			AttachStderr: true,
			Env:          []string{"SLUG_URL=http://" + shelfURL() + "/" + q.Get("app_id") + ".tgz"},
		}
		if req.Type == "web" {
			config.Env = append(config.Env, "SD_NAME="+q.Get("app_id"))
		}
		schedReq := &sampi.ScheduleReq{
			HostJobs: make(map[string][]*sampi.Job),
		}
	outer:
		for {
			for host := range state {
				schedReq.HostJobs[host] = append(schedReq.HostJobs[host], &sampi.Job{ID: prefix + randomID(), TCPPorts: 1, Config: config})
				diff--
				if diff == 0 {
					break outer
				}
			}
		}

		res, err := scheduler.Schedule(schedReq)
		if err != nil || !res.Success {
			log.Println("schedule error", err)
			return 500, nil, nil, err
		}
	} else if diff < 0 {
		for _, job := range jobs[:-diff] {
			host, err := lornec.New(job.Attributes["host_id"])
			if err != nil {
				log.Println("error connecting to", job.Attributes["host_id"], err)
				continue
			}
			if err := host.StopJob(job.ID); err != nil {
				log.Println("error stopping", job.ID, "on", job.Attributes["host_id"], err)
			}
		}
	}

	return 200, nil, req, nil
}

// GET /apps/{app_id}/jobs/{job_id}/logs
func getJobLog(w http.ResponseWriter, req *http.Request) {
	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	q := req.URL.Query()
	jobID := q.Get("job_id")
	if prefix := q.Get("app_id") + "-"; !strings.HasPrefix(jobID, prefix) {
		jobID = prefix + jobID
	}
	var job *sampi.Job
	var host sampi.Host
outer:
	for _, host = range state {
		for _, job = range host.Jobs {
			if job.ID == jobID {
				break outer
			}
		}
		job = nil
	}
	if job == nil {
		w.WriteHeader(404)
		return
	}

	attachReq := &lorne.AttachReq{
		JobID: job.ID,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagLogs,
	}
	err, errChan := lorneAttach(host.ID, attachReq, w, nil)
	if err != nil {
		w.WriteHeader(500)
		log.Println("attach error", err)
		return
	}
	if err := <-errChan; err != nil {
		log.Println("attach failed", err)
	}
}

type NewJob struct {
	Cmd     []string          `json:"cmd"`
	Env     map[string]string `json:"env"`
	Attach  bool              `json:"attach"`
	TTY     bool              `json:"tty"`
	Columns int               `json:"tty_columns"`
	Lines   int               `json:"tty_lines"`
}

// POST /apps/{app_id}/jobs
func runJob(w http.ResponseWriter, req *http.Request) {
	var jobReq NewJob
	if err := json.NewDecoder(req.Body).Decode(&jobReq); err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}
	// pick a random host
	var hostID string
	for hostID = range state {
		break
	}
	if hostID == "" {
		w.WriteHeader(500)
		log.Println("no hosts found")
		return
	}

	env := make([]string, 0, len(jobReq.Env))
	for k, v := range jobReq.Env {
		env = append(env, k+"="+v)
	}

	q := req.URL.Query()
	job := &sampi.Job{
		ID: q.Get("app_id") + "-run." + randomID(),
		Config: &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          jobReq.Cmd,
			AttachStdin:  true,
			AttachStdout: true,
			AttachStderr: true,
			StdinOnce:    true,
			Env:          append(env, "SLUG_URL=http://"+shelfURL()+"/"+q.Get("app_id")+".tgz"),
		},
	}
	if jobReq.TTY {
		job.Config.Tty = true
	}
	if jobReq.Attach {
		job.Config.AttachStdin = true
		job.Config.StdinOnce = true
		job.Config.OpenStdin = true
	}

	outR, outW := io.Pipe()
	inR, inW := io.Pipe()
	defer outR.Close()
	defer inW.Close()
	var errChan <-chan error
	if jobReq.Attach {
		attachReq := &lorne.AttachReq{
			JobID:  job.ID,
			Flags:  lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
			Height: jobReq.Lines,
			Width:  jobReq.Columns,
		}
		err, errChan = lorneAttach(hostID, attachReq, outW, inR)
		if err != nil {
			w.WriteHeader(500)
			log.Println("attach failed", err)
			return
		}
	}

	res, err := scheduler.Schedule(&sampi.ScheduleReq{HostJobs: map[string][]*sampi.Job{hostID: {job}}})
	if err != nil || !res.Success {
		w.WriteHeader(500)
		log.Println("schedule failed", err)
		return
	}

	if jobReq.Attach {
		w.Header().Set("Content-Type", "application/vnd.flynn.hijack")
		w.Header().Set("Content-Length", "0")
		w.WriteHeader(200)
		conn, bufrw, err := w.(http.Hijacker).Hijack()
		if err != nil {
			return
		}
		bufrw.Flush()
		go func() {
			buf := make([]byte, bufrw.Reader.Buffered())
			bufrw.Read(buf)
			inW.Write(buf)
			io.Copy(inW, conn)
			inW.Close()
		}()
		go io.Copy(conn, outR)
		<-errChan
		conn.Close()
		return
	}
	w.WriteHeader(200)
}

func lorneAttach(host string, req *lorne.AttachReq, out io.Writer, in io.Reader) (error, <-chan error) {
	services, err := disc.Services("flynn-lorne-attach." + host)
	if err != nil {
		return err, nil
	}
	addrs := services.OnlineAddrs()
	if len(addrs) == 0 {
		return err, nil
	}
	conn, err := net.Dial("tcp", addrs[0])
	if err != nil {
		return err, nil
	}
	err = gob.NewEncoder(conn).Encode(req)
	if err != nil {
		conn.Close()
		return err, nil
	}

	errChan := make(chan error)

	attach := func() {
		defer conn.Close()
		inErr := make(chan error, 1)
		if in != nil {
			go func() {
				io.Copy(conn, in)
			}()
		} else {
			close(inErr)
		}
		_, outErr := io.Copy(out, conn)
		if outErr != nil {
			errChan <- outErr
			return
		}
		errChan <- <-inErr
	}

	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		conn.Close()
		return err, nil
	}
	switch attachState[0] {
	case lorne.AttachError:
		errBytes, err := ioutil.ReadAll(conn)
		conn.Close()
		if err != nil {
			return err, nil
		}
		return errors.New(string(errBytes)), nil
	case lorne.AttachWaiting:
		go func() {
			if _, err := conn.Read(attachState); err != nil {
				conn.Close()
				errChan <- err
				return
			}
			if attachState[0] == lorne.AttachError {
				errBytes, err := ioutil.ReadAll(conn)
				conn.Close()
				if err != nil {
					errChan <- err
					return
				}
				errChan <- errors.New(string(errBytes))
				return
			}
			attach()
		}()
		return nil, errChan
	default:
		go attach()
		return nil, errChan
	}
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<commit_msg>Use new lorne attach client<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/url"
	"strings"

	"github.com/flynn/go-discover/discover"
	lornec "github.com/flynn/lorne/client"
	"github.com/flynn/lorne/types"
	sampic "github.com/flynn/sampi/client"
	"github.com/flynn/sampi/types"
	strowgerc "github.com/flynn/strowger/client"
	"github.com/flynn/strowger/types"
	"github.com/titanous/go-dockerclient"
	"github.com/titanous/go-tigertonic"
)

func main() {
	var err error
	scheduler, err = sampic.New()
	if err != nil {
		log.Fatal(err)
	}
	disc, err = discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	router, err = strowgerc.New()
	if err != nil {
		log.Fatal(err)
	}

	mux := tigertonic.NewTrieServeMux()
	mux.Handle("PUT", "/apps/{app_id}/domains/{domain}", tigertonic.Marshaled(addDomain))
	mux.Handle("POST", "/apps/{app_id}/formation/{formation_id}", tigertonic.Marshaled(changeFormation))
	mux.Handle("GET", "/apps/{app_id}/jobs", tigertonic.Marshaled(getJobs))
	mux.HandleFunc("GET", "/apps/{app_id}/jobs/{job_id}/logs", getJobLog)
	mux.HandleFunc("POST", "/apps/{app_id}/jobs", runJob)
	http.ListenAndServe("127.0.0.1:1200", tigertonic.Logged(mux, nil))
}

var scheduler *sampic.Client
var disc *discover.Client
var router *strowgerc.Client

type Job struct {
	ID   string `json:"id"`
	Type string `json:"type"`
}

func addDomain(u *url.URL, h http.Header, data *struct{}) (int, http.Header, struct{}, error) {
	q := u.Query()
	if err := router.AddFrontend(&strowger.Config{Service: q.Get("app_id"), HTTPDomain: q.Get("domain")}); err != nil {
		return 500, nil, struct{}{}, err
	}
	return 200, nil, struct{}{}, nil
}

// GET /apps/{app_id}/jobs
func getJobs(u *url.URL, h http.Header) (int, http.Header, []Job, error) {
	state, err := scheduler.State()
	if err != nil {
		return 500, nil, nil, err
	}

	q := u.Query()
	prefix := q.Get("app_id") + "-"
	jobs := make([]Job, 0)
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				typ := strings.Split(job.ID[len(prefix):], ".")[0]
				jobs = append(jobs, Job{ID: job.ID, Type: typ})
			}
		}
	}

	return 200, nil, jobs, nil
}

type Formation struct {
	Quantity int    `json:"quantity"`
	Type     string `json:"type"`
}

func shelfURL() string {
	set, _ := disc.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	return addrs[0]
}

// POST /apps/{app_id}/formation/{formation_id}
func changeFormation(u *url.URL, h http.Header, req *Formation) (int, http.Header, *Formation, error) {
	state, err := scheduler.State()
	if err != nil {
		log.Println("scheduler state error", err)
		return 500, nil, nil, err
	}

	q := u.Query()
	req.Type = q.Get("formation_id")
	prefix := q.Get("app_id") + "-" + req.Type + "."
	var jobs []*sampi.Job
	for _, host := range state {
		for _, job := range host.Jobs {
			if strings.HasPrefix(job.ID, prefix) {
				if job.Attributes == nil {
					job.Attributes = make(map[string]string)
				}
				job.Attributes["host_id"] = host.ID
				jobs = append(jobs, job)
			}
		}
	}

	if req.Quantity < 0 {
		req.Quantity = 0
	}
	diff := req.Quantity - len(jobs)
	log.Printf("have %d %s, diff %d", len(jobs), req.Type, diff)
	if diff > 0 {
		config := &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          []string{"start", req.Type},
			AttachStdout: true,
			AttachStderr: true,
			Env:          []string{"SLUG_URL=http://" + shelfURL() + "/" + q.Get("app_id") + ".tgz"},
		}
		if req.Type == "web" {
			config.Env = append(config.Env, "SD_NAME="+q.Get("app_id"))
		}
		schedReq := &sampi.ScheduleReq{
			HostJobs: make(map[string][]*sampi.Job),
		}
	outer:
		for {
			for host := range state {
				schedReq.HostJobs[host] = append(schedReq.HostJobs[host], &sampi.Job{ID: prefix + randomID(), TCPPorts: 1, Config: config})
				diff--
				if diff == 0 {
					break outer
				}
			}
		}

		res, err := scheduler.Schedule(schedReq)
		if err != nil || !res.Success {
			log.Println("schedule error", err)
			return 500, nil, nil, err
		}
	} else if diff < 0 {
		for _, job := range jobs[:-diff] {
			host, err := lornec.New(job.Attributes["host_id"])
			if err != nil {
				log.Println("error connecting to", job.Attributes["host_id"], err)
				continue
			}
			if err := host.StopJob(job.ID); err != nil {
				log.Println("error stopping", job.ID, "on", job.Attributes["host_id"], err)
			}
		}
	}

	return 200, nil, req, nil
}

// GET /apps/{app_id}/jobs/{job_id}/logs
func getJobLog(w http.ResponseWriter, req *http.Request) {
	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	q := req.URL.Query()
	jobID := q.Get("job_id")
	if prefix := q.Get("app_id") + "-"; !strings.HasPrefix(jobID, prefix) {
		jobID = prefix + jobID
	}
	var job *sampi.Job
	var host sampi.Host
outer:
	for _, host = range state {
		for _, job = range host.Jobs {
			if job.ID == jobID {
				break outer
			}
		}
		job = nil
	}
	if job == nil {
		w.WriteHeader(404)
		return
	}

	attachReq := &lorne.AttachReq{
		JobID: job.ID,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagLogs,
	}

	client, err := lornec.New(host.ID)
	if err != nil {
		w.WriteHeader(500)
		log.Println("lorne connect failed", err)
		return
	}
	attachConn, _, err := client.Attach(attachReq, false)
	if err != nil {
		w.WriteHeader(500)
		log.Println("attach failed", err)
		return
	}
	defer attachConn.Close()
	io.Copy(w, attachConn)
}

type NewJob struct {
	Cmd     []string          `json:"cmd"`
	Env     map[string]string `json:"env"`
	Attach  bool              `json:"attach"`
	TTY     bool              `json:"tty"`
	Columns int               `json:"tty_columns"`
	Lines   int               `json:"tty_lines"`
}

// POST /apps/{app_id}/jobs
func runJob(w http.ResponseWriter, req *http.Request) {
	var jobReq NewJob
	if err := json.NewDecoder(req.Body).Decode(&jobReq); err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}

	state, err := scheduler.State()
	if err != nil {
		w.WriteHeader(500)
		log.Println(err)
		return
	}
	// pick a random host
	var hostID string
	for hostID = range state {
		break
	}
	if hostID == "" {
		w.WriteHeader(500)
		log.Println("no hosts found")
		return
	}

	env := make([]string, 0, len(jobReq.Env))
	for k, v := range jobReq.Env {
		env = append(env, k+"="+v)
	}

	q := req.URL.Query()
	job := &sampi.Job{
		ID: q.Get("app_id") + "-run." + randomID(),
		Config: &docker.Config{
			Image:        "flynn/slugrunner",
			Cmd:          jobReq.Cmd,
			AttachStdin:  true,
			AttachStdout: true,
			AttachStderr: true,
			StdinOnce:    true,
			Env:          append(env, "SLUG_URL=http://"+shelfURL()+"/"+q.Get("app_id")+".tgz"),
		},
	}
	if jobReq.TTY {
		job.Config.Tty = true
	}
	if jobReq.Attach {
		job.Config.AttachStdin = true
		job.Config.StdinOnce = true
		job.Config.OpenStdin = true
	}

	var attachConn lornec.ReadWriteCloser
	var attachWait func() error
	if jobReq.Attach {
		attachReq := &lorne.AttachReq{
			JobID:  job.ID,
			Flags:  lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
			Height: jobReq.Lines,
			Width:  jobReq.Columns,
		}
		client, err := lornec.New(hostID)
		if err != nil {
			w.WriteHeader(500)
			log.Println("lorne connect failed", err)
			return
		}
		attachConn, attachWait, err = client.Attach(attachReq, true)
		if err != nil {
			w.WriteHeader(500)
			log.Println("attach failed", err)
			return
		}
		defer attachConn.Close()
	}

	res, err := scheduler.Schedule(&sampi.ScheduleReq{HostJobs: map[string][]*sampi.Job{hostID: {job}}})
	if err != nil || !res.Success {
		w.WriteHeader(500)
		log.Println("schedule failed", err)
		return
	}

	if jobReq.Attach {
		if err := attachWait(); err != nil {
			w.WriteHeader(500)
			log.Println("attach wait failed", err)
			return
		}
		w.Header().Set("Content-Type", "application/vnd.flynn.hijack")
		w.Header().Set("Content-Length", "0")
		w.WriteHeader(200)
		conn, _, err := w.(http.Hijacker).Hijack()
		if err != nil {
			return
		}
		defer conn.Close()

		done := make(chan struct{})
		copy := func(to lornec.ReadWriteCloser, from io.Reader) {
			io.Copy(to, from)
			to.CloseWrite()
			done <- struct{}{}
		}
		go copy(conn.(lornec.ReadWriteCloser), attachConn)
		go copy(attachConn, conn)
		<-done
		<-done

		return
	}
	w.WriteHeader(200)
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<|endoftext|>


<commit_before>// This is a client that writes out to a file, and optionally rolls the file

package main

import (
	"../../nsq"
	"../../util"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	filenamePattern  = "%s.%s.%d-%02d-%02d_%02d.log" // topic.host.YYY-MM-DD_HH.log
	hostIdentifier   = flag.String("host-identifier", "", "value to output in log filename in place of hostname. <SHORT_HOST> and <HOSTNAME> are valid replacement tokens")
	outputDir        = flag.String("output-dir", "/tmp", "directory to write output files to")
	topic            = flag.String("topic-name", "", "nsq topic")
	channel          = flag.String("channel-name", "nsq_to_file", "nsq channel")
	buffer           = flag.Int("buffer", 1000, "number of messages to buffer in channel and disk before sync/ack")
	verbose          = flag.Bool("verbose", false, "verbose logging")
	nsqAddresses     = util.StringArray{}
	lookupdAddresses = util.StringArray{}
)

func init() {
	flag.Var(&nsqAddresses, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdAddresses, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type FileLogger struct {
	out      *os.File
	filename string
	logChan  chan *Message
}

type Message struct {
	*nsq.Message
	returnChannel chan *nsq.FinishedMessage
}

type SyncMsg struct {
	m             *nsq.FinishedMessage
	returnChannel chan *nsq.FinishedMessage
}

func (l *FileLogger) HandleMessage(m *nsq.Message, responseChannel chan *nsq.FinishedMessage) {
	l.logChan <- &Message{m, responseChannel}
}

func router(r *nsq.Reader, f *FileLogger, termChan chan os.Signal, hupChan chan os.Signal) {
	pos := 0
	output := make([]*SyncMsg, *buffer)
	sync := false
	ticker := time.Tick(time.Duration(30) * time.Second)

	for {
		select {
		case <-termChan:
			r.Stop()
			sync = true
		case <-hupChan:
			f.out.Close()
			f.out = nil
			updateFile(f)
			if pos != 0 {
				sync = true
			}
		case <-ticker:
			if pos != 0 || f.out != nil {
				updateFile(f)
				sync = true
			}
		case m := <-f.logChan:
			if updateFile(f) {
				sync = true
			}
			f.out.Write(m.Body)
			f.out.WriteString("\n")
			x := &nsq.FinishedMessage{m.Id, 0, true}
			output[pos] = &SyncMsg{x, m.returnChannel}
			pos++
		}

		// in the case where you have N connections, flush after the
		// smallest buffer size for a single connection (otherwise the async handler will wait to finish message)
		// and you will starve your connection
		if sync || pos >= *buffer || pos >= r.ConnectionBufferSize() {
			if pos > 0 {
				log.Printf("syncing %d records to disk", pos)
				f.out.Sync()
				for pos > 0 {
					pos--
					m := output[pos]
					m.returnChannel <- m.m
					output[pos] = nil
				}
			}
			sync = false
		}
	}
}

func updateFile(f *FileLogger) bool {
	t := time.Now()
	hostname, _ := os.Hostname()
	shortHostname := strings.Split(hostname, ".")[0]
	identifier := shortHostname
	if len(*hostIdentifier) != 0 {
		identifier = strings.Replace(*hostIdentifier, "<SHORT_HOST>", shortHostname, -1)
		identifier = strings.Replace(identifier, "<HOSTNAME>", hostname, -1)
	}
	filename := fmt.Sprintf(filenamePattern, *topic, identifier, t.Year(), t.Month(), t.Day(), t.Hour())

	if filename != f.filename || f.out == nil {
		log.Printf("old %s new %s", f.filename, filename)
		// roll it
		if f.out != nil {
			f.out.Close()
		}
		os.MkdirAll(*outputDir, 777)
		log.Printf("opening %s/%s", *outputDir, filename)
		newfile, err := os.OpenFile(fmt.Sprintf("%s/%s", *outputDir, filename), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		f.out = newfile
		if err != nil {
			log.Fatal(err)
		}
		f.filename = filename
		return true
	}

	return false
}

func main() {
	flag.Parse()

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic-name and --channel-name are required")
	}

	if *buffer < 0 {
		log.Fatalf("--buffer must be > 0")
	}

	if len(nsqAddresses) == 0 && len(lookupdAddresses) == 0 {
		log.Fatalf("--nsqd-tcp-address or --lookupd-http-address required.")
	}
	if len(nsqAddresses) != 0 && len(lookupdAddresses) != 0 {
		log.Fatalf("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	hupChan := make(chan os.Signal, 1)
	termChan := make(chan os.Signal, 1)
	signal.Notify(hupChan, syscall.SIGHUP)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	f := &FileLogger{
		logChan: make(chan *Message, *buffer),
	}

	r, _ := nsq.NewReader(*topic, *channel)
	r.BufferSize = *buffer * 2
	r.VerboseLogging = *verbose

	r.AddAsyncHandler(f)
	go router(r, f, termChan, hupChan)

	for _, addrString := range nsqAddresses {
		err := r.ConnectToNSQ(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	for _, addrString := range lookupdAddresses {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	<-r.ExitChan
}
<commit_msg>nsq_to_file: refactor finished message handling (fatally exit if file ops fail)<commit_after>// This is a client that writes out to a file, and optionally rolls the file

package main

import (
	"../../nsq"
	"../../util"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	filenamePattern  = "%s.%s.%d-%02d-%02d_%02d.log" // topic.host.YYY-MM-DD_HH.log
	hostIdentifier   = flag.String("host-identifier", "", "value to output in log filename in place of hostname. <SHORT_HOST> and <HOSTNAME> are valid replacement tokens")
	outputDir        = flag.String("output-dir", "/tmp", "directory to write output files to")
	topic            = flag.String("topic-name", "", "nsq topic")
	channel          = flag.String("channel-name", "nsq_to_file", "nsq channel")
	buffer           = flag.Int("buffer", 1000, "number of messages to buffer in channel and disk before sync/ack")
	verbose          = flag.Bool("verbose", false, "verbose logging")
	nsqAddresses     = util.StringArray{}
	lookupdAddresses = util.StringArray{}
)

func init() {
	flag.Var(&nsqAddresses, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdAddresses, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type FileLogger struct {
	out      *os.File
	filename string
	logChan  chan *Message
}

type Message struct {
	*nsq.Message
	returnChannel chan *nsq.FinishedMessage
}

type SyncMsg struct {
	m             *nsq.FinishedMessage
	returnChannel chan *nsq.FinishedMessage
}

func (l *FileLogger) HandleMessage(m *nsq.Message, responseChannel chan *nsq.FinishedMessage) {
	l.logChan <- &Message{m, responseChannel}
}

func router(r *nsq.Reader, f *FileLogger, termChan chan os.Signal, hupChan chan os.Signal) {
	pos := 0
	output := make([]*Message, *buffer)
	sync := false
	ticker := time.Tick(time.Duration(30) * time.Second)

	for {
		select {
		case <-termChan:
			r.Stop()
			sync = true
		case <-hupChan:
			f.out.Close()
			f.out = nil
			updateFile(f)
			if pos != 0 {
				sync = true
			}
		case <-ticker:
			if pos != 0 || f.out != nil {
				updateFile(f)
				sync = true
			}
		case m := <-f.logChan:
			if updateFile(f) {
				sync = true
			}
			_, err := f.out.Write(m.Body)
			if err != nil {
				log.Fatalf("ERROR: writing message to disk - %s", err.Error())
			}
			_, err = f.out.WriteString("\n")
			if err != nil {
				log.Fatalf("ERROR: writing newline to disk - %s", err.Error())
			}
			output[pos] = m
			pos++
		}

		if sync || pos >= *buffer {
			if pos > 0 {
				log.Printf("syncing %d records to disk", pos)
				err := f.out.Sync()
				if err != nil {
					log.Fatalf("ERROR: failed syncing messages - %s", err.Error())
				}
				for pos > 0 {
					pos--
					m := output[pos]
					m.returnChannel <- &nsq.FinishedMessage{m.Id, 0, true}
					output[pos] = nil
				}
			}
			sync = false
		}
	}
}

func updateFile(f *FileLogger) bool {
	t := time.Now()
	hostname, _ := os.Hostname()
	shortHostname := strings.Split(hostname, ".")[0]
	identifier := shortHostname
	if len(*hostIdentifier) != 0 {
		identifier = strings.Replace(*hostIdentifier, "<SHORT_HOST>", shortHostname, -1)
		identifier = strings.Replace(identifier, "<HOSTNAME>", hostname, -1)
	}
	filename := fmt.Sprintf(filenamePattern, *topic, identifier, t.Year(), t.Month(), t.Day(), t.Hour())

	if filename != f.filename || f.out == nil {
		log.Printf("old %s new %s", f.filename, filename)
		// roll it
		if f.out != nil {
			f.out.Close()
		}
		os.MkdirAll(*outputDir, 777)
		log.Printf("opening %s/%s", *outputDir, filename)
		newfile, err := os.OpenFile(fmt.Sprintf("%s/%s", *outputDir, filename), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		f.out = newfile
		if err != nil {
			log.Fatal(err)
		}
		f.filename = filename
		return true
	}

	return false
}

func main() {
	flag.Parse()

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic-name and --channel-name are required")
	}

	if *buffer < 0 {
		log.Fatalf("--buffer must be > 0")
	}

	if len(nsqAddresses) == 0 && len(lookupdAddresses) == 0 {
		log.Fatalf("--nsqd-tcp-address or --lookupd-http-address required.")
	}
	if len(nsqAddresses) != 0 && len(lookupdAddresses) != 0 {
		log.Fatalf("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	hupChan := make(chan os.Signal, 1)
	termChan := make(chan os.Signal, 1)
	signal.Notify(hupChan, syscall.SIGHUP)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	f := &FileLogger{
		logChan: make(chan *Message, 1),
	}

	r, _ := nsq.NewReader(*topic, *channel)
	r.BufferSize = *buffer
	r.VerboseLogging = *verbose

	r.AddAsyncHandler(f)
	go router(r, f, termChan, hupChan)

	for _, addrString := range nsqAddresses {
		err := r.ConnectToNSQ(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	for _, addrString := range lookupdAddresses {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	<-r.ExitChan
}
<|endoftext|>


<commit_before>package main

import (
    "gitee.com/johng/gf/g"
    "gitee.com/johng/gf/g/net/ghttp"
)

func main() {
    s := g.Server()
    s.BindHandler("/ws", func(r *ghttp.Request) {
        conn, _ := r.WebSocket()
        for {
            msgType, msg, err := conn.ReadMessage()
            if err != nil {
                return
            }
            if err = conn.WriteMessage(msgType, msg); err != nil {
                return
            }
        }
    })
    s.SetPort(8199)
    s.Run()
}

<commit_msg>改进websocket示例代码<commit_after>package main

import (
    "gitee.com/johng/gf/g"
    "gitee.com/johng/gf/g/net/ghttp"
)

func main() {
    s := g.Server()
    s.BindHandler("/ws", func(r *ghttp.Request) {
        ws, _ := r.WebSocket()
        for {
            msgType, msg, err := ws.ReadMessage()
            if err != nil {
                return
            }
            if err = ws.WriteMessage(msgType, msg); err != nil {
                return
            }
        }
    })
    s.SetPort(8199)
    s.Run()
}

<|endoftext|>


<commit_before>package postgres

import "github.com/lfq7413/tomato/types"

const postgresSchemaCollectionName = "_SCHEMA"

// PostgresAdapter postgres 数据库适配器
type PostgresAdapter struct {
	collectionPrefix string
	collectionList   []string
}

// NewPostgresAdapter ...
func NewPostgresAdapter(collectionPrefix string) *PostgresAdapter {
	return &PostgresAdapter{
		collectionPrefix: collectionPrefix,
		collectionList:   []string{},
	}
}

// ClassExists ...
func (p *PostgresAdapter) ClassExists(name string) bool {
	return false
}

// SetClassLevelPermissions ...
func (p *PostgresAdapter) SetClassLevelPermissions(className string, CLPs types.M) error {
	return nil
}

// CreateClass ...
func (p *PostgresAdapter) CreateClass(className string, schema types.M) (types.M, error) {
	return nil, nil
}

// AddFieldIfNotExists ...
func (p *PostgresAdapter) AddFieldIfNotExists(className, fieldName string, fieldType types.M) error {
	return nil
}

// DeleteClass ...
func (p *PostgresAdapter) DeleteClass(className string) (types.M, error) {
	return nil, nil
}

// DeleteAllClasses ...
func (p *PostgresAdapter) DeleteAllClasses() error {
	return nil
}

// DeleteFields ...
func (p *PostgresAdapter) DeleteFields(className string, schema types.M, fieldNames []string) error {
	return nil
}

// CreateObject ...
func (p *PostgresAdapter) CreateObject(className string, schema, object types.M) error {
	return nil
}

// GetAllClasses ...
func (p *PostgresAdapter) GetAllClasses() ([]types.M, error) {
	return nil, nil
}

// GetClass ...
func (p *PostgresAdapter) GetClass(className string) (types.M, error) {
	return nil, nil
}

// DeleteObjectsByQuery ...
func (p *PostgresAdapter) DeleteObjectsByQuery(className string, schema, query types.M) error {
	return nil
}

// Find ...
func (p *PostgresAdapter) Find(className string, schema, query, options types.M) ([]types.M, error) {
	return nil, nil
}

// Count ...
func (p *PostgresAdapter) Count(className string, schema, query types.M) (int, error) {
	return 0, nil
}

// UpdateObjectsByQuery ...
func (p *PostgresAdapter) UpdateObjectsByQuery(className string, schema, query, update types.M) error {
	return nil
}

// FindOneAndUpdate ...
func (p *PostgresAdapter) FindOneAndUpdate(className string, schema, query, update types.M) (types.M, error) {
	return nil, nil
}

// UpsertOneObject ...
func (p *PostgresAdapter) UpsertOneObject(className string, schema, query, update types.M) error {
	return nil
}

// EnsureUniqueness ...
func (p *PostgresAdapter) EnsureUniqueness(className string, schema types.M, fieldNames []string) error {
	return nil
}

// PerformInitialization ...
func (p *PostgresAdapter) PerformInitialization(options types.M) error {
	return nil
}
<commit_msg>添加要实现的函数定义<commit_after>package postgres

import "github.com/lfq7413/tomato/types"

const postgresSchemaCollectionName = "_SCHEMA"

const postgresRelationDoesNotExistError = "42P01"
const postgresDuplicateRelationError = "42P07"
const postgresDuplicateColumnError = "42701"
const postgresUniqueIndexViolationError = "23505"
const postgresTransactionAbortedError = "25P02"

// PostgresAdapter postgres 数据库适配器
type PostgresAdapter struct {
	collectionPrefix string
	collectionList   []string
}

// NewPostgresAdapter ...
func NewPostgresAdapter(collectionPrefix string) *PostgresAdapter {
	return &PostgresAdapter{
		collectionPrefix: collectionPrefix,
		collectionList:   []string{},
	}
}

// ClassExists ...
func (p *PostgresAdapter) ClassExists(name string) bool {
	return false
}

// SetClassLevelPermissions ...
func (p *PostgresAdapter) SetClassLevelPermissions(className string, CLPs types.M) error {
	return nil
}

// CreateClass ...
func (p *PostgresAdapter) CreateClass(className string, schema types.M) (types.M, error) {
	return nil, nil
}

// AddFieldIfNotExists ...
func (p *PostgresAdapter) AddFieldIfNotExists(className, fieldName string, fieldType types.M) error {
	return nil
}

// DeleteClass ...
func (p *PostgresAdapter) DeleteClass(className string) (types.M, error) {
	return nil, nil
}

// DeleteAllClasses ...
func (p *PostgresAdapter) DeleteAllClasses() error {
	return nil
}

// DeleteFields ...
func (p *PostgresAdapter) DeleteFields(className string, schema types.M, fieldNames []string) error {
	return nil
}

// CreateObject ...
func (p *PostgresAdapter) CreateObject(className string, schema, object types.M) error {
	return nil
}

// GetAllClasses ...
func (p *PostgresAdapter) GetAllClasses() ([]types.M, error) {
	return nil, nil
}

// GetClass ...
func (p *PostgresAdapter) GetClass(className string) (types.M, error) {
	return nil, nil
}

// DeleteObjectsByQuery ...
func (p *PostgresAdapter) DeleteObjectsByQuery(className string, schema, query types.M) error {
	return nil
}

// Find ...
func (p *PostgresAdapter) Find(className string, schema, query, options types.M) ([]types.M, error) {
	return nil, nil
}

// Count ...
func (p *PostgresAdapter) Count(className string, schema, query types.M) (int, error) {
	return 0, nil
}

// UpdateObjectsByQuery ...
func (p *PostgresAdapter) UpdateObjectsByQuery(className string, schema, query, update types.M) error {
	return nil
}

// FindOneAndUpdate ...
func (p *PostgresAdapter) FindOneAndUpdate(className string, schema, query, update types.M) (types.M, error) {
	return nil, nil
}

// UpsertOneObject ...
func (p *PostgresAdapter) UpsertOneObject(className string, schema, query, update types.M) error {
	return nil
}

// EnsureUniqueness ...
func (p *PostgresAdapter) EnsureUniqueness(className string, schema types.M, fieldNames []string) error {
	return nil
}

// PerformInitialization ...
func (p *PostgresAdapter) PerformInitialization(options types.M) error {
	return nil
}

var parseToPosgresComparator = map[string]string{
	"$gt":  ">",
	"$lt":  "<",
	"$gte": ">=",
	"$lte": "<=",
}

func parseTypeToPostgresType(t types.M) (string, error) {
	// TODO
	return "", nil
}

func toPostgresValue(value interface{}) interface{} {
	// TODO
	return nil
}

func transformValue(value interface{}) interface{} {
	// TODO
	return nil
}

var emptyCLPS = types.M{
	"find":     types.M{},
	"get":      types.M{},
	"create":   types.M{},
	"update":   types.M{},
	"delete":   types.M{},
	"addField": types.M{},
}

var defaultCLPS = types.M{
	"find":     types.M{"*": true},
	"get":      types.M{"*": true},
	"create":   types.M{"*": true},
	"update":   types.M{"*": true},
	"delete":   types.M{"*": true},
	"addField": types.M{"*": true},
}

func toParseSchema(schema types.M) types.M {
	// TODO
	return nil
}

func toPostgresSchema(schema types.M) types.M {
	// TODO
	return nil
}

func handleDotFields(object types.M) types.M {
	// TODO
	return nil
}

func validateKeys(object interface{}) error {
	// TODO
	return nil
}

func joinTablesForSchema(schema types.M) []string {
	// TODO
	return nil
}

func buildWhereClause(schema, query types.M, index int) (types.M, error) {
	// TODO
	// toPostgresSchema
	// removeWhiteSpace
	// processRegexPattern
	// toPostgresValue
	// transformValue
	return nil, nil
}

func removeWhiteSpace(s string) string {
	// TODO
	return ""
}

func processRegexPattern(s string) string {
	// TODO
	// literalizeRegexPart
	return ""
}

func createLiteralRegex(s string) string {
	// TODO
	return ""
}

func literalizeRegexPart(s string) string {
	// TODO
	// createLiteralRegex
	return ""
}
<|endoftext|>


<commit_before>package bearychat

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
)

const (
	DEFAULT_RTM_API_BASE = "https://rtm.bearychat.com"
)

// RTMClient is used to interactive with BearyChat's RTM api
// and websocket message protocol.
type RTMClient struct {
	// rtm token
	Token string

	// rtm api base, defaults to `https://rtm.bearychat.com`
	APIBase string

	// services
	CurrentTeam *RTMCurrentTeamService
	User        *RTMUserService
	Channel     *RTMChannelService

	httpClient *http.Client
}

type rtmOptSetter func(*RTMClient) error

// enabled services
var services = []rtmOptSetter{
	newRTMCurrentTeamService,
	newRTMUserService,
	newRTMChannelService,
}

// NewRTMClient creates a rtm client.
//
//      client, _ := NewRTMClient(
//              "rtm-token",
//              WithRTMAPIBase("https://rtm.bearychat.com"),
//      )
func NewRTMClient(token string, setters ...rtmOptSetter) (*RTMClient, error) {
	c := &RTMClient{
		Token:   token,
		APIBase: DEFAULT_RTM_API_BASE,

		httpClient: http.DefaultClient,
	}

	for _, setter := range services {
		if err := setter(c); err != nil {
			return nil, err
		}
	}

	for _, setter := range setters {
		if err := setter(c); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// WithRTMAPIBase can be used to set rtm client's base api.
func WithRTMAPIBase(apiBase string) rtmOptSetter {
	return func(c *RTMClient) error {
		c.APIBase = apiBase
		return nil
	}
}

// WithRTMHTTPClient sets http client.
func WithRTMHTTPClient(httpClient *http.Client) rtmOptSetter {
	return func(c *RTMClient) error {
		c.httpClient = httpClient
		return nil
	}
}

// Do performs an api request.
func (c RTMClient) Do(resource, method string, in, result interface{}) (*http.Response, error) {
	uri, err := addTokenToResourceUri(
		fmt.Sprintf("%s/%s", c.APIBase, resource),
		c.Token,
	)
	if err != nil {
		return nil, err
	}

	// build payload (if any)
	var buf io.ReadWriter
	if in != nil {
		buf = new(bytes.Buffer)
		err := json.NewEncoder(buf).Encode(in)
		if err != nil {
			return nil, err
		}
	}

	// build request
	req, err := http.NewRequest(method, uri, buf)
	if err != nil {
		return nil, err
	}
	if in != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("Accept", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	// parse response
	defer resp.Body.Close()
	response := new(RTMAPIResponse)
	if err := json.NewDecoder(resp.Body).Decode(response); err != nil {
		return resp, err
	}

	// request failed
	if resp.StatusCode/100 != 2 || response.Code != 0 {
		return resp, response
	}

	// parse result (if any)
	if result != nil {
		return resp, json.Unmarshal(response.Result, result)
	}

	return resp, nil
}

func (c RTMClient) Get(resource string, result interface{}) (*http.Response, error) {
	return c.Do(resource, "GET", nil, result)
}

func (c RTMClient) Post(resource string, in, result interface{}) (*http.Response, error) {
	return c.Do(resource, "POST", in, result)
}

// Start performs rtm.start
func (c RTMClient) Start() (*User, string, error) {
	userAndWSHost := new(struct {
		User   *User  `json:"user"`
		WSHost string `json:"ws_host"`
	})
	_, err := c.Post("start", nil, userAndWSHost)

	return userAndWSHost.User, userAndWSHost.WSHost, err
}

// RTM api request response
type RTMAPIResponse struct {
	Code        int             `json:"code"`
	Result      json.RawMessage `json:"result,omitempty"`
	ErrorReason string          `json:"error,omitempty"`
}

func (r *RTMAPIResponse) Error() string {
	return r.ErrorReason
}

func addTokenToResourceUri(resource, token string) (string, error) {
	uri, err := url.Parse(resource)
	if err != nil {
		return "", err
	}

	q := uri.Query()
	q.Set("token", token)
	uri.RawQuery = q.Encode()

	return uri.String(), nil
}
<commit_msg>feat(rtm): implement `rtm.message` api<commit_after>package bearychat

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
)

const (
	DEFAULT_RTM_API_BASE = "https://rtm.bearychat.com"
)

// RTMClient is used to interactive with BearyChat's RTM api
// and websocket message protocol.
type RTMClient struct {
	// rtm token
	Token string

	// rtm api base, defaults to `https://rtm.bearychat.com`
	APIBase string

	// services
	CurrentTeam *RTMCurrentTeamService
	User        *RTMUserService
	Channel     *RTMChannelService

	httpClient *http.Client
}

type rtmOptSetter func(*RTMClient) error

// enabled services
var services = []rtmOptSetter{
	newRTMCurrentTeamService,
	newRTMUserService,
	newRTMChannelService,
}

// NewRTMClient creates a rtm client.
//
//      client, _ := NewRTMClient(
//              "rtm-token",
//              WithRTMAPIBase("https://rtm.bearychat.com"),
//      )
func NewRTMClient(token string, setters ...rtmOptSetter) (*RTMClient, error) {
	c := &RTMClient{
		Token:   token,
		APIBase: DEFAULT_RTM_API_BASE,

		httpClient: http.DefaultClient,
	}

	for _, setter := range services {
		if err := setter(c); err != nil {
			return nil, err
		}
	}

	for _, setter := range setters {
		if err := setter(c); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// WithRTMAPIBase can be used to set rtm client's base api.
func WithRTMAPIBase(apiBase string) rtmOptSetter {
	return func(c *RTMClient) error {
		c.APIBase = apiBase
		return nil
	}
}

// WithRTMHTTPClient sets http client.
func WithRTMHTTPClient(httpClient *http.Client) rtmOptSetter {
	return func(c *RTMClient) error {
		c.httpClient = httpClient
		return nil
	}
}

// Do performs an api request.
func (c RTMClient) Do(resource, method string, in, result interface{}) (*http.Response, error) {
	uri, err := addTokenToResourceUri(
		fmt.Sprintf("%s/%s", c.APIBase, resource),
		c.Token,
	)
	if err != nil {
		return nil, err
	}

	// build payload (if any)
	var buf io.ReadWriter
	if in != nil {
		buf = new(bytes.Buffer)
		err := json.NewEncoder(buf).Encode(in)
		if err != nil {
			return nil, err
		}
	}

	// build request
	req, err := http.NewRequest(method, uri, buf)
	if err != nil {
		return nil, err
	}
	if in != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("Accept", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	// parse response
	defer resp.Body.Close()
	response := new(RTMAPIResponse)
	if err := json.NewDecoder(resp.Body).Decode(response); err != nil {
		return resp, err
	}

	// request failed
	if resp.StatusCode/100 != 2 || response.Code != 0 {
		return resp, response
	}

	// parse result (if any)
	if result != nil {
		return resp, json.Unmarshal(response.Result, result)
	}

	return resp, nil
}

func (c RTMClient) Get(resource string, result interface{}) (*http.Response, error) {
	return c.Do(resource, "GET", nil, result)
}

func (c RTMClient) Post(resource string, in, result interface{}) (*http.Response, error) {
	return c.Do(resource, "POST", in, result)
}

// Start performs rtm.start
func (c RTMClient) Start() (*User, string, error) {
	userAndWSHost := new(struct {
		User   *User  `json:"user"`
		WSHost string `json:"ws_host"`
	})
	_, err := c.Post("start", nil, userAndWSHost)

	return userAndWSHost.User, userAndWSHost.WSHost, err
}

// Incoming performs rtm.message
func (c RTMClient) Incoming(m RTMIncoming) error {
	_, err := c.Post("message", m, nil)

	return err
}

// RTM api request response
type RTMAPIResponse struct {
	Code        int             `json:"code"`
	Result      json.RawMessage `json:"result,omitempty"`
	ErrorReason string          `json:"error,omitempty"`
}

func (r *RTMAPIResponse) Error() string {
	return r.ErrorReason
}

func addTokenToResourceUri(resource, token string) (string, error) {
	uri, err := url.Parse(resource)
	if err != nil {
		return "", err
	}

	q := uri.Query()
	q.Set("token", token)
	uri.RawQuery = q.Encode()

	return uri.String(), nil
}

// RTMIncoming represents message sent vai `rtm.message` api
type RTMIncoming struct {
	Text        string               `json:"text"`
	VChannelId  string               `json:"vchannel"`
	Markdown    bool                 `json:"markdown,omitempty"`
	Attachments []IncomingAttachment `json:"attachments,omitempty"`
}
<|endoftext|>


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

package snowstorm

import (
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/choices"

	sbcon "github.com/ava-labs/avalanchego/snow/consensus/snowball"
)

// DirectedFactory implements Factory by returning a directed struct
type DirectedFactory struct{}

// New implements Factory
func (DirectedFactory) New() Consensus { return &Directed{} }

// Directed is an implementation of a multi-color, non-transitive, snowball
// instance
type Directed struct {
	common

	// Key: Transaction ID
	// Value: Node that represents this transaction in the conflict graph
	txs map[[32]byte]*directedTx

	// Key: UTXO ID
	// Value: IDs of transactions that consume the UTXO specified in the key
	utxos map[[32]byte]ids.Set
}

type directedTx struct {
	snowball

	// pendingAccept identifies if this transaction has been marked as accepted
	// once its transitive dependencies have also been accepted
	pendingAccept bool

	// ins is the set of txIDs that this tx conflicts with that are less
	// preferred than this tx
	ins ids.Set

	// outs is the set of txIDs that this tx conflicts with that are more
	// preferred than this tx
	outs ids.Set

	// tx is the actual transaction this node represents
	tx Tx
}

// Initialize implements the Consensus interface
func (dg *Directed) Initialize(
	ctx *snow.Context,
	params sbcon.Parameters,
) error {
	dg.txs = make(map[[32]byte]*directedTx)
	dg.utxos = make(map[[32]byte]ids.Set)

	return dg.common.Initialize(ctx, params)
}

// IsVirtuous implements the Consensus interface
func (dg *Directed) IsVirtuous(tx Tx) bool {
	txID := tx.ID()
	// If the tx is currently processing, we should just return if was
	// registered as rogue or not.
	if node, exists := dg.txs[txID.Key()]; exists {
		return !node.rogue
	}

	// The tx isn't processing, so we need to check to see if it conflicts with
	// any of the other txs that are currently processing.
	for _, utxoID := range tx.InputIDs() {
		if _, exists := dg.utxos[utxoID.Key()]; exists {
			// A currently processing tx names the same input as the provided
			// tx, so the provided tx would be rogue.
			return false
		}
	}

	// This tx is virtuous as far as this consensus instance knows.
	return true
}

// Conflicts implements the Consensus interface
func (dg *Directed) Conflicts(tx Tx) ids.Set {
	var conflicts ids.Set = nil
	if node, exists := dg.txs[tx.ID().Key()]; exists {
		// If the tx is currently processing, the conflicting txs are just the
		// union of the inbound conflicts and the outbound conflicts.
		// Only bother to call Union, which will do a memory allocation, if ins or outs are non-empty.
		if node.ins.Len() > 0 || node.outs.Len() > 0 {
			conflicts.Union(node.ins)
			conflicts.Union(node.outs)
		}
	} else {
		// If the tx isn't currently processing, the conflicting txs are the
		// union of all the txs that spend an input that this tx spends.
		for _, inputID := range tx.InputIDs() {
			if spends, exists := dg.utxos[inputID.Key()]; exists {
				conflicts.Union(spends)
			}
		}
	}
	return conflicts
}

// Add implements the Consensus interface
func (dg *Directed) Add(tx Tx) error {
	if shouldVote, err := dg.shouldVote(dg, tx); !shouldVote || err != nil {
		return err
	}

	txID := tx.ID()
	txNode := &directedTx{tx: tx}

	// For each UTXO consumed by the tx:
	// * Add edges between this tx and txs that consume this UTXO
	// * Mark this tx as attempting to consume this UTXO
	for _, inputID := range tx.InputIDs() {
		inputIDKey := inputID.Key()

		// Get the set of txs that are currently processing that also consume
		// this UTXO
		spenders := dg.utxos[inputIDKey]

		// Add all the txs that spend this UTXO to this txs conflicts. These
		// conflicting txs must be preferred over this tx. We know this because
		// this tx currently has a bias of 0 and the tie goes to the tx whose
		// bias was updated first.
		txNode.outs.Union(spenders)

		// Update txs conflicting with tx to account for its issuance
		for conflictIDKey := range spenders {

			// Get the node that contains this conflicting tx
			conflict := dg.txs[conflictIDKey]

			// This conflicting tx can't be virtuous anymore. So, we attempt to
			// remove it from all of the virtuous sets.
			delete(dg.virtuous, conflictIDKey)
			delete(dg.virtuousVoting, conflictIDKey)

			// This tx should be set to rogue if it wasn't rogue before.
			conflict.rogue = true

			// This conflicting tx is preferred over the tx being inserted, as
			// described above. So we add the conflict to the inbound set.
			conflict.ins.Add(txID)
		}

		// Add this tx to list of txs consuming the current UTXO
		spenders.Add(txID)

		// Because this isn't a pointer, we should re-map the set.
		dg.utxos[inputIDKey] = spenders
	}

	// Mark this transaction as rogue if had any conflicts registered above
	txNode.rogue = txNode.outs.Len() != 0

	if !txNode.rogue {
		// If this tx is currently virtuous, add it to the virtuous sets
		dg.virtuous.Add(txID)
		dg.virtuousVoting.Add(txID)

		// If a tx is virtuous, it must be preferred.
		dg.preferences.Add(txID)
	}

	// Add this tx to the set of currently processing txs
	dg.txs[txID.Key()] = txNode

	// If a tx that this tx depends on is rejected, this tx should also be
	// rejected.
	dg.registerRejector(dg, tx)
	return nil
}

// Issued implements the Consensus interface
func (dg *Directed) Issued(tx Tx) bool {
	// If the tx is either Accepted or Rejected, then it must have been issued
	// previously.
	if tx.Status().Decided() {
		return true
	}

	// If the tx is currently processing, then it must have been issued.
	_, ok := dg.txs[tx.ID().Key()]
	return ok
}

// RecordPoll implements the Consensus interface
func (dg *Directed) RecordPoll(votes ids.Bag) (bool, error) {
	// Increase the vote ID. This is only updated here and is used to reset the
	// confidence values of transactions lazily.
	dg.currentVote++

	// This flag tracks if the Avalanche instance needs to recompute its
	// frontiers. Frontiers only need to be recalculated if preferences change
	// or if a tx was accepted.
	changed := false

	// We only want to iterate over txs that received alpha votes
	votes.SetThreshold(dg.params.Alpha)
	// Get the set of IDs that meet this alpha threshold
	metThreshold := votes.Threshold()
	for txIDKey := range metThreshold {

		// Get the node this tx represents
		txNode, exist := dg.txs[txIDKey]
		if !exist {
			// This tx may have already been accepted because of tx
			// dependencies. If this is the case, we can just drop the vote.
			continue
		}

		txNode.RecordSuccessfulPoll(dg.currentVote)

		dg.ctx.Log.Verbo("Updated TxID=%v to have consensus state=%s",
			txIDKey, &txNode.snowball)

		// If the tx should be accepted, then we should defer its acceptance
		// until its dependencies are decided. If this tx was already marked to
		// be accepted, we shouldn't register it again.
		if !txNode.pendingAccept &&
			txNode.Finalized(dg.params.BetaVirtuous, dg.params.BetaRogue) {
			// Mark that this tx is pending acceptance so acceptance is only
			// registered once.
			txNode.pendingAccept = true

			dg.registerAcceptor(dg, txNode.tx)
			if dg.errs.Errored() {
				return changed, dg.errs.Err
			}
		}

		if txNode.tx.Status() != choices.Accepted {
			// If this tx wasn't accepted, then this instance is only changed if
			// preferences changed.
			changed = dg.redirectEdges(txNode) || changed
		} else {
			// By accepting a tx, the state of this instance has changed.
			changed = true
		}
	}
	return changed, dg.errs.Err
}

func (dg *Directed) String() string {
	nodes := make([]*snowballNode, 0, len(dg.txs))
	for _, txNode := range dg.txs {
		nodes = append(nodes, &snowballNode{
			txID:               txNode.tx.ID(),
			numSuccessfulPolls: txNode.numSuccessfulPolls,
			confidence:         txNode.Confidence(dg.currentVote),
		})
	}
	return ConsensusString("DG", nodes)
}

// accept the named txID and remove it from the graph
func (dg *Directed) accept(txID ids.ID) error {
	txKey := txID.Key()
	txNode := dg.txs[txKey]
	// We are accepting the tx, so we should remove the node from the graph.
	delete(dg.txs, txKey)

	// This tx is consuming all the UTXOs from its inputs, so we can prune them
	// all from memory
	for _, inputID := range txNode.tx.InputIDs() {
		delete(dg.utxos, inputID.Key())
	}

	// This tx is now accepted, so it shouldn't be part of the virtuous set or
	// the preferred set. Its status as Accepted implies these descriptions.
	dg.virtuous.Remove(txID)
	dg.preferences.Remove(txID)

	// Reject all the txs that conflicted with this tx.
	if err := dg.reject(txNode.ins); err != nil {
		return err
	}
	// While it is typically true that a tx this is being accepted is preferred,
	// it is possible for this to not be the case. So this is handled for
	// completeness.
	if err := dg.reject(txNode.outs); err != nil {
		return err
	}
	return dg.acceptTx(txNode.tx)
}

// reject all the named txIDs and remove them from the graph
func (dg *Directed) reject(conflictIDs ids.Set) error {
	for conflictKey := range conflictIDs {
		conflict := dg.txs[conflictKey]
		// This tx is no longer an option for consuming the UTXOs from its
		// inputs, so we should remove their reference to this tx.
		for _, inputID := range conflict.tx.InputIDs() {
			inputIDKey := inputID.Key()
			txIDs, exists := dg.utxos[inputIDKey]
			if !exists {
				// This UTXO may no longer exist because it was removed due to
				// the acceptance of a tx. If that is the case, there is nothing
				// left to remove from memory.
				continue
			}
			delete(txIDs, conflictKey)
			if txIDs.Len() == 0 {
				// If this tx was the last tx consuming this UTXO, we should
				// prune the UTXO from memory entirely.
				delete(dg.utxos, inputIDKey)
			} else {
				// If this UTXO still has txs consuming it, then we should make
				// sure this update is written back to the UTXOs map.
				dg.utxos[inputIDKey] = txIDs
			}
		}

		// We are rejecting the tx, so we should remove it from the graph
		delete(dg.txs, conflictKey)

		// While it's statistically unlikely that something being rejected is
		// preferred, it is handled for completion.
		delete(dg.preferences, conflictKey)

		// remove the edge between this node and all its neighbors
		dg.removeConflict(conflictKey, conflict.ins)
		dg.removeConflict(conflictKey, conflict.outs)

		if err := dg.rejectTx(conflict.tx); err != nil {
			return err
		}
	}
	return nil
}

// redirectEdges attempts to turn outbound edges into inbound edges if the
// preferences have changed
func (dg *Directed) redirectEdges(tx *directedTx) bool {
	changed := false
	for conflictIDKey := range tx.outs {
		changed = dg.redirectEdge(tx, ids.NewID(conflictIDKey)) || changed
	}
	return changed
}

// Change the direction of this edge if needed. Returns true if the direction
// was switched.
// TODO replace
func (dg *Directed) redirectEdge(txNode *directedTx, conflictID ids.ID) bool {
	conflict := dg.txs[conflictID.Key()]
	if txNode.numSuccessfulPolls <= conflict.numSuccessfulPolls {
		return false
	}

	// Because this tx has a higher preference than the conflicting tx, we must
	// ensure that the edge is directed towards this tx.
	nodeID := txNode.tx.ID()

	// Change the edge direction according to the conflict tx
	conflict.ins.Remove(nodeID)
	conflict.outs.Add(nodeID)
	dg.preferences.Remove(conflictID) // This conflict has an outbound edge

	// Change the edge direction according to this tx
	txNode.ins.Add(conflictID)
	txNode.outs.Remove(conflictID)
	if txNode.outs.Len() == 0 {
		// If this tx doesn't have any outbound edges, it's preferred
		dg.preferences.Add(nodeID)
	}
	return true
}

func (dg *Directed) removeConflict(txIDKey [32]byte, neighborIDs ids.Set) {
	for neighborIDKey := range neighborIDs {
		neighbor, exists := dg.txs[neighborIDKey]
		if !exists {
			// If the neighbor doesn't exist, they may have already been
			// rejected, so this mapping can be skipped.
			continue
		}

		// Remove any edge to this tx.
		delete(neighbor.ins, txIDKey)
		delete(neighbor.outs, txIDKey)

		if neighbor.outs.Len() == 0 {
			// If this tx should now be preferred, make sure its status is
			// updated.
			dg.preferences.Add(ids.NewID(neighborIDKey))
		}
	}
}
<commit_msg>remove verbo log<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package snowstorm

import (
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/choices"

	sbcon "github.com/ava-labs/avalanchego/snow/consensus/snowball"
)

// DirectedFactory implements Factory by returning a directed struct
type DirectedFactory struct{}

// New implements Factory
func (DirectedFactory) New() Consensus { return &Directed{} }

// Directed is an implementation of a multi-color, non-transitive, snowball
// instance
type Directed struct {
	common

	// Key: Transaction ID
	// Value: Node that represents this transaction in the conflict graph
	txs map[[32]byte]*directedTx

	// Key: UTXO ID
	// Value: IDs of transactions that consume the UTXO specified in the key
	utxos map[[32]byte]ids.Set
}

type directedTx struct {
	snowball

	// pendingAccept identifies if this transaction has been marked as accepted
	// once its transitive dependencies have also been accepted
	pendingAccept bool

	// ins is the set of txIDs that this tx conflicts with that are less
	// preferred than this tx
	ins ids.Set

	// outs is the set of txIDs that this tx conflicts with that are more
	// preferred than this tx
	outs ids.Set

	// tx is the actual transaction this node represents
	tx Tx
}

// Initialize implements the Consensus interface
func (dg *Directed) Initialize(
	ctx *snow.Context,
	params sbcon.Parameters,
) error {
	dg.txs = make(map[[32]byte]*directedTx)
	dg.utxos = make(map[[32]byte]ids.Set)

	return dg.common.Initialize(ctx, params)
}

// IsVirtuous implements the Consensus interface
func (dg *Directed) IsVirtuous(tx Tx) bool {
	txID := tx.ID()
	// If the tx is currently processing, we should just return if was
	// registered as rogue or not.
	if node, exists := dg.txs[txID.Key()]; exists {
		return !node.rogue
	}

	// The tx isn't processing, so we need to check to see if it conflicts with
	// any of the other txs that are currently processing.
	for _, utxoID := range tx.InputIDs() {
		if _, exists := dg.utxos[utxoID.Key()]; exists {
			// A currently processing tx names the same input as the provided
			// tx, so the provided tx would be rogue.
			return false
		}
	}

	// This tx is virtuous as far as this consensus instance knows.
	return true
}

// Conflicts implements the Consensus interface
func (dg *Directed) Conflicts(tx Tx) ids.Set {
	var conflicts ids.Set = nil
	if node, exists := dg.txs[tx.ID().Key()]; exists {
		// If the tx is currently processing, the conflicting txs are just the
		// union of the inbound conflicts and the outbound conflicts.
		// Only bother to call Union, which will do a memory allocation, if ins or outs are non-empty.
		if node.ins.Len() > 0 || node.outs.Len() > 0 {
			conflicts.Union(node.ins)
			conflicts.Union(node.outs)
		}
	} else {
		// If the tx isn't currently processing, the conflicting txs are the
		// union of all the txs that spend an input that this tx spends.
		for _, inputID := range tx.InputIDs() {
			if spends, exists := dg.utxos[inputID.Key()]; exists {
				conflicts.Union(spends)
			}
		}
	}
	return conflicts
}

// Add implements the Consensus interface
func (dg *Directed) Add(tx Tx) error {
	if shouldVote, err := dg.shouldVote(dg, tx); !shouldVote || err != nil {
		return err
	}

	txID := tx.ID()
	txNode := &directedTx{tx: tx}

	// For each UTXO consumed by the tx:
	// * Add edges between this tx and txs that consume this UTXO
	// * Mark this tx as attempting to consume this UTXO
	for _, inputID := range tx.InputIDs() {
		inputIDKey := inputID.Key()

		// Get the set of txs that are currently processing that also consume
		// this UTXO
		spenders := dg.utxos[inputIDKey]

		// Add all the txs that spend this UTXO to this txs conflicts. These
		// conflicting txs must be preferred over this tx. We know this because
		// this tx currently has a bias of 0 and the tie goes to the tx whose
		// bias was updated first.
		txNode.outs.Union(spenders)

		// Update txs conflicting with tx to account for its issuance
		for conflictIDKey := range spenders {

			// Get the node that contains this conflicting tx
			conflict := dg.txs[conflictIDKey]

			// This conflicting tx can't be virtuous anymore. So, we attempt to
			// remove it from all of the virtuous sets.
			delete(dg.virtuous, conflictIDKey)
			delete(dg.virtuousVoting, conflictIDKey)

			// This tx should be set to rogue if it wasn't rogue before.
			conflict.rogue = true

			// This conflicting tx is preferred over the tx being inserted, as
			// described above. So we add the conflict to the inbound set.
			conflict.ins.Add(txID)
		}

		// Add this tx to list of txs consuming the current UTXO
		spenders.Add(txID)

		// Because this isn't a pointer, we should re-map the set.
		dg.utxos[inputIDKey] = spenders
	}

	// Mark this transaction as rogue if had any conflicts registered above
	txNode.rogue = txNode.outs.Len() != 0

	if !txNode.rogue {
		// If this tx is currently virtuous, add it to the virtuous sets
		dg.virtuous.Add(txID)
		dg.virtuousVoting.Add(txID)

		// If a tx is virtuous, it must be preferred.
		dg.preferences.Add(txID)
	}

	// Add this tx to the set of currently processing txs
	dg.txs[txID.Key()] = txNode

	// If a tx that this tx depends on is rejected, this tx should also be
	// rejected.
	dg.registerRejector(dg, tx)
	return nil
}

// Issued implements the Consensus interface
func (dg *Directed) Issued(tx Tx) bool {
	// If the tx is either Accepted or Rejected, then it must have been issued
	// previously.
	if tx.Status().Decided() {
		return true
	}

	// If the tx is currently processing, then it must have been issued.
	_, ok := dg.txs[tx.ID().Key()]
	return ok
}

// RecordPoll implements the Consensus interface
func (dg *Directed) RecordPoll(votes ids.Bag) (bool, error) {
	// Increase the vote ID. This is only updated here and is used to reset the
	// confidence values of transactions lazily.
	dg.currentVote++

	// This flag tracks if the Avalanche instance needs to recompute its
	// frontiers. Frontiers only need to be recalculated if preferences change
	// or if a tx was accepted.
	changed := false

	// We only want to iterate over txs that received alpha votes
	votes.SetThreshold(dg.params.Alpha)
	// Get the set of IDs that meet this alpha threshold
	metThreshold := votes.Threshold()
	for txIDKey := range metThreshold {

		// Get the node this tx represents
		txNode, exist := dg.txs[txIDKey]
		if !exist {
			// This tx may have already been accepted because of tx
			// dependencies. If this is the case, we can just drop the vote.
			continue
		}

		txNode.RecordSuccessfulPoll(dg.currentVote)

		// If the tx should be accepted, then we should defer its acceptance
		// until its dependencies are decided. If this tx was already marked to
		// be accepted, we shouldn't register it again.
		if !txNode.pendingAccept &&
			txNode.Finalized(dg.params.BetaVirtuous, dg.params.BetaRogue) {
			// Mark that this tx is pending acceptance so acceptance is only
			// registered once.
			txNode.pendingAccept = true

			dg.registerAcceptor(dg, txNode.tx)
			if dg.errs.Errored() {
				return changed, dg.errs.Err
			}
		}

		if txNode.tx.Status() != choices.Accepted {
			// If this tx wasn't accepted, then this instance is only changed if
			// preferences changed.
			changed = dg.redirectEdges(txNode) || changed
		} else {
			// By accepting a tx, the state of this instance has changed.
			changed = true
		}
	}
	return changed, dg.errs.Err
}

func (dg *Directed) String() string {
	nodes := make([]*snowballNode, 0, len(dg.txs))
	for _, txNode := range dg.txs {
		nodes = append(nodes, &snowballNode{
			txID:               txNode.tx.ID(),
			numSuccessfulPolls: txNode.numSuccessfulPolls,
			confidence:         txNode.Confidence(dg.currentVote),
		})
	}
	return ConsensusString("DG", nodes)
}

// accept the named txID and remove it from the graph
func (dg *Directed) accept(txID ids.ID) error {
	txKey := txID.Key()
	txNode := dg.txs[txKey]
	// We are accepting the tx, so we should remove the node from the graph.
	delete(dg.txs, txKey)

	// This tx is consuming all the UTXOs from its inputs, so we can prune them
	// all from memory
	for _, inputID := range txNode.tx.InputIDs() {
		delete(dg.utxos, inputID.Key())
	}

	// This tx is now accepted, so it shouldn't be part of the virtuous set or
	// the preferred set. Its status as Accepted implies these descriptions.
	dg.virtuous.Remove(txID)
	dg.preferences.Remove(txID)

	// Reject all the txs that conflicted with this tx.
	if err := dg.reject(txNode.ins); err != nil {
		return err
	}
	// While it is typically true that a tx this is being accepted is preferred,
	// it is possible for this to not be the case. So this is handled for
	// completeness.
	if err := dg.reject(txNode.outs); err != nil {
		return err
	}
	return dg.acceptTx(txNode.tx)
}

// reject all the named txIDs and remove them from the graph
func (dg *Directed) reject(conflictIDs ids.Set) error {
	for conflictKey := range conflictIDs {
		conflict := dg.txs[conflictKey]
		// This tx is no longer an option for consuming the UTXOs from its
		// inputs, so we should remove their reference to this tx.
		for _, inputID := range conflict.tx.InputIDs() {
			inputIDKey := inputID.Key()
			txIDs, exists := dg.utxos[inputIDKey]
			if !exists {
				// This UTXO may no longer exist because it was removed due to
				// the acceptance of a tx. If that is the case, there is nothing
				// left to remove from memory.
				continue
			}
			delete(txIDs, conflictKey)
			if txIDs.Len() == 0 {
				// If this tx was the last tx consuming this UTXO, we should
				// prune the UTXO from memory entirely.
				delete(dg.utxos, inputIDKey)
			} else {
				// If this UTXO still has txs consuming it, then we should make
				// sure this update is written back to the UTXOs map.
				dg.utxos[inputIDKey] = txIDs
			}
		}

		// We are rejecting the tx, so we should remove it from the graph
		delete(dg.txs, conflictKey)

		// While it's statistically unlikely that something being rejected is
		// preferred, it is handled for completion.
		delete(dg.preferences, conflictKey)

		// remove the edge between this node and all its neighbors
		dg.removeConflict(conflictKey, conflict.ins)
		dg.removeConflict(conflictKey, conflict.outs)

		if err := dg.rejectTx(conflict.tx); err != nil {
			return err
		}
	}
	return nil
}

// redirectEdges attempts to turn outbound edges into inbound edges if the
// preferences have changed
func (dg *Directed) redirectEdges(tx *directedTx) bool {
	changed := false
	for conflictIDKey := range tx.outs {
		changed = dg.redirectEdge(tx, ids.NewID(conflictIDKey)) || changed
	}
	return changed
}

// Change the direction of this edge if needed. Returns true if the direction
// was switched.
// TODO replace
func (dg *Directed) redirectEdge(txNode *directedTx, conflictID ids.ID) bool {
	conflict := dg.txs[conflictID.Key()]
	if txNode.numSuccessfulPolls <= conflict.numSuccessfulPolls {
		return false
	}

	// Because this tx has a higher preference than the conflicting tx, we must
	// ensure that the edge is directed towards this tx.
	nodeID := txNode.tx.ID()

	// Change the edge direction according to the conflict tx
	conflict.ins.Remove(nodeID)
	conflict.outs.Add(nodeID)
	dg.preferences.Remove(conflictID) // This conflict has an outbound edge

	// Change the edge direction according to this tx
	txNode.ins.Add(conflictID)
	txNode.outs.Remove(conflictID)
	if txNode.outs.Len() == 0 {
		// If this tx doesn't have any outbound edges, it's preferred
		dg.preferences.Add(nodeID)
	}
	return true
}

func (dg *Directed) removeConflict(txIDKey [32]byte, neighborIDs ids.Set) {
	for neighborIDKey := range neighborIDs {
		neighbor, exists := dg.txs[neighborIDKey]
		if !exists {
			// If the neighbor doesn't exist, they may have already been
			// rejected, so this mapping can be skipped.
			continue
		}

		// Remove any edge to this tx.
		delete(neighbor.ins, txIDKey)
		delete(neighbor.outs, txIDKey)

		if neighbor.outs.Len() == 0 {
			// If this tx should now be preferred, make sure its status is
			// updated.
			dg.preferences.Add(ids.NewID(neighborIDKey))
		}
	}
}
<|endoftext|>