<commit_before>package machinery

import (
	"fmt"
	"os"
	"testing"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/errors"
	"github.com/RichardKnop/machinery/v1/signatures"
)

var (
	task0, task1, task2, task3 signatures.TaskSignature
)

func init() {
	task0 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
		},
	}

	task1 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
		},
	}

	task2 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 5,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 6,
			},
		},
	}

	task3 = signatures.TaskSignature{
		Name: "multiply",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 4,
			},
		},
	}
}

func TestIntegration(t *testing.T) {
	brokerURL := os.Getenv("AMQP_URL")
	if brokerURL == "" {
		return
	}

	server1 := setup(brokerURL, "amqp")
	worker1 := server1.NewWorker("test_worker")
	go worker1.Launch()
	_testSendTask(server1, t)
	_testSendGroup(server1, t)
	_testSendChain(server1, t)
	worker1.Quit()

	memcacheURL := os.Getenv("MEMCACHE_URL")
	if memcacheURL == "" {
		return
	}

	server2 := setup(brokerURL, fmt.Sprintf("memcache://%v", memcacheURL))
	worker2 := server2.NewWorker("test_worker")
	go worker2.Launch()
	_testSendTask(server2, t)
	_testSendGroup(server2, t)
	_testSendChain(server2, t)
	worker2.Quit()
}

func _testSendTask(server *Server, t *testing.T) {
	asyncResult, err := server.SendTask(&task0)
	if err != nil {
		t.Error(err)
	}

	result, err := asyncResult.Get()
	if err != nil {
		t.Error(err)
	}

	if result.Interface() != int64(2) {
		t.Errorf(
			"result = %v(%v), want int64(2)",
			result.Type().String(),
			result.Interface(),
		)
	}
}

func _testSendGroup(server *Server, t *testing.T) {
	group := NewGroup(&task1, &task2, &task3)
	asyncResults, err := server.SendGroup(group)
	if err != nil {
		t.Error(err)
	}

	expectedResults := []int64{2, 11, 4}

	for i, asyncResult := range asyncResults {
		result, err := asyncResult.Get()
		if err != nil {
			t.Error(err)
		}

		if result.Interface() != int64(expectedResults[i]) {
			t.Errorf(
				"result = %v(%v), want int64(%v)",
				result.Type().String(),
				result.Interface(),
				expectedResults[i],
			)
		}
	}
}

func _testSendChain(server *Server, t *testing.T) {
	chain := NewChain(&task1, &task2, &task3)
	chainAsyncResult, err := server.SendChain(chain)
	if err != nil {
		t.Error(err)
	}

	result, err := chainAsyncResult.Get()
	if err != nil {
		t.Error(err)
	}

	if result.Interface() != int64(52) {
		t.Errorf(
			"result = %v(%v), want int64(52)",
			result.Type().String(),
			result.Interface(),
		)
	}
}

func setup(brokerURL, backend string) *Server {
	cnf := config.Config{
		Broker:        brokerURL,
		ResultBackend: backend,
		Exchange:      "test_exchange",
		ExchangeType:  "direct",
		DefaultQueue:  "test_queue",
		BindingKey:    "test_task",
	}

	server, err := NewServer(&cnf)
	errors.Fail(err, "Could not initialize server")

	tasks := map[string]interface{}{
		"add": func(args ...int64) (int64, error) {
			sum := int64(0)
			for _, arg := range args {
				sum += arg
			}
			return sum, nil
		},
		"multiply": func(args ...int64) (int64, error) {
			sum := int64(1)
			for _, arg := range args {
				sum *= arg
			}
			return sum, nil
		},
	}
	server.RegisterTasks(tasks)

	return server
}
<commit_msg>Fixed integration for task group. Async results can arrive in non deterministic order so the test must take that into account.<commit_after>package machinery

import (
	"fmt"
	"os"
	"testing"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/errors"
	"github.com/RichardKnop/machinery/v1/signatures"
)

var (
	task0, task1, task2, task3 signatures.TaskSignature
)

func init() {
	task0 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
		},
	}

	task1 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 1,
			},
		},
	}

	task2 = signatures.TaskSignature{
		Name: "add",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 5,
			},
			signatures.TaskArg{
				Type:  "int64",
				Value: 6,
			},
		},
	}

	task3 = signatures.TaskSignature{
		Name: "multiply",
		Args: []signatures.TaskArg{
			signatures.TaskArg{
				Type:  "int64",
				Value: 4,
			},
		},
	}
}

func TestIntegration(t *testing.T) {
	brokerURL := os.Getenv("AMQP_URL")
	if brokerURL == "" {
		return
	}

	server1 := setup(brokerURL, "amqp")
	worker1 := server1.NewWorker("test_worker")
	go worker1.Launch()
	_testSendTask(server1, t)
	_testSendGroup(server1, t)
	_testSendChain(server1, t)
	worker1.Quit()

	memcacheURL := os.Getenv("MEMCACHE_URL")
	if memcacheURL == "" {
		return
	}

	server2 := setup(brokerURL, fmt.Sprintf("memcache://%v", memcacheURL))
	worker2 := server2.NewWorker("test_worker")
	go worker2.Launch()
	_testSendTask(server2, t)
	_testSendGroup(server2, t)
	_testSendChain(server2, t)
	worker2.Quit()
}

func _testSendTask(server *Server, t *testing.T) {
	asyncResult, err := server.SendTask(&task0)
	if err != nil {
		t.Error(err)
	}

	result, err := asyncResult.Get()
	if err != nil {
		t.Error(err)
	}

	if result.Interface() != int64(2) {
		t.Errorf(
			"result = %v(%v), want int64(2)",
			result.Type().String(),
			result.Interface(),
		)
	}
}

func _testSendGroup(server *Server, t *testing.T) {
	group := NewGroup(&task1, &task2, &task3)
	asyncResults, err := server.SendGroup(group)
	if err != nil {
		t.Error(err)
	}

	expectedResults := []int64{2, 11, 4}

	isInSlice := func(lookup interface{}, list []int64) bool {
		for _, item := range list {
			if item == lookup {
				return true
			}
		}
		return false
	}

	for _, asyncResult := range asyncResults {
		result, err := asyncResult.Get()
		if err != nil {
			t.Error(err)
		}

		if !isInSlice(result.Interface(), expectedResults) {
			t.Errorf(
				"result = %v(%v), not in slice of expected results: []int64{2, 11, 4}",
				result.Type().String(),
				result.Interface(),
			)
		}
	}
}

func _testSendChain(server *Server, t *testing.T) {
	chain := NewChain(&task1, &task2, &task3)
	chainAsyncResult, err := server.SendChain(chain)
	if err != nil {
		t.Error(err)
	}

	result, err := chainAsyncResult.Get()
	if err != nil {
		t.Error(err)
	}

	if result.Interface() != int64(52) {
		t.Errorf(
			"result = %v(%v), want int64(52)",
			result.Type().String(),
			result.Interface(),
		)
	}
}

func setup(brokerURL, backend string) *Server {
	cnf := config.Config{
		Broker:        brokerURL,
		ResultBackend: backend,
		Exchange:      "test_exchange",
		ExchangeType:  "direct",
		DefaultQueue:  "test_queue",
		BindingKey:    "test_task",
	}

	server, err := NewServer(&cnf)
	errors.Fail(err, "Could not initialize server")

	tasks := map[string]interface{}{
		"add": func(args ...int64) (int64, error) {
			sum := int64(0)
			for _, arg := range args {
				sum += arg
			}
			return sum, nil
		},
		"multiply": func(args ...int64) (int64, error) {
			sum := int64(1)
			for _, arg := range args {
				sum *= arg
			}
			return sum, nil
		},
	}
	server.RegisterTasks(tasks)

	return server
}
<|endoftext|>


<commit_before>package replication

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"github.com/juju/errors"
)

type BinlogParser struct {
	format *FormatDescriptionEvent

	tables map[uint64]*TableMapEvent

	// for rawMode, we only parse FormatDescriptionEvent and RotateEvent
	rawMode bool
}

func NewBinlogParser() *BinlogParser {
	p := new(BinlogParser)

	p.tables = make(map[uint64]*TableMapEvent)

	return p
}

func (p *BinlogParser) Reset() {
	p.tables = make(map[uint64]*TableMapEvent)
	p.format = nil
}

type OnEventFunc func(*BinlogEvent) error

func (p *BinlogParser) ParseFile(name string, offset int64, onEvent OnEventFunc) error {
	f, err := os.Open(name)
	if err != nil {
		return errors.Trace(err)
	}
	defer f.Close()

	b := make([]byte, 4)
	if _, err = f.Read(b); err != nil {
		return errors.Trace(err)
	} else if !bytes.Equal(b, BinLogFileHeader) {
		return errors.Errorf("%s is not a valid binlog file, head 4 bytes must fe'bin' ", name)
	}

	if offset < 4 {
		offset = 4
	}

	if _, err = f.Seek(offset, os.SEEK_SET); err != nil {
		return errors.Errorf("seek %s to %d error %v", name, offset, err)
	}

	return p.ParseReader(f, onEvent)
}

