<commit_before>package lock

import (
	"fmt"
	"testing"
	"time"

	"github.com/coreos/etcd/server"
	"github.com/coreos/etcd/tests"
	"github.com/coreos/etcd/third_party/github.com/stretchr/testify/assert"
)

// Ensure that a lock can be acquired and released.
func TestModLockAcquireAndRelease(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "", 10)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Release lock.
		body, status, err = testReleaseLock(s, "foo", "2", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock can be acquired and another process is blocked until released.
func TestModLockBlockUntilAcquire(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		c := make(chan bool)

		// Acquire lock #1.
		go func() {
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "2")
			c <- true
		}()
		<-c

		// Acquire lock #2.
		waiting := true
		go func() {
			c <- true
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "4")
			waiting = false
		}()
		<-c

		time.Sleep(1 * time.Second)

		// Check that we have the lock #1.
		body, status, err := testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we are still waiting for lock #2.
		assert.Equal(t, waiting, true)

		// Release lock #1.
		_, status, err = testReleaseLock(s, "foo", "2", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)

		// Check that we have lock #2.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "4")

		// Release lock #2.
		_, status, err = testReleaseLock(s, "foo", "4", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)

		// Check that we have no lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock will be released after the TTL.
func TestModLockExpireAndRelease(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		c := make(chan bool)

		// Acquire lock #1.
		go func() {
			body, status, err := testAcquireLock(s, "foo", "", 2)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "2")
			c <- true
		}()
		<-c

		// Acquire lock #2.
		go func() {
			c <- true
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "4")
		}()
		<-c

		time.Sleep(1 * time.Second)

		// Check that we have the lock #1.
		body, status, err := testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Wait for lock #1 TTL.
		time.Sleep(2 * time.Second)

		// Check that we have lock #2.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "4")
	})
}

// Ensure that a lock can be renewed.
func TestModLockRenew(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "", 3)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		time.Sleep(2 * time.Second)

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Renew lock.
		body, status, err = testRenewLock(s, "foo", "2", "", 3)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		time.Sleep(2 * time.Second)

		// Check that we still have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		time.Sleep(2 * time.Second)

		// Check that lock was released.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock can be acquired with a value and released by value.
