<commit_before>package mesos

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/CiscoCloud/mesos-consul/config"
	"github.com/CiscoCloud/mesos-consul/consul"
	"github.com/CiscoCloud/mesos-consul/registry"

	consulapi "github.com/hashicorp/consul/api"
)

type CacheEntry struct {
	service      *consulapi.AgentServiceRegistration
	isRegistered bool
}

type Mesos struct {
	Registry registry.Registry
	Masters  *[]MesosHost
	Lock     sync.Mutex
}

func New(c *config.Config) *Mesos {
	m := new(Mesos)

	if c.Zk == "" {
		return nil
	}

	if consul.IsEnabled() {
		m.Registry = consul.New()
	}

	if m.Registry == nil {
		log.Fatal("[ERROR] No registry specified")
	}

	if m.Registry.CacheCreate() {
		m.LoadCache()
	}

	m.zkDetector(c.Zk)

	return m
}

func (m *Mesos) Refresh() error {
	sj, err := m.loadState()
	if err != nil {
		log.Print("[ERROR] No master")
		return err
	}

	if sj.Leader == "" {
		return errors.New("Empty master")
	}

	m.parseState(sj)

	return nil
}

func (m *Mesos) loadState() (StateJSON, error) {
	var err error
	var sj StateJSON

	defer func() {
		if rec := recover(); rec != nil {
			err = errors.New("can't connect to Mesos")
		}
	}()

	ip, port := m.getLeader()
	if ip == "" {
		return sj, errors.New("No master in zookeeper")
	}

	log.Printf("[INFO] Zookeeper leader: %s:%s", ip, port)

	log.Print("[INFO] reloading from master ", ip)
	sj = m.loadFromMaster(ip, port)

	if rip := leaderIP(sj.Leader); rip != ip {
		log.Print("[WARN] master changed to ", rip)
		sj = m.loadFromMaster(rip, port)
	}

	return sj, err
}

func (m *Mesos) loadFromMaster(ip string, port string) (sj StateJSON) {
	url := "http://" + ip + ":" + port + "/master/state.json"

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	err = json.Unmarshal(body, &sj)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	return sj
}

func (m *Mesos) parseState(sj StateJSON) {
	log.Print("[INFO] Running parseState")

	m.RegisterHosts(sj)
	log.Print("[DEBUG] Done running RegisterHosts")

	for _, fw := range sj.Frameworks {
		for _, task := range fw.Tasks {
			host, err := sj.GetFollowerById(task.FollowerId)
			if err == nil && task.State == "TASK_RUNNING" {
				m.registerTask(&task, host)

			}
		}
	}

	// Remove completed tasks
	m.Registry.Deregister()
}

func yankPorts(ports string) []int {
	rhs := strings.Split(ports, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []int{}

	mports := strings.Split(lhs, ",")
	for _, mport := range mports {
		pz := strings.Split(strings.TrimSpace(mport), "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, t)
		}
	}

	return yports
}
<commit_msg>Load cache in the proper location<commit_after>package mesos

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/CiscoCloud/mesos-consul/config"
	"github.com/CiscoCloud/mesos-consul/consul"
	"github.com/CiscoCloud/mesos-consul/registry"

	consulapi "github.com/hashicorp/consul/api"
)

type CacheEntry struct {
	service      *consulapi.AgentServiceRegistration
	isRegistered bool
}

type Mesos struct {
	Registry registry.Registry
	Masters  *[]MesosHost
	Lock     sync.Mutex
}

func New(c *config.Config) *Mesos {
	m := new(Mesos)

	if c.Zk == "" {
		return nil
	}

	if consul.IsEnabled() {
		m.Registry = consul.New()
	}

	if m.Registry == nil {
		log.Fatal("[ERROR] No registry specified")
	}

	m.zkDetector(c.Zk)

	return m
}

func (m *Mesos) Refresh() error {
	sj, err := m.loadState()
	if err != nil {
		log.Print("[ERROR] No master")
		return err
	}

	if sj.Leader == "" {
		return errors.New("Empty master")
	}

	if m.Registry.CacheCreate() {
		m.LoadCache()
	}

	m.parseState(sj)

	return nil
}

func (m *Mesos) loadState() (StateJSON, error) {
	var err error
	var sj StateJSON

	defer func() {
		if rec := recover(); rec != nil {
			err = errors.New("can't connect to Mesos")
		}
	}()

	ip, port := m.getLeader()
	if ip == "" {
		return sj, errors.New("No master in zookeeper")
	}

	log.Printf("[INFO] Zookeeper leader: %s:%s", ip, port)

	log.Print("[INFO] reloading from master ", ip)
	sj = m.loadFromMaster(ip, port)

	if rip := leaderIP(sj.Leader); rip != ip {
		log.Print("[WARN] master changed to ", rip)
		sj = m.loadFromMaster(rip, port)
	}

	return sj, err
}

func (m *Mesos) loadFromMaster(ip string, port string) (sj StateJSON) {
	url := "http://" + ip + ":" + port + "/master/state.json"

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	err = json.Unmarshal(body, &sj)
	if err != nil {
		log.Fatal("[ERROR] ", err)
	}

	return sj
}

func (m *Mesos) parseState(sj StateJSON) {
	log.Print("[INFO] Running parseState")

	m.RegisterHosts(sj)
	log.Print("[DEBUG] Done running RegisterHosts")

	for _, fw := range sj.Frameworks {
		for _, task := range fw.Tasks {
			host, err := sj.GetFollowerById(task.FollowerId)
			if err == nil && task.State == "TASK_RUNNING" {
				m.registerTask(&task, host)

			}
		}
	}

	// Remove completed tasks
	m.Registry.Deregister()
}

func yankPorts(ports string) []int {
	rhs := strings.Split(ports, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []int{}

	mports := strings.Split(lhs, ",")
	for _, mport := range mports {
		pz := strings.Split(strings.TrimSpace(mport), "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, t)
		}
	}

	return yports
}
<|endoftext|>


<commit_before>package mgmt

import (
	"net/http"

	"github.com/nyaxt/otaru/webui"

	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/mux"
	"github.com/rs/cors"
)

type Server struct {
	rtr     *mux.Router
	apirtr  *mux.Router
	httpsrv *http.Server
}

func NewServer() *Server {
	rtr := mux.NewRouter()

	rtr.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte("OK"))
	})

	apirtr := rtr.PathPrefix("/api").Subrouter()

	rtr.NotFoundHandler = http.FileServer(
		&assetfs.AssetFS{Asset: webui.Asset, AssetDir: webui.AssetDir, Prefix: "dist"})

	c := cors.New(cors.Options{
		AllowedOrigins: []string{"http://localhost:9000"}, // gulp devsrv
	})

	httpsrv := &http.Server{
		Addr:    ":10246",
		Handler: c.Handler(rtr),
	}
	return &Server{rtr: rtr, apirtr: apirtr, httpsrv: httpsrv}
}

func (srv *Server) APIRouter() *mux.Router { return srv.apirtr }

func (srv *Server) Run() error {
	if err := srv.httpsrv.ListenAndServe(); err != nil {
		return err
	}
	return nil
}
<commit_msg>fix 404 handler<commit_after>package mgmt

import (
	"net/http"
	"os"

	"github.com/nyaxt/otaru/webui"

	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/mux"
	"github.com/rs/cors"
)

type Server struct {
	rtr     *mux.Router
	apirtr  *mux.Router
	httpsrv *http.Server
}

func NewServer() *Server {
	rtr := mux.NewRouter()

	rtr.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte("OK"))
	})
	assetDirWrap := func(path string) ([]string, error) {
		children, err := webui.AssetDir(path)
		if err != nil {
			return nil, os.ErrNotExist
		}
		return children, nil
	}
	rtr.NotFoundHandler = http.FileServer(
		&assetfs.AssetFS{Asset: webui.Asset, AssetDir: assetDirWrap, Prefix: "dist"})

	apirtr := rtr.PathPrefix("/api").Subrouter()
	apirtr.NotFoundHandler = http.NotFoundHandler()

	c := cors.New(cors.Options{
		AllowedOrigins: []string{"http://localhost:9000"}, // gulp devsrv
	})

	httpsrv := &http.Server{
		Addr:    ":10246",
		Handler: c.Handler(rtr),
	}
	return &Server{rtr: rtr, apirtr: apirtr, httpsrv: httpsrv}
}

func (srv *Server) APIRouter() *mux.Router { return srv.apirtr }