func (p *BinlogParser) ParseReader(r io.Reader, onEvent OnEventFunc) error {
	p.Reset()

	var err error
	var n int64

	for {
		headBuf := make([]byte, EventHeaderSize)

		if _, err = io.ReadFull(r, headBuf); err == io.EOF {
			return nil
		} else if err != nil {
			return errors.Trace(err)
		}

		var h *EventHeader
		h, err = p.parseHeader(headBuf)
		if err != nil {
			return errors.Trace(err)
		}

		if h.EventSize <= uint32(EventHeaderSize) {
			return errors.Errorf("invalid event header, event size is %d, too small", h.EventSize)

		}

		var buf bytes.Buffer
		if n, err = io.CopyN(&buf, r, int64(h.EventSize)-int64(EventHeaderSize)); err != nil {
			return errors.Errorf("get event body err %v, need %d - %d, but got %d", err, h.EventSize, EventHeaderSize, n)
		}

		data := buf.Bytes()
		rawData := data

		eventLen := int(h.EventSize) - EventHeaderSize

		if len(data) != eventLen {
			return errors.Errorf("invalid data size %d in event %s, less event length %d", len(data), h.EventType, eventLen)
		}

		var e Event
		e, err = p.parseEvent(h, data)
		if err != nil {
			break
		}

		if err = onEvent(&BinlogEvent{rawData, h, e}); err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func (p *BinlogParser) SetRawMode(mode bool) {
	p.rawMode = mode
}

func (p *BinlogParser) parseHeader(data []byte) (*EventHeader, error) {
	h := new(EventHeader)
	err := h.Decode(data)
	if err != nil {
		return nil, err
	}

	return h, nil
}

func (p *BinlogParser) parseEvent(h *EventHeader, data []byte) (Event, error) {
	var e Event

	if h.EventType == FORMAT_DESCRIPTION_EVENT {
		p.format = &FormatDescriptionEvent{}
		e = p.format
	} else {
		if p.format != nil && p.format.ChecksumAlgorithm == BINLOG_CHECKSUM_ALG_CRC32 {
			data = data[0 : len(data)-4]
		}

		if h.EventType == ROTATE_EVENT {
			e = &RotateEvent{}
		} else if !p.rawMode {
			switch h.EventType {
			case QUERY_EVENT:
				e = &QueryEvent{}
			case XID_EVENT:
				e = &XIDEvent{}
			case TABLE_MAP_EVENT:
				te := &TableMapEvent{}
				if p.format.EventTypeHeaderLengths[TABLE_MAP_EVENT-1] == 6 {
					te.tableIDSize = 4
				} else {
					te.tableIDSize = 6
				}
				e = te
			case WRITE_ROWS_EVENTv0,
				UPDATE_ROWS_EVENTv0,
				DELETE_ROWS_EVENTv0,
				WRITE_ROWS_EVENTv1,
				DELETE_ROWS_EVENTv1,
				UPDATE_ROWS_EVENTv1,
				WRITE_ROWS_EVENTv2,
				UPDATE_ROWS_EVENTv2,
				DELETE_ROWS_EVENTv2:
				e = p.newRowsEvent(h)
			case ROWS_QUERY_EVENT:
				e = &RowsQueryEvent{}
			case GTID_EVENT:
				e = &GTIDEvent{}
			case BEGIN_LOAD_QUERY_EVENT:
				e = &BeginLoadQueryEvent{}
			case EXECUTE_LOAD_QUERY_EVENT:
				e = &ExecuteLoadQueryEvent{}
			case MARIADB_ANNOTATE_ROWS_EVENT:
				e = &MariadbAnnotaeRowsEvent{}
			case MARIADB_BINLOG_CHECKPOINT_EVENT:
				e = &MariadbBinlogCheckPointEvent{}
			case MARIADB_GTID_LIST_EVENT:
				e = &MariadbGTIDListEvent{}
			case MARIADB_GTID_EVENT:
				ee := &MariadbGTIDEvent{}
				ee.GTID.ServerID = h.ServerID
				e = ee
			default:
				e = &GenericEvent{}
			}
		} else {
			e = &GenericEvent{}
		}
	}

	if err := e.Decode(data); err != nil {
		return nil, &EventError{h, err.Error(), data}
	}

	if te, ok := e.(*TableMapEvent); ok {
		p.tables[te.TableID] = te
	}

	if re, ok := e.(*RowsEvent); ok {
		if (re.Flags & RowsEventStmtEndFlag) > 0 {
			// Refer https://github.com/alibaba/canal/blob/38cc81b7dab29b51371096fb6763ca3a8432ffee/dbsync/src/main/java/com/taobao/tddl/dbsync/binlog/event/RowsLogEvent.java#L176
			p.tables = make(map[uint64]*TableMapEvent)
		}
	}

	return e, nil
}

func (p *BinlogParser) parse(data []byte) (*BinlogEvent, error) {
	rawData := data

	h, err := p.parseHeader(data)

	if err != nil {
		return nil, err
	}

	data = data[EventHeaderSize:]
	eventLen := int(h.EventSize) - EventHeaderSize

	if len(data) != eventLen {
		return nil, fmt.Errorf("invalid data size %d in event %s, less event length %d", len(data), h.EventType, eventLen)
	}

	e, err := p.parseEvent(h, data)
	if err != nil {
		return nil, err
	}

	return &BinlogEvent{rawData, h, e}, nil
}

func (p *BinlogParser) newRowsEvent(h *EventHeader) *RowsEvent {
	e := &RowsEvent{}
	if p.format.EventTypeHeaderLengths[h.EventType-1] == 6 {
		e.tableIDSize = 4
	} else {
		e.tableIDSize = 6
	}

	e.needBitmap2 = false
	e.tables = p.tables

	switch h.EventType {
	case WRITE_ROWS_EVENTv0:
		e.Version = 0
	case UPDATE_ROWS_EVENTv0:
		e.Version = 0
	case DELETE_ROWS_EVENTv0:
		e.Version = 0
	case WRITE_ROWS_EVENTv1:
		e.Version = 1
	case DELETE_ROWS_EVENTv1:
		e.Version = 1
	case UPDATE_ROWS_EVENTv1:
		e.Version = 1
		e.needBitmap2 = true
	case WRITE_ROWS_EVENTv2:
		e.Version = 2
	case UPDATE_ROWS_EVENTv2:
		e.Version = 2
		e.needBitmap2 = true
	case DELETE_ROWS_EVENTv2:
		e.Version = 2
	}

	return e
}
<commit_msg>replication: tables can only be cleared with RowsEventStmtEndFlag (#79)<commit_after>package replication

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"github.com/juju/errors"
)

type BinlogParser struct {
	format *FormatDescriptionEvent

	tables map[uint64]*TableMapEvent

	// for rawMode, we only parse FormatDescriptionEvent and RotateEvent
	rawMode bool
}

func NewBinlogParser() *BinlogParser {
	p := new(BinlogParser)

	p.tables = make(map[uint64]*TableMapEvent)

	return p
}

func (p *BinlogParser) Reset() {
	p.format = nil
}

type OnEventFunc func(*BinlogEvent) error

func (p *BinlogParser) ParseFile(name string, offset int64, onEvent OnEventFunc) error {
	f, err := os.Open(name)
	if err != nil {
		return errors.Trace(err)
	}
	defer f.Close()

	b := make([]byte, 4)
	if _, err = f.Read(b); err != nil {
		return errors.Trace(err)
	} else if !bytes.Equal(b, BinLogFileHeader) {
		return errors.Errorf("%s is not a valid binlog file, head 4 bytes must fe'bin' ", name)
	}

	if offset < 4 {
		offset = 4
	}

	if _, err = f.Seek(offset, os.SEEK_SET); err != nil {
		return errors.Errorf("seek %s to %d error %v", name, offset, err)
	}

	return p.parseReader(f, onEvent)
}

func (p *BinlogParser) parseReader(r io.Reader, onEvent OnEventFunc) error {
	p.Reset()

	var err error
	var n int64

	for {
		headBuf := make([]byte, EventHeaderSize)

		if _, err = io.ReadFull(r, headBuf); err == io.EOF {
			return nil
		} else if err != nil {
			return errors.Trace(err)
		}

		var h *EventHeader
		h, err = p.parseHeader(headBuf)
		if err != nil {
			return errors.Trace(err)
		}

		if h.EventSize <= uint32(EventHeaderSize) {
			return errors.Errorf("invalid event header, event size is %d, too small", h.EventSize)

		}

		var buf bytes.Buffer
		if n, err = io.CopyN(&buf, r, int64(h.EventSize)-int64(EventHeaderSize)); err != nil {
			return errors.Errorf("get event body err %v, need %d - %d, but got %d", err, h.EventSize, EventHeaderSize, n)
		}

		data := buf.Bytes()
		rawData := data

		eventLen := int(h.EventSize) - EventHeaderSize

		if len(data) != eventLen {
			return errors.Errorf("invalid data size %d in event %s, less event length %d", len(data), h.EventType, eventLen)
		}

		var e Event
		e, err = p.parseEvent(h, data)
		if err != nil {
			break
		}

		if err = onEvent(&BinlogEvent{rawData, h, e}); err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func (p *BinlogParser) SetRawMode(mode bool) {
	p.rawMode = mode
}

func (p *BinlogParser) parseHeader(data []byte) (*EventHeader, error) {
	h := new(EventHeader)
	err := h.Decode(data)
	if err != nil {
		return nil, err
	}

	return h, nil
}

func (p *BinlogParser) parseEvent(h *EventHeader, data []byte) (Event, error) {
	var e Event

	if h.EventType == FORMAT_DESCRIPTION_EVENT {
		p.format = &FormatDescriptionEvent{}
		e = p.format
	} else {
		if p.format != nil && p.format.ChecksumAlgorithm == BINLOG_CHECKSUM_ALG_CRC32 {
			data = data[0 : len(data)-4]
		}

		if h.EventType == ROTATE_EVENT {
			e = &RotateEvent{}
		} else if !p.rawMode {
			switch h.EventType {
			case QUERY_EVENT:
				e = &QueryEvent{}
			case XID_EVENT:
				e = &XIDEvent{}
			case TABLE_MAP_EVENT:
				te := &TableMapEvent{}
				if p.format.EventTypeHeaderLengths[TABLE_MAP_EVENT-1] == 6 {
					te.tableIDSize = 4
				} else {
					te.tableIDSize = 6
				}
				e = te
			case WRITE_ROWS_EVENTv0,
				UPDATE_ROWS_EVENTv0,
				DELETE_ROWS_EVENTv0,
				WRITE_ROWS_EVENTv1,
				DELETE_ROWS_EVENTv1,
				UPDATE_ROWS_EVENTv1,
				WRITE_ROWS_EVENTv2,
				UPDATE_ROWS_EVENTv2,
				DELETE_ROWS_EVENTv2:
				e = p.newRowsEvent(h)
			case ROWS_QUERY_EVENT:
				e = &RowsQueryEvent{}
			case GTID_EVENT:
				e = &GTIDEvent{}
			case BEGIN_LOAD_QUERY_EVENT:
				e = &BeginLoadQueryEvent{}
			case EXECUTE_LOAD_QUERY_EVENT:
				e = &ExecuteLoadQueryEvent{}
			case MARIADB_ANNOTATE_ROWS_EVENT:
				e = &MariadbAnnotaeRowsEvent{}
			case MARIADB_BINLOG_CHECKPOINT_EVENT:
				e = &MariadbBinlogCheckPointEvent{}
			case MARIADB_GTID_LIST_EVENT:
				e = &MariadbGTIDListEvent{}
			case MARIADB_GTID_EVENT:
				ee := &MariadbGTIDEvent{}
				ee.GTID.ServerID = h.ServerID
				e = ee
			default:
				e = &GenericEvent{}
			}
		} else {
			e = &GenericEvent{}
		}
	}

	if err := e.Decode(data); err != nil {
		return nil, &EventError{h, err.Error(), data}
	}

	if te, ok := e.(*TableMapEvent); ok {
		p.tables[te.TableID] = te
	}

	if re, ok := e.(*RowsEvent); ok {
		if (re.Flags & RowsEventStmtEndFlag) > 0 {
			// Refer https://github.com/alibaba/canal/blob/38cc81b7dab29b51371096fb6763ca3a8432ffee/dbsync/src/main/java/com/taobao/tddl/dbsync/binlog/event/RowsLogEvent.java#L176
			p.tables = make(map[uint64]*TableMapEvent)
		}
	}

	return e, nil
}

func (p *BinlogParser) parse(data []byte) (*BinlogEvent, error) {
	rawData := data

	h, err := p.parseHeader(data)

	if err != nil {
		return nil, err
	}

	data = data[EventHeaderSize:]
	eventLen := int(h.EventSize) - EventHeaderSize

	if len(data) != eventLen {
		return nil, fmt.Errorf("invalid data size %d in event %s, less event length %d", len(data), h.EventType, eventLen)
	}

	e, err := p.parseEvent(h, data)
	if err != nil {
		return nil, err
	}

	return &BinlogEvent{rawData, h, e}, nil
}

func (p *BinlogParser) newRowsEvent(h *EventHeader) *RowsEvent {
	e := &RowsEvent{}
	if p.format.EventTypeHeaderLengths[h.EventType-1] == 6 {
		e.tableIDSize = 4
	} else {
		e.tableIDSize = 6
	}

	e.needBitmap2 = false
	e.tables = p.tables

	switch h.EventType {
	case WRITE_ROWS_EVENTv0:
		e.Version = 0
	case UPDATE_ROWS_EVENTv0:
		e.Version = 0
	case DELETE_ROWS_EVENTv0:
		e.Version = 0
	case WRITE_ROWS_EVENTv1:
		e.Version = 1
	case DELETE_ROWS_EVENTv1:
		e.Version = 1
	case UPDATE_ROWS_EVENTv1:
		e.Version = 1
		e.needBitmap2 = true
	case WRITE_ROWS_EVENTv2:
		e.Version = 2
	case UPDATE_ROWS_EVENTv2:
		e.Version = 2
		e.needBitmap2 = true
	case DELETE_ROWS_EVENTv2:
		e.Version = 2
	}

	return e
}
<|endoftext|>


<commit_before>package services

import (
	"fmt"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/off-sync/platform-proxy-aws/interfaces"
	"github.com/off-sync/platform-proxy-domain/services"
)

// ServiceRepository implements the ServiceRepository interface using an AWS ECS
// cluster as its backend.
type ServiceRepository struct {
	// AWS ECS API
	api interfaces.AwsEcsAPI

	// Configuration
	serverContainerName string
	dockerLabelPort     string
	defaultPort         int
}

// Default values for the ServiceRepository struct.
const (
	DefaultServerContainerName = "server"
	DefaultDockerLabelPort     = "com.off-sync.platform.proxy.port"
	DefaultDefaultPort         = 8080
)

// ServiceRepositoryOption defines the type used to further configure a
// ServiceRepository.
type ServiceRepositoryOption func(*ServiceRepository) error

// NewServiceRepository creates a new service repository based on the provided
// AWS ECS API.
func NewServiceRepository(api interfaces.AwsEcsAPI, options ...ServiceRepositoryOption) (*ServiceRepository, error) {
	r := &ServiceRepository{
		api:                 api,
		serverContainerName: DefaultServerContainerName,
		dockerLabelPort:     DefaultDockerLabelPort,
		defaultPort:         DefaultDefaultPort,
	}

	for _, opt := range options {
		err := opt(r)
		if err != nil {
			return nil, err
		}
	}

	return r, nil
}

// WithServerContainerName configures a service repository with the provided
// server container name.
func WithServerContainerName(name string) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.serverContainerName = name
		return nil
	}
}

// WithDockerLabelPort configures a service repository with the provided
// docker label for the port.
func WithDockerLabelPort(label string) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.dockerLabelPort = label
		return nil
	}
}