func TestModLockAcquireAndReleaseByValue(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "XXX", 10)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we have the lock.
		body, status, err = testGetLockValue(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "XXX")

		// Release lock.
		body, status, err = testReleaseLock(s, "foo", "", "XXX")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		// Check that we released the lock.
		body, status, err = testGetLockValue(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

func testAcquireLock(s *server.Server, key string, value string, ttl int) (string, int, error) {
	resp, err := tests.PostForm(fmt.Sprintf("%s/mod/v2/lock/%s?value=%s&ttl=%d", s.URL(), key, value, ttl), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testGetLockIndex(s *server.Server, key string) (string, int, error) {
	resp, err := tests.Get(fmt.Sprintf("%s/mod/v2/lock/%s?field=index", s.URL(), key))
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testGetLockValue(s *server.Server, key string) (string, int, error) {
	resp, err := tests.Get(fmt.Sprintf("%s/mod/v2/lock/%s", s.URL(), key))
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testReleaseLock(s *server.Server, key string, index string, value string) (string, int, error) {
	resp, err := tests.DeleteForm(fmt.Sprintf("%s/mod/v2/lock/%s?index=%s&value=%s", s.URL(), key, index, value), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testRenewLock(s *server.Server, key string, index string, value string, ttl int) (string, int, error) {
	resp, err := tests.PutForm(fmt.Sprintf("%s/mod/v2/lock/%s?index=%s&value=%s&ttl=%d", s.URL(), key, index, value, ttl), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}
<commit_msg>Add lock timeout test<commit_after>package lock

import (
	"fmt"
	"testing"
	"time"

	"github.com/coreos/etcd/server"
	"github.com/coreos/etcd/tests"
	"github.com/coreos/etcd/third_party/github.com/stretchr/testify/assert"
)

// Ensure that a lock can be acquired and released.
func TestModLockAcquireAndRelease(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "", 10)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Release lock.
		body, status, err = testReleaseLock(s, "foo", "2", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock can be acquired and another process is blocked until released.
func TestModLockBlockUntilAcquire(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		c := make(chan bool)

		// Acquire lock #1.
		go func() {
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "2")
			c <- true
		}()
		<-c

		// Acquire lock #2.
		waiting := true
		go func() {
			c <- true
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "4")
			waiting = false
		}()
		<-c

		time.Sleep(1 * time.Second)

		// Check that we have the lock #1.
		body, status, err := testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we are still waiting for lock #2.
		assert.Equal(t, waiting, true)

		// Release lock #1.
		_, status, err = testReleaseLock(s, "foo", "2", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)

		// Check that we have lock #2.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "4")

		// Release lock #2.
		_, status, err = testReleaseLock(s, "foo", "4", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)

		// Check that we have no lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock will be released after the TTL.
func TestModLockExpireAndRelease(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		c := make(chan bool)

		// Acquire lock #1.
		go func() {
			body, status, err := testAcquireLock(s, "foo", "", 2)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "2")
			c <- true
		}()
		<-c

		// Acquire lock #2.
		go func() {
			c <- true
			body, status, err := testAcquireLock(s, "foo", "", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "4")
		}()
		<-c

		time.Sleep(1 * time.Second)

		// Check that we have the lock #1.
		body, status, err := testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Wait for lock #1 TTL.
		time.Sleep(2 * time.Second)

		// Check that we have lock #2.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "4")
	})
}

// Ensure that a lock can be renewed.
func TestModLockRenew(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "", 3)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		time.Sleep(2 * time.Second)

		// Check that we have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Renew lock.
		body, status, err = testRenewLock(s, "foo", "2", "", 3)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		time.Sleep(2 * time.Second)

		// Check that we still have the lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		time.Sleep(2 * time.Second)

		// Check that lock was released.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock can be acquired with a value and released by value.
func TestModLockAcquireAndReleaseByValue(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		// Acquire lock.
		body, status, err := testAcquireLock(s, "foo", "XXX", 10)
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we have the lock.
		body, status, err = testGetLockValue(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "XXX")

		// Release lock.
		body, status, err = testReleaseLock(s, "foo", "", "XXX")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")

		// Check that we released the lock.
		body, status, err = testGetLockValue(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "")
	})
}

// Ensure that a lock honours the timeout option
func TestModLockAcquireTimeout(t *testing.T) {
	tests.RunServer(func(s *server.Server) {
		c := make(chan bool)

		// Acquire lock #1.
		go func() {
			body, status, err := testAcquireLock(s, "foo", "first", 10)
			assert.NoError(t, err)
			assert.Equal(t, status, 200)
			assert.Equal(t, body, "2")
			c <- true
		}()
		<-c

		// Attempt to acquire lock #2, timing out after 1s.
		waiting := true
		go func() {
			c <- true
			_, status, err := testAcquireLockWithTimeout(s, "foo", "second", 10, 1)
			assert.NoError(t, err)
			assert.Equal(t, status, 500)
			waiting = false
		}()
		<-c

		time.Sleep(5 * time.Second)

		// Check that we have the lock #1.
		body, status, err := testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
		assert.Equal(t, body, "2")

		// Check that we are not still waiting for lock #2.
		assert.Equal(t, waiting, false)

		// Release lock #1.
		_, status, err = testReleaseLock(s, "foo", "2", "")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)

		// Check that we have no lock.
		body, status, err = testGetLockIndex(s, "foo")
		assert.NoError(t, err)
		assert.Equal(t, status, 200)
	})
}

func testAcquireLock(s *server.Server, key string, value string, ttl int) (string, int, error) {
	resp, err := tests.PostForm(fmt.Sprintf("%s/mod/v2/lock/%s?value=%s&ttl=%d", s.URL(), key, value, ttl), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testAcquireLockWithTimeout(s *server.Server, key string, value string, ttl int, timeout int) (string, int, error) {
	resp, err := tests.PostForm(fmt.Sprintf("%s/mod/v2/lock/%s?value=%s&ttl=%d&timeout=%d", s.URL(), key, value, ttl, timeout), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testGetLockIndex(s *server.Server, key string) (string, int, error) {
	resp, err := tests.Get(fmt.Sprintf("%s/mod/v2/lock/%s?field=index", s.URL(), key))
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testGetLockValue(s *server.Server, key string) (string, int, error) {
	resp, err := tests.Get(fmt.Sprintf("%s/mod/v2/lock/%s", s.URL(), key))
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testReleaseLock(s *server.Server, key string, index string, value string) (string, int, error) {
	resp, err := tests.DeleteForm(fmt.Sprintf("%s/mod/v2/lock/%s?index=%s&value=%s", s.URL(), key, index, value), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}

func testRenewLock(s *server.Server, key string, index string, value string, ttl int) (string, int, error) {
	resp, err := tests.PutForm(fmt.Sprintf("%s/mod/v2/lock/%s?index=%s&value=%s&ttl=%d", s.URL(), key, index, value, ttl), nil)
	ret := tests.ReadBody(resp)
	return string(ret), resp.StatusCode, err
}
<|endoftext|>


<commit_before>package connector

import (
	"fmt"
	"github.com/op/go-logging"
	"strings"
	"sync"

	"github.com/bcicen/ctop/connector/collector"
	"github.com/bcicen/ctop/connector/manager"
	"github.com/bcicen/ctop/container"
	api "github.com/fsouza/go-dockerclient"
)

func init() { enabled["docker"] = NewDocker }

type Docker struct {
	client       *api.Client
	containers   map[string]*container.Container
	needsRefresh chan string // container IDs requiring refresh
	closed       chan struct{}
	lock         sync.RWMutex
}

func NewDocker() (Connector, error) {
	// init docker client
	client, err := api.NewClientFromEnv()
	if err != nil {
		return nil, err
	}
	cm := &Docker{
		client:       client,
		containers:   make(map[string]*container.Container),
		needsRefresh: make(chan string, 60),
		closed:       make(chan struct{}),
		lock:         sync.RWMutex{},
	}

	// query info as pre-flight healthcheck
	info, err := client.Info()
	if err != nil {
		return nil, err
	}

	log.Debugf("docker-connector ID: %s", info.ID)
	log.Debugf("docker-connector Driver: %s", info.Driver)
	log.Debugf("docker-connector Images: %d", info.Images)
	log.Debugf("docker-connector Name: %s", info.Name)
	log.Debugf("docker-connector ServerVersion: %s", info.ServerVersion)

	go cm.Loop()
	cm.refreshAll()
	go cm.watchEvents()
	return cm, nil
}

// Docker implements Connector
func (cm *Docker) Wait() struct{} { return <-cm.closed }

// Docker events watcher
func (cm *Docker) watchEvents() {
	log.Info("docker event listener starting")
	events := make(chan *api.APIEvents)
	cm.client.AddEventListener(events)

	for e := range events {
		if e.Type != "container" {
			continue
		}

		actionName := e.Action
		// Action may have additional param: "exec_create: redis-cli ping" or "health_status: healthy"
		// We need to strip to have only action name
		sepIdx := strings.Index(actionName, ": ")
		if sepIdx != -1 {
			actionName = actionName[:sepIdx]
		}

		switch actionName {
		case "start", "die", "pause", "unpause", "health_status":
			if log.IsEnabledFor(logging.DEBUG) {
				log.Debugf("handling docker event: action=%s id=%s", e.Action, e.ID)
			}
			cm.needsRefresh <- e.ID
		case "destroy":
			if log.IsEnabledFor(logging.DEBUG) {
				log.Debugf("handling docker event: action=destroy id=%s", e.ID)
			}
			cm.delByID(e.ID)
		}
	}
	log.Info("docker event listener exited")
	close(cm.closed)
}

func portsFormat(ports map[api.Port][]api.PortBinding) string {
	var exposed []string
	var published []string

	for k, v := range ports {
		if len(v) == 0 {
			exposed = append(exposed, string(k))
			continue
		}
		for _, binding := range v {
			s := fmt.Sprintf("%s:%s -> %s", binding.HostIP, binding.HostPort, k)
			published = append(published, s)
		}
	}

	return strings.Join(append(exposed, published...), "\n")
}

func ipsFormat(networks map[string]api.ContainerNetwork) string {
	var ips []string

	for k, v := range networks {
		s := fmt.Sprintf("%s:%s", k, v.IPAddress)
		ips = append(ips, s)
	}

	return strings.Join(ips, "\n")
}

func (cm *Docker) refresh(c *container.Container) {
	insp := cm.inspect(c.Id)
	// remove container if no longer exists
	if insp == nil {
		cm.delByID(c.Id)
		return
	}
	c.SetMeta("name", shortName(insp.Name))
	c.SetMeta("image", insp.Config.Image)
	c.SetMeta("IPs", ipsFormat(insp.NetworkSettings.Networks))
	c.SetMeta("ports", portsFormat(insp.NetworkSettings.Ports))
	c.SetMeta("created", insp.Created.Format("Mon Jan 2 15:04:05 2006"))
	c.SetMeta("health", insp.State.Health.Status)
	for _, env := range insp.Config.Env {
		c.SetMeta("[ENV-VAR]", env)
	}
	c.SetState(insp.State.Status)
}

func (cm *Docker) inspect(id string) *api.Container {
	c, err := cm.client.InspectContainer(id)
	if err != nil {
		if _, ok := err.(*api.NoSuchContainer); !ok {
			log.Errorf("%s (%T)", err.Error(), err)
		}
	}
	return c
}

// Mark all container IDs for refresh
func (cm *Docker) refreshAll() {
	opts := api.ListContainersOptions{All: true}
	allContainers, err := cm.client.ListContainers(opts)
	if err != nil {
		log.Errorf("%s (%T)", err.Error(), err)
		return
	}

	for _, i := range allContainers {
		c := cm.MustGet(i.ID)
		c.SetMeta("name", shortName(i.Names[0]))
		c.SetState(i.State)
		cm.needsRefresh <- c.Id
	}
}

func (cm *Docker) Loop() {
	for {
		select {
		case id := <-cm.needsRefresh:
			c := cm.MustGet(id)
			cm.refresh(c)
		case <-cm.closed:
			return
		}
	}
}

// MustGet gets a single container, creating one anew if not existing
func (cm *Docker) MustGet(id string) *container.Container {
	c, ok := cm.Get(id)
	// append container struct for new containers
	if !ok {
		// create collector
		collector := collector.NewDocker(cm.client, id)
		// create manager
		manager := manager.NewDocker(cm.client, id)
		// create container
		c = container.New(id, collector, manager)
		cm.lock.Lock()
		cm.containers[id] = c
		cm.lock.Unlock()
	}
	return c
}

// Docker implements Connector
func (cm *Docker) Get(id string) (*container.Container, bool) {
	cm.lock.Lock()
	c, ok := cm.containers[id]
	cm.lock.Unlock()
	return c, ok
}

// Remove containers by ID
func (cm *Docker) delByID(id string) {
	cm.lock.Lock()
	delete(cm.containers, id)
	cm.lock.Unlock()
	log.Infof("removed dead container: %s", id)
}

// Docker implements Connector
func (cm *Docker) All() (containers container.Containers) {
	cm.lock.Lock()
	for _, c := range cm.containers {
		containers = append(containers, c)
	}

	containers.Sort()
	containers.Filter()
	cm.lock.Unlock()
	return containers
}

// use primary container name
func shortName(name string) string {
	return strings.Replace(name, "/", "", 1)
}
<commit_msg>Improve docker events handling<commit_after>package connector

import (
	"fmt"
	"github.com/op/go-logging"
	"strings"
	"sync"

	"github.com/bcicen/ctop/connector/collector"
	"github.com/bcicen/ctop/connector/manager"
	"github.com/bcicen/ctop/container"
	api "github.com/fsouza/go-dockerclient"
)

func init() { enabled["docker"] = NewDocker }

var actionToStatus = map[string]string{
	"create":  "created",
	"start":   "running",
	"kill":    "exited",
	"stop":    "exited",
	"pause":   "paused",
	"unpause": "running",
}

type Docker struct {
	client       *api.Client
	containers   map[string]*container.Container
	needsRefresh chan string // container IDs requiring refresh
	closed       chan struct{}
	lock         sync.RWMutex
}

func NewDocker() (Connector, error) {
	// init docker client
	client, err := api.NewClientFromEnv()
	if err != nil {
		return nil, err
	}
	cm := &Docker{
		client:       client,
		containers:   make(map[string]*container.Container),
		needsRefresh: make(chan string, 60),
		closed:       make(chan struct{}),
		lock:         sync.RWMutex{},
	}

	// query info as pre-flight healthcheck
	info, err := client.Info()
	if err != nil {
		return nil, err
	}

	log.Debugf("docker-connector ID: %s", info.ID)
	log.Debugf("docker-connector Driver: %s", info.Driver)
	log.Debugf("docker-connector Images: %d", info.Images)
	log.Debugf("docker-connector Name: %s", info.Name)
	log.Debugf("docker-connector ServerVersion: %s", info.ServerVersion)

	go cm.Loop()
	cm.refreshAll()
	go cm.watchEvents()
	return cm, nil
}

// Docker implements Connector
func (cm *Docker) Wait() struct{} { return <-cm.closed }

// Docker events watcher
func (cm *Docker) watchEvents() {
	log.Info("docker event listener starting")
	events := make(chan *api.APIEvents)
	cm.client.AddEventListener(events)

	for e := range events {
		if e.Type != "container" {
			continue
		}

		actionName := e.Action
		// Action may have additional param: "exec_create: redis-cli ping" or "health_status: healthy"
		// We need to strip to have only action name
		sepIdx := strings.Index(actionName, ": ")
		if sepIdx != -1 {
			actionName = actionName[:sepIdx]
		}

		switch actionName {
		// most frequent event is a health checks
		case "health_status":
			healthStatus := e.Action[sepIdx+2:]
			if log.IsEnabledFor(logging.DEBUG) {
				log.Debugf("handling docker event: action=health_status id=%s %s", e.ID, healthStatus)
			}
			cm.needsRefresh <- e.ID
		case "create":
			if log.IsEnabledFor(logging.DEBUG) {
				log.Debugf("handling docker event: action=create id=%s", e.ID)
			}
			cm.needsRefresh <- e.ID
		case "destroy":
			if log.IsEnabledFor(logging.DEBUG) {
				log.Debugf("handling docker event: action=destroy id=%s", e.ID)
			}
			cm.delByID(e.ID)
		default:
			// check if this action changes status e.g. start -> running
			status := actionToStatus[actionName]
			if status != "" {
				if log.IsEnabledFor(logging.DEBUG) {
					log.Debugf("handling docker event: action=%s id=%s %s", actionName, e.ID, status)
				}
				cm.needsRefresh <- e.ID
			}
		}
	}
	log.Info("docker event listener exited")
	close(cm.closed)
}

func portsFormat(ports map[api.Port][]api.PortBinding) string {
	var exposed []string
	var published []string

	for k, v := range ports {
		if len(v) == 0 {
			exposed = append(exposed, string(k))
			continue
		}
		for _, binding := range v {
			s := fmt.Sprintf("%s:%s -> %s", binding.HostIP, binding.HostPort, k)
			published = append(published, s)
		}
	}

	return strings.Join(append(exposed, published...), "\n")
}

func ipsFormat(networks map[string]api.ContainerNetwork) string {
	var ips []string

	for k, v := range networks {
		s := fmt.Sprintf("%s:%s", k, v.IPAddress)
		ips = append(ips, s)
	}

	return strings.Join(ips, "\n")
}

func (cm *Docker) refresh(c *container.Container) {
	insp := cm.inspect(c.Id)
	// remove container if no longer exists
	if insp == nil {
		cm.delByID(c.Id)
		return
	}
	c.SetMeta("name", shortName(insp.Name))
	c.SetMeta("image", insp.Config.Image)
	c.SetMeta("IPs", ipsFormat(insp.NetworkSettings.Networks))
	c.SetMeta("ports", portsFormat(insp.NetworkSettings.Ports))
	c.SetMeta("created", insp.Created.Format("Mon Jan 2 15:04:05 2006"))
	c.SetMeta("health", insp.State.Health.Status)
	for _, env := range insp.Config.Env {
		c.SetMeta("[ENV-VAR]", env)
	}
	c.SetState(insp.State.Status)
}

func (cm *Docker) inspect(id string) *api.Container {
	c, err := cm.client.InspectContainer(id)
	if err != nil {
		if _, ok := err.(*api.NoSuchContainer); !ok {
			log.Errorf("%s (%T)", err.Error(), err)
		}
	}
	return c
}

// Mark all container IDs for refresh
func (cm *Docker) refreshAll() {
	opts := api.ListContainersOptions{All: true}
	allContainers, err := cm.client.ListContainers(opts)
	if err != nil {
		log.Errorf("%s (%T)", err.Error(), err)
		return
	}

	for _, i := range allContainers {
		c := cm.MustGet(i.ID)
		c.SetMeta("name", shortName(i.Names[0]))
		c.SetState(i.State)
		cm.needsRefresh <- c.Id
	}
}

func (cm *Docker) Loop() {
	for {
		select {
		case id := <-cm.needsRefresh:
			c := cm.MustGet(id)
			cm.refresh(c)
		case <-cm.closed:
			return
		}
	}
}

// MustGet gets a single container, creating one anew if not existing
func (cm *Docker) MustGet(id string) *container.Container {
	c, ok := cm.Get(id)
	// append container struct for new containers
	if !ok {
		// create collector
		collector := collector.NewDocker(cm.client, id)
		// create manager
		manager := manager.NewDocker(cm.client, id)
		// create container
		c = container.New(id, collector, manager)
		cm.lock.Lock()
		cm.containers[id] = c
		cm.lock.Unlock()
	}
	return c
}

// Docker implements Connector
func (cm *Docker) Get(id string) (*container.Container, bool) {
	cm.lock.Lock()
	c, ok := cm.containers[id]
	cm.lock.Unlock()
	return c, ok
}

// Remove containers by ID
func (cm *Docker) delByID(id string) {
	cm.lock.Lock()
	delete(cm.containers, id)
	cm.lock.Unlock()
	log.Infof("removed dead container: %s", id)
}

// Docker implements Connector
func (cm *Docker) All() (containers container.Containers) {
	cm.lock.Lock()
	for _, c := range cm.containers {
		containers = append(containers, c)
	}

	containers.Sort()
	containers.Filter()
	cm.lock.Unlock()
	return containers
}

// use primary container name
func shortName(name string) string {
	return strings.Replace(name, "/", "", 1)
}
<|endoftext|>


<commit_before>package sdbot

import (
	"bytes"
	"fmt"
	"strings"
	"time"
)

// The colour codes for tty to be used with colourize.
const (
	reset   string = "\x1b[0m"
	bold    string = "\x1b[1m"
	red     string = "\x1b[31m"
	green   string = "\x1b[32m"
	yellow  string = "\x1b[33m"
	blue    string = "\x1b[34m"
	black   string = "\x1b[30m"
	bgWhite string = "\x1b[47m"
)

// PrettyLogger logs everything with pretty colours. Looks good with the
// Solarized terminal theme.
// Don't like how it looks? Make your own!
type PrettyLogger struct {
	AnyLogger
}

func (lo *PrettyLogger) formatMessage(s string, level int) string {
	switch level {
	case levelDebug:
		fallthrough
	case levelWarn:
		return formatDebug(s)
	case levelInfo:
		return formatInfo(s)
	case levelError:
		return formatError(s)
	case levelIncoming:
		return formatIncoming(s)
	case levelOutgoing:
		return formatOutgoing(s)
	default:
		return formatGeneral(s)
	}
}

func formatDebug(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), colourize("!!", yellow), s)
}

func formatInfo(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), "II", s)
}

func formatError(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), colourize("!!", red), s)
}

func formatGeneral(s string) string {
	return fmt.Sprintf("%s %s\n", timestamp(), s)
}

func formatIncoming(s string) string {
	split := strings.Split(s, "\n")
	if len(split) < 2 {
		return formatGeneral(s)
	}
	room := split[0]
	rest := split[1]

	parts := strings.Split(rest, "|")
	prefix := colourize(">>", green)

	if room == "" {
		// Private messages
		if len(parts) > 1 {
			parts = parts[1:]
			return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, colourize(parts[0], blue), strings.Join(parts[1:], "|"))
		}
		return fmt.Sprintf("%s %s %s\n", timestamp(), prefix, parts[0])
	}

	room = colourize(room[1:], bold)

	if len(parts) == 0 {
		return fmt.Sprintf("%s %s %s\n", timestamp(), prefix, room)
	}

	if len(parts) == 1 {
		// Raw server messages/ban messages
		return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, room, colourize(parts[0], bold))
	}

	cmd := colourize(parts[1], blue)
	params := strings.Join(parts[2:], "|")

	return fmt.Sprintf("%s %s %s|%s|%s\n", timestamp(), prefix, room, cmd, params)
}

func formatOutgoing(s string) string {
	split := strings.Split(s, "|")
	room := colourize(split[0], bold)
	rest := split[1]
	prefix := colourize("<<", red)

	return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, room, rest)
}