func (srv *Server) Run() error {
	if err := srv.httpsrv.ListenAndServe(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package cluster

import (
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/rqlite/rqlite/command"
	"github.com/rqlite/rqlite/command/encoding"
)

const shortWait = 1 * time.Second
const longWait = 5 * time.Second

func Test_ServiceExecute(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// Ready for Execute tests now.
	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		return nil, errors.New("execute failed")
	}
	_, err := c.Execute(executeRequestFromString("some SQL"), s.Addr(), nil, longWait)
	if err == nil {
		t.Fatalf("client failed to report error")
	}
	if err.Error() != "execute failed" {
		t.Fatalf("incorrect error message received, got: %s", err.Error())
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		result := &command.ExecuteResult{
			LastInsertId: 1234,
			RowsAffected: 5678,
		}
		return []*command.ExecuteResult{result}, nil
	}
	res, err := c.Execute(executeRequestFromString("some SQL"), s.Addr(), nil, longWait)
	if err != nil {
		t.Fatalf("failed to execute query: %s", err.Error())
	}
	if exp, got := `[{"last_insert_id":1234,"rows_affected":5678}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for execute, expected %s, got %s", exp, got)
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		result := &command.ExecuteResult{
			Error: "no such table",
		}
		return []*command.ExecuteResult{result}, nil
	}
	res, err = c.Execute(executeRequestFromString("some SQL"), s.Addr(), nil, longWait)
	if err != nil {
		t.Fatalf("failed to execute: %s", err.Error())
	}
	if exp, got := `[{"error":"no such table"}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for execute, expected %s, got %s", exp, got)
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		time.Sleep(longWait)
		return nil, nil
	}
	_, err = c.Execute(executeRequestFromString("some SQL"), s.Addr(), nil, shortWait)
	if err == nil {
		t.Fatalf("failed to receive expected error")
	}
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("failed to receive expected error, got: %T %s", err, err)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

func Test_ServiceQuery(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// Ready for Query tests now.
	db.queryFn = func(er *command.QueryRequest) ([]*command.QueryRows, error) {
		if er.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL query received")
		}
		return nil, errors.New("query failed")
	}
	_, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), nil, longWait)
	if err == nil {
		t.Fatalf("client failed to report error")
	}
	if err.Error() != "query failed" {
		t.Fatalf("incorrect error message received, got: %s", err.Error())
	}

	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		if qr.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL statement received")
		}
		rows := &command.QueryRows{
			Columns: []string{"c1", "c2"},
			Types:   []string{"t1", "t2"},
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), nil, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := `[{"columns":["c1","c2"],"types":["t1","t2"]}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		if qr.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL statement received")
		}
		rows := &command.QueryRows{
			Error: "no such table",
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err = c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), nil, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := `[{"error":"no such table"}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	db.queryFn = func(er *command.QueryRequest) ([]*command.QueryRows, error) {
		time.Sleep(longWait)
		return nil, nil
	}
	_, err = c.Query(queryRequestFromString("some SQL"), s.Addr(), nil, shortWait)
	if err == nil {
		t.Fatalf("failed to receive expected error")
	}
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("failed to receive expected error, got: %T %s", err, err)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

// Test_ServiceQueryLarge ensures that query responses larger than 64K are
// encoded and decoded correctly.
func Test_ServiceQueryLarge(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	var b strings.Builder
	for i := 0; i < 100000; i++ {
		b.WriteString("bar")
	}
	if b.Len() < 64000 {
		t.Fatalf("failed to generate a large enough string for test")
	}

	// Ready for Query tests now.
	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		parameter := &command.Parameter{
			Value: &command.Parameter_S{
				S: b.String(),
			},
		}
		value := &command.Values{
			Parameters: []*command.Parameter{parameter},
		}

		rows := &command.QueryRows{
			Columns: []string{"c1"},
			Types:   []string{"t1"},
			Values:  []*command.Values{value},
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), nil, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := fmt.Sprintf(`[{"columns":["c1"],"types":["t1"],"values":[["%s"]]}]`, b.String()), asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

// Test_BinaryEncoding_Backwards ensures that software earlier than v6.6.2
// can communicate with v6.6.2+ releases. v6.6.2 increased the maximum size
// of cluster responses.
func Test_BinaryEncoding_Backwards(t *testing.T) {
	// Simulate a 6.6.2 encoding a response size less than 16 bits.
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b[0:], uint32(1234))

	// Can older software decode it OK?
	if binary.LittleEndian.Uint16(b[0:]) != 1234 {
		t.Fatal("failed to read correct value")
	}
}

func executeRequestFromString(s string) *command.ExecuteRequest {
	return executeRequestFromStrings([]string{s})
}

// queryRequestFromStrings converts a slice of strings into a command.ExecuteRequest
func executeRequestFromStrings(s []string) *command.ExecuteRequest {
	stmts := make([]*command.Statement, len(s))
	for i := range s {
		stmts[i] = &command.Statement{
			Sql: s[i],
		}
	}
	return &command.ExecuteRequest{
		Request: &command.Request{
			Statements:  stmts,
			Transaction: false,
		},
		Timings: false,
	}
}

func queryRequestFromString(s string) *command.QueryRequest {
	return queryRequestFromStrings([]string{s})
}

// queryRequestFromStrings converts a slice of strings into a command.QueryRequest
func queryRequestFromStrings(s []string) *command.QueryRequest {
	stmts := make([]*command.Statement, len(s))
	for i := range s {
		stmts[i] = &command.Statement{
			Sql: s[i],
		}
	}
	return &command.QueryRequest{
		Request: &command.Request{
			Statements:  stmts,
			Transaction: false,
		},
		Timings: false,
	}
}

func asJSON(v interface{}) string {
	b, err := encoding.JSONMarshal(v)
	if err != nil {
		panic(fmt.Sprintf("failed to JSON marshal value: %s", err.Error()))
	}
	return string(b)
}
<commit_msg>More NO_CREDS use<commit_after>package cluster

import (
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/rqlite/rqlite/command"
	"github.com/rqlite/rqlite/command/encoding"
)

const shortWait = 1 * time.Second
const longWait = 5 * time.Second

var (
	NO_CREDS = (*Credentials)(nil)
)

func Test_ServiceExecute(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// Ready for Execute tests now.
	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		return nil, errors.New("execute failed")
	}
	_, err := c.Execute(executeRequestFromString("some SQL"), s.Addr(), NO_CREDS, longWait)
	if err == nil {
		t.Fatalf("client failed to report error")
	}
	if err.Error() != "execute failed" {
		t.Fatalf("incorrect error message received, got: %s", err.Error())
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		result := &command.ExecuteResult{
			LastInsertId: 1234,
			RowsAffected: 5678,
		}
		return []*command.ExecuteResult{result}, nil
	}
	res, err := c.Execute(executeRequestFromString("some SQL"), s.Addr(), NO_CREDS, longWait)
	if err != nil {
		t.Fatalf("failed to execute query: %s", err.Error())
	}
	if exp, got := `[{"last_insert_id":1234,"rows_affected":5678}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for execute, expected %s, got %s", exp, got)
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		if er.Request.Statements[0].Sql != "some SQL" {
			t.Fatalf("incorrect SQL statement received")
		}
		result := &command.ExecuteResult{
			Error: "no such table",
		}
		return []*command.ExecuteResult{result}, nil
	}
	res, err = c.Execute(executeRequestFromString("some SQL"), s.Addr(), NO_CREDS, longWait)
	if err != nil {
		t.Fatalf("failed to execute: %s", err.Error())
	}
	if exp, got := `[{"error":"no such table"}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for execute, expected %s, got %s", exp, got)
	}

	db.executeFn = func(er *command.ExecuteRequest) ([]*command.ExecuteResult, error) {
		time.Sleep(longWait)
		return nil, nil
	}
	_, err = c.Execute(executeRequestFromString("some SQL"), s.Addr(), NO_CREDS, shortWait)
	if err == nil {
		t.Fatalf("failed to receive expected error")
	}
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("failed to receive expected error, got: %T %s", err, err)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

func Test_ServiceQuery(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	// Ready for Query tests now.
	db.queryFn = func(er *command.QueryRequest) ([]*command.QueryRows, error) {
		if er.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL query received")
		}
		return nil, errors.New("query failed")
	}
	_, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), NO_CREDS, longWait)
	if err == nil {
		t.Fatalf("client failed to report error")
	}
	if err.Error() != "query failed" {
		t.Fatalf("incorrect error message received, got: %s", err.Error())
	}

	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		if qr.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL statement received")
		}
		rows := &command.QueryRows{
			Columns: []string{"c1", "c2"},
			Types:   []string{"t1", "t2"},
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), NO_CREDS, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := `[{"columns":["c1","c2"],"types":["t1","t2"]}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		if qr.Request.Statements[0].Sql != "SELECT * FROM foo" {
			t.Fatalf("incorrect SQL statement received")
		}
		rows := &command.QueryRows{
			Error: "no such table",
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err = c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), NO_CREDS, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := `[{"error":"no such table"}]`, asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	db.queryFn = func(er *command.QueryRequest) ([]*command.QueryRows, error) {
		time.Sleep(longWait)
		return nil, nil
	}
	_, err = c.Query(queryRequestFromString("some SQL"), s.Addr(), NO_CREDS, shortWait)
	if err == nil {
		t.Fatalf("failed to receive expected error")
	}
	if !errors.Is(err, os.ErrDeadlineExceeded) {
		t.Fatalf("failed to receive expected error, got: %T %s", err, err)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

// Test_ServiceQueryLarge ensures that query responses larger than 64K are
// encoded and decoded correctly.
func Test_ServiceQueryLarge(t *testing.T) {
	ln, mux := mustNewMux()
	go mux.Serve()
	tn := mux.Listen(1) // Could be any byte value.
	db := mustNewMockDatabase()
	cred := mustNewMockCredentialStore()
	s := New(tn, db, cred)
	if s == nil {
		t.Fatalf("failed to create cluster service")
	}

	c := NewClient(mustNewDialer(1, false, false), 30*time.Second)

	if err := s.Open(); err != nil {
		t.Fatalf("failed to open cluster service: %s", err.Error())
	}

	var b strings.Builder
	for i := 0; i < 100000; i++ {
		b.WriteString("bar")
	}
	if b.Len() < 64000 {
		t.Fatalf("failed to generate a large enough string for test")
	}

	// Ready for Query tests now.
	db.queryFn = func(qr *command.QueryRequest) ([]*command.QueryRows, error) {
		parameter := &command.Parameter{
			Value: &command.Parameter_S{
				S: b.String(),
			},
		}
		value := &command.Values{
			Parameters: []*command.Parameter{parameter},
		}

		rows := &command.QueryRows{
			Columns: []string{"c1"},
			Types:   []string{"t1"},
			Values:  []*command.Values{value},
		}
		return []*command.QueryRows{rows}, nil
	}
	res, err := c.Query(queryRequestFromString("SELECT * FROM foo"), s.Addr(), NO_CREDS, longWait)
	if err != nil {
		t.Fatalf("failed to query: %s", err.Error())
	}
	if exp, got := fmt.Sprintf(`[{"columns":["c1"],"types":["t1"],"values":[["%s"]]}]`, b.String()), asJSON(res); exp != got {
		t.Fatalf("unexpected results for query, expected %s, got %s", exp, got)
	}

	// Clean up resources.
	if err := ln.Close(); err != nil {
		t.Fatalf("failed to close Mux's listener: %s", err)
	}
	if err := s.Close(); err != nil {
		t.Fatalf("failed to close cluster service")
	}
}

// Test_BinaryEncoding_Backwards ensures that software earlier than v6.6.2
// can communicate with v6.6.2+ releases. v6.6.2 increased the maximum size
// of cluster responses.
func Test_BinaryEncoding_Backwards(t *testing.T) {
	// Simulate a 6.6.2 encoding a response size less than 16 bits.
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b[0:], uint32(1234))

	// Can older software decode it OK?
	if binary.LittleEndian.Uint16(b[0:]) != 1234 {
		t.Fatal("failed to read correct value")
	}
}

func executeRequestFromString(s string) *command.ExecuteRequest {
	return executeRequestFromStrings([]string{s})
}

// queryRequestFromStrings converts a slice of strings into a command.ExecuteRequest
func executeRequestFromStrings(s []string) *command.ExecuteRequest {
	stmts := make([]*command.Statement, len(s))
	for i := range s {
		stmts[i] = &command.Statement{
			Sql: s[i],
		}
	}
	return &command.ExecuteRequest{
		Request: &command.Request{
			Statements:  stmts,
			Transaction: false,
		},
		Timings: false,
	}
}

func queryRequestFromString(s string) *command.QueryRequest {
	return queryRequestFromStrings([]string{s})
}

// queryRequestFromStrings converts a slice of strings into a command.QueryRequest
func queryRequestFromStrings(s []string) *command.QueryRequest {
	stmts := make([]*command.Statement, len(s))
	for i := range s {
		stmts[i] = &command.Statement{
			Sql: s[i],
		}
	}
	return &command.QueryRequest{
		Request: &command.Request{
			Statements:  stmts,
			Transaction: false,
		},
		Timings: false,
	}
}

func asJSON(v interface{}) string {
	b, err := encoding.JSONMarshal(v)
	if err != nil {
		panic(fmt.Sprintf("failed to JSON marshal value: %s", err.Error()))
	}
	return string(b)
}
<|endoftext|>


<commit_before>package miggle

import (
	"fmt"
	"github.com/gorilla/mux"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
)

func init() {
	http.Handle("/", Router())
}

func Router() *mux.Router {
	var res *Resource = NewResource()
	r := mux.NewRouter()

	r.HandleFunc("/test", res.RequestHandler(new(EndPoint)))
	r.HandleFunc("/get/status/{status:[0-9]+}", res.RequestHandler(new(EndPointStatus)))

	return r
}

type EndPoint struct{}

func (this *EndPoint) Init(request *http.Request) error {
	return nil
}

func (this *EndPoint) Get(request *http.Request) (int, interface{}, http.Header) {
	return 200, "data", nil
}

func TestBasicGet(t *testing.T) {
	r, err := http.NewRequest("GET", "/test", nil)
	if err != nil {
		t.Error(err)
	}
	w := httptest.NewRecorder()
	Router().ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Error("Response status code does not equal")
	}

	body, _ := ioutil.ReadAll(w.Body)

	if string(body) != "data" {
		t.Error("Response body does not equal")
	}

}

type EndPointStatus struct{}

func (this *EndPointStatus) Init(request *http.Request) error {
	return nil
}

func (this *EndPointStatus) Get(request *http.Request) (int, interface{}, http.Header) {
	fmt.Println(request.RequestURI)
	vars := mux.Vars(request)
	status, err := strconv.Atoi(vars["status"])

	if err != nil {
		return 400, err, nil

	}
	return status, "data", nil
}

func TestAdvanceGet(t *testing.T) {
	statuses := []int{100, 200, 300, 400, 404, 500}
	for _, status := range statuses {
		r, err := http.NewRequest("GET", fmt.Sprintf("/get/status/%d", status), nil)
		if err != nil {
			t.Error(err)
		}

		w := httptest.NewRecorder()
		Router().ServeHTTP(w, r)
		if w.Code != status {
			t.Error("Response status code does not equal")
		}
		body, _ := ioutil.ReadAll(w.Body)

		if string(body) != "data" {
			t.Error("Response body does not equal")
		}
	}

	r, err := http.NewRequest("GET", "/get/status/string", nil)
	if err != nil {
		t.Error(err)
	}

	w := httptest.NewRecorder()
	Router().ServeHTTP(w, r)
	if w.Code != 404 {
		t.Error("Response status code does not equal")
	}

}
<commit_msg>Add proper content-type tests<commit_after>package miggle

import (
	"fmt"
	"github.com/gorilla/mux"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
)

func init() {
	http.Handle("/", Router())
}

func Router() *mux.Router {
	var res *Resource = NewResource()
	r := mux.NewRouter()

	r.HandleFunc("/test", res.RequestHandler(new(EndPoint)))
	r.HandleFunc("/get/status/{status:[0-9]+}", res.RequestHandler(new(EndPointStatus)))

	return r
}

type EndPoint struct{}

func (this *EndPoint) Init(request *http.Request) error {
	return nil
}

func (this *EndPoint) Get(request *http.Request) (int, interface{}, http.Header) {
	return 200, "data", nil
}

func TestBasicGet(t *testing.T) {
	r, err := http.NewRequest("GET", "/test", nil)
	if err != nil {
		t.Error(err)
	}
	w := httptest.NewRecorder()
	Router().ServeHTTP(w, r)
	if w.Code != http.StatusOK {
		t.Error("Response status code does not equal")
	}

	body, _ := ioutil.ReadAll(w.Body)

	if string(body) != "data" {
		t.Error("Response body does not equal")
	}

	if header, ok := w.HeaderMap["Content-Type"]; ok {
		if header[0] != "plain/text" {
			t.Error("Response header Content-Type: %s does not exists", header[0])
		}
	} else {
		t.Error("Response header Content-Type: does not exists", header[0])
	}
}

type EndPointStatus struct {
	Status int
}

func (this *EndPointStatus) Init(request *http.Request) error {
	vars := mux.Vars(request)
	if id, ok := vars["status"]; ok {
		status, err := strconv.Atoi(id)

		if err != nil {
			return err
		}
		this.Status = status
	}
	return nil
}

func (this *EndPointStatus) Get(request *http.Request) (int, interface{}, http.Header) {
	return this.Status, "data", http.Header{"Content-Type": {"application/json; charset=utf-8"}}
}

func TestAdvanceGet(t *testing.T) {
	statuses := []int{100, 200, 300, 400, 404, 500}
	for _, status := range statuses {
		r, err := http.NewRequest("GET", fmt.Sprintf("/get/status/%d", status), nil)
		if err != nil {
			t.Error(err)
		}

		w := httptest.NewRecorder()
		Router().ServeHTTP(w, r)
		if w.Code != status {
			t.Error("Response status code does not equal")
		}
		body, _ := ioutil.ReadAll(w.Body)

		if string(body) != "data" {
			t.Error("Response body does not equal")
		}

		if header, ok := w.HeaderMap["Content-Type"]; ok {
			if header[0] != "application/json; charset=utf-8" {
				t.Error("Response header Content-Type: %s does not exists", header[0])
			}
		} else {
			t.Error("Response header Content-Type' does not exists", header)
		}
	}

	r, err := http.NewRequest("GET", "/get/status/string", nil)
	if err != nil {
		t.Error(err)
	}

	w := httptest.NewRecorder()
	Router().ServeHTTP(w, r)
	if w.Code != 404 {
		t.Error("Response status code does not equal")
	}

}
<|endoftext|>


<commit_before>package main

import (
	"container/heap"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path"
	"strconv"
	"strings"
)

const (
	maxMatchLength = 258
	minMatchLength = 4
	hashBits       = 15
	hashSize       = 1 << hashBits
	hashMask       = (1 << hashBits) - 1
)

type compressionLevel struct {
	good, lazy, nice, chain int
}

var levels = []compressionLevel{
	{}, // 0
	{3, 0, 8, 4},
	{3, 0, 16, 8},
	{3, 0, 32, 32},
	{4, 4, 16, 16},
	{8, 16, 32, 32},
	{8, 16, 128, 128},
	{8, 32, 128, 256},
	{32, 128, 258, 1024},
	{32, 258, 258, 4096},
}

type dictator struct {
	// Pseudo deflate variables, we need those to perform deflate like matching, to identify strings that are emmited as is
	compressionLevel
	window   []byte
	hashHead [32768]int
	hashPrev [16384]int
	// Accumulate characters emmited as is
	stringBuf [16384]byte
	stringLen int
	// Count all the strings here
	table map[string]int
}

func (d *dictator) init(level int) (err error) {
	if level < 4 || level > 9 {
		return fmt.Errorf("Only supposts levels [4, 9], got %d", level)
	}

	d.compressionLevel = levels[level]
	d.stringLen = 0
	d.table = make(map[string]int)
	for i := range d.hashHead {
		d.hashHead[i] = -1
	}
	return nil
}

// Try to find a match starting at index whose length is greater than prevSize.
// We only look at chainCount possibilities before giving up.
func (d *dictator) findMatch(pos int, prevHead int, prevLength int, lookahead int) (length, offset int, ok bool) {

	minMatchLook := maxMatchLength
	if lookahead < minMatchLook {
		minMatchLook = lookahead
	}

	win := d.window

	// We quit when we get a match that's at least nice long
	nice := len(win) - pos

	if d.nice < nice {
		nice = d.nice
	}

	// If we've got a match that's good enough, only look in 1/4 the chain.
	tries := d.chain
	length = prevLength
	if length >= d.good {
		tries >>= 2
	}

	w0 := win[pos]
	w1 := win[pos+1]
	wEnd := win[pos+length]

	for i := prevHead; tries > 0; tries-- {
		if w0 == win[i] && w1 == win[i+1] && wEnd == win[i+length] {
			n := 2
			for pos+n < len(win) && win[i+n] == win[pos+n] {
				n++
			}

			if n > length && (n > 3) {
				length = n
				offset = pos - i
				ok = true
				if n >= nice {
					// The match is good enough that we don't try to find a better one.
					break
				}
				wEnd = win[pos+n]
			}
		}
		if i = d.hashPrev[i]; i < 0 {
			break
		}
	}
	return
}

func (d *dictator) findUncompressable(in []byte) {
	d.window = in
	pos := 0
	length := minMatchLength - 1

	for {
		lookahead := len(in) - pos
		if lookahead <= minMatchLength {
			break
		}

		hash := ((int(in[pos]) << 10) ^ (int(in[pos]) << 5) ^ (int(in[pos]))) & hashMask
		hashHead := d.hashHead[hash]
		d.hashPrev[pos] = hashHead
		d.hashHead[hash] = pos

		prevLength := length
		if hashHead >= 0 && prevLength < d.nice {
			if newLength, _, ok := d.findMatch(pos, hashHead, minMatchLength-1, lookahead); ok {
				length = newLength
			}
		}

		// Now deflate would output the previous match, therefore if accumulated enough uncompressed bytes, "flush" them
		if prevLength >= minMatchLength && length <= prevLength {
			if d.stringLen >= minMatchLength {
				key := string(d.stringBuf[:d.stringLen])
				d.table[key]++
				d.stringLen = 0
			}
			newPos := pos + prevLength - 1
			if newPos >= len(in) {
				break
			}
			pos++
			for pos < newPos {
				hash := ((int(in[pos]) << 10) ^ (int(in[pos]) << 5) ^ (int(in[pos]))) & hashMask
				hashHead := d.hashHead[hash]
				d.hashPrev[pos] = hashHead
				d.hashHead[hash] = pos
				pos++
			}
			length = minMatchLength - 1
			// Or the previous literal
		} else if pos > 0 {
			d.stringBuf[d.stringLen] = d.window[pos-1]
			d.stringLen++
			pos++
		} else {
			pos++
		}
	}
	if d.stringLen > minMatchLength {
		//fmt.Println(string(d.stringBuf[:d.stringLen]))
		d.table[string(d.stringBuf[:d.stringLen])]++
		d.stringLen = 0
	}
}

// An Item is something we manage in a priority queue.
type scoredString struct {
	value string // The value of the item; arbitrary.
	score int    // The priority of the item in the queue.
}

// A PriorityQueue implements heap.Interface and holds scoredStrings
type PriorityQueue []*scoredString

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].score > pq[j].score
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *PriorityQueue) Push(x interface{}) {
	item := x.(*scoredString)
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	*pq = old[0 : n-1]
	return item
}

func PrintUsage() {
	fmt.Printf("Usage: %s dictionary_size input_directory output_file\n", path.Base(os.Args[0]))
}

func main() {
	var window [16384]byte
	var d dictator
	var table map[string]int
	table = make(map[string]int)
	pq := make(PriorityQueue, 0)
	var dictionary string

	if len(os.Args) != 4 {
		PrintUsage()
		return
	}

	dictLen, _ := strconv.Atoi(os.Args[1])
	path := os.Args[2]

	files, _ := ioutil.ReadDir(path)

	percent := float64(0)
	for num, f := range files {
		file, err := os.Open(path + "/" + f.Name()) // For read access.
		if err != nil {
			continue
		}
		count, err := file.Read(window[:len(window)])
		d.init(4)
		// Create a table of all uncompressable strings in the fime
		d.findUncompressable(window[:count])
		// Merge with the main table
		for k, _ := range d.table {
			table[k]++
		}

		file.Close()

		if newPercent := float64(num) / float64(len(files)) * 100; (newPercent - percent) >= 1 {
			percent = math.Floor(newPercent)
			fmt.Printf("\r%.2f%% ", newPercent)
		}
	}
	fmt.Println("\r100%%")
	fmt.Println("Total uncompressible strings found: ", len(table))
	// If a string appeares in less than 1% of the files, it is probably useless
	threshold := int(math.Ceil(float64(len(files)) * 0.01))
	// Remove unique strings, score others and put into a heap
	heap.Init(&pq)
	for i, v := range table {
		if v < threshold {
			delete(table, i)
		} else {
			item := &scoredString{i, (v * (len(i) - 3)) / len(i)}
			heap.Push(&pq, item)
		}
	}
	fmt.Println("Uncompressible strings with frequency greater than ", threshold, ": ", len(table))

	// Start poping strings from the heap. We want the highest scoring closer to the end, so they are encoded with smaller distance value
	for (pq.Len() > 0) && (len(dictionary) < dictLen) {
		item := heap.Pop(&pq).(*scoredString)
		// Ignore strings that already made it to the dictionary, append others in front
		if !strings.Contains(dictionary, item.value) {
			dictionary = item.value + dictionary
		}
	}
	// Truncate
	if len(dictionary) > dictLen {
		dictionary = dictionary[len(dictionary)-dictLen:]
	}
	// Write
	ioutil.WriteFile(os.Args[3], []byte(dictionary), 0644)
}
<commit_msg>Add flag to configure window size<commit_after>package main

import (
	"container/heap"
	"flag"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path"
	"strconv"
	"strings"
)

var windowSize = flag.Int("windowsize", 16384, "Window size")

const (
	maxMatchLength = 258
	minMatchLength = 4
	hashBits       = 15
	hashSize       = 1 << hashBits
	hashMask       = (1 << hashBits) - 1
)

type compressionLevel struct {
	good, lazy, nice, chain int
}

var levels = []compressionLevel{
	{}, // 0
	{3, 0, 8, 4},
	{3, 0, 16, 8},
	{3, 0, 32, 32},
	{4, 4, 16, 16},
	{8, 16, 32, 32},
	{8, 16, 128, 128},
	{8, 32, 128, 256},
	{32, 128, 258, 1024},
	{32, 258, 258, 4096},
}

type dictator struct {
	// Pseudo deflate variables, we need those to perform deflate like matching, to identify strings that are emmited as is
	compressionLevel
	window   []byte
	hashHead [32768]int
	hashPrev []int
	// Accumulate characters emitted as is
	stringBuf []byte
	stringLen int
	// Count all the strings here
	table map[string]int
}

func NewDictator(windowSize int) *dictator {
	dictator := new(dictator)
	dictator.hashPrev = make([]int, windowSize);
	dictator.stringBuf = make([]byte, windowSize);
	return dictator
}

func (d *dictator) init(level int) (err error) {
	if level < 4 || level > 9 {
		return fmt.Errorf("Only supposts levels [4, 9], got %d", level)
	}

	d.compressionLevel = levels[level]
	d.stringLen = 0
	d.table = make(map[string]int)
	for i := range d.hashHead {
		d.hashHead[i] = -1
	}
	return nil
}

// Try to find a match starting at index whose length is greater than prevSize.
// We only look at chainCount possibilities before giving up.
func (d *dictator) findMatch(pos int, prevHead int, prevLength int, lookahead int) (length, offset int, ok bool) {

	minMatchLook := maxMatchLength
	if lookahead < minMatchLook {
		minMatchLook = lookahead
	}

	win := d.window

	// We quit when we get a match that's at least nice long
	nice := len(win) - pos

	if d.nice < nice {
		nice = d.nice
	}

	// If we've got a match that's good enough, only look in 1/4 the chain.
	tries := d.chain
	length = prevLength
	if length >= d.good {
		tries >>= 2
	}

	w0 := win[pos]
	w1 := win[pos+1]
	wEnd := win[pos+length]

	for i := prevHead; tries > 0; tries-- {
		if w0 == win[i] && w1 == win[i+1] && wEnd == win[i+length] {
			n := 2
			for pos+n < len(win) && win[i+n] == win[pos+n] {
				n++
			}

			if n > length && (n > 3) {
				length = n
				offset = pos - i
				ok = true
				if n >= nice {
					// The match is good enough that we don't try to find a better one.
					break
				}
				wEnd = win[pos+n]
			}
		}
		if i = d.hashPrev[i]; i < 0 {
			break
		}
	}
	return
}

func (d *dictator) findUncompressable(in []byte) {
	d.window = in
	pos := 0
	length := minMatchLength - 1

	for {
		lookahead := len(in) - pos
		if lookahead <= minMatchLength {
			break
		}

		hash := ((int(in[pos]) << 10) ^ (int(in[pos]) << 5) ^ (int(in[pos]))) & hashMask
		hashHead := d.hashHead[hash]
		d.hashPrev[pos] = hashHead
		d.hashHead[hash] = pos

		prevLength := length
		if hashHead >= 0 && prevLength < d.nice {
			if newLength, _, ok := d.findMatch(pos, hashHead, minMatchLength-1, lookahead); ok {
				length = newLength
			}
		}

		// Now deflate would output the previous match, therefore if accumulated enough uncompressed bytes, "flush" them
		if prevLength >= minMatchLength && length <= prevLength {
			if d.stringLen >= minMatchLength {
				key := string(d.stringBuf[:d.stringLen])
				d.table[key]++
				d.stringLen = 0
			}
			newPos := pos + prevLength - 1
			if newPos >= len(in) {
				break
			}
			pos++
			for pos < newPos {
				hash := ((int(in[pos]) << 10) ^ (int(in[pos]) << 5) ^ (int(in[pos]))) & hashMask
				hashHead := d.hashHead[hash]
				d.hashPrev[pos] = hashHead
				d.hashHead[hash] = pos
				pos++
			}
			length = minMatchLength - 1
			// Or the previous literal
		} else if pos > 0 {
			d.stringBuf[d.stringLen] = d.window[pos-1]
			d.stringLen++
			pos++
		} else {
			pos++
		}
	}
	if d.stringLen > minMatchLength {
		//fmt.Println(string(d.stringBuf[:d.stringLen]))
		d.table[string(d.stringBuf[:d.stringLen])]++
		d.stringLen = 0
	}
}

// An Item is something we manage in a priority queue.
type scoredString struct {
	value string // The value of the item; arbitrary.
	score int    // The priority of the item in the queue.
}

// A PriorityQueue implements heap.Interface and holds scoredStrings
type PriorityQueue []*scoredString

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].score > pq[j].score
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *PriorityQueue) Push(x interface{}) {
	item := x.(*scoredString)
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	*pq = old[0 : n-1]
	return item
}

func PrintUsage() {
	fmt.Printf("Usage: %s [options] dictionary_size input_directory output_file\n", path.Base(os.Args[0]))
	flag.PrintDefaults()
}

func main() {
	flag.Parse()
	window := make([]byte, *windowSize)
	d := NewDictator(*windowSize)
	var table map[string]int
	table = make(map[string]int)
	pq := make(PriorityQueue, 0)
	var dictionary string

	if len(flag.Args()) != 3 {
		PrintUsage()
		return
	}

	dictLen, _ := strconv.Atoi(flag.Arg(0))
	path := flag.Arg(1)

	files, _ := ioutil.ReadDir(path)

	percent := float64(0)
	for num, f := range files {
		file, err := os.Open(path + "/" + f.Name()) // For read access.
		if err != nil {
			continue
		}
		count, err := file.Read(window[:len(window)])
		d.init(4)
		// Create a table of all uncompressable strings in the fime
		d.findUncompressable(window[:count])
		// Merge with the main table
		for k, _ := range d.table {
			table[k]++
		}

		file.Close()

		if newPercent := float64(num) / float64(len(files)) * 100; (newPercent - percent) >= 1 {
			percent = math.Floor(newPercent)
			fmt.Printf("\r%.2f%% ", newPercent)
		}
	}
	fmt.Println("\r100%%")
	fmt.Println("Total uncompressible strings found: ", len(table))
	// If a string appeares in less than 1% of the files, it is probably useless
	threshold := int(math.Ceil(float64(len(files)) * 0.01))
	// Remove unique strings, score others and put into a heap
	heap.Init(&pq)
	for i, v := range table {
		if v < threshold {
			delete(table, i)
		} else {
			item := &scoredString{i, (v * (len(i) - 3)) / len(i)}
			heap.Push(&pq, item)
		}
	}
	fmt.Println("Uncompressible strings with frequency greater than ", threshold, ": ", len(table))

	// Start poping strings from the heap. We want the highest scoring closer to the end, so they are encoded with smaller distance value
	for (pq.Len() > 0) && (len(dictionary) < dictLen) {
		item := heap.Pop(&pq).(*scoredString)
		// Ignore strings that already made it to the dictionary, append others in front
		if !strings.Contains(dictionary, item.value) {
			dictionary = item.value + dictionary
		}
	}
	// Truncate
	if len(dictionary) > dictLen {
		dictionary = dictionary[len(dictionary)-dictLen:]
	}
	// Write
	ioutil.WriteFile(flag.Arg(2), []byte(dictionary), 0644)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/spf13/pflag"

	"golang.org/x/crypto/ssh/terminal"

	kflag "k8s.io/apiserver/pkg/util/flag"
	"k8s.io/client-go/pkg/apis/clientauthentication/v1alpha1"
	"k8s.io/cloud-provider-openstack/pkg/identity/keystone"
)

const errRespTemplate string = `{
	"apiVersion": "client.authentication.k8s.io/v1alpha1",
	"kind": "ExecCredential",
	"spec": {
		"response": {
			"code": 401,
			"header": {},
		},
	}
}`

const respTemplate string = `{
	"apiVersion": "client.authentication.k8s.io/v1alpha1",
	"kind": "ExecCredential",
	"status": {
		"token": "%v",
		"expirationTimestamp": "%v"
	}
}`

func promptForString(field string, r io.Reader, show bool) (result string, err error) {
	// We have to print output to Stderr, because Stdout is redirected and not shown to the user.
	fmt.Fprintf(os.Stderr, "Please enter %s: ", field)

	if show {
		_, err = fmt.Fscan(r, &result)
	} else {
		var data []byte
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			data, err = terminal.ReadPassword(int(os.Stdin.Fd()))
			result = string(data)
			fmt.Fprintf(os.Stderr, "\n")
		} else {
			return "", fmt.Errorf("error reading input for %s", field)
		}
	}
	return result, err
}

// prompt pulls keystone auth url, domain, username and password from stdin,
// if they are not specified initially (i.e. equal "").
func prompt(url string, domain string, user string, project string, password string) (gophercloud.AuthOptions, error) {
	var err error
	var options gophercloud.AuthOptions

	if url == "" {
		url, err = promptForString("Keystone Auth URL", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if domain == "" {
		domain, err = promptForString("domain name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if user == "" {
		user, err = promptForString("user name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if project == "" {
		project, err = promptForString("project name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if password == "" {
		password, err = promptForString("password", nil, false)
		if err != nil {
			return options, err
		}
	}

	options = gophercloud.AuthOptions{
		IdentityEndpoint: url,
		Username:         user,
		TenantName:       project,
		Password:         password,
		DomainName:       domain,
	}

	return options, nil
}

// KuberneteExecInfo holds information passed to the plugin by the transport. This
// contains runtime specific information, such as if the session is interactive,
// auth API version and kind of request.
type KuberneteExecInfo struct {
	APIVersion string                      `json:"apiVersion"`
	Kind       string                      `json:"kind"`
	Spec       v1alpha1.ExecCredentialSpec `json:"spec"`
}

func validateKebernetesExecInfo(kei KuberneteExecInfo) error {
	if kei.APIVersion != v1alpha1.SchemeGroupVersion.String() {
		return fmt.Errorf("unsupported API version: %v", kei.APIVersion)
	}

	if kei.Kind != "ExecCredential" {
		return fmt.Errorf("incorrect request kind: %v", kei.Kind)
	}

	return nil
}

func main() {
	var url string
	var domain string
	var user string
	var project string
	var password string
	var options gophercloud.AuthOptions
	var err error

	pflag.StringVar(&url, "keystone-url", os.Getenv("OS_AUTH_URL"), "URL for the OpenStack Keystone API")
	pflag.StringVar(&domain, "domain-name", os.Getenv("OS_DOMAIN_NAME"), "Keystone domain name")
	pflag.StringVar(&user, "user-name", os.Getenv("OS_USERNAME"), "User name")
	pflag.StringVar(&project, "project-name", os.Getenv("OS_PROJECT_NAME"), "Keystone project name")
	pflag.StringVar(&password, "password", os.Getenv("OS_PASSWORD"), "Password")
	kflag.InitFlags()

	keiData := os.Getenv("KUBERNETES_EXEC_INFO")
	if keiData == "" {
		fmt.Fprintln(os.Stderr, "KUBERNETES_EXEC_INFO env variable must be set.")
		os.Exit(1)
	}

	kei := KuberneteExecInfo{}
	err = json.Unmarshal([]byte(keiData), &kei)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to parse KUBERNETES_EXEC_INFO value.")
		os.Exit(1)
	}

	err = validateKebernetesExecInfo(kei)
	if err != nil {
		fmt.Fprintf(os.Stderr, "An error occurred: %v\n", err)
		os.Exit(1)
	}

	// Generate Gophercloud Auth Options based on input data from stdin
	// if "intaractive" is set to true, or from env variables otherwise.
	if !kei.Spec.Interactive {
		options, err = openstack.AuthOptionsFromEnv()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read openstack env vars: %s", err)
			os.Exit(1)
		}
	} else {
		options, err = prompt(url, domain, user, project, password)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read data from console: %s", err)
			os.Exit(1)
		}
	}

	token, err := keystone.GetToken(options)
	if err != nil {
		if _, ok := err.(gophercloud.ErrDefault401); ok {
			fmt.Println(errRespTemplate)
			os.Exit(0)
		}
		fmt.Fprintf(os.Stderr, "An error occurred: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf(respTemplate, token.ID, token.ExpiresAt.Format(time.RFC3339Nano))
}
<commit_msg>fix nits in client-keystone-auth<commit_after>/*
Copyright 2018 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 main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/gophercloud/gophercloud"
	"github.com/gophercloud/gophercloud/openstack"
	"github.com/spf13/pflag"

	"golang.org/x/crypto/ssh/terminal"

	kflag "k8s.io/apiserver/pkg/util/flag"
	"k8s.io/client-go/pkg/apis/clientauthentication/v1alpha1"
	"k8s.io/cloud-provider-openstack/pkg/identity/keystone"
)

const errRespTemplate string = `{
	"apiVersion": "client.authentication.k8s.io/v1alpha1",
	"kind": "ExecCredential",
	"spec": {
		"response": {
			"code": 401,
			"header": {},
		},
	}
}`

const respTemplate string = `{
	"apiVersion": "client.authentication.k8s.io/v1alpha1",
	"kind": "ExecCredential",
	"status": {
		"token": "%v",
		"expirationTimestamp": "%v"
	}
}`

func promptForString(field string, r io.Reader, show bool) (result string, err error) {
	// We have to print output to Stderr, because Stdout is redirected and not shown to the user.
	fmt.Fprintf(os.Stderr, "Please enter %s: ", field)

	if show {
		_, err = fmt.Fscan(r, &result)
	} else {
		var data []byte
		if terminal.IsTerminal(int(os.Stdin.Fd())) {
			data, err = terminal.ReadPassword(int(os.Stdin.Fd()))
			result = string(data)
			fmt.Fprintf(os.Stderr, "\n")
		} else {
			return "", fmt.Errorf("error reading input for %s", field)
		}
	}
	return result, err
}

// prompt pulls keystone auth url, domain, project, username and password from stdin,
// if they are not specified initially (i.e. equal "").
func prompt(url string, domain string, user string, project string, password string) (gophercloud.AuthOptions, error) {
	var err error
	var options gophercloud.AuthOptions

	if url == "" {
		url, err = promptForString("Keystone Auth URL", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if domain == "" {
		domain, err = promptForString("domain name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if user == "" {
		user, err = promptForString("user name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if project == "" {
		project, err = promptForString("project name", os.Stdin, true)
		if err != nil {
			return options, err
		}
	}

	if password == "" {
		password, err = promptForString("password", nil, false)
		if err != nil {
			return options, err
		}
	}

	options = gophercloud.AuthOptions{
		IdentityEndpoint: url,
		Username:         user,
		TenantName:       project,
		Password:         password,
		DomainName:       domain,
	}

	return options, nil
}

// KuberneteExecInfo holds information passed to the plugin by the transport. This
// contains runtime specific information, such as if the session is interactive,
// auth API version and kind of request.
type KuberneteExecInfo struct {
	APIVersion string                      `json:"apiVersion"`
	Kind       string                      `json:"kind"`
	Spec       v1alpha1.ExecCredentialSpec `json:"spec"`
}

func validateKebernetesExecInfo(kei KuberneteExecInfo) error {
	if kei.APIVersion != v1alpha1.SchemeGroupVersion.String() {
		return fmt.Errorf("unsupported API version: %v", kei.APIVersion)
	}

	if kei.Kind != "ExecCredential" {
		return fmt.Errorf("incorrect request kind: %v", kei.Kind)
	}

	return nil
}

func main() {
	var url string
	var domain string
	var user string
	var project string
	var password string
	var options gophercloud.AuthOptions
	var err error

	pflag.StringVar(&url, "keystone-url", os.Getenv("OS_AUTH_URL"), "URL for the OpenStack Keystone API")
	pflag.StringVar(&domain, "domain-name", os.Getenv("OS_DOMAIN_NAME"), "Keystone domain name")
	pflag.StringVar(&user, "user-name", os.Getenv("OS_USERNAME"), "User name")
	pflag.StringVar(&project, "project-name", os.Getenv("OS_PROJECT_NAME"), "Keystone project name")
	pflag.StringVar(&password, "password", os.Getenv("OS_PASSWORD"), "Password")
	kflag.InitFlags()

	keiData := os.Getenv("KUBERNETES_EXEC_INFO")
	if keiData == "" {
		fmt.Fprintln(os.Stderr, "KUBERNETES_EXEC_INFO env variable must be set.")
		os.Exit(1)
	}

	kei := KuberneteExecInfo{}
	err = json.Unmarshal([]byte(keiData), &kei)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to parse KUBERNETES_EXEC_INFO value.")
		os.Exit(1)
	}

	err = validateKebernetesExecInfo(kei)
	if err != nil {
		fmt.Fprintf(os.Stderr, "An error occurred: %v\n", err)
		os.Exit(1)
	}

	// Generate Gophercloud Auth Options based on input data from stdin
	// if "intaractive" is set to true, or from env variables otherwise.
	if !kei.Spec.Interactive {
		options, err = openstack.AuthOptionsFromEnv()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read openstack env vars: %s", err)
			os.Exit(1)
		}
	} else {
		options, err = prompt(url, domain, user, project, password)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read data from console: %s", err)
			os.Exit(1)
		}
	}

	token, err := keystone.GetToken(options)
	if err != nil {
		if _, ok := err.(gophercloud.ErrDefault401); ok {
			fmt.Println(errRespTemplate)
			os.Exit(0)
		}
		fmt.Fprintf(os.Stderr, "An error occurred: %v\n", err)
		os.Exit(1)
	}

	out := fmt.Sprintf(respTemplate, token.ID, token.ExpiresAt.Format(time.RFC3339Nano))
	fmt.Println(out)
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"os"

	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager/lagerflags"
)

type FileServerConfig struct {
	ServerAddress   string `json:"server_address,omitempty"`
	StaticDirectory string `json:"static_directory,omitempty"`
	DropsondePort   int    `json:"dropsonde_port,omitempty"`
	ConsulCluster   string `json:"consul_cluster,omitempty"`

	debugserver.DebugServerConfig
	lagerflags.LagerConfig
}

func DefaultFileServerConfig() FileServerConfig {
	return FileServerConfig{
		ServerAddress: "0.0.0.0:8080",
		DropsondePort: 3457,
		LagerConfig:   lagerflags.DefaultLagerConfig(),
	}
}

func NewFileServerConfig(configPath string) (FileServerConfig, error) {
	fileServerConfig := DefaultFileServerConfig()

	configFile, err := os.Open(configPath)
	if err != nil {
		return FileServerConfig{}, err
	}

	decoder := json.NewDecoder(configFile)
	err = decoder.Decode(&fileServerConfig)
	if err != nil {
		return FileServerConfig{}, err
	}

	return fileServerConfig, nil
}
<commit_msg>Close config file after reading.<commit_after>package config

import (
	"encoding/json"
	"os"

	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager/lagerflags"
)

type FileServerConfig struct {
	ServerAddress   string `json:"server_address,omitempty"`
	StaticDirectory string `json:"static_directory,omitempty"`
	DropsondePort   int    `json:"dropsonde_port,omitempty"`
	ConsulCluster   string `json:"consul_cluster,omitempty"`

	debugserver.DebugServerConfig
	lagerflags.LagerConfig
}

func DefaultFileServerConfig() FileServerConfig {
	return FileServerConfig{
		ServerAddress: "0.0.0.0:8080",
		DropsondePort: 3457,
		LagerConfig:   lagerflags.DefaultLagerConfig(),
	}
}

func NewFileServerConfig(configPath string) (FileServerConfig, error) {
	fileServerConfig := DefaultFileServerConfig()

	configFile, err := os.Open(configPath)
	if err != nil {
		return FileServerConfig{}, err
	}

	defer configFile.Close()

	decoder := json.NewDecoder(configFile)
	err = decoder.Decode(&fileServerConfig)
	if err != nil {
		return FileServerConfig{}, err
	}

	return fileServerConfig, nil
}
<|endoftext|>


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

package main

import (
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer/consumererror"
	"go.opentelemetry.io/collector/service/defaultcomponents"

	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/alibabacloudlogserviceexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsemfexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awskinesisexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsprometheusremotewriteexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsxrayexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/azuremonitorexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/carbonexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/datadogexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/dynatraceexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/elasticexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/f5cloudexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/honeycombexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/jaegerthrifthttpexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/loadbalancingexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/logzioexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/lokiexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/newrelicexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sapmexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sentryexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/signalfxexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/splunkhecexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/stackdriverexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sumologicexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/httpforwarder"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/observer/hostobserver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/observer/k8sobserver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/groupbyattrsprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/groupbytraceprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/k8sprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/metricstransformprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/resourcedetectionprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/routingprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/spanmetricsprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/awsecscontainermetricsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/awsxrayreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/carbonreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/collectdreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/dockerstatsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/filelogreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jmxreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/k8sclusterreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/kubeletstatsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/prometheusexecreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/receivercreator"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/redisreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/sapmreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/signalfxreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/simpleprometheusreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/splunkhecreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/statsdreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/wavefrontreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/windowsperfcountersreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/zookeeperreceiver"
)

func components() (component.Factories, error) {
	var errs []error
	factories, err := defaultcomponents.Components()
	if err != nil {
		return component.Factories{}, err
	}

	extensions := []component.ExtensionFactory{
		hostobserver.NewFactory(),
		httpforwarder.NewFactory(),
		k8sobserver.NewFactory(),
	}

	for _, ext := range factories.Extensions {
		extensions = append(extensions, ext)
	}

	factories.Extensions, err = component.MakeExtensionFactoryMap(extensions...)
	if err != nil {
		errs = append(errs, err)
	}

	receivers := []component.ReceiverFactory{
		awsecscontainermetricsreceiver.NewFactory(),
		awsxrayreceiver.NewFactory(),
		carbonreceiver.NewFactory(),
		collectdreceiver.NewFactory(),
		dockerstatsreceiver.NewFactory(),
		filelogreceiver.NewFactory(),
		jmxreceiver.NewFactory(),
		k8sclusterreceiver.NewFactory(),
		kubeletstatsreceiver.NewFactory(),
		prometheusexecreceiver.NewFactory(),
		receivercreator.NewFactory(),
		redisreceiver.NewFactory(),
		sapmreceiver.NewFactory(),
		signalfxreceiver.NewFactory(),
		simpleprometheusreceiver.NewFactory(),
		splunkhecreceiver.NewFactory(),
		statsdreceiver.NewFactory(),
		wavefrontreceiver.NewFactory(),
		windowsperfcountersreceiver.NewFactory(),
		zookeeperreceiver.NewFactory(),
	}

	receivers = append(receivers, extraReceivers()...)

	for _, rcv := range factories.Receivers {
		receivers = append(receivers, rcv)
	}
	factories.Receivers, err = component.MakeReceiverFactoryMap(receivers...)
	if err != nil {
		errs = append(errs, err)
	}

	exporters := []component.ExporterFactory{
		alibabacloudlogserviceexporter.NewFactory(),
		awsemfexporter.NewFactory(),
		awskinesisexporter.NewFactory(),
		awsprometheusremotewriteexporter.NewFactory(),
		awsxrayexporter.NewFactory(),
		azuremonitorexporter.NewFactory(),
		carbonexporter.NewFactory(),
		datadogexporter.NewFactory(),
		dynatraceexporter.NewFactory(),
		elasticexporter.NewFactory(),
		f5cloudexporter.NewFactory(),
		honeycombexporter.NewFactory(),
		jaegerthrifthttpexporter.NewFactory(),
		loadbalancingexporter.NewFactory(),
		logzioexporter.NewFactory(),
		lokiexporter.NewFactory(),
		newrelicexporter.NewFactory(),
		sapmexporter.NewFactory(),
		sentryexporter.NewFactory(),
		signalfxexporter.NewFactory(),
		splunkhecexporter.NewFactory(),
		stackdriverexporter.NewFactory(),
		sumologicexporter.NewFactory(),
	}
	for _, exp := range factories.Exporters {
		exporters = append(exporters, exp)
	}
	factories.Exporters, err = component.MakeExporterFactoryMap(exporters...)
	if err != nil {
		errs = append(errs, err)
	}

	processors := []component.ProcessorFactory{
		groupbyattrsprocessor.NewFactory(),
		groupbytraceprocessor.NewFactory(),
		k8sprocessor.NewFactory(),
		metricstransformprocessor.NewFactory(),
		resourcedetectionprocessor.NewFactory(),
		routingprocessor.NewFactory(),
		tailsamplingprocessor.NewFactory(),
		spanmetricsprocessor.NewFactory(),
	}
	for _, pr := range factories.Processors {
		processors = append(processors, pr)
	}
	factories.Processors, err = component.MakeProcessorFactoryMap(processors...)
	if err != nil {
		errs = append(errs, err)
	}

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

package main

import (
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/service/defaultcomponents"

	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/alibabacloudlogserviceexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsemfexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awskinesisexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsprometheusremotewriteexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/awsxrayexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/azuremonitorexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/carbonexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/datadogexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/dynatraceexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/elasticexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/f5cloudexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/honeycombexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/jaegerthrifthttpexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/loadbalancingexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/logzioexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/lokiexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/newrelicexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sapmexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sentryexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/signalfxexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/splunkhecexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/stackdriverexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/sumologicexporter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/httpforwarder"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/observer/hostobserver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/observer/k8sobserver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/groupbyattrsprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/groupbytraceprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/k8sprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/metricstransformprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/resourcedetectionprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/routingprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/spanmetricsprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/awsecscontainermetricsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/awsxrayreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/carbonreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/collectdreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/dockerstatsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/filelogreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jmxreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/k8sclusterreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/kubeletstatsreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/prometheusexecreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/receivercreator"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/redisreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/sapmreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/signalfxreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/simpleprometheusreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/splunkhecreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/statsdreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/wavefrontreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/windowsperfcountersreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/zookeeperreceiver"
)

func components() (component.Factories, error) {
	factories, err := defaultcomponents.Components()
	if err != nil {
		return component.Factories{}, err
	}

	extensions := []component.ExtensionFactory{
		hostobserver.NewFactory(),
		httpforwarder.NewFactory(),
		k8sobserver.NewFactory(),
	}

	for _, ext := range factories.Extensions {
		extensions = append(extensions, ext)
	}

	factories.Extensions, err = component.MakeExtensionFactoryMap(extensions...)
	if err != nil {
		return component.Factories{}, err
	}

	receivers := []component.ReceiverFactory{
		awsecscontainermetricsreceiver.NewFactory(),
		awsxrayreceiver.NewFactory(),
		carbonreceiver.NewFactory(),
		collectdreceiver.NewFactory(),
		dockerstatsreceiver.NewFactory(),
		filelogreceiver.NewFactory(),
		jmxreceiver.NewFactory(),
		k8sclusterreceiver.NewFactory(),
		kubeletstatsreceiver.NewFactory(),
		prometheusexecreceiver.NewFactory(),
		receivercreator.NewFactory(),
		redisreceiver.NewFactory(),
		sapmreceiver.NewFactory(),
		signalfxreceiver.NewFactory(),
		simpleprometheusreceiver.NewFactory(),
		splunkhecreceiver.NewFactory(),
		statsdreceiver.NewFactory(),
		wavefrontreceiver.NewFactory(),
		windowsperfcountersreceiver.NewFactory(),
		zookeeperreceiver.NewFactory(),
	}

	receivers = append(receivers, extraReceivers()...)

	for _, rcv := range factories.Receivers {
		receivers = append(receivers, rcv)
	}
	factories.Receivers, err = component.MakeReceiverFactoryMap(receivers...)
	if err != nil {
		return component.Factories{}, err
	}

	exporters := []component.ExporterFactory{
		alibabacloudlogserviceexporter.NewFactory(),
		awsemfexporter.NewFactory(),
		awskinesisexporter.NewFactory(),
		awsprometheusremotewriteexporter.NewFactory(),
		awsxrayexporter.NewFactory(),
		azuremonitorexporter.NewFactory(),
		carbonexporter.NewFactory(),
		datadogexporter.NewFactory(),
		dynatraceexporter.NewFactory(),
		elasticexporter.NewFactory(),
		f5cloudexporter.NewFactory(),
		honeycombexporter.NewFactory(),
		jaegerthrifthttpexporter.NewFactory(),
		loadbalancingexporter.NewFactory(),
		logzioexporter.NewFactory(),
		lokiexporter.NewFactory(),
		newrelicexporter.NewFactory(),
		sapmexporter.NewFactory(),
		sentryexporter.NewFactory(),
		signalfxexporter.NewFactory(),
		splunkhecexporter.NewFactory(),
		stackdriverexporter.NewFactory(),
		sumologicexporter.NewFactory(),
	}
	for _, exp := range factories.Exporters {
		exporters = append(exporters, exp)
	}
	factories.Exporters, err = component.MakeExporterFactoryMap(exporters...)
	if err != nil {
		return component.Factories{}, err
	}

	processors := []component.ProcessorFactory{
		groupbyattrsprocessor.NewFactory(),
		groupbytraceprocessor.NewFactory(),
		k8sprocessor.NewFactory(),
		metricstransformprocessor.NewFactory(),
		resourcedetectionprocessor.NewFactory(),
		routingprocessor.NewFactory(),
		tailsamplingprocessor.NewFactory(),
		spanmetricsprocessor.NewFactory(),
	}
	for _, pr := range factories.Processors {
		processors = append(processors, pr)
	}
	factories.Processors, err = component.MakeProcessorFactoryMap(processors...)
	if err != nil {
		return component.Factories{}, err
	}

	return factories, nil
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   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/>
*/

package instance

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/percona/percona-agent/agent"

	"strconv"

	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/mrms"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
)

type empty struct{}

type Manager struct {
	logger    *pct.Logger
	configDir string
	api       pct.APIConnector
	// --
	status         *pct.Status
	repo           *Repo
	stopChan       chan empty
	mrm            mrms.Monitor
	mrmChans       map[string]<-chan bool
	mrmsGlobalChan chan string
	agentConfig    *agent.Config
}

func NewManager(logger *pct.Logger, configDir string, api pct.APIConnector, mrm mrms.Monitor) *Manager {
	repo := NewRepo(pct.NewLogger(logger.LogChan(), "instance-repo"), configDir, api)
	m := &Manager{
		logger:    logger,
		configDir: configDir,
		api:       api,
		// --
		status:         pct.NewStatus([]string{"instance", "instance-repo", "instance-mrms"}),
		repo:           repo,
		mrm:            mrm,
		mrmChans:       make(map[string]<-chan bool),
		mrmsGlobalChan: make(chan string, 100), // monitor up to 100 instances
	}
	return m
}

/////////////////////////////////////////////////////////////////////////////
// Interface
/////////////////////////////////////////////////////////////////////////////

// @goroutine[0]
func (m *Manager) Start() error {
	m.status.Update("instance", "Starting")
	if err := m.repo.Init(); err != nil {
		return err
	}
	m.logger.Info("Started")
	m.status.Update("instance", "Running")

	mrmsGlobalChan, err := m.mrm.GlobalSubscribe()
	if err != nil {
		return err
	}

	for _, instance := range m.GetMySQLInstances() {
		ch, err := m.mrm.Add(instance.DSN)
		if err != nil {
			m.logger.Error("Cannot add instance to the monitor:", err)
			continue
		}
		m.pushInstanceInfo(instance)
		// Store the channel to be able to remove it from mrms
		m.mrmChans[instance.DSN] = ch
	}
	go m.monitorInstancesRestart(mrmsGlobalChan)
	return nil
}

// @goroutine[0]
func (m *Manager) Stop() error {
	// Can't stop the instance manager.
	return nil
}

// @goroutine[0]
func (m *Manager) Handle(cmd *proto.Cmd) *proto.Reply {
	m.status.UpdateRe("instance", "Handling", cmd)
	defer m.status.Update("instance", "Running")

	it := &proto.ServiceInstance{}
	if err := json.Unmarshal(cmd.Data, it); err != nil {
		return cmd.Reply(nil, err)
	}

	switch cmd.Cmd {
	case "Add":
		err := m.repo.Add(it.Service, it.InstanceId, it.Instance, true) // true = write to disk
		if err != nil {
			return cmd.Reply(nil, err)
		}
		if it.Service == "mysql" {
			// Get the instance as type proto.MySQLInstance instead of proto.ServiceInstance
			// because we need the dsn field
			iit := &proto.MySQLInstance{}
			err := m.repo.Get(it.Service, it.InstanceId, iit)
			if err != nil {
				return cmd.Reply(nil, err)
			}
			if err != nil {
				return cmd.Reply(nil, err)
			}
			ch, err := m.mrm.Add(iit.DSN)
			if err != nil {
				m.mrmChans[iit.DSN] = ch
			}
			err = m.pushInstanceInfo(iit)
			if err != nil {
				m.logger.Error(err)
			}
		}
		return cmd.Reply(nil, nil)
	case "Remove":
		if it.Service == "mysql" {
			// Get the instance as type proto.MySQLInstance instead of proto.ServiceInstance
			// because we need the dsn field
			iit := &proto.MySQLInstance{}
			err := m.repo.Get(it.Service, it.InstanceId, iit)
			// Don't return an error. This is just a remove from mrms
			if err != nil {
				m.logger.Error(err)
			} else {
				m.mrm.Remove(iit.DSN, m.mrmChans[iit.DSN])
			}
		}
		err := m.repo.Remove(it.Service, it.InstanceId)
		return cmd.Reply(nil, err)
	case "GetInfo":
		info, err := m.handleGetInfo(it.Service, it.Instance)
		return cmd.Reply(info, err)
	default:
		return cmd.Reply(nil, pct.UnknownCmdError{Cmd: cmd.Cmd})
	}
}

func (m *Manager) Status() map[string]string {
	m.status.Update("instance-repo", strings.Join(m.repo.List(), " "))
	return m.status.All()
}

func (m *Manager) GetConfig() ([]proto.AgentConfig, []error) {
	return nil, nil
}

func (m *Manager) Repo() *Repo {
	return m.repo
}

/////////////////////////////////////////////////////////////////////////////
// Implementation
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) handleGetInfo(service string, data []byte) (interface{}, error) {
	switch service {
	case "mysql":
		it := &proto.MySQLInstance{}
		if err := json.Unmarshal(data, it); err != nil {
			return nil, errors.New("instance.Repo:json.Unmarshal:" + err.Error())
		}
		if it.DSN == "" {
			return nil, fmt.Errorf("MySQL instance DSN is not set")
		}
		if err := GetMySQLInfo(it); err != nil {
			return nil, err
		}
		return it, nil
	default:
		return nil, fmt.Errorf("Don't know how to get info for %s service", service)
	}
}

func GetMySQLInfo(it *proto.MySQLInstance) error {
	conn := mysql.NewConnection(it.DSN)
	if err := conn.Connect(1); err != nil {
		return err
	}
	defer conn.Close()
	sql := "SELECT /* percona-agent */" +
		" CONCAT_WS('.', @@hostname, IF(@@port='3306',NULL,@@port)) AS Hostname," +
		" @@version_comment AS Distro," +
		" @@version AS Version"
	err := conn.DB().QueryRow(sql).Scan(
		&it.Hostname,
		&it.Distro,
		&it.Version,
	)
	if err != nil {
		return err
	}
	return nil
}

func (m *Manager) GetMySQLInstances() []*proto.MySQLInstance {
	m.logger.Debug("getMySQLInstances:call")
	defer m.logger.Debug("getMySQLInstances:return")

	var instances []*proto.MySQLInstance
	for _, name := range m.Repo().List() {
		parts := strings.Split(name, "-") // mysql-1 or server-12
		if len(parts) != 2 {
			m.logger.Error("Invalid instance name: %s: expected 2 parts, got %d", name, len(parts))
			continue
		}
		if parts[0] == "mysql" {
			id, err := strconv.ParseInt(parts[1], 10, 64)
			if err != nil {
				m.logger.Error("Invalid instance ID: %s: %s", name, err)
				continue
			}
			it := &proto.MySQLInstance{}
			if err := m.Repo().Get(parts[0], uint(id), it); err != nil {
				m.logger.Error("Failed to get instance %s: %s", name, err)
				continue
			}
			instances = append(instances, it)
		}
	}
	return instances
}

func (m *Manager) monitorInstancesRestart(ch chan string) {
	m.logger.Debug("monitorInstancesRestart:call")
	defer func() {
		if err := recover(); err != nil {
			m.logger.Error("MySQL connection crashed: ", err)
			m.status.Update("instance-mrms", "Crashed")
		} else {
			m.status.Update("instance-mrms", "Stopped")
		}
		m.logger.Debug("monitorInstancesRestart:return")
	}()

	ch, err := m.mrm.GlobalSubscribe()
	if err != nil {
		m.logger.Error(fmt.Sprintf("Failed to get MySQL restart monitor global channel: %s", err))
		return
	}

	for {
		m.status.Update("instance-mrms", "Idle")
		select {
		case dsn := <-ch:
			safeDSN := mysql.HideDSNPassword(dsn)
			m.logger.Debug("mrms:restart:" + safeDSN)
			m.status.Update("instance-mrms", "Updating "+safeDSN)

			// Get the updated instances list. It should be updated every time since
			// the Add method can add new instances to the list.
			for _, instance := range m.GetMySQLInstances() {
				if instance.DSN != dsn {
					continue
				}

				m.status.Update("instance-mrms", "Getting info "+safeDSN)
				if err := GetMySQLInfo(instance); err != nil {
					m.logger.Warn(fmt.Sprintf("Failed to get MySQL info %s: %s", safeDSN, err))
					continue
				}

				m.status.Update("instance-mrms", "Updating info "+safeDSN)
				err := m.pushInstanceInfo(instance)
				if err != nil {
					m.logger.Warn(err)
				}
			}
		}
	}
}

func (m *Manager) pushInstanceInfo(instance *proto.MySQLInstance) error {
	uri := fmt.Sprintf("%s/%s/%d", m.api.EntryLink("instances"), "mysql", instance.Id)
	data, err := json.Marshal(instance)
	if err != nil {
		m.logger.Error(err)
		return err
	}
	resp, body, err := m.api.Put(m.api.ApiKey(), uri, data)
	if err != nil {
		return err
	}
	// Sometimes the API returns only a status code for an error, without a message
	// so body = nil and in that case string(body) will fail.
	if body == nil {
		body = []byte{}
	}
	if resp != nil && resp.StatusCode != 200 {
		return fmt.Errorf("Failed to PUT: %d, %s", resp.StatusCode, string(body))
	}
	return nil
}
<commit_msg>PCT-562-B PR#123<commit_after>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   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/>
*/

package instance

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/percona/percona-agent/agent"

	"strconv"

	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/mrms"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
)

type empty struct{}

type Manager struct {
	logger    *pct.Logger
	configDir string
	api       pct.APIConnector
	// --
	status         *pct.Status
	repo           *Repo
	stopChan       chan empty
	mrm            mrms.Monitor
	mrmChans       map[string]<-chan bool
	mrmsGlobalChan chan string
	agentConfig    *agent.Config
}

func NewManager(logger *pct.Logger, configDir string, api pct.APIConnector, mrm mrms.Monitor) *Manager {
	repo := NewRepo(pct.NewLogger(logger.LogChan(), "instance-repo"), configDir, api)
	m := &Manager{
		logger:    logger,
		configDir: configDir,
		api:       api,
		// --
		status:         pct.NewStatus([]string{"instance", "instance-repo", "instance-mrms"}),
		repo:           repo,
		mrm:            mrm,
		mrmChans:       make(map[string]<-chan bool),
		mrmsGlobalChan: make(chan string, 100), // monitor up to 100 instances
	}
	return m
}

/////////////////////////////////////////////////////////////////////////////
// Interface
/////////////////////////////////////////////////////////////////////////////

// @goroutine[0]
func (m *Manager) Start() error {
	m.status.Update("instance", "Starting")
	if err := m.repo.Init(); err != nil {
		return err
	}
	m.logger.Info("Started")
	m.status.Update("instance", "Running")

	mrmsGlobalChan, err := m.mrm.GlobalSubscribe()
	if err != nil {
		return err
	}

	for _, instance := range m.GetMySQLInstances() {
		ch, err := m.mrm.Add(instance.DSN)
		if err != nil {
			m.logger.Error("Cannot add instance to the monitor:", err)
			continue
		}
		m.pushInstanceInfo(instance)
		// Store the channel to be able to remove it from mrms
		m.mrmChans[instance.DSN] = ch
	}
	go m.monitorInstancesRestart(mrmsGlobalChan)
	return nil
}

// @goroutine[0]
func (m *Manager) Stop() error {
	// Can't stop the instance manager.
	return nil
}

// @goroutine[0]
func (m *Manager) Handle(cmd *proto.Cmd) *proto.Reply {
	m.status.UpdateRe("instance", "Handling", cmd)
	defer m.status.Update("instance", "Running")

	it := &proto.ServiceInstance{}
	if err := json.Unmarshal(cmd.Data, it); err != nil {
		return cmd.Reply(nil, err)
	}

	switch cmd.Cmd {
	case "Add":
		err := m.repo.Add(it.Service, it.InstanceId, it.Instance, true) // true = write to disk
		if err != nil {
			return cmd.Reply(nil, err)
		}
		if it.Service == "mysql" {
			// Get the instance as type proto.MySQLInstance instead of proto.ServiceInstance
			// because we need the dsn field
			// We only return errors for repo.Add, not for mrm so all returns within this block
			// will return nil, nil
			iit := &proto.MySQLInstance{}
			err := m.repo.Get(it.Service, it.InstanceId, iit)
			if err != nil {
				m.logger.Error(err)
				return cmd.Reply(nil, nil)
			}
			ch, err := m.mrm.Add(iit.DSN)
			if err != nil {
				m.logger.Error(err)
				return cmd.Reply(nil, nil)
			}
			m.mrmChans[iit.DSN] = ch
			err = m.pushInstanceInfo(iit)
			if err != nil {
				m.logger.Error(err)
				return cmd.Reply(nil, nil)
			}
		}
		return cmd.Reply(nil, nil)
	case "Remove":
		if it.Service == "mysql" {
			// Get the instance as type proto.MySQLInstance instead of proto.ServiceInstance
			// because we need the dsn field
			iit := &proto.MySQLInstance{}
			err := m.repo.Get(it.Service, it.InstanceId, iit)
			// Don't return an error. This is just a remove from mrms
			if err != nil {
				m.logger.Error(err)
			} else {
				m.mrm.Remove(iit.DSN, m.mrmChans[iit.DSN])
			}
		}
		err := m.repo.Remove(it.Service, it.InstanceId)
		return cmd.Reply(nil, err)
	case "GetInfo":
		info, err := m.handleGetInfo(it.Service, it.Instance)
		return cmd.Reply(info, err)
	default:
		return cmd.Reply(nil, pct.UnknownCmdError{Cmd: cmd.Cmd})
	}
}

func (m *Manager) Status() map[string]string {
	m.status.Update("instance-repo", strings.Join(m.repo.List(), " "))
	return m.status.All()
}

func (m *Manager) GetConfig() ([]proto.AgentConfig, []error) {
	return nil, nil
}

func (m *Manager) Repo() *Repo {
	return m.repo
}

/////////////////////////////////////////////////////////////////////////////
// Implementation
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) handleGetInfo(service string, data []byte) (interface{}, error) {
	switch service {
	case "mysql":
		it := &proto.MySQLInstance{}
		if err := json.Unmarshal(data, it); err != nil {
			return nil, errors.New("instance.Repo:json.Unmarshal:" + err.Error())
		}
		if it.DSN == "" {
			return nil, fmt.Errorf("MySQL instance DSN is not set")
		}
		if err := GetMySQLInfo(it); err != nil {
			return nil, err
		}
		return it, nil
	default:
		return nil, fmt.Errorf("Don't know how to get info for %s service", service)
	}
}

func GetMySQLInfo(it *proto.MySQLInstance) error {
	conn := mysql.NewConnection(it.DSN)
	if err := conn.Connect(1); err != nil {
		return err
	}
	defer conn.Close()
	sql := "SELECT /* percona-agent */" +
		" CONCAT_WS('.', @@hostname, IF(@@port='3306',NULL,@@port)) AS Hostname," +
		" @@version_comment AS Distro," +
		" @@version AS Version"
	err := conn.DB().QueryRow(sql).Scan(
		&it.Hostname,
		&it.Distro,
		&it.Version,
	)
	if err != nil {
		return err
	}
	return nil
}

func (m *Manager) GetMySQLInstances() []*proto.MySQLInstance {
	m.logger.Debug("getMySQLInstances:call")
	defer m.logger.Debug("getMySQLInstances:return")

	var instances []*proto.MySQLInstance
	for _, name := range m.Repo().List() {
		parts := strings.Split(name, "-") // mysql-1 or server-12
		if len(parts) != 2 {
			m.logger.Error("Invalid instance name: %s: expected 2 parts, got %d", name, len(parts))
			continue
		}
		if parts[0] == "mysql" {
			id, err := strconv.ParseInt(parts[1], 10, 64)
			if err != nil {
				m.logger.Error("Invalid instance ID: %s: %s", name, err)
				continue
			}
			it := &proto.MySQLInstance{}
			if err := m.Repo().Get(parts[0], uint(id), it); err != nil {
				m.logger.Error("Failed to get instance %s: %s", name, err)
				continue
			}
			instances = append(instances, it)
		}
	}
	return instances
}

func (m *Manager) monitorInstancesRestart(ch chan string) {
	m.logger.Debug("monitorInstancesRestart:call")
	defer func() {
		if err := recover(); err != nil {
			m.logger.Error("MySQL connection crashed: ", err)
			m.status.Update("instance-mrms", "Crashed")
		} else {
			m.status.Update("instance-mrms", "Stopped")
		}
		m.logger.Debug("monitorInstancesRestart:return")
	}()

	ch, err := m.mrm.GlobalSubscribe()
	if err != nil {
		m.logger.Error(fmt.Sprintf("Failed to get MySQL restart monitor global channel: %s", err))
		return
	}

	for {
		m.status.Update("instance-mrms", "Idle")
		select {
		case dsn := <-ch:
			safeDSN := mysql.HideDSNPassword(dsn)
			m.logger.Debug("mrms:restart:" + safeDSN)
			m.status.Update("instance-mrms", "Updating "+safeDSN)

			// Get the updated instances list. It should be updated every time since
			// the Add method can add new instances to the list.
			for _, instance := range m.GetMySQLInstances() {
				if instance.DSN != dsn {
					continue
				}
				m.status.Update("instance-mrms", "Updating info "+safeDSN)
				err := m.pushInstanceInfo(instance)
				if err != nil {
					m.logger.Warn(err)
				}
			}
		}
	}
}

func (m *Manager) pushInstanceInfo(instance *proto.MySQLInstance) error {
	safeDSN := mysql.HideDSNPassword(instance.DSN)
	m.status.Update("instance-mrms", "Getting info "+safeDSN)
	if err := GetMySQLInfo(instance); err != nil {
		m.logger.Warn(fmt.Sprintf("Failed to get MySQL info %s: %s", safeDSN, err))
		return err
	}

	uri := fmt.Sprintf("%s/%s/%d", m.api.EntryLink("instances"), "mysql", instance.Id)
	data, err := json.Marshal(instance)
	if err != nil {
		m.logger.Error(err)
		return err
	}
	resp, body, err := m.api.Put(m.api.ApiKey(), uri, data)
	if err != nil {
		return err
	}
	// Sometimes the API returns only a status code for an error, without a message
	// so body = nil and in that case string(body) will fail.
	if body == nil {
		body = []byte{}
	}
	if resp != nil && resp.StatusCode != 200 {
		return fmt.Errorf("Failed to PUT: %d, %s", resp.StatusCode, string(body))
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build docker_integration

package docker

import (
	"bytes"
	"testing"
)

func TestIntegrationPullCreateStartLogs(t *testing.T) {
	imageName := pullImage(t)
	client, err := NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	hostConfig := HostConfig{PublishAllPorts: true}
	createOpts := CreateContainerOptions{
		Config: &Config{
			Image: imageName,
			Cmd:   []string{"cat", "/home/gopher/file.txt"},
			User:  "gopher",
		},
		HostConfig: &hostConfig,
	}
	container, err := client.CreateContainer(createOpts)
	if err != nil {
		t.Fatal(err)
	}
	err = client.StartContainer(container.ID, &hostConfig)
	if err != nil {
		t.Fatal(err)
	}
	status, err := client.WaitContainer(container.ID)
	if err != nil {
		t.Error(err)
	}
	if status != 0 {
		t.Errorf("WaitContainer(%q): wrong status. Want 0. Got %d", container.ID, status)
	}
	var stdout, stderr bytes.Buffer
	logsOpts := LogsOptions{
		Container:    container.ID,
		OutputStream: &stdout,
		ErrorStream:  &stderr,
		Stdout:       true,
		Stderr:       true,
	}
	err = client.Logs(logsOpts)
	if err != nil {
		t.Error(err)
	}
	if stderr.String() != "" {
		t.Errorf("Got unexpected stderr from logs: %q", stderr.String())
	}
	expected := `Welcome to reality, wake up and rejoice