// WithDefaultPort configures a service repository with the provided
// default port.
func WithDefaultPort(port int) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.defaultPort = port
		return nil
	}
}

// ListServices returns all service names contained in this repository.
func (r *ServiceRepository) ListServices() ([]string, error) {
	return r.api.ListServices()
}

// DescribeService returns the service with the specified name. If no service
// exists with that name an ErrUnknownService is returned.
func (r *ServiceRepository) DescribeService(name string) (*services.Service, error) {
	service, err := r.api.DescribeService(name)
	if err != nil {
		return nil, err
	}

	serverURL, err := r.getTaskDefinitionServerURL(aws.StringValue(service.TaskDefinition))
	if err != nil {
		return nil, err
	}

	return services.NewService(name, serverURL)
}

func (r *ServiceRepository) getTaskDefinitionServerURL(taskDefArn string) (string, error) {
	tdef, err := r.api.DescribeTaskDefinition(taskDefArn)
	if err != nil {
		return "", err
	}

	for _, cdef := range tdef.ContainerDefinitions {
		if aws.StringValue(cdef.Name) != r.serverContainerName {
			// not the server
			continue
		}

		port := r.defaultPort

		portLabel, found := cdef.DockerLabels[r.dockerLabelPort]
		if found {
			port, err = strconv.Atoi(*portLabel)
			if err != nil {
				return "", fmt.Errorf("invalid port: %s", *portLabel)
			}
		}

		return fmt.Sprintf("http://%s:%d", *cdef.Hostname, port), nil
	}

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

package services

import (
	"fmt"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/off-sync/platform-proxy-aws/interfaces"
	"github.com/off-sync/platform-proxy-domain/services"
)

// ServiceRepository implements the ServiceRepository interface using an AWS ECS
// cluster as its backend.
type ServiceRepository struct {
	// AWS ECS API
	api interfaces.AwsEcsAPI

	// Configuration
	serverContainerName string
	dockerLabelPort     string
	defaultPort         int
}

// Default values for the ServiceRepository struct.
const (
	DefaultServerContainerName = "server"
	DefaultDockerLabelPort     = "com.off-sync.platform.proxy.port"
	DefaultDefaultPort         = 8080
)

// ServiceRepositoryOption defines the type used to further configure a
// ServiceRepository.
type ServiceRepositoryOption func(*ServiceRepository) error

// NewServiceRepository creates a new service repository based on the provided
// AWS ECS API.
func NewServiceRepository(api interfaces.AwsEcsAPI, options ...ServiceRepositoryOption) (*ServiceRepository, error) {
	r := &ServiceRepository{
		api:                 api,
		serverContainerName: DefaultServerContainerName,
		dockerLabelPort:     DefaultDockerLabelPort,
		defaultPort:         DefaultDefaultPort,
	}

	for _, opt := range options {
		err := opt(r)
		if err != nil {
			return nil, err
		}
	}

	return r, nil
}

// WithServerContainerName configures a service repository with the provided
// server container name.
func WithServerContainerName(name string) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.serverContainerName = name
		return nil
	}
}

// WithDockerLabelPort configures a service repository with the provided
// docker label for the port.
func WithDockerLabelPort(label string) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.dockerLabelPort = label
		return nil
	}
}

// WithDefaultPort configures a service repository with the provided
// default port.
func WithDefaultPort(port int) ServiceRepositoryOption {
	return func(r *ServiceRepository) error {
		r.defaultPort = port
		return nil
	}
}

// ListServices returns all service names contained in this repository.
func (r *ServiceRepository) ListServices() ([]string, error) {
	return r.api.ListServices()
}

// DescribeService returns the service with the specified name. If no service
// exists with that name an ErrUnknownService is returned.
func (r *ServiceRepository) DescribeService(name string) (*services.Service, error) {
	service, err := r.api.DescribeService(name)
	if err != nil {
		return nil, err
	}

	serverURL, err := r.getTaskDefinitionServerURL(aws.StringValue(service.TaskDefinition))
	if err != nil {
		return nil, err
	}

	return services.NewService(name, serverURL)
}

func (r *ServiceRepository) getTaskDefinitionServerURL(taskDefArn string) (string, error) {
	tdef, err := r.api.DescribeTaskDefinition(taskDefArn)
	if err != nil {
		return "", err
	}

	for _, cdef := range tdef.ContainerDefinitions {
		if aws.StringValue(cdef.Name) != r.serverContainerName {
			// not the server
			continue
		}

		port := r.defaultPort

		portLabel, found := cdef.DockerLabels[r.dockerLabelPort]
		if found {
			port, err = strconv.Atoi(*portLabel)
			if err != nil {
				return "", fmt.Errorf("invalid port: %s", *portLabel)
			}
		}

		return fmt.Sprintf("http://%s:%d", *cdef.Hostname, port), nil
	}

	return "", fmt.Errorf("no server container found for task definition: %s", taskDefArn)
}
<|endoftext|>


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

package gcmautoscaling

import (
	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"

	"github.com/GoogleCloudPlatform/heapster/extpoints"
	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	"github.com/GoogleCloudPlatform/heapster/sinks/gcm"
	"github.com/golang/glog"
)

var (
	LabelHostname = sink_api.LabelDescriptor{
		Key:         "hostname",
		Description: "Hostname where the container ran",
	}
	LabelGCEResourceID = sink_api.LabelDescriptor{
		Key:         "compute.googleapis.com/resource_id",
		Description: "Resource id for nodes specific for GCE.",
	}
	LabelGCEResourceType = sink_api.LabelDescriptor{
		Key:         "compute.googleapis.com/resource_type",
		Description: "Resource types for nodes specific for GCE.",
	}
)

var autoscalingLabels = []sink_api.LabelDescriptor{
	LabelHostname,
	LabelGCEResourceID,
	LabelGCEResourceType,
}

type utilizationMetric struct {
	name        string
	description string
}

var autoscalingMetrics = map[string]utilizationMetric{
	"cpu/usage": {
		name:        "cpu/node_utilization",
		description: "Cpu utilization as a percentage of node capacity",
	},
	"memory/usage": {
		name:        "memory/node_utilization",
		description: "Memory utilization as a percentage of memory capacity",
	},
}

type gcmAutocalingSink struct {
	core *gcm.GcmCore
	// For given hostname remember its capacity in milicores.
	cpuCapacity map[string]int64
	// Memory capacity in bytes.
	memCapacity map[string]int64
}

// Adds the specified metrics or updates them if they already exist.
func (self gcmAutocalingSink) Register(_ []sink_api.MetricDescriptor) error {
	for _, metric := range autoscalingMetrics {
		if err := self.core.Register(metric.name, metric.description, sink_api.MetricGauge.String(), sink_api.ValueDouble.String(), autoscalingLabels); err != nil {
			return err
		}
	}
	return nil
}

// Stores events into the backend.
func (self gcmAutocalingSink) StoreEvents([]kube_api.Event) error {
	// No-op, Google Cloud Monitoring doesn't store events
	return nil
}

func isNode(metric *sink_api.Point) bool {
	return metric.Labels[sink_api.LabelContainerName.Key] == "/"
}

func (self gcmAutocalingSink) updateMachineCapacity(input []sink_api.Timeseries) {
	for _, entry := range input {
		metric := entry.Point
		if !isNode(metric) {
			continue
		}

		host := metric.Labels[sink_api.LabelHostname.Key]
		value, ok := metric.Value.(int64)
		if !ok {
			continue
		}
		if metric.Name == "cpu/limit" {
			self.cpuCapacity[host] = value
		} else if metric.Name == "memory/limit" {
			self.memCapacity[host] = value
		}
	}
}