func timestamp() string {
	return strings.Join([]string{"[", time.Now().Format(time.RFC3339), "]"}, "")
}

func colourize(s string, codes ...string) string {
	var buffer bytes.Buffer
	for _, code := range codes {
		buffer.WriteString(code)
	}

	return strings.Join([]string{buffer.String(), s, reset}, "")
}
<commit_msg>pretty_logger: add fatal level logging<commit_after>package sdbot

import (
	"bytes"
	"fmt"
	"strings"
	"time"
)

// The colour codes for tty to be used with colourize.
const (
	reset   string = "\x1b[0m"
	bold    string = "\x1b[1m"
	red     string = "\x1b[31m"
	green   string = "\x1b[32m"
	yellow  string = "\x1b[33m"
	blue    string = "\x1b[34m"
	black   string = "\x1b[30m"
	bgWhite string = "\x1b[47m"
)

// PrettyLogger logs everything with pretty colours. Looks good with the
// Solarized terminal theme.
// Don't like how it looks? Make your own!
type PrettyLogger struct {
	AnyLogger
}

func (lo *PrettyLogger) formatMessage(s string, level int) string {
	switch level {
	case levelDebug:
		fallthrough
	case levelWarn:
		return formatDebug(s)
	case levelInfo:
		return formatInfo(s)
	case levelFatal:
		fallthrough
	case levelError:
		return formatError(s)
	case levelIncoming:
		return formatIncoming(s)
	case levelOutgoing:
		return formatOutgoing(s)
	default:
		return formatGeneral(s)
	}
}

func formatDebug(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), colourize("!!", yellow), s)
}

func formatInfo(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), "II", s)
}

func formatError(s string) string {
	return fmt.Sprintf("%s %s %s\n", timestamp(), colourize("!!", red), s)
}

func formatGeneral(s string) string {
	return fmt.Sprintf("%s %s\n", timestamp(), s)
}

func formatIncoming(s string) string {
	split := strings.Split(s, "\n")
	if len(split) < 2 {
		return formatGeneral(s)
	}
	room := split[0]
	rest := split[1]

	parts := strings.Split(rest, "|")
	prefix := colourize(">>", green)

	if room == "" {
		// Private messages
		if len(parts) > 1 {
			parts = parts[1:]
			return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, colourize(parts[0], blue), strings.Join(parts[1:], "|"))
		}
		return fmt.Sprintf("%s %s %s\n", timestamp(), prefix, parts[0])
	}

	room = colourize(room[1:], bold)

	if len(parts) == 0 {
		return fmt.Sprintf("%s %s %s\n", timestamp(), prefix, room)
	}

	if len(parts) == 1 {
		// Raw server messages/ban messages
		return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, room, colourize(parts[0], bold))
	}

	cmd := colourize(parts[1], blue)
	params := strings.Join(parts[2:], "|")

	return fmt.Sprintf("%s %s %s|%s|%s\n", timestamp(), prefix, room, cmd, params)
}

func formatOutgoing(s string) string {
	split := strings.Split(s, "|")
	room := colourize(split[0], bold)
	rest := split[1]
	prefix := colourize("<<", red)

	return fmt.Sprintf("%s %s %s|%s\n", timestamp(), prefix, room, rest)
}

func timestamp() string {
	return strings.Join([]string{"[", time.Now().Format(time.RFC3339), "]"}, "")
}

func colourize(s string, codes ...string) string {
	var buffer bytes.Buffer
	for _, code := range codes {
		buffer.WriteString(code)
	}

	return strings.Join([]string{buffer.String(), s, reset}, "")
}
<|endoftext|>


<commit_before>package kinetic

import (
	"errors"
	"syscall"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestProducerStop(t *testing.T) {
	producer, _ := new(Producer).Init()
	producer.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running producer", t, func() {
		go producer.produce()

		Convey("It should stop producing if sent an interrupt signal", func() {
			producer.interrupts <- syscall.SIGINT

			// Wait for it to stop
			for {
				if !producer.IsProducing() {
					break
				}
			}

			So(producer.IsProducing(), ShouldEqual, false)
		})
	})

	producer.Close()
}

func TestProducerError(t *testing.T) {
	producer, _ := new(Producer).Init()
	producer.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running producer", t, func() {
		go producer.produce()

		Convey("It should handle errors successfully", func() {
			producer.errors <- errors.New("All your base are belong to us!")
			// Let the error propagate
			<-time.After(1 * time.Second)
			So(producer.getErrCount(), ShouldEqual, 1)
			So(producer.IsProducing(), ShouldEqual, true)
		})
	})

	producer.Close()
}

func TestProducerMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4)
	producer, _ := new(Producer).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")

	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				if err != nil {
					t.Fatalf(err.Error())
				}

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}

	listener.Close()
	producer.Close()
}
<commit_msg>give sig_int more time for circle box<commit_after>package kinetic

import (
	"errors"
	"syscall"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestProducerStop(t *testing.T) {
	producer, _ := new(Producer).Init()
	producer.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running producer", t, func() {
		go producer.produce()

		Convey("It should stop producing if sent an interrupt signal", func() {
			producer.interrupts <- syscall.SIGINT

			// Wait for it to stop
			for {
				if !producer.IsProducing() {
					break
				}
			}

			So(producer.IsProducing(), ShouldEqual, false)
		})
	})

	producer.Close()
}

func TestProducerError(t *testing.T) {
	producer, _ := new(Producer).Init()
	producer.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running producer", t, func() {
		go producer.produce()

		Convey("It should handle errors successfully", func() {
			producer.errors <- errors.New("All your base are belong to us!")
			// Let the error propagate
			<-time.After(3 * time.Second)
			So(producer.getErrCount(), ShouldEqual, 1)
			So(producer.IsProducing(), ShouldEqual, true)
		})
	})

	producer.Close()
}

func TestProducerMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4)
	producer, _ := new(Producer).InitC("your-stream", "0", "LATEST", "accesskey", "secretkey", "us-east-1", 4)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")

	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				if err != nil {
					t.Fatalf(err.Error())
				}

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}

	listener.Close()
	producer.Close()
}
<|endoftext|>


<commit_before>package mpb_test

import (
	"context"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"

	"github.com/vbauerster/mpb/v7"
	"github.com/vbauerster/mpb/v7/decor"
)