Welcome to reality, you've made the right choice
Welcome to reality, and let them hear your voice, shout it out!
`
	if stdout.String() != expected {
		t.Errorf("Got wrong stdout from logs.\nWant:\n%#v.\n\nGot:\n%#v.", expected, stdout.String())
	}
}

func pullImage(t *testing.T) string {
	t.Helper()
	imageName := "fsouza/go-dockerclient-integration:latest"
	var buf bytes.Buffer
	pullOpts := PullImageOptions{
		Repository:   imageName,
		OutputStream: &buf,
	}
	client, err := NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	err = client.PullImage(pullOpts, AuthConfiguration{})
	if err != nil {
		t.Logf("Pull output: %s", buf.String())
		t.Fatal(err)
	}
	return imageName
}
<commit_msg>integration_test: pullImage is not really a test helper<commit_after>// Copyright 2015 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build docker_integration

package docker

import (
	"bytes"
	"testing"
)

func TestIntegrationPullCreateStartLogs(t *testing.T) {
	imageName := pullImage(t)
	client, err := NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	hostConfig := HostConfig{PublishAllPorts: true}
	createOpts := CreateContainerOptions{
		Config: &Config{
			Image: imageName,
			Cmd:   []string{"cat", "/home/gopher/file.txt"},
			User:  "gopher",
		},
		HostConfig: &hostConfig,
	}
	container, err := client.CreateContainer(createOpts)
	if err != nil {
		t.Fatal(err)
	}
	err = client.StartContainer(container.ID, &hostConfig)
	if err != nil {
		t.Fatal(err)
	}
	status, err := client.WaitContainer(container.ID)
	if err != nil {
		t.Error(err)
	}
	if status != 0 {
		t.Errorf("WaitContainer(%q): wrong status. Want 0. Got %d", container.ID, status)
	}
	var stdout, stderr bytes.Buffer
	logsOpts := LogsOptions{
		Container:    container.ID,
		OutputStream: &stdout,
		ErrorStream:  &stderr,
		Stdout:       true,
		Stderr:       true,
	}
	err = client.Logs(logsOpts)
	if err != nil {
		t.Error(err)
	}
	if stderr.String() != "" {
		t.Errorf("Got unexpected stderr from logs: %q", stderr.String())
	}
	expected := `Welcome to reality, wake up and rejoice
Welcome to reality, you've made the right choice
Welcome to reality, and let them hear your voice, shout it out!
`
	if stdout.String() != expected {
		t.Errorf("Got wrong stdout from logs.\nWant:\n%#v.\n\nGot:\n%#v.", expected, stdout.String())
	}
}

func pullImage(t *testing.T) string {
	imageName := "fsouza/go-dockerclient-integration:latest"
	var buf bytes.Buffer
	pullOpts := PullImageOptions{
		Repository:   imageName,
		OutputStream: &buf,
	}
	client, err := NewClientFromEnv()
	if err != nil {
		t.Fatal(err)
	}
	err = client.PullImage(pullOpts, AuthConfiguration{})
	if err != nil {
		t.Logf("Pull output: %s", buf.String())
		t.Fatal(err)
	}
	return imageName
}
<|endoftext|>