// Pushes the specified metric values in input. The metrics must already exist.
func (self gcmAutocalingSink) StoreTimeseries(input []sink_api.Timeseries) error {
	self.updateMachineCapacity(input)

	// Build a map of metrics by name.
	metrics := make(map[string][]gcm.Timeseries)
	for _, entry := range input {
		metric := entry.Point
		// We want to export only node metrics.
		if !isNode(metric) {
			continue
		}

		labels := make(map[string]string, 3)
		labels[LabelHostname.Key] = metric.Labels[sink_api.LabelHostname.Key]
		labels[LabelGCEResourceID.Key] = metric.Labels[sink_api.LabelHostID.Key]
		labels[LabelGCEResourceType.Key] = "instance"
		metric.Labels = labels

		var ts *gcm.Timeseries
		var err error
		var newVal float64
		if metric.Name == "cpu/usage" {
			ts, err = self.core.GetEquivalentRateMetric(metric)
			if err != nil {
				return err
			}
			capacity := self.cpuCapacity[labels[LabelHostname.Key]]
			if ts == nil || capacity < 1 {
				continue
			}
			newVal = *ts.Point.DoubleValue / float64(capacity)
		} else if metric.Name == "mem/usage" {
			ts, err = self.core.GetMetric(metric)
			if err != nil {
				return err
			}
			capacity := self.memCapacity[labels[LabelHostname.Key]]
			if ts == nil || capacity < 1 {
				continue
			}
			newVal = float64(*ts.Point.Int64Value) / float64(capacity)
		} else {
			continue
		}

		name := gcm.FullMetricName(autoscalingMetrics[metric.Name].name)
		ts.TimeseriesDescriptor.Metric = name
		ts.Point.DoubleValue = &newVal

		metrics[name] = append(metrics[name], *ts)
	}

	return self.core.StoreTimeseries(metrics)
}

func (self gcmAutocalingSink) DebugInfo() string {
	return "Sink Type: GCM Autoscaling"
}

func (self gcmAutocalingSink) Name() string {
	return "Google Cloud Monitoring Sink for Autoscaling"
}

func init() {
	extpoints.SinkFactories.Register(CreateGCMScalingSink, "gcmautoscaling")
}

func CreateGCMScalingSink(_ string, _ map[string][]string) ([]sink_api.ExternalSink, error) {
	core, err := gcm.NewCore()
	sink := gcmAutocalingSink{
		core:        core,
		cpuCapacity: make(map[string]int64),
		memCapacity: make(map[string]int64),
	}
	glog.Infof("created GCM Autocaling sink")
	return []sink_api.ExternalSink{sink}, err
}
<commit_msg>Added exporting nodes reservation metrics for Autoscaling sink<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcmautoscaling

import (
	"github.com/GoogleCloudPlatform/heapster/extpoints"
	"github.com/GoogleCloudPlatform/heapster/sinks/gcm"
	"github.com/golang/glog"

	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	kube_api "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
)

var (
	LabelHostname = sink_api.LabelDescriptor{
		Key:         "hostname",
		Description: "Hostname where the container ran",
	}
	LabelGCEResourceID = sink_api.LabelDescriptor{
		Key:         "compute.googleapis.com/resource_id",
		Description: "Resource id for nodes specific for GCE.",
	}
	LabelGCEResourceType = sink_api.LabelDescriptor{
		Key:         "compute.googleapis.com/resource_type",
		Description: "Resource types for nodes specific for GCE.",
	}
)

var autoscalingLabels = []sink_api.LabelDescriptor{
	LabelHostname,
	LabelGCEResourceID,
	LabelGCEResourceType,
}

type utilizationMetric struct {
	name        string
	description string
}

var autoscalingMetrics = map[string]utilizationMetric{
	"cpu/usage": {
		name:        "cpu/node_utilization",
		description: "Cpu utilization as a share of node capacity",
	},
	"cpu/limit": {
		name:        "cpu/node_reservation",
		description: "Share of cpu that is reserved on the node",
	},
	"memory/usage": {
		name:        "memory/node_utilization",
		description: "Memory utilization as a share of memory capacity",
	},
	"memory/limit": {
		name:        "memory/node_reservation",
		description: "Share of memory that is reserved on the node",
	},
}

type gcmAutocalingSink struct {
	core *gcm.GcmCore
	// For given hostname remember its capacity in milicores.
	cpuCapacity map[string]int64
	// For given hostname remember amount of reserved cpu in milicores.
	cpuReservation map[string]int64
	// Memory capacity in bytes.
	memCapacity map[string]int64
	// For given hostname remember amount of reserved memory in bytes.
	memReservation map[string]int64
}

// Adds the specified metrics or updates them if they already exist.
func (self gcmAutocalingSink) Register(_ []sink_api.MetricDescriptor) error {
	for _, metric := range autoscalingMetrics {
		if err := self.core.Register(metric.name, metric.description, sink_api.MetricGauge.String(), sink_api.ValueDouble.String(), autoscalingLabels); err != nil {
			return err
		}
	}
	return nil
}

// Stores events into the backend.
func (self gcmAutocalingSink) StoreEvents([]kube_api.Event) error {
	// No-op, Google Cloud Monitoring doesn't store events
	return nil
}

func isNode(metric *sink_api.Point) bool {
	return metric.Labels[sink_api.LabelContainerName.Key] == "/"
}

func isPodContainer(metric *sink_api.Point) bool {
	return len(metric.Labels[sink_api.LabelPodName.Key]) > 0
}

func (self *gcmAutocalingSink) updateMachineCapacityAndReservation(input []sink_api.Timeseries) {
	self.cpuCapacity = make(map[string]int64)
	self.cpuReservation = make(map[string]int64)
	self.memCapacity = make(map[string]int64)
	self.memReservation = make(map[string]int64)
	for _, entry := range input {
		metric := entry.Point
		if metric.Name != "cpu/limit" && metric.Name != "memory/limit" {
			continue
		}
		host := metric.Labels[sink_api.LabelHostname.Key]
		value, ok := metric.Value.(int64)
		if !ok || value < 1 {
			continue
		}

		if isNode(metric) {
			if metric.Name == "cpu/limit" {
				self.cpuCapacity[host] = value
			} else if metric.Name == "memory/limit" {
				self.memCapacity[host] = value
			}
		} else if isPodContainer(metric) {
			if metric.Name == "cpu/limit" {
				self.cpuReservation[host] += value
			} else if metric.Name == "memory/limit" {
				self.memReservation[host] += value
			}
		}
	}
}

func getLabels(metric *sink_api.Point) map[string]string {
	return map[string]string{
		LabelHostname.Key:        metric.Labels[sink_api.LabelHostname.Key],
		LabelGCEResourceID.Key:   metric.Labels[sink_api.LabelHostID.Key],
		LabelGCEResourceType.Key: "instance",
	}
}

func (self *gcmAutocalingSink) getNewValue(metric *sink_api.Point, ts *gcm.Timeseries) *float64 {
	host := metric.Labels[sink_api.LabelHostname.Key]

	var val float64
	switch metric.Name {
	case "cpu/usage":
		capacity, ok := self.cpuCapacity[host]
		if !ok || capacity < 1 || ts.Point.DoubleValue == nil {
			return nil
		}
		val = *ts.Point.DoubleValue / float64(capacity)
	case "cpu/limit":
		reserved, ok := self.cpuReservation[host]
		capacity, ok2 := self.cpuCapacity[host]
		if !ok || !ok2 || capacity < 1 {
			return nil
		}
		val = float64(reserved) / float64(capacity)
	case "memory/usage":
		capacity, ok := self.memCapacity[host]
		if !ok || capacity < 1 || ts.Point.Int64Value == nil {
			return nil
		}
		val = float64(*ts.Point.Int64Value) / float64(capacity)
	case "memory/limit":
		reserved, ok := self.memReservation[host]
		capacity, ok2 := self.memCapacity[host]
		if !ok || !ok2 || capacity < 1 {
			return nil
		}
		val = float64(reserved) / float64(capacity)
	default:
		return nil
	}
	return &val
}

// Pushes the specified metric values in input. The metrics must already exist.
func (self gcmAutocalingSink) StoreTimeseries(input []sink_api.Timeseries) error {
	self.updateMachineCapacityAndReservation(input)

	// Build a map of metrics by name.
	metrics := make(map[string][]gcm.Timeseries)
	for _, entry := range input {
		metric := entry.Point
		// We want to export only node metrics.
		if !isNode(metric) {
			continue
		}

		metric.Labels = getLabels(metric)

		var ts *gcm.Timeseries
		var err error
		if metric.Name == "cpu/usage" {
			ts, err = self.core.GetEquivalentRateMetric(metric)
		} else if metric.Name == "cpu/limit" || metric.Name == "memory/usage" || metric.Name == "memory/limit" {
			ts, err = self.core.GetMetric(metric)
		} else {
			continue
		}
		if err != nil || ts == nil {
			glog.Infof("Failed to create Timeseries for metric %v, host %v. Error %v.", autoscalingMetrics[metric.Name].name, metric.Labels[sink_api.LabelHostname.Key], err)
			continue
		}

		val := self.getNewValue(metric, ts)
		if val == nil {
			glog.Infof("Failed to compute new value for metric %v, host %v.", autoscalingMetrics[metric.Name].name, metric.Labels[sink_api.LabelHostname.Key])
			continue
		}
		ts.Point.DoubleValue = val
		name := gcm.FullMetricName(autoscalingMetrics[metric.Name].name)
		ts.TimeseriesDescriptor.Metric = name

		metrics[name] = append(metrics[name], *ts)
	}

	return self.core.StoreTimeseries(metrics)
}

func (self gcmAutocalingSink) DebugInfo() string {
	return "Sink Type: GCM Autoscaling"
}

func (self gcmAutocalingSink) Name() string {
	return "Google Cloud Monitoring Sink for Autoscaling"
}

func init() {
	extpoints.SinkFactories.Register(CreateGCMScalingSink, "gcmautoscaling")
}

func CreateGCMScalingSink(_ string, _ map[string][]string) ([]sink_api.ExternalSink, error) {
	core, err := gcm.NewCore()
	sink := gcmAutocalingSink{core: core}
	glog.Infof("created GCM Autocaling sink")
	return []sink_api.ExternalSink{sink}, err
}
<|endoftext|>


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

package avalanche

import (
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/consensus/snowstorm"
	"github.com/ava-labs/avalanchego/snow/engine/avalanche/vertex"
)

// Voter records chits received from [vdr] once its dependencies are met.
type voter struct {
	t         *Transitive
	vdr       ids.ShortID
	requestID uint32
	response  []ids.ID
	deps      ids.Set
}

func (v *voter) Dependencies() ids.Set { return v.deps }

// Mark that a dependency has been met.
func (v *voter) Fulfill(id ids.ID) {
	v.deps.Remove(id)
	v.Update()
}

// Abandon this attempt to record chits.
func (v *voter) Abandon(id ids.ID) { v.Fulfill(id) }