const (
	timeout = 200 * time.Millisecond
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func TestBarCount(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	b := p.AddBar(0, mpb.BarRemoveOnComplete())

	if count := p.BarCount(); count != 1 {
		t.Errorf("BarCount want: %d, got: %d\n", 1, count)
	}

	b.SetTotal(100, true)

	if count := p.BarCount(); count != 0 {
		t.Errorf("BarCount want: %d, got: %d\n", 0, count)
	}

	go p.Wait()

	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

func TestBarAbort(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))
	n := 2
	bars := make([]*mpb.Bar, n)
	for i := 0; i < n; i++ {
		b := p.AddBar(100)
		switch i {
		case n - 1:
			var abortCalledTimes int
			for j := 0; !b.Aborted(); j++ {
				if j >= 10 {
					b.Abort(true)
					abortCalledTimes++
				} else {
					b.Increment()
				}
			}
			if abortCalledTimes != 1 {
				t.Errorf("Expected abortCalledTimes: %d, got: %d\n", 1, abortCalledTimes)
			}
			count := p.BarCount()
			if count != 1 {
				t.Errorf("BarCount want: %d, got: %d\n", 1, count)
			}
		default:
			go func() {
				for !b.Completed() {
					b.Increment()
					time.Sleep(randomDuration(100 * time.Millisecond))
				}
			}()
		}
		bars[i] = b
	}

	bars[0].Abort(false)
	go p.Wait()
	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

func TestWithContext(t *testing.T) {
	shutdown := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	p := mpb.NewWithContext(ctx, mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	done := make(chan struct{})
	bar := p.AddBar(0) // never complete bar
	go func() {
		for !bar.Aborted() {
			time.Sleep(randomDuration(100 * time.Millisecond))
			cancel()
		}
		close(done)
	}()

	go func() {
		<-done
		p.Wait()
	}()

	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

// MaxWidthDistributor shouldn't stuck in the middle while removing or aborting a bar
func TestMaxWidthDistributor(t *testing.T) {

	makeWrapper := func(f func([]chan int), start, end chan struct{}) func([]chan int) {
		return func(column []chan int) {
			start <- struct{}{}
			f(column)
			<-end
		}
	}

	ready := make(chan struct{})
	start := make(chan struct{})
	end := make(chan struct{})
	mpb.MaxWidthDistributor = makeWrapper(mpb.MaxWidthDistributor, start, end)

	total := 100
	numBars := 6
	p := mpb.New(mpb.WithOutput(ioutil.Discard))
	for i := 0; i < numBars; i++ {
		bar := p.AddBar(int64(total),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(decor.EwmaETA(decor.ET_STYLE_GO, 60, decor.WCSyncSpace)),
		)
		go func() {
			<-ready
			for i := 0; i < total; i++ {
				start := time.Now()
				if id := bar.ID(); id > 1 && i >= 32 {
					if id&1 == 1 {
						bar.Abort(true)
					} else {
						bar.Abort(false)
					}
				}
				time.Sleep(randomDuration(100 * time.Millisecond))
				bar.IncrInt64(rand.Int63n(5) + 1)
				bar.DecoratorEwmaUpdate(time.Since(start))
			}
		}()
	}

	go func() {
		<-ready
		p.Wait()
		close(start)
	}()

	res := t.Run("maxWidthDistributor", func(t *testing.T) {
		close(ready)
		for v := range start {
			timer := time.NewTimer(100 * time.Millisecond)
			select {
			case end <- v:
				timer.Stop()
			case <-timer.C:
				t.FailNow()
			}
		}
	})

	if !res {
		t.Error("maxWidthDistributor stuck in the middle")
	}
}

func randomDuration(max time.Duration) time.Duration {
	return time.Duration(rand.Intn(10)+1) * max / 10
}
<commit_msg>cosmetics...<commit_after>package mpb_test

import (
	"context"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"

	"github.com/vbauerster/mpb/v7"
	"github.com/vbauerster/mpb/v7/decor"
)

const (
	timeout = 200 * time.Millisecond
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func TestBarCount(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	b := p.AddBar(0, mpb.BarRemoveOnComplete())

	if count := p.BarCount(); count != 1 {
		t.Errorf("BarCount want: %d, got: %d\n", 1, count)
	}

	b.SetTotal(100, true)

	if count := p.BarCount(); count != 0 {
		t.Errorf("BarCount want: %d, got: %d\n", 0, count)
	}

	go p.Wait()

	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

func TestBarAbort(t *testing.T) {
	shutdown := make(chan struct{})
	p := mpb.New(mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))
	n := 2
	bars := make([]*mpb.Bar, n)
	for i := 0; i < n; i++ {
		b := p.AddBar(100)
		switch i {
		case n - 1:
			var abortCalledTimes int
			for j := 0; !b.Aborted(); j++ {
				if j >= 10 {
					b.Abort(true)
					abortCalledTimes++
				} else {
					b.Increment()
				}
			}
			if abortCalledTimes != 1 {
				t.Errorf("Expected abortCalledTimes: %d, got: %d\n", 1, abortCalledTimes)
			}
			count := p.BarCount()
			if count != 1 {
				t.Errorf("BarCount want: %d, got: %d\n", 1, count)
			}
		default:
			go func() {
				for !b.Completed() {
					b.Increment()
					time.Sleep(randomDuration(100 * time.Millisecond))
				}
			}()
		}
		bars[i] = b
	}

	bars[0].Abort(false)

	go p.Wait()

	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

func TestWithContext(t *testing.T) {
	shutdown := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	p := mpb.NewWithContext(ctx, mpb.WithShutdownNotifier(shutdown), mpb.WithOutput(ioutil.Discard))

	done := make(chan struct{})
	bar := p.AddBar(0) // never complete bar
	go func() {
		for !bar.Aborted() {
			time.Sleep(randomDuration(100 * time.Millisecond))
			cancel()
		}
		close(done)
	}()

	go func() {
		<-done
		p.Wait()
	}()

	select {
	case <-shutdown:
	case <-time.After(timeout):
		t.Errorf("Progress didn't shutdown after %v", timeout)
	}
}

// MaxWidthDistributor shouldn't stuck in the middle while removing or aborting a bar
func TestMaxWidthDistributor(t *testing.T) {

	makeWrapper := func(f func([]chan int), start, end chan struct{}) func([]chan int) {
		return func(column []chan int) {
			start <- struct{}{}
			f(column)
			<-end
		}
	}

	ready := make(chan struct{})
	start := make(chan struct{})
	end := make(chan struct{})
	mpb.MaxWidthDistributor = makeWrapper(mpb.MaxWidthDistributor, start, end)

	total := 100
	numBars := 6
	p := mpb.New(mpb.WithOutput(ioutil.Discard))
	for i := 0; i < numBars; i++ {
		bar := p.AddBar(int64(total),
			mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0),
			mpb.PrependDecorators(decor.EwmaETA(decor.ET_STYLE_GO, 60, decor.WCSyncSpace)),
		)
		go func() {
			<-ready
			for i := 0; i < total; i++ {
				start := time.Now()
				if id := bar.ID(); id > 1 && i >= 32 {
					if id&1 == 1 {
						bar.Abort(true)
					} else {
						bar.Abort(false)
					}
				}
				time.Sleep(randomDuration(100 * time.Millisecond))
				bar.IncrInt64(rand.Int63n(5) + 1)
				bar.DecoratorEwmaUpdate(time.Since(start))
			}
		}()
	}

	go func() {
		<-ready
		p.Wait()
		close(start)
	}()

	res := t.Run("maxWidthDistributor", func(t *testing.T) {
		close(ready)
		for v := range start {
			timer := time.NewTimer(100 * time.Millisecond)
			select {
			case end <- v:
				timer.Stop()
			case <-timer.C:
				t.FailNow()
			}
		}
	})

	if !res {
		t.Error("maxWidthDistributor stuck in the middle")
	}
}

func randomDuration(max time.Duration) time.Duration {
	return time.Duration(rand.Intn(10)+1) * max / 10
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsNetworkInterface() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsNetworkInterfaceCreate,
		Read:   resourceAwsNetworkInterfaceRead,
		Update: resourceAwsNetworkInterfaceUpdate,
		Delete: resourceAwsNetworkInterfaceDelete,

		Schema: map[string]*schema.Schema{

			"subnet_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"private_ips": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},

			"security_groups": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},

			"source_dest_check": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"attachment": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"instance": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"device_index": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
						"attachment_id": &schema.Schema{
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
				Set: resourceAwsEniAttachmentHash,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsNetworkInterfaceCreate(d *schema.ResourceData, meta interface{}) error {

	conn := meta.(*AWSClient).ec2conn

	request := &ec2.CreateNetworkInterfaceInput{
		SubnetId: aws.String(d.Get("subnet_id").(string)),
	}

	security_groups := d.Get("security_groups").(*schema.Set).List()
	if len(security_groups) != 0 {
		request.Groups = expandStringList(security_groups)
	}

	private_ips := d.Get("private_ips").(*schema.Set).List()
	if len(private_ips) != 0 {
		request.PrivateIpAddresses = expandPrivateIPAddresses(private_ips)
	}

	log.Printf("[DEBUG] Creating network interface")
	resp, err := conn.CreateNetworkInterface(request)
	if err != nil {
		return fmt.Errorf("Error creating ENI: %s", err)
	}

	d.SetId(*resp.NetworkInterface.NetworkInterfaceId)
	log.Printf("[INFO] ENI ID: %s", d.Id())
	return resourceAwsNetworkInterfaceUpdate(d, meta)
}

func resourceAwsNetworkInterfaceRead(d *schema.ResourceData, meta interface{}) error {

	conn := meta.(*AWSClient).ec2conn
	describe_network_interfaces_request := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{aws.String(d.Id())},
	}
	describeResp, err := conn.DescribeNetworkInterfaces(describe_network_interfaces_request)

	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidNetworkInterfaceID.NotFound" {
			// The ENI is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving ENI: %s", err)
	}
	if len(describeResp.NetworkInterfaces) != 1 {
		return fmt.Errorf("Unable to find ENI: %#v", describeResp.NetworkInterfaces)
	}

	eni := describeResp.NetworkInterfaces[0]
	d.Set("subnet_id", eni.SubnetId)
	d.Set("private_ips", flattenNetworkInterfacesPrivateIPAddresses(eni.PrivateIpAddresses))
	d.Set("security_groups", flattenGroupIdentifiers(eni.Groups))
	d.Set("source_dest_check", eni.SourceDestCheck)

	// Tags
	d.Set("tags", tagsToMap(eni.TagSet))

	if eni.Attachment != nil {
		attachment := []map[string]interface{}{flattenAttachment(eni.Attachment)}
		d.Set("attachment", attachment)
	} else {
		d.Set("attachment", nil)
	}

	return nil
}

func networkInterfaceAttachmentRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {

		describe_network_interfaces_request := &ec2.DescribeNetworkInterfacesInput{
			NetworkInterfaceIds: []*string{aws.String(id)},
		}
		describeResp, err := conn.DescribeNetworkInterfaces(describe_network_interfaces_request)

		if err != nil {
			log.Printf("[ERROR] Could not find network interface %s. %s", id, err)
			return nil, "", err
		}

		eni := describeResp.NetworkInterfaces[0]
		hasAttachment := strconv.FormatBool(eni.Attachment != nil)
		log.Printf("[DEBUG] ENI %s has attachment state %s", id, hasAttachment)
		return eni, hasAttachment, nil
	}
}

func resourceAwsNetworkInterfaceDetach(oa *schema.Set, meta interface{}, eniId string) error {
	// if there was an old attachment, remove it
	if oa != nil && len(oa.List()) > 0 {
		old_attachment := oa.List()[0].(map[string]interface{})
		detach_request := &ec2.DetachNetworkInterfaceInput{
			AttachmentId: aws.String(old_attachment["attachment_id"].(string)),
			Force:        aws.Bool(true),
		}
		conn := meta.(*AWSClient).ec2conn
		_, detach_err := conn.DetachNetworkInterface(detach_request)
		if detach_err != nil {
			return fmt.Errorf("Error detaching ENI: %s", detach_err)
		}

		log.Printf("[DEBUG] Waiting for ENI (%s) to become dettached", eniId)
		stateConf := &resource.StateChangeConf{
			Pending: []string{"true"},
			Target:  "false",
			Refresh: networkInterfaceAttachmentRefreshFunc(conn, eniId),
			Timeout: 10 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for ENI (%s) to become dettached: %s", eniId, err)
		}
	}

	return nil
}

func convertToAwsStringSlice(s []interface{}) []*string {
	var b []*string
	for _, i := range s {
		b = append(b, aws.String(i.(string)))
	}
	return b
}

func resourceAwsNetworkInterfaceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn
	d.Partial(true)

	if d.HasChange("attachment") {
		oa, na := d.GetChange("attachment")

		detach_err := resourceAwsNetworkInterfaceDetach(oa.(*schema.Set), meta, d.Id())
		if detach_err != nil {
			return detach_err
		}

		// if there is a new attachment, attach it
		if na != nil && len(na.(*schema.Set).List()) > 0 {
			new_attachment := na.(*schema.Set).List()[0].(map[string]interface{})
			di := new_attachment["device_index"].(int)
			attach_request := &ec2.AttachNetworkInterfaceInput{
				DeviceIndex:        aws.Int64(int64(di)),
				InstanceId:         aws.String(new_attachment["instance"].(string)),
				NetworkInterfaceId: aws.String(d.Id()),
			}
			_, attach_err := conn.AttachNetworkInterface(attach_request)
			if attach_err != nil {
				return fmt.Errorf("Error attaching ENI: %s", attach_err)
			}
		}

		d.SetPartial("attachment")
	}

	if d.HasChange("private_ips") {
		o, n := d.GetChange("private_ips")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		// Unassign old IP addresses
		unassignIps := os.Difference(ns)
		if unassignIps.Len() != 0 {
			input := &ec2.UnassignPrivateIpAddressesInput{
				NetworkInterfaceId: aws.String(d.Id()),
				PrivateIpAddresses: convertToAwsStringSlice(unassignIps.List()),
			}
			_, err := conn.UnassignPrivateIpAddresses(input)
			if err != nil {
				return fmt.Errorf("Failure to unassign Private IPs: %s", err)
			}
		}

		// Assign new IP addresses
		assignIps := ns.Difference(os)
		if assignIps.Len() != 0 {
			input := &ec2.AssignPrivateIpAddressesInput{
				NetworkInterfaceId: aws.String(d.Id()),
				PrivateIpAddresses: convertToAwsStringSlice(assignIps.List()),
			}
			_, err := conn.AssignPrivateIpAddresses(input)
			if err != nil {
				return fmt.Errorf("Failure to assign Private IPs: %s", err)
			}
		}

		d.SetPartial("private_ips")
	}

	request := &ec2.ModifyNetworkInterfaceAttributeInput{
		NetworkInterfaceId: aws.String(d.Id()),
		SourceDestCheck:    &ec2.AttributeBooleanValue{Value: aws.Bool(d.Get("source_dest_check").(bool))},
	}

	_, err := conn.ModifyNetworkInterfaceAttribute(request)
	if err != nil {
		return fmt.Errorf("Failure updating ENI: %s", err)
	}

	d.SetPartial("source_dest_check")

	if d.HasChange("security_groups") {
		request := &ec2.ModifyNetworkInterfaceAttributeInput{
			NetworkInterfaceId: aws.String(d.Id()),
			Groups:             expandStringList(d.Get("security_groups").(*schema.Set).List()),
		}

		_, err := conn.ModifyNetworkInterfaceAttribute(request)
		if err != nil {
			return fmt.Errorf("Failure updating ENI: %s", err)
		}

		d.SetPartial("security_groups")
	}

	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	d.Partial(false)

	return resourceAwsNetworkInterfaceRead(d, meta)
}

func resourceAwsNetworkInterfaceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	log.Printf("[INFO] Deleting ENI: %s", d.Id())

	detach_err := resourceAwsNetworkInterfaceDetach(d.Get("attachment").(*schema.Set), meta, d.Id())
	if detach_err != nil {
		return detach_err
	}

	deleteEniOpts := ec2.DeleteNetworkInterfaceInput{
		NetworkInterfaceId: aws.String(d.Id()),
	}
	if _, err := conn.DeleteNetworkInterface(&deleteEniOpts); err != nil {
		return fmt.Errorf("Error deleting ENI: %s", err)
	}

	return nil
}

func resourceAwsEniAttachmentHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["instance"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["device_index"].(int)))
	return hashcode.String(buf.String())
}
<commit_msg>Reuse structure method for network interface ips<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsNetworkInterface() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsNetworkInterfaceCreate,
		Read:   resourceAwsNetworkInterfaceRead,
		Update: resourceAwsNetworkInterfaceUpdate,
		Delete: resourceAwsNetworkInterfaceDelete,

		Schema: map[string]*schema.Schema{

			"subnet_id": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"private_ips": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},

			"security_groups": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set:      schema.HashString,
			},

			"source_dest_check": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"attachment": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"instance": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"device_index": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
						"attachment_id": &schema.Schema{
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
				Set: resourceAwsEniAttachmentHash,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsNetworkInterfaceCreate(d *schema.ResourceData, meta interface{}) error {

	conn := meta.(*AWSClient).ec2conn

	request := &ec2.CreateNetworkInterfaceInput{
		SubnetId: aws.String(d.Get("subnet_id").(string)),
	}

	security_groups := d.Get("security_groups").(*schema.Set).List()
	if len(security_groups) != 0 {
		request.Groups = expandStringList(security_groups)
	}

	private_ips := d.Get("private_ips").(*schema.Set).List()
	if len(private_ips) != 0 {
		request.PrivateIpAddresses = expandPrivateIPAddresses(private_ips)
	}

	log.Printf("[DEBUG] Creating network interface")
	resp, err := conn.CreateNetworkInterface(request)
	if err != nil {
		return fmt.Errorf("Error creating ENI: %s", err)
	}

	d.SetId(*resp.NetworkInterface.NetworkInterfaceId)
	log.Printf("[INFO] ENI ID: %s", d.Id())
	return resourceAwsNetworkInterfaceUpdate(d, meta)
}

func resourceAwsNetworkInterfaceRead(d *schema.ResourceData, meta interface{}) error {

	conn := meta.(*AWSClient).ec2conn
	describe_network_interfaces_request := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{aws.String(d.Id())},
	}
	describeResp, err := conn.DescribeNetworkInterfaces(describe_network_interfaces_request)

	if err != nil {
		if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidNetworkInterfaceID.NotFound" {
			// The ENI is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving ENI: %s", err)
	}
	if len(describeResp.NetworkInterfaces) != 1 {
		return fmt.Errorf("Unable to find ENI: %#v", describeResp.NetworkInterfaces)
	}

	eni := describeResp.NetworkInterfaces[0]
	d.Set("subnet_id", eni.SubnetId)
	d.Set("private_ips", flattenNetworkInterfacesPrivateIPAddresses(eni.PrivateIpAddresses))
	d.Set("security_groups", flattenGroupIdentifiers(eni.Groups))
	d.Set("source_dest_check", eni.SourceDestCheck)

	// Tags
	d.Set("tags", tagsToMap(eni.TagSet))

	if eni.Attachment != nil {
		attachment := []map[string]interface{}{flattenAttachment(eni.Attachment)}
		d.Set("attachment", attachment)
	} else {
		d.Set("attachment", nil)
	}

	return nil
}

func networkInterfaceAttachmentRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {

		describe_network_interfaces_request := &ec2.DescribeNetworkInterfacesInput{
			NetworkInterfaceIds: []*string{aws.String(id)},
		}
		describeResp, err := conn.DescribeNetworkInterfaces(describe_network_interfaces_request)

		if err != nil {
			log.Printf("[ERROR] Could not find network interface %s. %s", id, err)
			return nil, "", err
		}

		eni := describeResp.NetworkInterfaces[0]
		hasAttachment := strconv.FormatBool(eni.Attachment != nil)
		log.Printf("[DEBUG] ENI %s has attachment state %s", id, hasAttachment)
		return eni, hasAttachment, nil
	}
}

func resourceAwsNetworkInterfaceDetach(oa *schema.Set, meta interface{}, eniId string) error {
	// if there was an old attachment, remove it
	if oa != nil && len(oa.List()) > 0 {
		old_attachment := oa.List()[0].(map[string]interface{})
		detach_request := &ec2.DetachNetworkInterfaceInput{
			AttachmentId: aws.String(old_attachment["attachment_id"].(string)),
			Force:        aws.Bool(true),
		}
		conn := meta.(*AWSClient).ec2conn
		_, detach_err := conn.DetachNetworkInterface(detach_request)
		if detach_err != nil {
			return fmt.Errorf("Error detaching ENI: %s", detach_err)
		}

		log.Printf("[DEBUG] Waiting for ENI (%s) to become dettached", eniId)
		stateConf := &resource.StateChangeConf{
			Pending: []string{"true"},
			Target:  "false",
			Refresh: networkInterfaceAttachmentRefreshFunc(conn, eniId),
			Timeout: 10 * time.Minute,
		}
		if _, err := stateConf.WaitForState(); err != nil {
			return fmt.Errorf(
				"Error waiting for ENI (%s) to become dettached: %s", eniId, err)
		}
	}

	return nil
}

func resourceAwsNetworkInterfaceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn
	d.Partial(true)

	if d.HasChange("attachment") {
		oa, na := d.GetChange("attachment")

		detach_err := resourceAwsNetworkInterfaceDetach(oa.(*schema.Set), meta, d.Id())
		if detach_err != nil {
			return detach_err
		}

		// if there is a new attachment, attach it
		if na != nil && len(na.(*schema.Set).List()) > 0 {
			new_attachment := na.(*schema.Set).List()[0].(map[string]interface{})
			di := new_attachment["device_index"].(int)
			attach_request := &ec2.AttachNetworkInterfaceInput{
				DeviceIndex:        aws.Int64(int64(di)),
				InstanceId:         aws.String(new_attachment["instance"].(string)),
				NetworkInterfaceId: aws.String(d.Id()),
			}
			_, attach_err := conn.AttachNetworkInterface(attach_request)
			if attach_err != nil {
				return fmt.Errorf("Error attaching ENI: %s", attach_err)
			}
		}

		d.SetPartial("attachment")
	}

	if d.HasChange("private_ips") {
		o, n := d.GetChange("private_ips")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		// Unassign old IP addresses
		unassignIps := os.Difference(ns)
		if unassignIps.Len() != 0 {
			input := &ec2.UnassignPrivateIpAddressesInput{
				NetworkInterfaceId: aws.String(d.Id()),
				PrivateIpAddresses: expandStringList(unassignIps.List()),
			}
			_, err := conn.UnassignPrivateIpAddresses(input)
			if err != nil {
				return fmt.Errorf("Failure to unassign Private IPs: %s", err)
			}
		}

		// Assign new IP addresses
		assignIps := ns.Difference(os)
		if assignIps.Len() != 0 {
			input := &ec2.AssignPrivateIpAddressesInput{
				NetworkInterfaceId: aws.String(d.Id()),
				PrivateIpAddresses: expandStringList(assignIps.List()),
			}
			_, err := conn.AssignPrivateIpAddresses(input)
			if err != nil {
				return fmt.Errorf("Failure to assign Private IPs: %s", err)
			}
		}

		d.SetPartial("private_ips")
	}

	request := &ec2.ModifyNetworkInterfaceAttributeInput{
		NetworkInterfaceId: aws.String(d.Id()),
		SourceDestCheck:    &ec2.AttributeBooleanValue{Value: aws.Bool(d.Get("source_dest_check").(bool))},
	}

	_, err := conn.ModifyNetworkInterfaceAttribute(request)
	if err != nil {
		return fmt.Errorf("Failure updating ENI: %s", err)
	}

	d.SetPartial("source_dest_check")

	if d.HasChange("security_groups") {
		request := &ec2.ModifyNetworkInterfaceAttributeInput{
			NetworkInterfaceId: aws.String(d.Id()),
			Groups:             expandStringList(d.Get("security_groups").(*schema.Set).List()),
		}

		_, err := conn.ModifyNetworkInterfaceAttribute(request)
		if err != nil {
			return fmt.Errorf("Failure updating ENI: %s", err)
		}

		d.SetPartial("security_groups")
	}

	if err := setTags(conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	d.Partial(false)

	return resourceAwsNetworkInterfaceRead(d, meta)
}

func resourceAwsNetworkInterfaceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ec2conn

	log.Printf("[INFO] Deleting ENI: %s", d.Id())

	detach_err := resourceAwsNetworkInterfaceDetach(d.Get("attachment").(*schema.Set), meta, d.Id())
	if detach_err != nil {
		return detach_err
	}

	deleteEniOpts := ec2.DeleteNetworkInterfaceInput{
		NetworkInterfaceId: aws.String(d.Id()),
	}
	if _, err := conn.DeleteNetworkInterface(&deleteEniOpts); err != nil {
		return fmt.Errorf("Error deleting ENI: %s", err)
	}

	return nil
}

func resourceAwsEniAttachmentHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["instance"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["device_index"].(int)))
	return hashcode.String(buf.String())
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/aws-sdk-go-base/tfawserr"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
	"github.com/terraform-providers/terraform-provider-aws/aws/internal/hashcode"
)

func resourceAwsLakeFormationDataLakeSettings() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationDataLakeSettingsCreate,
		Update: resourceAwsLakeFormationDataLakeSettingsCreate,
		Read:   resourceAwsLakeFormationDataLakeSettingsRead,
		Delete: resourceAwsLakeFormationDataLakeSettingsDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"admins": {
				Type:     schema.TypeSet,
				Computed: true,
				Optional: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validateArn,
				},
			},
			"catalog_id": {
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},
			"create_database_default_permissions": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				MaxItems: 3,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"permissions": {
							Type:     schema.TypeSet,
							Optional: true,
							Computed: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.StringInSlice(lakeformation.Permission_Values(), false),
							},
						},
						"principal": {
							Type:         schema.TypeString,
							Optional:     true,
							Computed:     true,
							ValidateFunc: validatePrincipal,
						},
					},
				},
			},
			"create_table_default_permissions": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				MaxItems: 3,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"permissions": {
							Type:     schema.TypeSet,
							Optional: true,
							Computed: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.StringInSlice(lakeformation.Permission_Values(), false),
							},
						},
						"principal": {
							Type:         schema.TypeString,
							Optional:     true,
							Computed:     true,
							ValidateFunc: validatePrincipal,
						},
					},
				},
			},
			"trusted_resource_owners": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validateAwsAccountId,
				},
			},
		},
	}
}

func resourceAwsLakeFormationDataLakeSettingsCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.PutDataLakeSettingsInput{}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	settings := &lakeformation.DataLakeSettings{}

	if v, ok := d.GetOk("create_database_default_permissions"); ok {
		settings.CreateDatabaseDefaultPermissions = expandDataLakeSettingsCreateDefaultPermissions(v.([]interface{}))
	}

	if v, ok := d.GetOk("create_table_default_permissions"); ok {
		settings.CreateTableDefaultPermissions = expandDataLakeSettingsCreateDefaultPermissions(v.([]interface{}))
	}

	if v, ok := d.GetOk("admins"); ok {
		settings.DataLakeAdmins = expandDataLakeSettingsAdmins(v.(*schema.Set))
	}

	if v, ok := d.GetOk("trusted_resource_owners"); ok {
		settings.TrustedResourceOwners = expandStringList(v.([]interface{}))
	}

	input.DataLakeSettings = settings

	var output *lakeformation.PutDataLakeSettingsOutput
	err := resource.Retry(2*time.Minute, func() *resource.RetryError {
		var err error
		output, err = conn.PutDataLakeSettings(input)
		if err != nil {
			if isAWSErr(err, lakeformation.ErrCodeInvalidInputException, "Invalid principal") {
				return resource.RetryableError(err)
			}
			if isAWSErr(err, lakeformation.ErrCodeConcurrentModificationException, "") {
				return resource.RetryableError(err)
			}

			return resource.NonRetryableError(fmt.Errorf("error creating Lake Formation data lake settings: %w", err))
		}
		return nil
	})

	if isResourceTimeoutError(err) {
		output, err = conn.PutDataLakeSettings(input)
	}

	if err != nil {
		return fmt.Errorf("error creating Lake Formation data lake settings: %w", err)
	}

	if output == nil {
		return fmt.Errorf("error creating Lake Formation data lake settings: empty response")
	}

	d.SetId(fmt.Sprintf("%d", hashcode.String(input.String())))

	return resourceAwsLakeFormationDataLakeSettingsRead(d, meta)
}

func resourceAwsLakeFormationDataLakeSettingsRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.GetDataLakeSettingsInput{}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	output, err := conn.GetDataLakeSettings(input)

	if !d.IsNewResource() && tfawserr.ErrCodeEquals(err, lakeformation.ErrCodeEntityNotFoundException) {
		log.Printf("[WARN] Lake Formation data lake settings (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("reading Lake Formation data lake settings (%s): %w", d.Id(), err)
	}

	if output == nil || output.DataLakeSettings == nil {
		return fmt.Errorf("reading Lake Formation data lake settings (%s): empty response", d.Id())
	}

	settings := output.DataLakeSettings

	d.Set("create_database_default_permissions", flattenDataLakeSettingsCreateDefaultPermissions(settings.CreateDatabaseDefaultPermissions))
	d.Set("create_table_default_permissions", flattenDataLakeSettingsCreateDefaultPermissions(settings.CreateTableDefaultPermissions))
	d.Set("admins", flattenDataLakeSettingsAdmins(settings.DataLakeAdmins))
	d.Set("trusted_resource_owners", flattenStringList(settings.TrustedResourceOwners))

	return nil
}

func resourceAwsLakeFormationDataLakeSettingsDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.PutDataLakeSettingsInput{
		DataLakeSettings: &lakeformation.DataLakeSettings{
			CreateDatabaseDefaultPermissions: make([]*lakeformation.PrincipalPermissions, 0),
			CreateTableDefaultPermissions:    make([]*lakeformation.PrincipalPermissions, 0),
			DataLakeAdmins:                   make([]*lakeformation.DataLakePrincipal, 0),
			TrustedResourceOwners:            make([]*string, 0),
		},
	}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	_, err := conn.PutDataLakeSettings(input)

	if tfawserr.ErrCodeEquals(err, lakeformation.ErrCodeEntityNotFoundException) {
		log.Printf("[WARN] Lake Formation data lake settings (%s) not found, removing from state", d.Id())
		return nil
	}

	if err != nil {
		return fmt.Errorf("deleting Lake Formation data lake settings (%s): %w", d.Id(), err)
	}

	return nil
}

func expandDataLakeSettingsCreateDefaultPermissions(tfMaps []interface{}) []*lakeformation.PrincipalPermissions {
	apiObjects := make([]*lakeformation.PrincipalPermissions, 0, len(tfMaps))

	for _, tfMap := range tfMaps {
		apiObjects = append(apiObjects, expandDataLakeSettingsCreateDefaultPermission(tfMap.(map[string]interface{})))
	}

	return apiObjects
}

func expandDataLakeSettingsCreateDefaultPermission(tfMap map[string]interface{}) *lakeformation.PrincipalPermissions {
	apiObject := &lakeformation.PrincipalPermissions{
		Permissions: expandStringSet(tfMap["permissions"].(*schema.Set)),
		Principal: &lakeformation.DataLakePrincipal{
			DataLakePrincipalIdentifier: aws.String(tfMap["principal"].(string)),
		},
	}

	return apiObject
}

func flattenDataLakeSettingsCreateDefaultPermissions(apiObjects []*lakeformation.PrincipalPermissions) []map[string]interface{} {
	if apiObjects == nil {
		return nil
	}

	tfMaps := make([]map[string]interface{}, len(apiObjects))
	for i, v := range apiObjects {
		tfMaps[i] = flattenDataLakeSettingsCreateDefaultPermission(v)
	}

	return tfMaps
}

func flattenDataLakeSettingsCreateDefaultPermission(apiObject *lakeformation.PrincipalPermissions) map[string]interface{} {
	tfMap := make(map[string]interface{})

	if apiObject == nil {
		return tfMap
	}

	if apiObject.Permissions != nil {
		tfMap["permissions"] = flattenStringSet(apiObject.Permissions)
	}

	if v := aws.StringValue(apiObject.Principal.DataLakePrincipalIdentifier); v != "" {
		tfMap["principal"] = v
	}

	return tfMap
}

func expandDataLakeSettingsAdmins(tfSet *schema.Set) []*lakeformation.DataLakePrincipal {
	tfSlice := tfSet.List()
	apiObjects := make([]*lakeformation.DataLakePrincipal, 0, len(tfSlice))

	for _, tfItem := range tfSlice {
		val, ok := tfItem.(string)
		if ok && val != "" {
			apiObjects = append(apiObjects, &lakeformation.DataLakePrincipal{
				DataLakePrincipalIdentifier: aws.String(tfItem.(string)),
			})
		}
	}

	return apiObjects
}

func flattenDataLakeSettingsAdmins(apiObjects []*lakeformation.DataLakePrincipal) *schema.Set {
	if apiObjects == nil {
		return nil
	}

	tfSlice := make([]*string, 0, len(apiObjects))

	for _, apiObject := range apiObjects {
		tfSlice = append(tfSlice, apiObject.DataLakePrincipalIdentifier)
	}

	return flattenStringSet(tfSlice)
}
<commit_msg>lakeformation_data_lake_settings: Switch to interface<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/aws-sdk-go-base/tfawserr"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
	"github.com/terraform-providers/terraform-provider-aws/aws/internal/hashcode"
)

func resourceAwsLakeFormationDataLakeSettings() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationDataLakeSettingsCreate,
		Update: resourceAwsLakeFormationDataLakeSettingsCreate,
		Read:   resourceAwsLakeFormationDataLakeSettingsRead,
		Delete: resourceAwsLakeFormationDataLakeSettingsDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"admins": {
				Type:     schema.TypeSet,
				Computed: true,
				Optional: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validateArn,
				},
			},
			"catalog_id": {
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},
			"create_database_default_permissions": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				MaxItems: 3,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"permissions": {
							Type:     schema.TypeSet,
							Optional: true,
							Computed: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.StringInSlice(lakeformation.Permission_Values(), false),
							},
						},
						"principal": {
							Type:         schema.TypeString,
							Optional:     true,
							Computed:     true,
							ValidateFunc: validatePrincipal,
						},
					},
				},
			},
			"create_table_default_permissions": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				MaxItems: 3,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"permissions": {
							Type:     schema.TypeSet,
							Optional: true,
							Computed: true,
							Elem: &schema.Schema{
								Type:         schema.TypeString,
								ValidateFunc: validation.StringInSlice(lakeformation.Permission_Values(), false),
							},
						},
						"principal": {
							Type:         schema.TypeString,
							Optional:     true,
							Computed:     true,
							ValidateFunc: validatePrincipal,
						},
					},
				},
			},
			"trusted_resource_owners": {
				Type:     schema.TypeList,
				Computed: true,
				Optional: true,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validateAwsAccountId,
				},
			},
		},
	}
}