func (v *voter) Update() {
	if v.deps.Len() != 0 || v.t.errs.Errored() {
		return
	}

	results, finished := v.t.polls.Vote(v.requestID, v.vdr, v.response)
	if !finished {
		return
	}
	results, err := v.bubbleVotes(results)
	if err != nil {
		v.t.errs.Add(err)
		return
	}

	v.t.Ctx.Log.Debug("Finishing poll with:\n%s", &results)
	if err := v.t.Consensus.RecordPoll(results); err != nil {
		v.t.errs.Add(err)
		return
	}

	orphans := v.t.Consensus.Orphans()
	txs := make([]snowstorm.Tx, 0, orphans.Len())
	for orphanID := range orphans {
		if tx, err := v.t.VM.Get(orphanID); err == nil {
			txs = append(txs, tx)
		} else {
			v.t.Ctx.Log.Warn("Failed to fetch %s during attempted re-issuance", orphanID)
		}
	}
	if len(txs) > 0 {
		v.t.Ctx.Log.Debug("Re-issuing %d transactions", len(txs))
	}
	if err := v.t.batch(txs, true /*=force*/, false /*empty*/); err != nil {
		v.t.errs.Add(err)
		return
	}

	if v.t.Consensus.Quiesce() {
		v.t.Ctx.Log.Debug("Avalanche engine can quiesce")
		return
	}

	v.t.Ctx.Log.Debug("Avalanche engine can't quiesce")
	v.t.errs.Add(v.t.repoll())
}

func (v *voter) bubbleVotes(votes ids.UniqueBag) (ids.UniqueBag, error) {
	vertexHeap := vertex.NewHeap()
	for vote := range votes {
		vtx, err := v.t.Manager.Get(vote)
		if err != nil {
			continue
		}
		vertexHeap.Push(vtx)
	}

	for vertexHeap.Len() > 0 {
		vtx := vertexHeap.Pop()
		vtxID := vtx.ID()
		set := votes.GetSet(vtxID)
		status := vtx.Status()

		if !status.Fetched() {
			v.t.Ctx.Log.Debug("Dropping %d vote(s) for %s because the vertex is unknown",
				set.Len(), vtxID)
			votes.RemoveSet(vtxID)
			continue
		}

		if status.Decided() {
			v.t.Ctx.Log.Debug("Dropping %d vote(s) for %s because the vertex is decided as %s",
				set.Len(), vtxID, status)
			votes.RemoveSet(vtxID)
			continue
		}

		if !v.t.Consensus.VertexIssued(vtx) {
			v.t.Ctx.Log.Debug("Bubbling %d vote(s) for %s because the vertex isn't issued",
				set.Len(), vtxID)
			votes.RemoveSet(vtxID) // Remove votes for this vertex because it hasn't been issued

			parents, err := vtx.Parents()
			if err != nil {
				return votes, err
			}
			for _, parentVtx := range parents {
				votes.UnionSet(parentVtx.ID(), set)
				vertexHeap.Push(parentVtx)
			}
		}
	}

	return votes, nil
}
<commit_msg>revert log level changes<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package avalanche

import (
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/consensus/snowstorm"
	"github.com/ava-labs/avalanchego/snow/engine/avalanche/vertex"
)

// Voter records chits received from [vdr] once its dependencies are met.
type voter struct {
	t         *Transitive
	vdr       ids.ShortID
	requestID uint32
	response  []ids.ID
	deps      ids.Set
}

func (v *voter) Dependencies() ids.Set { return v.deps }

// Mark that a dependency has been met.
func (v *voter) Fulfill(id ids.ID) {
	v.deps.Remove(id)
	v.Update()
}

// Abandon this attempt to record chits.
func (v *voter) Abandon(id ids.ID) { v.Fulfill(id) }

func (v *voter) Update() {
	if v.deps.Len() != 0 || v.t.errs.Errored() {
		return
	}

	results, finished := v.t.polls.Vote(v.requestID, v.vdr, v.response)
	if !finished {
		return
	}
	results, err := v.bubbleVotes(results)
	if err != nil {
		v.t.errs.Add(err)
		return
	}

	v.t.Ctx.Log.Debug("Finishing poll with:\n%s", &results)
	if err := v.t.Consensus.RecordPoll(results); err != nil {
		v.t.errs.Add(err)
		return
	}

	orphans := v.t.Consensus.Orphans()
	txs := make([]snowstorm.Tx, 0, orphans.Len())
	for orphanID := range orphans {
		if tx, err := v.t.VM.Get(orphanID); err == nil {
			txs = append(txs, tx)
		} else {
			v.t.Ctx.Log.Warn("Failed to fetch %s during attempted re-issuance", orphanID)
		}
	}
	if len(txs) > 0 {
		v.t.Ctx.Log.Debug("Re-issuing %d transactions", len(txs))
	}
	if err := v.t.batch(txs, true /*=force*/, false /*empty*/); err != nil {
		v.t.errs.Add(err)
		return
	}

	if v.t.Consensus.Quiesce() {
		v.t.Ctx.Log.Debug("Avalanche engine can quiesce")
		return
	}

	v.t.Ctx.Log.Debug("Avalanche engine can't quiesce")
	v.t.errs.Add(v.t.repoll())
}

func (v *voter) bubbleVotes(votes ids.UniqueBag) (ids.UniqueBag, error) {
	vertexHeap := vertex.NewHeap()
	for vote := range votes {
		vtx, err := v.t.Manager.Get(vote)
		if err != nil {
			continue
		}
		vertexHeap.Push(vtx)
	}

	for vertexHeap.Len() > 0 {
		vtx := vertexHeap.Pop()
		vtxID := vtx.ID()
		set := votes.GetSet(vtxID)
		status := vtx.Status()

		if !status.Fetched() {
			v.t.Ctx.Log.Verbo("Dropping %d vote(s) for %s because the vertex is unknown",
				set.Len(), vtxID)
			votes.RemoveSet(vtxID)
			continue
		}

		if status.Decided() {
			v.t.Ctx.Log.Verbo("Dropping %d vote(s) for %s because the vertex is decided as %s",
				set.Len(), vtxID, status)
			votes.RemoveSet(vtxID)
			continue
		}

		if !v.t.Consensus.VertexIssued(vtx) {
			v.t.Ctx.Log.Verbo("Bubbling %d vote(s) for %s because the vertex isn't issued",
				set.Len(), vtxID)
			votes.RemoveSet(vtxID) // Remove votes for this vertex because it hasn't been issued

			parents, err := vtx.Parents()
			if err != nil {
				return votes, err
			}
			for _, parentVtx := range parents {
				votes.UnionSet(parentVtx.ID(), set)
				vertexHeap.Push(parentVtx)
			}
		}
	}

	return votes, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package ensurer

import (
	"errors"
	"fmt"
	"strconv"

	flowcontrolv1beta2 "k8s.io/api/flowcontrol/v1beta2"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"

	"github.com/google/go-cmp/cmp"
)

const (
	fieldManager = "api-priority-and-fairness-config-producer-v1"
)

// ensureStrategy provides a strategy for ensuring apf bootstrap configurationWrapper.
// We have two types of configurationWrapper objects:
// - mandatory: the mandatory configurationWrapper objects are about ensuring that the P&F
//   system itself won't crash; we have to be sure there's 'catch-all' place for
//   everything to go. Any changes made by the cluster operators to these
//   configurationWrapper objects will be stomped by the apiserver.
//
// - suggested: additional configurationWrapper objects for initial behavior.
//   the cluster operators have an option to edit or delete these configurationWrapper objects.
type ensureStrategy interface {
	// Name of the strategy, for now we have two: 'mandatory' and 'suggested'.
	// This comes handy in logging.
	Name() string

	// ShouldUpdate accepts the current and the bootstrap configuration and determines
	// whether an update is necessary.
	// current is the existing in-cluster configuration object.
	// bootstrap is the configuration the kube-apiserver maintains in-memory.
	//
	// ok: true if auto update is required, otherwise false
	// object: the new object represents the new configuration to be stored in-cluster.
	// err: err is set when the function runs into an error and can not
	// determine if auto update is needed.
	ShouldUpdate(current, bootstrap configurationObject) (object runtime.Object, ok bool, err error)
}

// this internal interface provides abstraction for dealing with the `Spec`
// of both 'FlowSchema' and 'PriorityLevelConfiguration' objects.
// Since the ensure logic for both types is common, we use a few internal interfaces
// to abstract out the differences of these two types.
type specCopier interface {
	// HasSpecChanged returns true if the spec of both the bootstrap and
	// the current configuration object is same, otherwise false.
	HasSpecChanged(bootstrap, current runtime.Object) (bool, error)

	// CopySpec makes a deep copy the spec of the bootstrap object
	// and copies it to that of the current object.
	// CopySpec assumes that the current object is safe to mutate, so it
	// rests with the caller to make a deep copy of the current.
	CopySpec(bootstrap, current runtime.Object) error
}

// this internal interface provides abstraction for CRUD operation
// related to both 'FlowSchema' and 'PriorityLevelConfiguration' objects.
// Since the ensure logic for both types is common, we use a few internal interfaces
// to abstract out the differences of these two types.
type configurationClient interface {
	Create(object runtime.Object) (runtime.Object, error)
	Update(object runtime.Object) (runtime.Object, error)
	Get(name string) (configurationObject, error)
	Delete(name string) error
}

type configurationWrapper interface {
	// TypeName returns the type of the configuration that this interface deals with.
	// We use it to log the type name of the configuration object being ensured.
	// It is either 'PriorityLevelConfiguration' or 'FlowSchema'
	TypeName() string

	configurationClient
	specCopier
}

// A convenient wrapper interface that is used by the ensure logic.
type configurationObject interface {
	metav1.Object
	runtime.Object
}

func newSuggestedEnsureStrategy(copier specCopier) ensureStrategy {
	return &strategy{
		copier:               copier,
		alwaysAutoUpdateSpec: false,
		name:                 "suggested",
	}
}

func newMandatoryEnsureStrategy(copier specCopier) ensureStrategy {
	return &strategy{
		copier:               copier,
		alwaysAutoUpdateSpec: true,
		name:                 "mandatory",
	}
}

// auto-update strategy for the configuration objects
type strategy struct {
	copier               specCopier
	alwaysAutoUpdateSpec bool
	name                 string
}

func (s *strategy) Name() string {
	return s.name
}

func (s *strategy) ShouldUpdate(current, bootstrap configurationObject) (runtime.Object, bool, error) {
	if current == nil || bootstrap == nil {
		return nil, false, nil
	}

	autoUpdateSpec := s.alwaysAutoUpdateSpec
	if !autoUpdateSpec {
		autoUpdateSpec = shouldUpdateSpec(current)
	}
	updateAnnotation := shouldUpdateAnnotation(current, autoUpdateSpec)

	var specChanged bool
	if autoUpdateSpec {
		changed, err := s.copier.HasSpecChanged(bootstrap, current)
		if err != nil {
			return nil, false, fmt.Errorf("failed to compare spec - %w", err)
		}
		specChanged = changed
	}

	if !(updateAnnotation || specChanged) {
		// the annotation key is up to date and the spec has not changed, no update is necessary
		return nil, false, nil
	}

	// if we are here, either we need to update the annotation key or the spec.
	copy, ok := current.DeepCopyObject().(configurationObject)
	if !ok {
		// we should never be here
		return nil, false, errors.New("incompatible object type")
	}

	if updateAnnotation {
		setAutoUpdateAnnotation(copy, autoUpdateSpec)
	}
	if specChanged {
		s.copier.CopySpec(bootstrap, copy)
	}

	return copy, true, nil
}

// shouldUpdateSpec inspects the auto-update annotation key and generation field to determine
// whether the configurationWrapper object should be auto-updated.
func shouldUpdateSpec(accessor metav1.Object) bool {
	value, _ := accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]
	if autoUpdate, err := strconv.ParseBool(value); err == nil {
		return autoUpdate
	}

	// We are here because of either a or b:
	// a. the annotation key is missing.
	// b. the annotation key is present but the value does not represent a boolean.
	// In either case, if the operator hasn't changed the spec, we can safely auto update.
	// Please note that we can't protect the changes made by the operator in the following scenario:
	// - The operator deletes and recreates the same object with a variant spec (generation resets to 1).
	if accessor.GetGeneration() == 1 {
		return true
	}
	return false
}