func resourceAwsLakeFormationDataLakeSettingsCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.PutDataLakeSettingsInput{}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	settings := &lakeformation.DataLakeSettings{}

	if v, ok := d.GetOk("create_database_default_permissions"); ok {
		settings.CreateDatabaseDefaultPermissions = expandDataLakeSettingsCreateDefaultPermissions(v.([]interface{}))
	}

	if v, ok := d.GetOk("create_table_default_permissions"); ok {
		settings.CreateTableDefaultPermissions = expandDataLakeSettingsCreateDefaultPermissions(v.([]interface{}))
	}

	if v, ok := d.GetOk("admins"); ok {
		settings.DataLakeAdmins = expandDataLakeSettingsAdmins(v.(*schema.Set))
	}

	if v, ok := d.GetOk("trusted_resource_owners"); ok {
		settings.TrustedResourceOwners = expandStringList(v.([]interface{}))
	}

	input.DataLakeSettings = settings

	var output *lakeformation.PutDataLakeSettingsOutput
	err := resource.Retry(2*time.Minute, func() *resource.RetryError {
		var err error
		output, err = conn.PutDataLakeSettings(input)
		if err != nil {
			if isAWSErr(err, lakeformation.ErrCodeInvalidInputException, "Invalid principal") {
				return resource.RetryableError(err)
			}
			if isAWSErr(err, lakeformation.ErrCodeConcurrentModificationException, "") {
				return resource.RetryableError(err)
			}

			return resource.NonRetryableError(fmt.Errorf("error creating Lake Formation data lake settings: %w", err))
		}
		return nil
	})

	if isResourceTimeoutError(err) {
		output, err = conn.PutDataLakeSettings(input)
	}

	if err != nil {
		return fmt.Errorf("error creating Lake Formation data lake settings: %w", err)
	}

	if output == nil {
		return fmt.Errorf("error creating Lake Formation data lake settings: empty response")
	}

	d.SetId(fmt.Sprintf("%d", hashcode.String(input.String())))

	return resourceAwsLakeFormationDataLakeSettingsRead(d, meta)
}

func resourceAwsLakeFormationDataLakeSettingsRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.GetDataLakeSettingsInput{}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	output, err := conn.GetDataLakeSettings(input)

	if !d.IsNewResource() && tfawserr.ErrCodeEquals(err, lakeformation.ErrCodeEntityNotFoundException) {
		log.Printf("[WARN] Lake Formation data lake settings (%s) not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("reading Lake Formation data lake settings (%s): %w", d.Id(), err)
	}

	if output == nil || output.DataLakeSettings == nil {
		return fmt.Errorf("reading Lake Formation data lake settings (%s): empty response", d.Id())
	}

	settings := output.DataLakeSettings

	d.Set("create_database_default_permissions", flattenDataLakeSettingsCreateDefaultPermissions(settings.CreateDatabaseDefaultPermissions))
	d.Set("create_table_default_permissions", flattenDataLakeSettingsCreateDefaultPermissions(settings.CreateTableDefaultPermissions))
	d.Set("admins", flattenDataLakeSettingsAdmins(settings.DataLakeAdmins))
	d.Set("trusted_resource_owners", flattenStringList(settings.TrustedResourceOwners))

	return nil
}

func resourceAwsLakeFormationDataLakeSettingsDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn

	input := &lakeformation.PutDataLakeSettingsInput{
		DataLakeSettings: &lakeformation.DataLakeSettings{
			CreateDatabaseDefaultPermissions: make([]*lakeformation.PrincipalPermissions, 0),
			CreateTableDefaultPermissions:    make([]*lakeformation.PrincipalPermissions, 0),
			DataLakeAdmins:                   make([]*lakeformation.DataLakePrincipal, 0),
			TrustedResourceOwners:            make([]*string, 0),
		},
	}

	if v, ok := d.GetOk("catalog_id"); ok {
		input.CatalogId = aws.String(v.(string))
	}

	_, err := conn.PutDataLakeSettings(input)

	if tfawserr.ErrCodeEquals(err, lakeformation.ErrCodeEntityNotFoundException) {
		log.Printf("[WARN] Lake Formation data lake settings (%s) not found, removing from state", d.Id())
		return nil
	}

	if err != nil {
		return fmt.Errorf("deleting Lake Formation data lake settings (%s): %w", d.Id(), err)
	}

	return nil
}

func expandDataLakeSettingsCreateDefaultPermissions(tfMaps []interface{}) []*lakeformation.PrincipalPermissions {
	apiObjects := make([]*lakeformation.PrincipalPermissions, 0, len(tfMaps))

	for _, tfMap := range tfMaps {
		apiObjects = append(apiObjects, expandDataLakeSettingsCreateDefaultPermission(tfMap.(map[string]interface{})))
	}

	return apiObjects
}

func expandDataLakeSettingsCreateDefaultPermission(tfMap map[string]interface{}) *lakeformation.PrincipalPermissions {
	apiObject := &lakeformation.PrincipalPermissions{
		Permissions: expandStringSet(tfMap["permissions"].(*schema.Set)),
		Principal: &lakeformation.DataLakePrincipal{
			DataLakePrincipalIdentifier: aws.String(tfMap["principal"].(string)),
		},
	}

	return apiObject
}

func flattenDataLakeSettingsCreateDefaultPermissions(apiObjects []*lakeformation.PrincipalPermissions) []map[string]interface{} {
	if apiObjects == nil {
		return nil
	}

	tfMaps := make([]map[string]interface{}, len(apiObjects))
	for i, v := range apiObjects {
		tfMaps[i] = flattenDataLakeSettingsCreateDefaultPermission(v)
	}

	return tfMaps
}

func flattenDataLakeSettingsCreateDefaultPermission(apiObject *lakeformation.PrincipalPermissions) map[string]interface{} {
	tfMap := make(map[string]interface{})

	if apiObject == nil {
		return tfMap
	}

	if apiObject.Permissions != nil {
		tfMap["permissions"] = flattenStringSet(apiObject.Permissions)
	}

	if v := aws.StringValue(apiObject.Principal.DataLakePrincipalIdentifier); v != "" {
		tfMap["principal"] = v
	}

	return tfMap
}

func expandDataLakeSettingsAdmins(tfSet *schema.Set) []*lakeformation.DataLakePrincipal {
	tfSlice := tfSet.List()
	apiObjects := make([]*lakeformation.DataLakePrincipal, 0, len(tfSlice))

	for _, tfItem := range tfSlice {
		val, ok := tfItem.(string)
		if ok && val != "" {
			apiObjects = append(apiObjects, &lakeformation.DataLakePrincipal{
				DataLakePrincipalIdentifier: aws.String(tfItem.(string)),
			})
		}
	}

	return apiObjects
}

func flattenDataLakeSettingsAdmins(apiObjects []*lakeformation.DataLakePrincipal) []interface{} {
	if apiObjects == nil {
		return nil
	}

	tfSlice := make([]interface{}, 0, len(apiObjects))

	for _, apiObject := range apiObjects {
		tfSlice = append(tfSlice, *apiObject.DataLakePrincipalIdentifier)
	}

	return tfSlice
}
<|endoftext|>


<commit_before>package operations

import (
	"github.com/qiniu/qshell/v2/iqshell/common/config"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/host"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/storage/bucket"
)

func getDownloadHostProvider(cfg *config.Config, downloadCfg *DownloadCfg) host.Provider {
	var hosts []*host.Host
	if downloadCfg.GetFileApi {
		hosts = getFileApiHosts(cfg, downloadCfg)
	} else {
		hosts = defaultDownloadHosts(cfg, downloadCfg)
	}

	log.DebugF("download Domain:%s", hosts)
	return host.NewListProvider(hosts)
}

func defaultDownloadHosts(cfg *config.Config, downloadCfg *DownloadCfg) []*host.Host {

	hosts := make([]*host.Host, 0)

	// 1. 从 download 配置获取
	if len(downloadCfg.CdnDomain) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.CdnDomain,
		})
	}

	// 2. 动态获取 bucket 绑定的 domain
	b := downloadCfg.Bucket
	log.DebugF("get domain of bucket:%s", b)
	if domains, e := bucket.AllDomainsOfBucket(b); e != nil {
		log.DebugF("get bucket:%s domain error:%v", b, e)
	} else {
		for _, domain := range domains {
			if data.NotEmpty(domain.Domain) {
				hosts = append(hosts, &host.Host{
					Host:   "",
					Domain: domain.Domain.Value(),
				})
			}
		}
	}

	return hosts
}

func getFileApiHosts(cfg *config.Config, downloadCfg *DownloadCfg) []*host.Host {
	hosts := make([]*host.Host, 0)

	// 1. 从 download 配置获取
	if len(downloadCfg.CdnDomain) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.CdnDomain,
		})
	}

	if len(downloadCfg.IoHost) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.IoHost,
		})
	}

	// 2. 从 region 中获取 ioHost
	for _, io := range cfg.Hosts.Io {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: io,
		})
	}

	// 3. 通过 uc query 查询 bucket 所在的 region，并从 region 获取 ioHost
	log.DebugF("get region of bucket:%s", downloadCfg.Bucket)
	if region, err := bucket.Region(downloadCfg.Bucket); err != nil {
		log.DebugF("get region of bucket:%s err:%v", downloadCfg.Bucket, err)
	} else {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: region.IovipHost,
		})
	}

	return hosts
}
<commit_msg>optimize download retry<commit_after>package operations

import (
	"github.com/qiniu/qshell/v2/iqshell/common/config"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/host"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/storage/bucket"
	"strings"
)

func getDownloadHostProvider(cfg *config.Config, downloadCfg *DownloadCfg) host.Provider {
	var hosts []*host.Host
	if downloadCfg.GetFileApi {
		hosts = getFileApiHosts(cfg, downloadCfg)
	} else {
		hosts = defaultDownloadHosts(cfg, downloadCfg)
	}

	hostStrings := make([]string, 0, len(hosts))
	for _, h := range hosts {
		hostStrings = append(hostStrings, h.GetServer())
	}
	log.DebugF("download Domain:[%s]", strings.Join(hostStrings, ","))
	return host.NewListProvider(hosts)
}

func defaultDownloadHosts(cfg *config.Config, downloadCfg *DownloadCfg) []*host.Host {

	hosts := make([]*host.Host, 0)

	// 1. 从 download 配置获取
	if len(downloadCfg.CdnDomain) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.CdnDomain,
		})
	}

	// 2. 动态获取 bucket 绑定的 domain
	b := downloadCfg.Bucket
	log.DebugF("get domain of bucket:%s", b)
	if domains, e := bucket.AllDomainsOfBucket(b); e != nil {
		log.DebugF("get bucket:%s domain error:%v", b, e)
	} else {
		for _, domain := range domains {
			if data.NotEmpty(domain.Domain) {
				hosts = append(hosts, &host.Host{
					Host:   "",
					Domain: domain.Domain.Value(),
				})
			}
		}
	}

	return hosts
}

func getFileApiHosts(cfg *config.Config, downloadCfg *DownloadCfg) []*host.Host {
	hosts := make([]*host.Host, 0)

	// 1. 从 download 配置获取
	if len(downloadCfg.CdnDomain) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.CdnDomain,
		})
	}

	if len(downloadCfg.IoHost) > 0 {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: downloadCfg.IoHost,
		})
	}

	// 2. 从 region 中获取 ioHost
	for _, io := range cfg.Hosts.Io {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: io,
		})
	}

	// 3. 通过 uc query 查询 bucket 所在的 region，并从 region 获取 ioHost
	log.DebugF("get region of bucket:%s", downloadCfg.Bucket)
	if region, err := bucket.Region(downloadCfg.Bucket); err != nil {
		log.DebugF("get region of bucket:%s err:%v", downloadCfg.Bucket, err)
	} else {
		hosts = append(hosts, &host.Host{
			Host:   "",
			Domain: region.IovipHost,
		})
	}

	return hosts
}
<|endoftext|>


<commit_before>package inter

import "github.com/catorpilor/leetcode/utils"

type Interval struct {
	Start, End int
}