// shouldUpdateAnnotation determines whether the current value of the auto-update annotation
// key matches the desired value.
func shouldUpdateAnnotation(accessor metav1.Object, desired bool) bool {
	if value, ok := accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]; ok {
		if current, err := strconv.ParseBool(value); err == nil && current == desired {
			return false
		}
	}

	return true
}

// setAutoUpdateAnnotation sets the auto-update annotation key to the specified value.
func setAutoUpdateAnnotation(accessor metav1.Object, autoUpdate bool) {
	if accessor.GetAnnotations() == nil {
		accessor.SetAnnotations(map[string]string{})
	}

	accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey] = strconv.FormatBool(autoUpdate)
}

// ensureConfiguration ensures the boostrap configurationWrapper on the cluster based on the specified strategy.
func ensureConfiguration(wrapper configurationWrapper, strategy ensureStrategy, bootstrap configurationObject) error {
	name := bootstrap.GetName()
	configurationType := strategy.Name()

	current, err := wrapper.Get(bootstrap.GetName())
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return fmt.Errorf("failed to retrieve %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
		}

		// we always re-create a missing configuration object
		if _, err := wrapper.Create(bootstrap); err != nil {
			return fmt.Errorf("cannot create %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
		}

		klog.V(2).InfoS(fmt.Sprintf("Successfully created %s", wrapper.TypeName()), "type", configurationType, "name", name)
		return nil
	}

	klog.V(5).InfoS(fmt.Sprintf("The %s already exists, checking whether it is up to date", wrapper.TypeName()), "type", configurationType, "name", name)
	newObject, update, err := strategy.ShouldUpdate(current, bootstrap)
	if err != nil {
		return fmt.Errorf("failed to determine whether auto-update is required for %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
	}
	if !update {
		if klog.V(5).Enabled() {
			// TODO: if we use structured logging here the diff gets escaped and very awkward to read in the log
			klog.Infof("No update required for the %s type=%s name=%q diff: %s", wrapper.TypeName(), configurationType, name, cmp.Diff(current, bootstrap))
		}
		return nil
	}

	if _, err := wrapper.Update(newObject); err != nil {
		return fmt.Errorf("failed to update the %s, will retry later type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
	}

	klog.V(2).Infof("Updated the %s type=%s name=%q diff: %s", wrapper.TypeName(), configurationType, name, cmp.Diff(current, newObject))
	return nil
}

func removeConfiguration(wrapper configurationWrapper, name string) error {
	current, err := wrapper.Get(name)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}

		return fmt.Errorf("failed to retrieve the %s, will retry later name=%q error=%w", wrapper.TypeName(), name, err)
	}

	value := current.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]
	autoUpdate, err := strconv.ParseBool(value)
	if err != nil {
		klog.ErrorS(err, fmt.Sprintf("Skipping deletion of the %s", wrapper.TypeName()), "name", name)

		// This may need manual intervention, in case the annotation value is malformed,
		// so don't return an error, that might trigger futile retry loop.
		return nil
	}
	if !autoUpdate {
		klog.V(5).InfoS(fmt.Sprintf("Skipping deletion of the %s", wrapper.TypeName()), "name", name)
		return nil
	}

	if err := wrapper.Delete(name); err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}

		return fmt.Errorf("failed to delete the %s, will retry later name=%q error=%w", wrapper.TypeName(), name, err)
	}

	klog.V(2).InfoS(fmt.Sprintf("Successfully deleted the %s", wrapper.TypeName()), "name", name)
	return nil
}

// getRemoveCandidate returns a list of configuration objects we should delete
// from the cluster given a set of bootstrap and current configuration.
// bootstrap: a set of hard coded configuration kube-apiserver maintains in-memory.
// current: a set of configuration objects that exist on the cluster
// Any object present in current is a candidate for removal if both a and b are true:
//  a. the object in current is missing from the bootstrap configuration
//  b. the object has the designated auto-update annotation key
// This function shares the common logic for both FlowSchema and PriorityLevelConfiguration
// type and hence it accepts metav1.Object only.
func getRemoveCandidate(bootstrap sets.String, current []metav1.Object) []string {
	if len(current) == 0 {
		return nil
	}

	candidates := make([]string, 0)
	for i := range current {
		object := current[i]
		if _, ok := object.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]; !ok {
			// the configuration object does not have the annotation key
			continue
		}

		if _, ok := bootstrap[object.GetName()]; !ok {
			candidates = append(candidates, object.GetName())
		}
	}
	return candidates
}
<commit_msg>Concurrentize pkg/registry/flowcontrol/ensurer/strategy.go<commit_after>/*
Copyright 2021 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package ensurer

import (
	"errors"
	"fmt"
	"strconv"

	flowcontrolv1beta2 "k8s.io/api/flowcontrol/v1beta2"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"

	"github.com/google/go-cmp/cmp"
)

const (
	fieldManager = "api-priority-and-fairness-config-producer-v1"
)

// ensureStrategy provides a strategy for ensuring apf bootstrap configurationWrapper.
// We have two types of configurationWrapper objects:
// - mandatory: the mandatory configurationWrapper objects are about ensuring that the P&F
//   system itself won't crash; we have to be sure there's 'catch-all' place for
//   everything to go. Any changes made by the cluster operators to these
//   configurationWrapper objects will be stomped by the apiserver.
//
// - suggested: additional configurationWrapper objects for initial behavior.
//   the cluster operators have an option to edit or delete these configurationWrapper objects.
type ensureStrategy interface {
	// Name of the strategy, for now we have two: 'mandatory' and 'suggested'.
	// This comes handy in logging.
	Name() string

	// ShouldUpdate accepts the current and the bootstrap configuration and determines
	// whether an update is necessary.
	// current is the existing in-cluster configuration object.
	// bootstrap is the configuration the kube-apiserver maintains in-memory.
	//
	// ok: true if auto update is required, otherwise false
	// object: the new object represents the new configuration to be stored in-cluster.
	// err: err is set when the function runs into an error and can not
	// determine if auto update is needed.
	ShouldUpdate(current, bootstrap configurationObject) (object runtime.Object, ok bool, err error)
}

// this internal interface provides abstraction for dealing with the `Spec`
// of both 'FlowSchema' and 'PriorityLevelConfiguration' objects.
// Since the ensure logic for both types is common, we use a few internal interfaces
// to abstract out the differences of these two types.
type specCopier interface {
	// HasSpecChanged returns true if the spec of both the bootstrap and
	// the current configuration object is same, otherwise false.
	HasSpecChanged(bootstrap, current runtime.Object) (bool, error)

	// CopySpec makes a deep copy the spec of the bootstrap object
	// and copies it to that of the current object.
	// CopySpec assumes that the current object is safe to mutate, so it
	// rests with the caller to make a deep copy of the current.
	CopySpec(bootstrap, current runtime.Object) error
}

// this internal interface provides abstraction for CRUD operation
// related to both 'FlowSchema' and 'PriorityLevelConfiguration' objects.
// Since the ensure logic for both types is common, we use a few internal interfaces
// to abstract out the differences of these two types.
type configurationClient interface {
	Create(object runtime.Object) (runtime.Object, error)
	Update(object runtime.Object) (runtime.Object, error)
	Get(name string) (configurationObject, error)
	Delete(name string) error
}

type configurationWrapper interface {
	// TypeName returns the type of the configuration that this interface deals with.
	// We use it to log the type name of the configuration object being ensured.
	// It is either 'PriorityLevelConfiguration' or 'FlowSchema'
	TypeName() string

	configurationClient
	specCopier
}

// A convenient wrapper interface that is used by the ensure logic.
type configurationObject interface {
	metav1.Object
	runtime.Object
}

func newSuggestedEnsureStrategy(copier specCopier) ensureStrategy {
	return &strategy{
		copier:               copier,
		alwaysAutoUpdateSpec: false,
		name:                 "suggested",
	}
}

func newMandatoryEnsureStrategy(copier specCopier) ensureStrategy {
	return &strategy{
		copier:               copier,
		alwaysAutoUpdateSpec: true,
		name:                 "mandatory",
	}
}

// auto-update strategy for the configuration objects
type strategy struct {
	copier               specCopier
	alwaysAutoUpdateSpec bool
	name                 string
}

func (s *strategy) Name() string {
	return s.name
}

func (s *strategy) ShouldUpdate(current, bootstrap configurationObject) (runtime.Object, bool, error) {
	if current == nil || bootstrap == nil {
		return nil, false, nil
	}

	autoUpdateSpec := s.alwaysAutoUpdateSpec
	if !autoUpdateSpec {
		autoUpdateSpec = shouldUpdateSpec(current)
	}
	updateAnnotation := shouldUpdateAnnotation(current, autoUpdateSpec)

	var specChanged bool
	if autoUpdateSpec {
		changed, err := s.copier.HasSpecChanged(bootstrap, current)
		if err != nil {
			return nil, false, fmt.Errorf("failed to compare spec - %w", err)
		}
		specChanged = changed
	}

	if !(updateAnnotation || specChanged) {
		// the annotation key is up to date and the spec has not changed, no update is necessary
		return nil, false, nil
	}

	// if we are here, either we need to update the annotation key or the spec.
	copy, ok := current.DeepCopyObject().(configurationObject)
	if !ok {
		// we should never be here
		return nil, false, errors.New("incompatible object type")
	}

	if updateAnnotation {
		setAutoUpdateAnnotation(copy, autoUpdateSpec)
	}
	if specChanged {
		s.copier.CopySpec(bootstrap, copy)
	}

	return copy, true, nil
}

// shouldUpdateSpec inspects the auto-update annotation key and generation field to determine
// whether the configurationWrapper object should be auto-updated.
func shouldUpdateSpec(accessor metav1.Object) bool {
	value, _ := accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]
	if autoUpdate, err := strconv.ParseBool(value); err == nil {
		return autoUpdate
	}

	// We are here because of either a or b:
	// a. the annotation key is missing.
	// b. the annotation key is present but the value does not represent a boolean.
	// In either case, if the operator hasn't changed the spec, we can safely auto update.
	// Please note that we can't protect the changes made by the operator in the following scenario:
	// - The operator deletes and recreates the same object with a variant spec (generation resets to 1).
	if accessor.GetGeneration() == 1 {
		return true
	}
	return false
}

// shouldUpdateAnnotation determines whether the current value of the auto-update annotation
// key matches the desired value.
func shouldUpdateAnnotation(accessor metav1.Object, desired bool) bool {
	if value, ok := accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]; ok {
		if current, err := strconv.ParseBool(value); err == nil && current == desired {
			return false
		}
	}

	return true
}

// setAutoUpdateAnnotation sets the auto-update annotation key to the specified value.
func setAutoUpdateAnnotation(accessor metav1.Object, autoUpdate bool) {
	if accessor.GetAnnotations() == nil {
		accessor.SetAnnotations(map[string]string{})
	}

	accessor.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey] = strconv.FormatBool(autoUpdate)
}

// ensureConfiguration ensures the boostrap configurationWrapper on the cluster based on the specified strategy.
func ensureConfiguration(wrapper configurationWrapper, strategy ensureStrategy, bootstrap configurationObject) error {
	name := bootstrap.GetName()
	configurationType := strategy.Name()

	var current configurationObject
	var err error
	for {
		current, err = wrapper.Get(bootstrap.GetName())
		if err == nil {
			break
		}
		if !apierrors.IsNotFound(err) {
			return fmt.Errorf("failed to retrieve %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
		}

		// we always re-create a missing configuration object
		if _, err = wrapper.Create(bootstrap); err == nil {
			klog.V(2).InfoS(fmt.Sprintf("Successfully created %s", wrapper.TypeName()), "type", configurationType, "name", name)
			return nil
		}

		if !apierrors.IsAlreadyExists(err) {
			return fmt.Errorf("cannot create %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
		}
		klog.V(5).InfoS(fmt.Sprintf("Something created the %s concurrently", wrapper.TypeName()), "type", configurationType, "name", name)
	}

	klog.V(5).InfoS(fmt.Sprintf("The %s already exists, checking whether it is up to date", wrapper.TypeName()), "type", configurationType, "name", name)
	newObject, update, err := strategy.ShouldUpdate(current, bootstrap)
	if err != nil {
		return fmt.Errorf("failed to determine whether auto-update is required for %s type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
	}
	if !update {
		if klog.V(5).Enabled() {
			// TODO: if we use structured logging here the diff gets escaped and very awkward to read in the log
			klog.Infof("No update required for the %s type=%s name=%q diff: %s", wrapper.TypeName(), configurationType, name, cmp.Diff(current, bootstrap))
		}
		return nil
	}

	if _, err = wrapper.Update(newObject); err == nil {
		klog.V(2).Infof("Updated the %s type=%s name=%q diff: %s", wrapper.TypeName(), configurationType, name, cmp.Diff(current, newObject))
		return nil
	}

	if apierrors.IsConflict(err) {
		klog.V(2).InfoS(fmt.Sprintf("Something updated the %s concurrently, I will check its spec later", wrapper.TypeName()), "type", configurationType, "name", name)
		return nil
	}

	return fmt.Errorf("failed to update the %s, will retry later type=%s name=%q error=%w", wrapper.TypeName(), configurationType, name, err)
}

func removeConfiguration(wrapper configurationWrapper, name string) error {
	current, err := wrapper.Get(name)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}

		return fmt.Errorf("failed to retrieve the %s, will retry later name=%q error=%w", wrapper.TypeName(), name, err)
	}

	value := current.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]
	autoUpdate, err := strconv.ParseBool(value)
	if err != nil {
		klog.ErrorS(err, fmt.Sprintf("Skipping deletion of the %s", wrapper.TypeName()), "name", name)

		// This may need manual intervention, in case the annotation value is malformed,
		// so don't return an error, that might trigger futile retry loop.
		return nil
	}
	if !autoUpdate {
		klog.V(5).InfoS(fmt.Sprintf("Skipping deletion of the %s", wrapper.TypeName()), "name", name)
		return nil
	}

	if err := wrapper.Delete(name); err != nil {
		if apierrors.IsNotFound(err) {
			klog.V(5).InfoS(fmt.Sprintf("Something concurrently deleted the %s", wrapper.TypeName()), "name", name)
			return nil
		}

		return fmt.Errorf("failed to delete the %s, will retry later name=%q error=%w", wrapper.TypeName(), name, err)
	}

	klog.V(2).InfoS(fmt.Sprintf("Successfully deleted the %s", wrapper.TypeName()), "name", name)
	return nil
}

// getRemoveCandidate returns a list of configuration objects we should delete
// from the cluster given a set of bootstrap and current configuration.
// bootstrap: a set of hard coded configuration kube-apiserver maintains in-memory.
// current: a set of configuration objects that exist on the cluster
// Any object present in current is a candidate for removal if both a and b are true:
//  a. the object in current is missing from the bootstrap configuration
//  b. the object has the designated auto-update annotation key
// This function shares the common logic for both FlowSchema and PriorityLevelConfiguration
// type and hence it accepts metav1.Object only.
func getRemoveCandidate(bootstrap sets.String, current []metav1.Object) []string {
	if len(current) == 0 {
		return nil
	}

	candidates := make([]string, 0)
	for i := range current {
		object := current[i]
		if _, ok := object.GetAnnotations()[flowcontrolv1beta2.AutoUpdateAnnotationKey]; !ok {
			// the configuration object does not have the annotation key
			continue
		}

		if _, ok := bootstrap[object.GetName()]; !ok {
			candidates = append(candidates, object.GetName())
		}
	}
	return candidates
}
<|endoftext|>


<commit_before>package migrations

import . "github.com/grafana/grafana/pkg/services/sqlstore/migrator"

func addTeamMigrations(mg *Migrator) {
	teamV1 := Table{
		Name: "team",
		Columns: []*Column{
			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
			{Name: "name", Type: DB_NVarchar, Length: 255, Nullable: false},
			{Name: "org_id", Type: DB_BigInt},
			{Name: "created", Type: DB_DateTime, Nullable: false},
			{Name: "updated", Type: DB_DateTime, Nullable: false},
		},
		Indices: []*Index{
			{Cols: []string{"org_id"}},
			{Cols: []string{"org_id", "name"}, Type: UniqueIndex},
		},
	}

	mg.AddMigration("create team table", NewAddTableMigration(teamV1))

	//-------  indexes ------------------
	mg.AddMigration("add index team.org_id", NewAddIndexMigration(teamV1, teamV1.Indices[0]))
	mg.AddMigration("add unique index team_org_id_name", NewAddIndexMigration(teamV1, teamV1.Indices[1]))

	teamMemberV1 := Table{
		Name: "team_member",
		Columns: []*Column{
			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
			{Name: "org_id", Type: DB_BigInt},
			{Name: "team_id", Type: DB_BigInt},
			{Name: "user_id", Type: DB_BigInt},
			{Name: "created", Type: DB_DateTime, Nullable: false},
			{Name: "updated", Type: DB_DateTime, Nullable: false},
		},
		Indices: []*Index{
			{Cols: []string{"org_id"}},
			{Cols: []string{"org_id", "team_id", "user_id"}, Type: UniqueIndex},
		},
	}

	mg.AddMigration("create team member table", NewAddTableMigration(teamMemberV1))

	//-------  indexes ------------------
	mg.AddMigration("add index team_member.org_id", NewAddIndexMigration(teamMemberV1, teamMemberV1.Indices[0]))
	mg.AddMigration("add unique index team_member_org_id_team_id_user_id", NewAddIndexMigration(teamMemberV1, teamMemberV1.Indices[1]))
}
<commit_msg>fix: reduced team name column length, fixes #10244<commit_after>package migrations

import . "github.com/grafana/grafana/pkg/services/sqlstore/migrator"

func addTeamMigrations(mg *Migrator) {
	teamV1 := Table{
		Name: "team",
		Columns: []*Column{
			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
			{Name: "name", Type: DB_NVarchar, Length: 190, Nullable: false},
			{Name: "org_id", Type: DB_BigInt},
			{Name: "created", Type: DB_DateTime, Nullable: false},
			{Name: "updated", Type: DB_DateTime, Nullable: false},
		},
		Indices: []*Index{
			{Cols: []string{"org_id"}},
			{Cols: []string{"org_id", "name"}, Type: UniqueIndex},
		},
	}

	mg.AddMigration("create team table", NewAddTableMigration(teamV1))

	//-------  indexes ------------------
	mg.AddMigration("add index team.org_id", NewAddIndexMigration(teamV1, teamV1.Indices[0]))
	mg.AddMigration("add unique index team_org_id_name", NewAddIndexMigration(teamV1, teamV1.Indices[1]))

	teamMemberV1 := Table{
		Name: "team_member",
		Columns: []*Column{
			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
			{Name: "org_id", Type: DB_BigInt},
			{Name: "team_id", Type: DB_BigInt},
			{Name: "user_id", Type: DB_BigInt},
			{Name: "created", Type: DB_DateTime, Nullable: false},
			{Name: "updated", Type: DB_DateTime, Nullable: false},
		},
		Indices: []*Index{
			{Cols: []string{"org_id"}},
			{Cols: []string{"org_id", "team_id", "user_id"}, Type: UniqueIndex},
		},
	}

	mg.AddMigration("create team member table", NewAddTableMigration(teamMemberV1))

	//-------  indexes ------------------
	mg.AddMigration("add index team_member.org_id", NewAddIndexMigration(teamMemberV1, teamMemberV1.Indices[0]))
	mg.AddMigration("add unique index team_member_org_id_team_id_user_id", NewAddIndexMigration(teamMemberV1, teamMemberV1.Indices[1]))
}
<|endoftext|>


<commit_before>package apple

import (
	"github.com/ivan1993spb/snake-server/engine"
	"github.com/ivan1993spb/snake-server/world"
)

type Apple struct {
	world    *world.World
	location engine.Location
}

type ErrCreateApple string

func (e ErrCreateApple) Error() string {
	return "cannot create apple: " + string(e)
}

// NewApple creates and locates new apple
func NewApple(world *world.World) (*Apple, error) {
	apple := &Apple{}

	location, err := world.CreateObjectRandomDot(apple)
	if err != nil {
		return nil, ErrCreateApple(err.Error())
	}

	apple.location = location
	apple.world = world

	return apple, nil
}

func (a *Apple) NutritionalValue(dot *engine.Dot) uint16 {
	if a.location.Equals(engine.Location{dot}) {
		// TODO: Handle error.
		a.world.DeleteObject(a, a.location)
		return 1
	}

	return 0
}
<commit_msg>Impliment Stringer interface in apple.Apple<commit_after>package apple

import (
	"fmt"

	"github.com/ivan1993spb/snake-server/engine"
	"github.com/ivan1993spb/snake-server/world"
)

type Apple struct {
	world    *world.World
	location engine.Location
}

type ErrCreateApple string

func (e ErrCreateApple) Error() string {
	return "cannot create apple: " + string(e)
}

// NewApple creates and locates new apple
func NewApple(world *world.World) (*Apple, error) {
	apple := &Apple{}

	location, err := world.CreateObjectRandomDot(apple)
	if err != nil {
		return nil, ErrCreateApple(err.Error())
	}

	apple.location = location
	apple.world = world

	return apple, nil
}

func (a *Apple) String() string {
	return fmt.Sprintf("apple %s", a.location)
}

func (a *Apple) NutritionalValue(dot *engine.Dot) uint16 {
	if a.location.Equals(engine.Location{dot}) {
		// TODO: Handle error.
		a.world.DeleteObject(a, a.location)
		return 1
	}

	return 0
}
<|endoftext|>


<commit_before>/**
 * Ant Worker Project
 *
 * Copyright (c) 2015 Epinion Online Research Team
 *
 * --------------------------------------------------------------------
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Affero General Public License
 * as published by the Free Software Foundation, either version 3
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public
 * License along with this program.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 * --------------------------------------------------------------------
 *
 * Author:
 *     Jerry Pham       <jerry@andjerry.com>
 *     Loi Nguyen       <loint@penlook.com>
 */