func Insert(intervals []Interval, newInterval Interval) []Interval {
	var ret []Interval
	var idx int
	n := len(intervals)
	// add all the intervals ending befor newInterval starts
	for idx < n && intervals[idx].End < newInterval.Start {
		ret = append(ret, intervals[idx])
		idx++
	}
	// mrege the overlapped ones
	for idx < n && intervals[idx].Start <= newInterval.End {
		newInterval = Interval{
			Start: utils.Min(intervals[idx].Start, newInterval.Start),
			End:   utils.Max(intervals[idx].End, newInterval.End),
		}
		idx++
	}
	ret = append(ret, newInterval)
	if idx < n {
		// add teh rest
		ret = append(ret, intervals[idx:]...)
	}
	return ret
}
<commit_msg>no more allocation<commit_after>package inter

import "github.com/catorpilor/leetcode/utils"

type Interval struct {
	Start, End int
}

func Insert(intervals []Interval, newInterval Interval) []Interval {
	var ret []Interval
	var idx int
	n := len(intervals)
	// add all the intervals ending befor newInterval starts
	for idx < n && intervals[idx].End < newInterval.Start {
		ret = append(ret, intervals[idx])
		idx++
	}
	// mrege the overlapped ones
	for idx < n && intervals[idx].Start <= newInterval.End {
		// newInterval = Interval{
		// 	Start: utils.Min(intervals[idx].Start, newInterval.Start),
		// 	End:   utils.Max(intervals[idx].End, newInterval.End),
		// }
		newInterval.Start = utils.Min(intervals[idx].Start, newInterval.Start)
		newInterval.End = utils.Max(intervals[idx].End, newInterval.End)
		idx++
	}
	ret = append(ret, newInterval)
	if idx < n {
		// add teh rest
		ret = append(ret, intervals[idx:]...)
	}
	return ret
}
<|endoftext|>


<commit_before>package skill

import (
	"github.com/fudanchii/sifr/irc"
	"strings"
)

type Skill func(c *irc.Client, m *irc.Message)

func forgeSkill(cmd string, hasArg bool, fn Skill, c *irc.Client) irc.MessageHandler {
	return func(msg *irc.Message) {
		if nocmd(msg.Body, cmd, hasArg) {
			return
		}
		fn(c, msg)
	}
}

func nocmd(txt, cmd string, hasArg bool) bool {
	offset := len(cmd)
	if hasArg {
		offset += 2
	}
	if len(txt) < offset {
		return true
	}
	return !strings.HasPrefix(txt, cmd)
}

func ActivateFor(c *irc.Client) {
	c.AddHandler("PRIVMSG", forgeSkill(".c", true, flipCoin, c))
	c.AddHandler("PRIVMSG", forgeSkill(".j", true, joinChannel, c))
	c.AddHandler("PRIVMSG", forgeSkill(".leave", true, leaveChannel, c))
}
<commit_msg>[WIP] Stub for command authorization<commit_after>package skill

import (
	"github.com/fudanchii/sifr/irc"
	"strings"
)

type archetype func(c *irc.Client, m *irc.Message)

type archMeta struct {
    cmd string
    hasArg bool
    authorized bool
    archFn archetype
}

func forgeSkill(archmeta archMeta, c *irc.Client) irc.MessageHandler {
	return func(msg *irc.Message) {
		if nocmd(msg.Body, archmeta.cmd, archmeta.hasArg) {
			return
		}
		archmeta.archFn(c, msg)
	}
}

func nocmd(txt, cmd string, hasArg bool) bool {
	offset := len(cmd)
	if hasArg {
		offset += 2
	}
	if len(txt) < offset {
		return true
	}
	return !strings.HasPrefix(txt, cmd)
}

func ActivateFor(c *irc.Client) {
	c.AddHandler("PRIVMSG", forgeSkill(".c", true, flipCoin, c))
	c.AddHandler("PRIVMSG", forgeSkill(".j", true, joinChannel, c))
	c.AddHandler("PRIVMSG", forgeSkill(".leave", true, leaveChannel, c))
}
<|endoftext|>


<commit_before>/*
A template for building workers to work with slack-jobs
https://github.com/phamhongviet/slack-jobs
*/
package main

import (
	"fmt"
	"github.com/benmanns/goworker"
	"strings"
)

// struct Job represent the job received from resque
type Job struct {
	Request   string
	User      string
	Channel   string
	Timestamp string
}

// ParseJob create a new Job from arguments
func ParseJob(args []interface{}) Job {
	var job Job
	for _, v := range args {
		v := v.(string)
		switch {
		case strings.HasPrefix(v, "request="):
			job.Request = strings.TrimPrefix(v, "request=")
		case strings.HasPrefix(v, "user="):
			job.User = strings.TrimPrefix(v, "user=")
		case strings.HasPrefix(v, "timestamp="):
			job.Timestamp = strings.TrimPrefix(v, "timestamp=")
		case strings.HasPrefix(v, "channel_name="):
			job.Channel = strings.TrimPrefix(v, "channel_name=")
		}
	}
	return job
}

// Parse has the same logic as ParseJob
// Parse can update an existing Job
func (job *Job) Parse(args []interface{}) {
	for _, v := range args {
		v := v.(string)
		switch {
		case strings.HasPrefix(v, "request="):
			job.Request = strings.TrimPrefix(v, "request=")
		case strings.HasPrefix(v, "user="):
			job.User = strings.TrimPrefix(v, "user=")
		case strings.HasPrefix(v, "timestamp="):
			job.Timestamp = strings.TrimPrefix(v, "timestamp=")
		case strings.HasPrefix(v, "channel_name="):
			job.Channel = strings.TrimPrefix(v, "channel_name=")
		}
	}
}

// register worker with class and handle function
func init() {
	goworker.Register("SlackOPS", slackOPSWorker)
}

// simple handle function
func slackOPSWorker(queue string, args ...interface{}) error {
	// parse job and print its argument
	job := ParseJob(args)
	fmt.Printf("JOB: %s\n", job)
	return nil
}

// let the worker do its job :)
func main() {
	if err := goworker.Work(); err != nil {
		fmt.Println("Error:", err)
	}
}
<commit_msg>print out job<commit_after>/*
A template for building workers to work with slack-jobs
https://github.com/phamhongviet/slack-jobs
*/
package main

import (
	"fmt"
	"github.com/benmanns/goworker"
	"strings"
)

// struct Job represent the job received from resque
type Job struct {
	Request   string
	User      string
	Channel   string
	Timestamp string
}

// ParseJob create a new Job from arguments
func ParseJob(args []interface{}) Job {
	var job Job
	for _, v := range args {
		v := v.(string)
		switch {
		case strings.HasPrefix(v, "request="):
			job.Request = strings.TrimPrefix(v, "request=")
		case strings.HasPrefix(v, "user="):
			job.User = strings.TrimPrefix(v, "user=")
		case strings.HasPrefix(v, "timestamp="):
			job.Timestamp = strings.TrimPrefix(v, "timestamp=")
		case strings.HasPrefix(v, "channel_name="):
			job.Channel = strings.TrimPrefix(v, "channel_name=")
		}
	}
	return job
}

// Parse has the same logic as ParseJob
// Parse can update an existing Job
func (job *Job) Parse(args []interface{}) {
	for _, v := range args {
		v := v.(string)
		switch {
		case strings.HasPrefix(v, "request="):
			job.Request = strings.TrimPrefix(v, "request=")
		case strings.HasPrefix(v, "user="):
			job.User = strings.TrimPrefix(v, "user=")
		case strings.HasPrefix(v, "timestamp="):
			job.Timestamp = strings.TrimPrefix(v, "timestamp=")
		case strings.HasPrefix(v, "channel_name="):
			job.Channel = strings.TrimPrefix(v, "channel_name=")
		}
	}
}

// register worker with class and handle function
func init() {
	goworker.Register("SlackOPS", slackOPSWorker)
}

// simple handle function
func slackOPSWorker(queue string, args ...interface{}) error {
	// parse job and print its argument
	job := ParseJob(args)
	fmt.Printf("%s, %s, %s, %s\n", job.Request, job.User, job.Channel, job.Timestamp)
	return nil
}

// let the worker do its job :)
func main() {
	if err := goworker.Work(); err != nil {
		fmt.Println("Error:", err)
	}
}
<|endoftext|>


<commit_before>package endpoint

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"
)

const (
	httpExpiresAfter       = time.Second * 30
	httpRequestTimeout     = time.Second * 5
	httpMaxIdleConnections = 20
)

type HTTPEndpointConn struct {
	mu     sync.Mutex
	ep     Endpoint
	ex     bool
	t      time.Time
	client *http.Client
}

func newHTTPEndpointConn(ep Endpoint) *HTTPEndpointConn {
	return &HTTPEndpointConn{
		ep: ep,
		t:  time.Now(),
	}
}

func (conn *HTTPEndpointConn) Expired() bool {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if !conn.ex {
		if time.Now().Sub(conn.t) > httpExpiresAfter {
			conn.ex = true
			conn.client = nil
		}
	}
	return conn.ex
}

func (conn *HTTPEndpointConn) Send(msg string) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if conn.ex {
		return errors.New("expired")
	}
	conn.t = time.Now()
	if conn.client == nil {
		conn.client = &http.Client{
			Transport: &http.Transport{
				MaxIdleConnsPerHost: httpMaxIdleConnections,
			},
			Timeout: httpRequestTimeout,
		}
	}
	req, err := http.NewRequest("POST", conn.ep.Original, bytes.NewBufferString(msg))
	if err != nil {
		return err
	}
	resp, err := conn.client.Do(req)
	if err != nil {
		return err
	}
	// close the connection to reuse it
	defer resp.Body.Close()
	// discard response
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		return err
	}
	// we only care about the 200 response
	if resp.StatusCode != 200 {
		return fmt.Errorf("invalid status: %s", resp.Status)
	}
	return nil
}
<commit_msg>Fixed #124: Added Content-Type json to HTTP headers (#125)<commit_after>package endpoint

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"
)

const (
	httpExpiresAfter       = time.Second * 30
	httpRequestTimeout     = time.Second * 5
	httpMaxIdleConnections = 20
)

type HTTPEndpointConn struct {
	mu     sync.Mutex
	ep     Endpoint
	ex     bool
	t      time.Time
	client *http.Client
}

func newHTTPEndpointConn(ep Endpoint) *HTTPEndpointConn {
	return &HTTPEndpointConn{
		ep: ep,
		t:  time.Now(),
	}
}

func (conn *HTTPEndpointConn) Expired() bool {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if !conn.ex {
		if time.Now().Sub(conn.t) > httpExpiresAfter {
			conn.ex = true
			conn.client = nil
		}
	}
	return conn.ex
}

func (conn *HTTPEndpointConn) Send(msg string) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if conn.ex {
		return errors.New("expired")
	}
	conn.t = time.Now()
	if conn.client == nil {
		conn.client = &http.Client{
			Transport: &http.Transport{
				MaxIdleConnsPerHost: httpMaxIdleConnections,
			},
			Timeout: httpRequestTimeout,
		}
	}
	req, err := http.NewRequest("POST", conn.ep.Original, bytes.NewBufferString(msg))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json")
	resp, err := conn.client.Do(req)
	if err != nil {
		return err
	}
	// close the connection to reuse it
	defer resp.Body.Close()
	// discard response
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		return err
	}
	// we only care about the 200 response
	if resp.StatusCode != 200 {
		return fmt.Errorf("invalid status: %s", resp.Status)
	}
	return nil
}
<|endoftext|>