package main

import (
	"fmt"
	"testing"
	"github.com/stretchr/testify/assert"
	"path/filepath"
	"os"
	"log"
	"os/exec"
	"io/ioutil"
	"net/http"
)

func system(sudo bool, cmd string, arg string) string {
	var out string
	if (sudo) {
		if out, err := exec.Command("/bin/sh", "-c", "sudo " + cmd + " " + arg).Output(); err != nil {
			fmt.Fprintln(os.Stderr, "There was an error running git rev-parse command: ", err)
			fmt.Fprintln(os.Stderr, "sudo " + cmd + " " + arg, err)
			fmt.Println(out)
			os.Exit(1)
		}
	} else {
		if out, err := exec.Command(cmd, arg).Output(); err != nil {
			fmt.Fprintln(os.Stderr, "There was an error running git rev-parse command: ", err)
			fmt.Println(out)
			fmt.Fprintln(os.Stderr, cmd, err)
			os.Exit(1)
		}
	}
	return string(out)
}

func cwd() string {
	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		log.Fatal(err)
	}
	return string(dir)
}

func setUp() {
	pwd := os.Getenv("PWD")
	os.Chdir(pwd)
	system(false, "go", "build")
	system(false, "cp", "-f " + pwd + "/ant-worker /usr/bin/ant-worker")
	system(true, "chmod", "777 /usr/bin/ant-worker")
	system(true, "ant-worker", "install")
	system(true, "ant-worker", "start")
}

func tearDown() {
	system(true, "ant-worker", "stop")
	system(true, "ant-worker", "uninstall")
}

func TestService(t *testing.T) {
	assert := assert.New(t)
	setUp()
	defer tearDown()
	response, err := http.Get("http://localhost:2345/test")
	assert.Equal(true, err == nil)
	if err == nil {
		defer response.Body.Close()
		contents, err := ioutil.ReadAll(response.Body)
		assert.Equal(true, err == nil)
		assert.Equal("{\"status\":\"hello world\"}\n", string(contents))
	}
}<commit_msg>Reverse copy function and remove cp command<commit_after>/**
 * Ant Worker Project
 *
 * Copyright (c) 2015 Epinion Online Research Team
 *
 * --------------------------------------------------------------------
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Affero General Public License
 * as published by the Free Software Foundation, either version 3
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public
 * License along with this program.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 * --------------------------------------------------------------------
 *
 * Author:
 *     Jerry Pham       <jerry@andjerry.com>
 *     Loi Nguyen       <loint@penlook.com>
 */

package main

import (
	"fmt"
	"testing"
	"github.com/stretchr/testify/assert"
	"path/filepath"
	"os"
	"log"
	"os/exec"
	"io"
	"io/ioutil"
	"net/http"
)

func system(sudo bool, cmd string, arg string) string {
	var out string
	if (sudo) {
		if out, err := exec.Command("/bin/sh", "-c", "sudo " + cmd + " " + arg).Output(); err != nil {
			fmt.Fprintln(os.Stderr, "There was an error running git rev-parse command: ", err)
			fmt.Println(out)
			os.Exit(1)
		}
	} else {
		if out, err := exec.Command(cmd, arg).Output(); err != nil {
			fmt.Fprintln(os.Stderr, "There was an error running git rev-parse command: ", err)
			fmt.Println(out)
			os.Exit(1)
		}
	}
	return string(out)
}

func cwd() string {
	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err != nil {
		log.Fatal(err)
	}
	return string(dir)
}

func copy(target string, dest string) {
	r, err := os.Open(target)
	if err != nil {
		panic(err)
	}
	defer r.Close()

     w, err := os.Create(dest)
     if err != nil {
         panic(err)
     }
     defer w.Close()

     // do the actual work
     _, err = io.Copy(w, r)
     if err != nil {
         panic(err)
     }
}

func setUp() {
	pwd := os.Getenv("PWD")
	os.Chdir(pwd)
	system(false, "go", "build")
	copy(pwd + "/ant-worker", "/usr/bin/ant-worker")
	system(true, "chmod", "777 /usr/bin/ant-worker")
	system(true, "ant-worker", "install")
	system(true, "ant-worker", "start")
}

func tearDown() {
	system(true, "ant-worker", "stop")
	system(true, "ant-worker", "uninstall")
}

func TestService(t *testing.T) {
	assert := assert.New(t)
	setUp()
	defer tearDown()
	response, err := http.Get("http://localhost:2345/test")
	assert.Equal(true, err == nil)
	if err == nil {
		defer response.Body.Close()
		contents, err := ioutil.ReadAll(response.Body)
		assert.Equal(true, err == nil)
		assert.Equal("{\"status\":\"hello world\"}\n", string(contents))
	}
}<|endoftext|>


<commit_before>package main

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

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

type RequestLogger struct{}

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter() *HostRouter {
	return &HostRouter{}
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var err error
	host := req.Host
	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.adminHandler(w, req)
}

func (r *HostRouter) adminHandler(w http.ResponseWriter, req *http.Request) {
	r.Lock()
	defer r.Unlock()

	if Registry.VHostsLen() == 0 {
		http.Error(w, "no backends available", http.StatusServiceUnavailable)
		return
	}

	// TODO: better status lines
	stats := Registry.Stats()
	for _, svc := range stats {
		if len(svc.VirtualHosts) == 0 {
			continue
		}
		fmt.Fprintf(w, "%v\n", svc.VirtualHosts)
		for _, b := range svc.Backends {
			js, _ := json.Marshal(b)
			fmt.Fprintf(w, "\t%s\n", string(js))
		}
	}

	fmt.Fprintf(w, "\n")
	return
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()

	log.Printf("HTTP server listening at %s", listenAddr)

	// Proxy acts as http handler:
	r.server = &http.Server{
		Addr:           listenAddr,
		Handler:        r,
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	var err error
	r.listener, err = net.Listen("tcp", listenAddr)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	r.Unlock()
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(r.listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer() {
	//FIXME: this global wg?
	defer wg.Done()
	httpRouter = NewHostRouter()
	httpRouter.Start(nil)
}

func sslRedirect(pr *ProxyRequest) bool {
	pr.Request.Header.Set("X-Request-Id", genId())

	if sslOnly && pr.Request.Header.Get("X-Forwarded-Proto") != "https" {
		//TODO: verify RequestURI
		redirLoc := "https://" + pr.Request.Host + pr.Request.RequestURI
		http.Redirect(pr.ResponseWriter, pr.Request, redirLoc, http.StatusMovedPermanently)
		return false
	}

	return true
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string

	duration := pr.FinishTime.Sub(pr.StartTime)

	if pr.Request != nil {
		id = pr.Request.Header.Get("X-Request-Id")
		method = pr.Request.Method
		clientIP = pr.Request.RemoteAddr
		url = pr.Request.Host + pr.Request.RequestURI
		agent = pr.Request.UserAgent()
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, duration, agent, err)
	return true
}
<commit_msg>Add the status code back into the shuttle http log<commit_after>package main

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

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

type RequestLogger struct{}

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter() *HostRouter {
	return &HostRouter{}
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var err error
	host := req.Host
	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.adminHandler(w, req)
}

func (r *HostRouter) adminHandler(w http.ResponseWriter, req *http.Request) {
	r.Lock()
	defer r.Unlock()

	if Registry.VHostsLen() == 0 {
		http.Error(w, "no backends available", http.StatusServiceUnavailable)
		return
	}

	// TODO: better status lines
	stats := Registry.Stats()
	for _, svc := range stats {
		if len(svc.VirtualHosts) == 0 {
			continue
		}
		fmt.Fprintf(w, "%v\n", svc.VirtualHosts)
		for _, b := range svc.Backends {
			js, _ := json.Marshal(b)
			fmt.Fprintf(w, "\t%s\n", string(js))
		}
	}

	fmt.Fprintf(w, "\n")
	return
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()

	log.Printf("HTTP server listening at %s", listenAddr)

	// Proxy acts as http handler:
	r.server = &http.Server{
		Addr:           listenAddr,
		Handler:        r,
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	var err error
	r.listener, err = net.Listen("tcp", listenAddr)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	r.Unlock()
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(r.listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer() {
	//FIXME: this global wg?
	defer wg.Done()
	httpRouter = NewHostRouter()
	httpRouter.Start(nil)
}

func sslRedirect(pr *ProxyRequest) bool {
	pr.Request.Header.Set("X-Request-Id", genId())

	if sslOnly && pr.Request.Header.Get("X-Forwarded-Proto") != "https" {
		//TODO: verify RequestURI
		redirLoc := "https://" + pr.Request.Host + pr.Request.RequestURI
		http.Redirect(pr.ResponseWriter, pr.Request, redirLoc, http.StatusMovedPermanently)
		return false
	}

	return true
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string
	var status int

	duration := pr.FinishTime.Sub(pr.StartTime)

	if pr.Request != nil {
		id = pr.Request.Header.Get("X-Request-Id")
		method = pr.Request.Method
		clientIP = pr.Request.RemoteAddr
		url = pr.Request.Host + pr.Request.RequestURI
		agent = pr.Request.UserAgent()
		status = pr.Response.StatusCode
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s status=%d duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, status, duration, agent, err)
	return true
}
<|endoftext|>