<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

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

package tchannel_test

// This file contains functions for tests to access internal tchannel state.
// Since it has a _test.go suffix, it is only compiled with tests in this package.

import (
	"bytes"
	"math/rand"
	"sync"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"
	"golang.org/x/net/context"
)

type swapper struct {
	t *testing.T
}

func (s *swapper) OnError(ctx context.Context, err error) {
	s.t.Errorf("OnError: %v", err)
}

func (*swapper) Handle(ctx context.Context, args *raw.Args) (*raw.Res, error) {
	return &raw.Res{
		Arg2: args.Arg3,
		Arg3: args.Arg2,
	}, nil
}

func TestFramesReleased(t *testing.T) {
	CheckStress(t)

	defer testutils.SetTimeout(t, 10*time.Second)()
	const (
		requestsPerGoroutine = 10
		numGoroutines        = 10
		maxRandArg           = 512 * 1024
	)

	var connections []*Connection
	pool := NewRecordingFramePool()
	WithVerifiedServer(t, &testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: pool,
		},
	}, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		clientCh, err := NewChannel("swap-client", nil)
		require.NoError(t, err)
		defer clientCh.Close()

		// Create an active connection that can be shared by the goroutines by calling Ping.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, clientCh.Ping(ctx, hostPort))

		var wg sync.WaitGroup
		worker := func() {
			for i := 0; i < requestsPerGoroutine; i++ {
				ctx, cancel := NewContext(time.Second * 5)
				defer cancel()

				require.NoError(t, clientCh.Ping(ctx, hostPort))

				arg2 := testutils.RandBytes(rand.Intn(maxRandArg))
				arg3 := testutils.RandBytes(rand.Intn(maxRandArg))
				resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
				if !assert.NoError(t, err, "error during sendRecv") {
					continue
				}

				// We expect the arguments to be swapped.
				if bytes.Compare(arg3, resArg2) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
				}
				if bytes.Compare(arg2, resArg3) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
				}
			}
			wg.Done()
		}

		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go worker()
		}

		wg.Wait()

		connections = append(connections, GetConnections(serverCh)...)
		connections = append(connections, GetConnections(clientCh)...)
	})

	// Wait a few milliseconds for the closing of channels to take effect.
	time.Sleep(10 * time.Millisecond)

	if unreleasedCount, isEmpty := pool.CheckEmpty(); isEmpty != "" || unreleasedCount > 0 {
		t.Errorf("Frame pool has %v unreleased frames, errors:\n%v", unreleasedCount, isEmpty)
	}

	// Check the message exchanges and make sure they are all empty.
	if exchangesLeft := CheckEmptyExchangesConns(connections); exchangesLeft != "" {
		t.Errorf("Found uncleared message exchanges:\n%v", exchangesLeft)
	}
}
<commit_msg>Add test to reproduce hasMoreFragments bug<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

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

package tchannel_test

// This file contains functions for tests to access internal tchannel state.
// Since it has a _test.go suffix, it is only compiled with tests in this package.

import (
	"bytes"
	"math/rand"
	"sync"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"
	"golang.org/x/net/context"
)

type swapper struct {
	t *testing.T
}

func (s *swapper) OnError(ctx context.Context, err error) {
	s.t.Errorf("OnError: %v", err)
}

func (*swapper) Handle(ctx context.Context, args *raw.Args) (*raw.Res, error) {
	return &raw.Res{
		Arg2: args.Arg3,
		Arg3: args.Arg2,
	}, nil
}

func TestFramesReleased(t *testing.T) {
	CheckStress(t)

	defer testutils.SetTimeout(t, 10*time.Second)()
	const (
		requestsPerGoroutine = 10
		numGoroutines        = 10
		maxRandArg           = 512 * 1024
	)

	var connections []*Connection
	pool := NewRecordingFramePool()
	WithVerifiedServer(t, &testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: pool,
		},
	}, func(serverCh *Channel, hostPort string) {
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		clientCh, err := NewChannel("swap-client", nil)
		require.NoError(t, err)
		defer clientCh.Close()

		// Create an active connection that can be shared by the goroutines by calling Ping.
		ctx, cancel := NewContext(time.Second)
		defer cancel()
		require.NoError(t, clientCh.Ping(ctx, hostPort))

		var wg sync.WaitGroup
		worker := func() {
			for i := 0; i < requestsPerGoroutine; i++ {
				ctx, cancel := NewContext(time.Second * 5)
				defer cancel()

				require.NoError(t, clientCh.Ping(ctx, hostPort))

				arg2 := testutils.RandBytes(rand.Intn(maxRandArg))
				arg3 := testutils.RandBytes(rand.Intn(maxRandArg))
				resArg2, resArg3, _, err := raw.Call(ctx, clientCh, hostPort, "swap-server", "swap", arg2, arg3)
				if !assert.NoError(t, err, "error during sendRecv") {
					continue
				}

				// We expect the arguments to be swapped.
				if bytes.Compare(arg3, resArg2) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg2, arg3)
				}
				if bytes.Compare(arg2, resArg3) != 0 {
					t.Errorf("returned arg2 does not match expected:\n  got %v\n want %v", resArg3, arg2)
				}
			}
			wg.Done()
		}

		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go worker()
		}

		wg.Wait()

		connections = append(connections, GetConnections(serverCh)...)
		connections = append(connections, GetConnections(clientCh)...)
	})

	// Wait a few milliseconds for the closing of channels to take effect.
	time.Sleep(10 * time.Millisecond)

	if unreleasedCount, isEmpty := pool.CheckEmpty(); isEmpty != "" || unreleasedCount > 0 {
		t.Errorf("Frame pool has %v unreleased frames, errors:\n%v", unreleasedCount, isEmpty)
	}

	// Check the message exchanges and make sure they are all empty.
	if exchangesLeft := CheckEmptyExchangesConns(connections); exchangesLeft != "" {
		t.Errorf("Found uncleared message exchanges:\n%v", exchangesLeft)
	}
}

type dirtyFramePool struct{}

func (p dirtyFramePool) Get() *Frame {
	f := NewFrame(MaxFramePayloadSize)
	for i := range f.Payload {
		f.Payload[i] = ^byte(0)
	}
	return f
}

func (p dirtyFramePool) Release(f *Frame) {}

func TestDirtyFrameRequests(t *testing.T) {

	argSizes := []int{50000, 100000, 150000}
	WithVerifiedServer(t, &testutils.ChannelOpts{
		ServiceName: "swap-server",
		DefaultConnectionOptions: ConnectionOptions{
			FramePool: dirtyFramePool{},
		},
	}, func(serverCh *Channel, hostPort string) {
		peerInfo := serverCh.PeerInfo()
		serverCh.Register(raw.Wrap(&swapper{t}), "swap")

		for _, arg2Size := range argSizes {
			for _, arg3Size := range argSizes {
				ctx, cancel := NewContext(time.Second)
				defer cancel()

				arg2, arg3 := testutils.RandBytes(arg2Size), testutils.RandBytes(arg3Size)
				res2, res3, _, err := raw.Call(ctx, serverCh, hostPort, peerInfo.ServiceName, "swap", arg2, arg3)
				if assert.NoError(t, err, "Call failed") {
					assert.Equal(t, arg2, res3, "Result arg3 wrong")
					assert.Equal(t, arg3, res2, "Result arg3 wrong")
				}
			}
		}
	})
}
<|endoftext|>


<commit_before>package test_with_docker

import (
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

type LocalDaemon struct {
	serviceTimeout time.Duration
}

func ldTrial() agentBuilderF {
	log.Println("Checking docker socket:")
	o, _ := exec.Command("ls", "-l", "/var/run/docker.sock").CombinedOutput()
	log.Print(string(o))
	ps := runCommand("docker", "ps")
	if ps.err != nil {
		log.Printf("docker ps failed:\n  Stdout:\n%s\n  Stderr:\n%s", ps.stdout, ps.stderr)
		return nil
	}
	return ldBuild
}

func ldBuild(c agentCfg) Agent {
	log.Println("Using local docker daemon")
	return &LocalDaemon{serviceTimeout: c.timeout}
}

func (ld *LocalDaemon) ComposeServices(dir string, svcs serviceMap) (*command, error) {
	ip, _ := ld.IP()

	return composeService(dir, ip, []string{}, svcs, ld.serviceTimeout)
}

// RebuildService forces the rebuild of a docker-compose service.
func (ld *LocalDaemon) RebuildService(dir, name string) error {
	return rebuildService(dir, name, []string{})
}

// Shutdown terminates the set of services started by ComposeServices
// If passed a nil (as ComposeServices returns in the event that all services
// were available, Shutdown is a no-op
func (ld *LocalDaemon) Shutdown(c *command) {
	if c != nil {
		dockerComposeDown(c)
	}
}

// ShutdownNow implements Agent for LocalDaemon
func (ld *LocalDaemon) ShutdownNow() {
	dockerComposeDown(nil)
}

func (ld *LocalDaemon) Cleanup() error {
	log.Println("Cleaning up...")
	o, _ := exec.Command("ls", "-l", "/var/run/docker.sock").CombinedOutput()
	log.Print(string(o))
	return nil
}

// InstallRegistryCertificate for a local daemon simply logs it's unwillingness to comply.
func (ld *LocalDaemon) InstallRegistryCertificate(registryName, composeDir, registryCertPath string) error {
	log.Print("Local daemon agent doesn't need registry certs installed, and would require sudo to do it. No-op.")
	log.Printf("If you have trouble with registry push in your tests, consider copying %q to /etc/docker/certs.d/%q/ca.crt", registryCertPath, registryName)
	return nil
}

// InstallFile puts a path found on the local machine to a path on the docker host.
func (ld *LocalDaemon) InstallFile(src string, dest string) error {
	destDir := filepath.Dir(dest)
	ld.Exec("mkdir", "-p", destDir)
	return ld.Exec("cp", src, dest)
}

// DifferingFiles compares specific files involved in docker
func (ld *LocalDaemon) DifferingFiles(pathPairs ...[]string) (differentPairs [][]string, err error) {
	localPaths, remotePaths := make([]string, 0, len(pathPairs)), make([]string, 0, len(pathPairs))

	for _, pair := range pathPairs {
		localPaths = append(localPaths, pair[0])
		remotePaths = append(remotePaths, pair[1])
	}

	localMD5s := localMD5s(localPaths...)
	remoteMD5s, err := ld.MD5s(remotePaths...)
	if err != nil {
		return
	}

	return fileDiffs(pathPairs, localMD5s, remoteMD5s), nil
}

// IP returns the IP address where the daemon is located.
// In order to access the services provided by a docker-compose on a
// docker-machine, we need to know the ip address. Some client test code
// needs to know the IP address prior to starting up the services, which is
// why this function is exposed
func (ld *LocalDaemon) IP() (net.IP, error) {
	return net.ParseIP(`127.0.0.1`), nil
}

// MD5s computes digests of a list of paths
// This can be used to compare to local digests and avoid copying files or
// restarting the daemon
func (ld *LocalDaemon) MD5s(paths ...string) (map[string]string, error) {
	dir, err := ioutil.TempDir("/tmp", "md5")
	if err != nil {
		return nil, err
	}
	var tgts []string
	md5 := make(map[string]string)
	for _, p := range paths {
		tgt := filepath.Join(dir, p)
		os.MkdirAll(filepath.Dir(tgt), os.ModeDir|os.ModePerm)
		ld.Exec("cp", p, tgt)
		tgts = append(tgts, tgt)
		md5[p] = tgt
	}
	ld.Exec("chmod", "-R", "a+rw", dir)
	tmd5s := localMD5s(tgts...)
	os.RemoveAll(dir)
	for p, t := range md5 {
		md5[p] = tmd5s[t]
	}
	return md5, nil
}

// RestartDaemon reboots the docker daemon
func (ld *LocalDaemon) RestartDaemon() error {
	rss := [][]string{
		[]string{"/etc/init.d/docker", "restart"},
		[]string{"service", "docker", "restart"},
		[]string{"systemctl", "restart", "docker", "docker.socket"},
	}
	var err error

	for _, rs := range rss {
		err = ld.Exec(rs...)
		if err == nil {
			break
		}
	}

	UntilReady(time.Second/5, 30*time.Second, func() (string, func() bool, func()) {
		return "Docker socket recreated",
			func() bool {
				err := exec.Command("ls", "-l", "/var/run/docker.sock").Run()
				return err == nil
			},
			func() {}
	})

	return err
}

// Exec executes commands as root on the daemon host
// It uses sudo. This is dangerous.
func (ld *LocalDaemon) Exec(args ...string) error {
	cmd := runCommand("sudo", args...)
	log.Println(cmd.String())
	return cmd.err
}

/*
 */
<commit_msg>Unquoting some strings<commit_after>package test_with_docker

import (
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

type LocalDaemon struct {
	serviceTimeout time.Duration
}

func ldTrial() agentBuilderF {
	log.Println("Checking docker socket:")
	o, _ := exec.Command("ls", "-l", "/var/run/docker.sock").CombinedOutput()
	log.Print(string(o))
	ps := runCommand("docker", "ps")
	if ps.err != nil {
		log.Printf("docker ps failed:\n  Stdout:\n%s\n  Stderr:\n%s", ps.stdout, ps.stderr)
		return nil
	}
	return ldBuild
}

func ldBuild(c agentCfg) Agent {
	log.Println("Using local docker daemon")
	return &LocalDaemon{serviceTimeout: c.timeout}
}

func (ld *LocalDaemon) ComposeServices(dir string, svcs serviceMap) (*command, error) {
	ip, _ := ld.IP()

	return composeService(dir, ip, []string{}, svcs, ld.serviceTimeout)
}

// RebuildService forces the rebuild of a docker-compose service.
func (ld *LocalDaemon) RebuildService(dir, name string) error {
	return rebuildService(dir, name, []string{})
}

// Shutdown terminates the set of services started by ComposeServices
// If passed a nil (as ComposeServices returns in the event that all services
// were available, Shutdown is a no-op
func (ld *LocalDaemon) Shutdown(c *command) {
	if c != nil {
		dockerComposeDown(c)
	}
}

// ShutdownNow implements Agent for LocalDaemon
func (ld *LocalDaemon) ShutdownNow() {
	dockerComposeDown(nil)
}

func (ld *LocalDaemon) Cleanup() error {
	log.Println("Cleaning up...")
	o, _ := exec.Command("ls", "-l", "/var/run/docker.sock").CombinedOutput()
	log.Print(string(o))
	return nil
}

// InstallRegistryCertificate for a local daemon simply logs it's unwillingness to comply.
func (ld *LocalDaemon) InstallRegistryCertificate(registryName, composeDir, registryCertPath string) error {
	log.Print("Local daemon agent doesn't need registry certs installed, and would require sudo to do it. No-op.")
	log.Printf("If you have trouble with registry push in your tests,\n consider copying %s to /etc/docker/certs.d/%s/ca.crt", registryCertPath, registryName)
	return nil
}

// InstallFile puts a path found on the local machine to a path on the docker host.
func (ld *LocalDaemon) InstallFile(src string, dest string) error {
	destDir := filepath.Dir(dest)
	ld.Exec("mkdir", "-p", destDir)
	return ld.Exec("cp", src, dest)
}

// DifferingFiles compares specific files involved in docker
func (ld *LocalDaemon) DifferingFiles(pathPairs ...[]string) (differentPairs [][]string, err error) {
	localPaths, remotePaths := make([]string, 0, len(pathPairs)), make([]string, 0, len(pathPairs))

	for _, pair := range pathPairs {
		localPaths = append(localPaths, pair[0])
		remotePaths = append(remotePaths, pair[1])
	}

	localMD5s := localMD5s(localPaths...)
	remoteMD5s, err := ld.MD5s(remotePaths...)
	if err != nil {
		return
	}

	return fileDiffs(pathPairs, localMD5s, remoteMD5s), nil
}

// IP returns the IP address where the daemon is located.
// In order to access the services provided by a docker-compose on a
// docker-machine, we need to know the ip address. Some client test code
// needs to know the IP address prior to starting up the services, which is
// why this function is exposed
func (ld *LocalDaemon) IP() (net.IP, error) {
	return net.ParseIP(`127.0.0.1`), nil
}

// MD5s computes digests of a list of paths
// This can be used to compare to local digests and avoid copying files or
// restarting the daemon
func (ld *LocalDaemon) MD5s(paths ...string) (map[string]string, error) {
	dir, err := ioutil.TempDir("/tmp", "md5")
	if err != nil {
		return nil, err
	}
	var tgts []string
	md5 := make(map[string]string)
	for _, p := range paths {
		tgt := filepath.Join(dir, p)
		os.MkdirAll(filepath.Dir(tgt), os.ModeDir|os.ModePerm)
		ld.Exec("cp", p, tgt)
		tgts = append(tgts, tgt)
		md5[p] = tgt
	}
	ld.Exec("chmod", "-R", "a+rw", dir)
	tmd5s := localMD5s(tgts...)
	os.RemoveAll(dir)
	for p, t := range md5 {
		md5[p] = tmd5s[t]
	}
	return md5, nil
}

// RestartDaemon reboots the docker daemon
func (ld *LocalDaemon) RestartDaemon() error {
	rss := [][]string{
		[]string{"/etc/init.d/docker", "restart"},
		[]string{"service", "docker", "restart"},
		[]string{"systemctl", "restart", "docker", "docker.socket"},
	}
	var err error

	for _, rs := range rss {
		err = ld.Exec(rs...)
		if err == nil {
			break
		}
	}

	UntilReady(time.Second/5, 30*time.Second, func() (string, func() bool, func()) {
		return "Docker socket recreated",
			func() bool {
				err := exec.Command("ls", "-l", "/var/run/docker.sock").Run()
				return err == nil
			},
			func() {}
	})

	return err
}

// Exec executes commands as root on the daemon host
// It uses sudo. This is dangerous.
func (ld *LocalDaemon) Exec(args ...string) error {
	cmd := runCommand("sudo", args...)
	log.Println(cmd.String())
	return cmd.err
}

/*
 */
<|endoftext|>


<commit_before>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	ctx         context.Context
	cancel      context.CancelFunc
	canWrite    *sync.Cond // Condition variable to implement TCP backpressure.
}

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler) (Connection, error) {
	// prepare key
	connKeyArg := NewConnKeyArg()
	connKey := rand.Uint32()
	SetConnKeyVal(unsafe.Pointer(connKeyArg), connKey)

	if tcpConnectionHandler == nil {
		return nil, errors.New("no registered TCP connection handlers found")
	}

	ctx, cancel := context.WithCancel(context.Background())

	conn := &tcpConn{
		pcb:         pcb,
		handler:     handler,
		network:     "tcp",
		localAddr:   ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:  ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:  connKeyArg,
		connKey:     connKey,
		closing:     false,
		localClosed: false,
		aborting:    false,
		ctx:         ctx,
		cancel:      cancel,
		canWrite:    sync.NewCond(&sync.Mutex{}),
	}

	// Associate conn with key and save to the global map.
	tcpConns.Store(connKey, conn)

	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(TCP_POLL_INTERVAL))

	// Unlocks lwip thread during connecting remote host, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := handler.Connect(conn, conn.RemoteAddr())
	lwipMutex.Lock()
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}
	if conn.isAborting() {
		return errors.New(fmt.Sprintf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr()))
	}
	// Unlocks lwip thread during sending data to remote, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := conn.handler.DidReceive(conn, data)
	lwipMutex.Lock()
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}
	C.tcp_recved(conn.pcb, C.u16_t(len(data)))
	return nil
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (int, error) {
	err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
	if err == C.ERR_OK {
		C.tcp_output(conn.pcb)
		return len(data), nil
	} else if err == C.ERR_MEM {
		return 0, nil
	}
	return 0, fmt.Errorf("lwip tcp_write failed with error code: %v", int(err))
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	totalWritten := 0
	conn.canWrite.L.Lock()
	for len(data) > 0 {
		if conn.isLocalClosed() {
			return 0, fmt.Errorf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr())
		}
		if conn.isAborting() {
			return 0, fmt.Errorf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr())
		}

		lwipMutex.Lock()
		toWrite := len(data)
		if toWrite > int(conn.pcb.snd_buf) {
			// Write at most the size of the LWIP buffer.
			toWrite = int(conn.pcb.snd_buf)
		}
		if toWrite > 0 {
			written, err := conn.tcpWrite(data[0:toWrite])
			totalWritten += written
			if err != nil {
				lwipMutex.Unlock()
				conn.canWrite.L.Unlock()
				return totalWritten, err
			}
			data = data[written:len(data)]
		}
		lwipMutex.Unlock()
		if len(data) == 0 {
			break // Don't block if all the data has been written.
		}
		conn.canWrite.Wait()
	}
	conn.canWrite.L.Unlock()

	return totalWritten, nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) CheckState() error {
	if !conn.isLocalClosed() {
		// Signal the writer to try writting.
		conn.canWrite.Broadcast()
		// Return and wait for the Sent() callback to be called, and then check again.
		return NewLWIPError(LWIP_ERR_OK)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
	}

	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	conn.canWrite.Broadcast()
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	conn.cancel()

	// TODO: may return ERR_MEM if no memory to allocate segments use for closing the conn,
	// should check and try again in Sent() for Poll() callbacks.
	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprintf("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
	conn.canWrite.Broadcast()
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	conn.Release()
	conn.cancel()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<commit_msg>make sure the cond var is unlocked before return<commit_after>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	ctx         context.Context
	cancel      context.CancelFunc
	canWrite    *sync.Cond // Condition variable to implement TCP backpressure.
}

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler) (Connection, error) {
	// prepare key
	connKeyArg := NewConnKeyArg()
	connKey := rand.Uint32()
	SetConnKeyVal(unsafe.Pointer(connKeyArg), connKey)

	if tcpConnectionHandler == nil {
		return nil, errors.New("no registered TCP connection handlers found")
	}

	ctx, cancel := context.WithCancel(context.Background())

	conn := &tcpConn{
		pcb:         pcb,
		handler:     handler,
		network:     "tcp",
		localAddr:   ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:  ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:  connKeyArg,
		connKey:     connKey,
		closing:     false,
		localClosed: false,
		aborting:    false,
		ctx:         ctx,
		cancel:      cancel,
		canWrite:    sync.NewCond(&sync.Mutex{}),
	}

	// Associate conn with key and save to the global map.
	tcpConns.Store(connKey, conn)

	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(TCP_POLL_INTERVAL))

	// Unlocks lwip thread during connecting remote host, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := handler.Connect(conn, conn.RemoteAddr())
	lwipMutex.Lock()
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}
	if conn.isAborting() {
		return errors.New(fmt.Sprintf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr()))
	}
	// Unlocks lwip thread during sending data to remote, gives other goroutines
	// chances to interact with the lwip thread. Assuming lwip thread has already
	// been locked.
	lwipMutex.Unlock()
	err := conn.handler.DidReceive(conn, data)
	lwipMutex.Lock()
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}
	C.tcp_recved(conn.pcb, C.u16_t(len(data)))
	return nil
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (int, error) {
	err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
	if err == C.ERR_OK {
		C.tcp_output(conn.pcb)
		return len(data), nil
	} else if err == C.ERR_MEM {
		return 0, nil
	}
	return 0, fmt.Errorf("lwip tcp_write failed with error code: %v", int(err))
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	totalWritten := 0

	conn.canWrite.L.Lock()
	defer conn.canWrite.L.Unlock()

	for len(data) > 0 {
		if conn.isLocalClosed() {
			return 0, fmt.Errorf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr())
		}
		if conn.isAborting() {
			return 0, fmt.Errorf("connection %v->%v is aborting", conn.LocalAddr(), conn.RemoteAddr())
		}

		lwipMutex.Lock()
		toWrite := len(data)
		if toWrite > int(conn.pcb.snd_buf) {
			// Write at most the size of the LWIP buffer.
			toWrite = int(conn.pcb.snd_buf)
		}
		if toWrite > 0 {
			written, err := conn.tcpWrite(data[0:toWrite])
			totalWritten += written
			if err != nil {
				lwipMutex.Unlock()
				return totalWritten, err
			}
			data = data[written:len(data)]
		}
		lwipMutex.Unlock()
		if len(data) == 0 {
			break // Don't block if all the data has been written.
		}
		conn.canWrite.Wait()
	}

	return totalWritten, nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) CheckState() error {
	if !conn.isLocalClosed() {
		// Signal the writer to try writting.
		conn.canWrite.Broadcast()
		// Return and wait for the Sent() callback to be called, and then check again.
		return NewLWIPError(LWIP_ERR_OK)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
	}

	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	conn.canWrite.Broadcast()
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	conn.cancel()

	// TODO: may return ERR_MEM if no memory to allocate segments use for closing the conn,
	// should check and try again in Sent() for Poll() callbacks.
	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprintf("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
	conn.canWrite.Broadcast()
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	conn.Release()
	conn.cancel()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<|endoftext|>


<commit_before>package core

// Topology is a topology which can add Sources, Boxes, and Sinks
// dynamically. Boxes and Sinks can also add inputs dynamically from running
// Sources or Boxes.
type Topology interface {
	// Name returns the name of this topology.
	Name() string

	// Context returns the Context tied to the topology. It isn't always safe
	// to modify fields of the Context. However, some fields of it can handle
	// concurrent access properly. See the document of Context for details.
	Context() *Context

	// AddSource adds a Source to the topology. It will asynchronously call
	// source's GenerateStream and returns after the source becomes ready.
	// GenerateStream could be called lazily to avoid unnecessary computation.
	// GenerateStream and Stop might also be called when this method returns
	// an error. The caller must not call GenerateStream or Stop of the source.
	AddSource(name string, s Source, config *SourceConfig) (SourceNode, error)

	// AddBox adds a Box to the topology. It returns BoxNode and the
	// caller can configure inputs or other settings of the Box node through it.
	//
	// Don't add the same instance of a Box more than once even if they have
	// different names. However, if a Box has a reference counter and
	// its initialization and termination are done exactly once at proper
	// timing, it can be added multiple times when a builder supports duplicated
	// registration of the same instance of a Box.
	AddBox(name string, b Box, config *BoxConfig) (BoxNode, error)

	// AddSink adds a Sink to the topology. It returns SinkNode and the
	// caller can configure inputs or other settings of the Sink node through it.
	AddSink(name string, s Sink, config *SinkConfig) (SinkNode, error)

	// Remove removes a node from the topology. It doesn't return an error when
	// the topology couldn't find a node having the name. The removed node is
	// stopped by the topology and Remove methods blocks until the node actually
	// stops.
	Remove(name string) error

	// Stop stops the topology. If the topology doesn't have a cycle, it stops
	// after all tuples generated from Sources at the time of the invocation
	// are written into Sinks. Stop method returns after processing all the
	// tuples.
	//
	// BUG: Currently Stop method doesn't work if the topology has a cycle.
	Stop() error

	// State returns the current state of the topology. The topology's state
	// isn't relevant to those nodes have.
	State() TopologyStateHolder

	// TODO: low priority: Pause, Resume

	// Node returns a node registered to the topology. It returns an error
	// when the topology doesn't have the node.
	Node(name string) (Node, error)

	// Nodes returns all nodes registered to the topology. The map returned
	// from this method can safely be modified.
	Nodes() map[string]Node

	// Source returns a source registered to the topology. It returns an error
	// when the topology doesn't have the source.
	Source(name string) (SourceNode, error)

	// Sources returns all sources registered to the topology. The map returned
	// from this method can safely be modified.
	Sources() map[string]SourceNode

	// Box returns a box registered to the topology. It returns an error when
	// the topology doesn't have the box.
	Box(name string) (BoxNode, error)

	// Boxes returns all boxes registered to the topology. The map returned
	// from this method can safely be modified.
	Boxes() map[string]BoxNode

	// Sink returns a sink registereed to the topology. It returns an error
	// when the topology doesn't have the sink.
	Sink(name string) (SinkNode, error)

	// Sinks returns all sinks registered to the topology. The map returned
	// from this method can safely be modified.
	Sinks() map[string]SinkNode
}

// SourceConfig has configuration parameters of a Source node.
type SourceConfig struct {
	// PausedOnStartup is a flag which indicates the initial state of the
	// source. If it is true, the source is paused. Otherwise, source runs just
	// after it is added to a topology.
	PausedOnStartup bool

	// RemoveOnStop is a flag which indicates the stop state of the
	// topology. If it is true, the source is removed.
	RemoveOnStop bool

	// Meta contains meta information of the source. This field won't be used
	// by core package and application can store any form of information
	// related to the source.
	Meta interface{}
}

// BoxConfig has configuration parameters of a Box node.
type BoxConfig struct {
	// TODO: parallelism

	// RemoveOnStop is a flag which indicates the stop state of the
	// topology. If it is true, the box is removed.
	RemoveOnStop bool

	// Meta contains meta information of the box. This field won't be used
	// by core package and application can store any form of information
	// related to the box.
	Meta interface{}
}

// SinkConfig has configuration parameters of a Sink node.
type SinkConfig struct {
	// RemoveOnStop is a flag which indicates the stop state of the
	// topology. If it is true, the sink is removed.
	RemoveOnStop bool

	// Meta contains meta information of the sink. This field won't be used
	// by core package and application can store any form of information
	// related to the sink.
	Meta interface{}
}
<commit_msg>Adjust comment width<commit_after>package core

// Topology is a topology which can add Sources, Boxes, and Sinks
// dynamically. Boxes and Sinks can also add inputs dynamically from running
// Sources or Boxes.
type Topology interface {
	// Name returns the name of this topology.
	Name() string

	// Context returns the Context tied to the topology. It isn't always safe
	// to modify fields of the Context. However, some fields of it can handle
	// concurrent access properly. See the document of Context for details.
	Context() *Context

	// AddSource adds a Source to the topology. It will asynchronously call
	// source's GenerateStream and returns after the source becomes ready.
	// GenerateStream could be called lazily to avoid unnecessary computation.
	// GenerateStream and Stop might also be called when this method returns
	// an error. The caller must not call GenerateStream or Stop of the source.
	AddSource(name string, s Source, config *SourceConfig) (SourceNode, error)

	// AddBox adds a Box to the topology. It returns BoxNode and the
	// caller can configure inputs or other settings of the Box node through it.
	//
	// Don't add the same instance of a Box more than once even if they have
	// different names. However, if a Box has a reference counter and
	// its initialization and termination are done exactly once at proper
	// timing, it can be added multiple times when a builder supports duplicated
	// registration of the same instance of a Box.
	AddBox(name string, b Box, config *BoxConfig) (BoxNode, error)

	// AddSink adds a Sink to the topology. It returns SinkNode and the
	// caller can configure inputs or other settings of the Sink node through it.
	AddSink(name string, s Sink, config *SinkConfig) (SinkNode, error)

	// Remove removes a node from the topology. It doesn't return an error when
	// the topology couldn't find a node having the name. The removed node is
	// stopped by the topology and Remove methods blocks until the node actually
	// stops.
	Remove(name string) error

	// Stop stops the topology. If the topology doesn't have a cycle, it stops
	// after all tuples generated from Sources at the time of the invocation
	// are written into Sinks. Stop method returns after processing all the
	// tuples.
	//
	// BUG: Currently Stop method doesn't work if the topology has a cycle.
	Stop() error

	// State returns the current state of the topology. The topology's state
	// isn't relevant to those nodes have.
	State() TopologyStateHolder

	// TODO: low priority: Pause, Resume

	// Node returns a node registered to the topology. It returns an error
	// when the topology doesn't have the node.
	Node(name string) (Node, error)

	// Nodes returns all nodes registered to the topology. The map returned
	// from this method can safely be modified.
	Nodes() map[string]Node

	// Source returns a source registered to the topology. It returns an error
	// when the topology doesn't have the source.
	Source(name string) (SourceNode, error)

	// Sources returns all sources registered to the topology. The map returned
	// from this method can safely be modified.
	Sources() map[string]SourceNode

	// Box returns a box registered to the topology. It returns an error when
	// the topology doesn't have the box.
	Box(name string) (BoxNode, error)

	// Boxes returns all boxes registered to the topology. The map returned
	// from this method can safely be modified.
	Boxes() map[string]BoxNode

	// Sink returns a sink registereed to the topology. It returns an error
	// when the topology doesn't have the sink.
	Sink(name string) (SinkNode, error)

	// Sinks returns all sinks registered to the topology. The map returned
	// from this method can safely be modified.
	Sinks() map[string]SinkNode
}

// SourceConfig has configuration parameters of a Source node.
type SourceConfig struct {
	// PausedOnStartup is a flag which indicates the initial state of the
	// source. If it is true, the source is paused. Otherwise, source runs just
	// after it is added to a topology.
	PausedOnStartup bool

	// RemoveOnStop is a flag which indicates the stop state of the topology.
	// If it is true, the source is removed.
	RemoveOnStop bool

	// Meta contains meta information of the source. This field won't be used
	// by core package and application can store any form of information
	// related to the source.
	Meta interface{}
}

// BoxConfig has configuration parameters of a Box node.
type BoxConfig struct {
	// TODO: parallelism

	// RemoveOnStop is a flag which indicates the stop state of the topology.
	// If it is true, the box is removed.
	RemoveOnStop bool

	// Meta contains meta information of the box. This field won't be used
	// by core package and application can store any form of information
	// related to the box.
	Meta interface{}
}

// SinkConfig has configuration parameters of a Sink node.
type SinkConfig struct {
	// RemoveOnStop is a flag which indicates the stop state of the topology.
	// If it is true, the sink is removed.
	RemoveOnStop bool

	// Meta contains meta information of the sink. This field won't be used
	// by core package and application can store any form of information
	// related to the sink.
	Meta interface{}
}
<|endoftext|>


<commit_before>package jpeg

//
// Original codes are bollowed from go-thumber.
// Copyright (c) 2014 pixiv Inc. All rights reserved.
//
// See: https://github.com/pixiv/go-thumber
//

/*
#include <stdlib.h>
#include <stdio.h>
#include <jpeglib.h>

// exported from golang
void sourceInit(struct jpeg_decompress_struct*);
void sourceSkip(struct jpeg_decompress_struct*, long);
boolean sourceFill(struct jpeg_decompress_struct*);
void sourceTerm(struct jpeg_decompress_struct*);

*/
import "C"

import (
	"io"
	"unsafe"
)

const readBufferSize = 16384

type sourceManager struct {
	magic       uint32
	pub         C.struct_jpeg_source_mgr
	buffer      [readBufferSize]byte
	src         io.Reader
	startOfFile bool
	currentSize int
}

func getSourceManager(dinfo *C.struct_jpeg_decompress_struct) (ret *sourceManager) {
	// unsafe upcast magic to get the sourceManager associated with a dinfo
	ret = (*sourceManager)(unsafe.Pointer(uintptr(unsafe.Pointer(dinfo.src)) - unsafe.Offsetof(sourceManager{}.pub)))
	// just in case this ever breaks in a future release for some reason,
	// check the magic
	if ret.magic != magic {
		panic("Invalid sourceManager magic; upcast failed.")
	}
	return
}

//export sourceInit
func sourceInit(dinfo *C.struct_jpeg_decompress_struct) {
	mgr := getSourceManager(dinfo)
	mgr.startOfFile = true
}

//export sourceSkip
func sourceSkip(dinfo *C.struct_jpeg_decompress_struct, bytes C.long) {
	mgr := getSourceManager(dinfo)
	if bytes > 0 {
		for bytes >= C.long(mgr.pub.bytes_in_buffer) {
			bytes -= C.long(mgr.pub.bytes_in_buffer)
			sourceFill(dinfo)
		}
	}
	mgr.pub.bytes_in_buffer -= C.size_t(bytes)
	if mgr.pub.bytes_in_buffer != 0 {
		mgr.pub.next_input_byte = (*C.JOCTET)(&mgr.buffer[mgr.currentSize-int(mgr.pub.bytes_in_buffer)])
	}
}

//export sourceTerm
func sourceTerm(dinfo *C.struct_jpeg_decompress_struct) {
	// do nothing
}

//export sourceFill
func sourceFill(dinfo *C.struct_jpeg_decompress_struct) C.boolean {
	mgr := getSourceManager(dinfo)
	bytes, err := mgr.src.Read(mgr.buffer[:])
	mgr.pub.bytes_in_buffer = C.size_t(bytes)
	mgr.currentSize = bytes
	mgr.pub.next_input_byte = (*C.JOCTET)(&mgr.buffer[0])
	if err == io.EOF {
		if bytes == 0 {
			if mgr.startOfFile {
				panic("input is empty")
			}
			// EOF and need more data. Fill in a fake EOI to get a partial image.
			mgr.buffer[0] = 0xff
			mgr.buffer[1] = C.JPEG_EOI
			mgr.pub.bytes_in_buffer = 2
		}
	} else if err != nil {
		panic(err)
	}
	mgr.startOfFile = false

	return C.TRUE
}

func makeSourceManager(src io.Reader, dinfo *C.struct_jpeg_decompress_struct) (mgr sourceManager) {
	mgr.magic = magic
	mgr.src = src
	mgr.pub.init_source = (*[0]byte)(C.sourceInit)
	mgr.pub.fill_input_buffer = (*[0]byte)(C.sourceFill)
	mgr.pub.skip_input_data = (*[0]byte)(C.sourceSkip)
	mgr.pub.resync_to_restart = (*[0]byte)(C.jpeg_resync_to_restart) // default implementation
	mgr.pub.term_source = (*[0]byte)(C.sourceTerm)
	mgr.pub.bytes_in_buffer = 0
	mgr.pub.next_input_byte = nil
	dinfo.src = &mgr.pub
	return
}
<commit_msg>Fix #8 linker error in golang 1.5 on darwin.<commit_after>package jpeg

//
// Original codes are bollowed from go-thumber.
// Copyright (c) 2014 pixiv Inc. All rights reserved.
//
// See: https://github.com/pixiv/go-thumber
//

/*
#include <stdlib.h>
#include <stdio.h>
#include <jpeglib.h>

// exported from golang
void sourceInit(struct jpeg_decompress_struct*);
void sourceSkip(struct jpeg_decompress_struct*, long);
boolean sourceFill(struct jpeg_decompress_struct*);
void sourceTerm(struct jpeg_decompress_struct*);

static void* get_jpeg_resync_to_restart() {
	return jpeg_resync_to_restart;
}

*/
import "C"

import (
	"io"
	"unsafe"
)

const readBufferSize = 16384

type sourceManager struct {
	magic       uint32
	pub         C.struct_jpeg_source_mgr
	buffer      [readBufferSize]byte
	src         io.Reader
	startOfFile bool
	currentSize int
}

func getSourceManager(dinfo *C.struct_jpeg_decompress_struct) (ret *sourceManager) {
	// unsafe upcast magic to get the sourceManager associated with a dinfo
	ret = (*sourceManager)(unsafe.Pointer(uintptr(unsafe.Pointer(dinfo.src)) - unsafe.Offsetof(sourceManager{}.pub)))
	// just in case this ever breaks in a future release for some reason,
	// check the magic
	if ret.magic != magic {
		panic("Invalid sourceManager magic; upcast failed.")
	}
	return
}

//export sourceInit
func sourceInit(dinfo *C.struct_jpeg_decompress_struct) {
	mgr := getSourceManager(dinfo)
	mgr.startOfFile = true
}

//export sourceSkip
func sourceSkip(dinfo *C.struct_jpeg_decompress_struct, bytes C.long) {
	mgr := getSourceManager(dinfo)
	if bytes > 0 {
		for bytes >= C.long(mgr.pub.bytes_in_buffer) {
			bytes -= C.long(mgr.pub.bytes_in_buffer)
			sourceFill(dinfo)
		}
	}
	mgr.pub.bytes_in_buffer -= C.size_t(bytes)
	if mgr.pub.bytes_in_buffer != 0 {
		mgr.pub.next_input_byte = (*C.JOCTET)(&mgr.buffer[mgr.currentSize-int(mgr.pub.bytes_in_buffer)])
	}
}

//export sourceTerm
func sourceTerm(dinfo *C.struct_jpeg_decompress_struct) {
	// do nothing
}

//export sourceFill
func sourceFill(dinfo *C.struct_jpeg_decompress_struct) C.boolean {
	mgr := getSourceManager(dinfo)
	bytes, err := mgr.src.Read(mgr.buffer[:])
	mgr.pub.bytes_in_buffer = C.size_t(bytes)
	mgr.currentSize = bytes
	mgr.pub.next_input_byte = (*C.JOCTET)(&mgr.buffer[0])
	if err == io.EOF {
		if bytes == 0 {
			if mgr.startOfFile {
				panic("input is empty")
			}
			// EOF and need more data. Fill in a fake EOI to get a partial image.
			mgr.buffer[0] = 0xff
			mgr.buffer[1] = C.JPEG_EOI
			mgr.pub.bytes_in_buffer = 2
		}
	} else if err != nil {
		panic(err)
	}
	mgr.startOfFile = false

	return C.TRUE
}

func makeSourceManager(src io.Reader, dinfo *C.struct_jpeg_decompress_struct) (mgr sourceManager) {
	mgr.magic = magic
	mgr.src = src
	mgr.pub.init_source = (*[0]byte)(C.sourceInit)
	mgr.pub.fill_input_buffer = (*[0]byte)(C.sourceFill)
	mgr.pub.skip_input_data = (*[0]byte)(C.sourceSkip)
	mgr.pub.resync_to_restart = (*[0]byte)(C.get_jpeg_resync_to_restart()) // default implementation
	mgr.pub.term_source = (*[0]byte)(C.sourceTerm)
	mgr.pub.bytes_in_buffer = 0
	mgr.pub.next_input_byte = nil
	dinfo.src = &mgr.pub
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package chaos

import (
	"fmt"
	"math"
	"math/rand"
	"strings"
	"sync"
	"time"

	"k8s.io/perf-tests/clusterloader2/api"
	"k8s.io/perf-tests/clusterloader2/pkg/framework/client"
	"k8s.io/perf-tests/clusterloader2/pkg/util"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog"
)

const (
	monitoringNamespace = "monitoring"
	prometheusLabel     = "prometheus=k8s"
)

// NodeKiller is a utility to simulate node failures.
type NodeKiller struct {
	config   api.NodeFailureConfig
	client   clientset.Interface
	provider string
	// killedNodes stores names of the nodes that have been killed by NodeKiller.
	killedNodes sets.String
	recorder    *eventRecorder
}

type nodeAction string

const (
	stopServices nodeAction = "stopService"
	rebootNode              = "rebootNode"
)

type event struct {
	time     time.Time
	action   nodeAction
	nodeName string
}

type eventRecorder struct {
	events []event
	mux    sync.Mutex
}

func newEventRecorder() *eventRecorder {
	return &eventRecorder{[]event{}, sync.Mutex{}}
}

func (r *eventRecorder) record(a nodeAction, nodeName string) {
	e := event{time.Now(), a, nodeName}
	r.mux.Lock()
	r.events = append(r.events, e)
	r.mux.Unlock()
}

// NewNodeKiller creates new NodeKiller.
func NewNodeKiller(config api.NodeFailureConfig, client clientset.Interface, provider string) (*NodeKiller, error) {
	if provider != "gce" && provider != "gke" {
		return nil, fmt.Errorf("provider %q is not supported by NodeKiller", provider)
	}
	return &NodeKiller{config, client, provider, sets.NewString(), newEventRecorder()}, nil
}

// Run starts NodeKiller until stopCh is closed.
func (k *NodeKiller) Run(stopCh <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	// wait.JitterUntil starts work immediately, so wait first.
	sleepInterrupt(wait.Jitter(time.Duration(k.config.Interval), k.config.JitterFactor), stopCh)
	wait.JitterUntil(func() {
		nodes, err := k.pickNodes()
		if err != nil {
			klog.Errorf("%s: Unable to pick nodes to kill: %v", k, err)
			return
		}
		k.kill(nodes, stopCh)
	}, time.Duration(k.config.Interval), k.config.JitterFactor, true, stopCh)
}

func (k *NodeKiller) pickNodes() ([]v1.Node, error) {
	allNodes, err := util.GetSchedulableUntainedNodes(k.client)
	if err != nil {
		return nil, err
	}

	prometheusPods, err := client.ListPodsWithOptions(k.client, monitoringNamespace, metav1.ListOptions{
		LabelSelector: prometheusLabel,
	})
	if err != nil {
		return nil, err
	}
	nodesHasPrometheusPod := sets.NewString()
	for i := range prometheusPods {
		if prometheusPods[i].Spec.NodeName != "" {
			nodesHasPrometheusPod.Insert(prometheusPods[i].Spec.NodeName)
			klog.Infof("%s: Node %s removed from killing. Runs pod %s", k, prometheusPods[i].Spec.NodeName, prometheusPods[i].Name)
		}
	}

	nodes := allNodes[:0]
	for _, node := range allNodes {
		if !nodesHasPrometheusPod.Has(node.Name) && !k.killedNodes.Has(node.Name) {
			nodes = append(nodes, node)
		}
	}
	rand.Shuffle(len(nodes), func(i, j int) {
		nodes[i], nodes[j] = nodes[j], nodes[i]
	})
	numNodes := int(math.Ceil(k.config.FailureRate * float64(len(nodes))))
	klog.Infof("%s: %d nodes available, wants to fail %d nodes", k, len(nodes), numNodes)
	if len(nodes) > numNodes {
		nodes = nodes[:numNodes]
	}
	for _, node := range nodes {
		klog.Infof("%s: Node %q schedule for failure", k, node.Name)
	}
	return nodes, nil
}

func (k *NodeKiller) kill(nodes []v1.Node, stopCh <-chan struct{}) {
	wg := sync.WaitGroup{}
	wg.Add(len(nodes))
	for _, node := range nodes {
		k.killedNodes.Insert(node.Name)
		node := node
		go func() {
			defer wg.Done()

			klog.Infof("%s: Stopping docker and kubelet on %q to simulate failure", k, node.Name)
			k.addStopServicesEvent(node.Name)
			err := util.SSH("sudo systemctl stop docker kubelet", &node, nil)
			if err != nil {
				klog.Errorf("%s: ERROR while stopping node %q: %v", k, node.Name, err)
				return
			}

			// Listening for interruptions on stopCh or wait for the simulated downtime
			sleepInterrupt(time.Duration(k.config.SimulatedDowntime), stopCh)

			klog.Infof("%s: Rebooting %q to repair the node", k, node.Name)
			// Scheduling a reboot in one second, then disconnecting.
			//
			// Bash command explanation:
			// 'nohup' - Making sure that end of SSH connection signal will not break sudo
			// 'sudo' - Elevated priviliages, required by 'shutdown'
			// 'shutdown' - Control machine power
			// '-r' - Making 'shutdown' to reboot, instead of power-off
			// '+1s' - Parameter to 'reboot', to wait 1 second before rebooting.
			// '> /dev/null 2> /dev/null < /dev/null' - File descriptor redirect, all three I/O to avoid ssh hanging,
			//                                          see https://web.archive.org/web/20090429074212/http://www.openssh.com/faq.html#3.10
			// '&' - Execute command in background, end without waiting for result
			k.addRebootEvent(node.Name)
			err = util.SSH("nohup sudo shutdown -r +1s > /dev/null 2> /dev/null < /dev/null &", &node, nil)
			if err != nil {
				klog.Errorf("%s: Error while rebooting node %q: %v", k, node.Name, err)
				return
			}
		}()
	}
	wg.Wait()
}

func (k *NodeKiller) addStopServicesEvent(nodeName string) {
	k.recorder.record(stopServices, nodeName)
}

func (k *NodeKiller) addRebootEvent(nodeName string) {
	k.recorder.record(rebootNode, nodeName)
}

// Summary logs NodeKiller execution
func (k *NodeKiller) Summary() string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("%s: Recorded following events\n", k))
	for _, e := range k.recorder.events {
		sb.WriteString(fmt.Sprintf("%s: At %v %v happend for node %s\n", k, e.time.Format(time.UnixDate), e.action, e.nodeName))
	}
	return sb.String()
}

func (k *NodeKiller) String() string {
	return "NodeKiller"
}

// Utility method to put the current routine to sleep or break the sleep if stopCh closes
// Note of warning: if stopCh is already closed the process may not sleep at all.
func sleepInterrupt(duration time.Duration, stopCh <-chan struct{}) {
	select {
	case <-stopCh:
		break
	case <-time.After(duration):
		break
	}
}
<commit_msg>Revert "[NodeKiller] Schedule reboot and disconnect from node"<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package chaos

import (
	"fmt"
	"math"
	"math/rand"
	"strings"
	"sync"
	"time"

	"k8s.io/perf-tests/clusterloader2/api"
	"k8s.io/perf-tests/clusterloader2/pkg/framework/client"
	"k8s.io/perf-tests/clusterloader2/pkg/util"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/klog"
)

const (
	monitoringNamespace = "monitoring"
	prometheusLabel     = "prometheus=k8s"
)

// NodeKiller is a utility to simulate node failures.
type NodeKiller struct {
	config   api.NodeFailureConfig
	client   clientset.Interface
	provider string
	// killedNodes stores names of the nodes that have been killed by NodeKiller.
	killedNodes sets.String
	recorder    *eventRecorder
}

type nodeAction string

const (
	stopServices nodeAction = "stopService"
	rebootNode              = "rebootNode"
)

type event struct {
	time     time.Time
	action   nodeAction
	nodeName string
}

type eventRecorder struct {
	events []event
	mux    sync.Mutex
}

func newEventRecorder() *eventRecorder {
	return &eventRecorder{[]event{}, sync.Mutex{}}
}

func (r *eventRecorder) record(a nodeAction, nodeName string) {
	e := event{time.Now(), a, nodeName}
	r.mux.Lock()
	r.events = append(r.events, e)
	r.mux.Unlock()
}

// NewNodeKiller creates new NodeKiller.
func NewNodeKiller(config api.NodeFailureConfig, client clientset.Interface, provider string) (*NodeKiller, error) {
	if provider != "gce" && provider != "gke" {
		return nil, fmt.Errorf("provider %q is not supported by NodeKiller", provider)
	}
	return &NodeKiller{config, client, provider, sets.NewString(), newEventRecorder()}, nil
}

// Run starts NodeKiller until stopCh is closed.
func (k *NodeKiller) Run(stopCh <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	// wait.JitterUntil starts work immediately, so wait first.
	sleepInterrupt(wait.Jitter(time.Duration(k.config.Interval), k.config.JitterFactor), stopCh)
	wait.JitterUntil(func() {
		nodes, err := k.pickNodes()
		if err != nil {
			klog.Errorf("%s: Unable to pick nodes to kill: %v", k, err)
			return
		}
		k.kill(nodes, stopCh)
	}, time.Duration(k.config.Interval), k.config.JitterFactor, true, stopCh)
}

func (k *NodeKiller) pickNodes() ([]v1.Node, error) {
	allNodes, err := util.GetSchedulableUntainedNodes(k.client)
	if err != nil {
		return nil, err
	}

	prometheusPods, err := client.ListPodsWithOptions(k.client, monitoringNamespace, metav1.ListOptions{
		LabelSelector: prometheusLabel,
	})
	if err != nil {
		return nil, err
	}
	nodesHasPrometheusPod := sets.NewString()
	for i := range prometheusPods {
		if prometheusPods[i].Spec.NodeName != "" {
			nodesHasPrometheusPod.Insert(prometheusPods[i].Spec.NodeName)
			klog.Infof("%s: Node %s removed from killing. Runs pod %s", k, prometheusPods[i].Spec.NodeName, prometheusPods[i].Name)
		}
	}

	nodes := allNodes[:0]
	for _, node := range allNodes {
		if !nodesHasPrometheusPod.Has(node.Name) && !k.killedNodes.Has(node.Name) {
			nodes = append(nodes, node)
		}
	}
	rand.Shuffle(len(nodes), func(i, j int) {
		nodes[i], nodes[j] = nodes[j], nodes[i]
	})
	numNodes := int(math.Ceil(k.config.FailureRate * float64(len(nodes))))
	klog.Infof("%s: %d nodes available, wants to fail %d nodes", k, len(nodes), numNodes)
	if len(nodes) > numNodes {
		nodes = nodes[:numNodes]
	}
	for _, node := range nodes {
		klog.Infof("%s: Node %q schedule for failure", k, node.Name)
	}
	return nodes, nil
}

func (k *NodeKiller) kill(nodes []v1.Node, stopCh <-chan struct{}) {
	wg := sync.WaitGroup{}
	wg.Add(len(nodes))
	for _, node := range nodes {
		k.killedNodes.Insert(node.Name)
		node := node
		go func() {
			defer wg.Done()

			klog.Infof("%s: Stopping docker and kubelet on %q to simulate failure", k, node.Name)
			k.addStopServicesEvent(node.Name)
			err := util.SSH("sudo systemctl stop docker kubelet", &node, nil)
			if err != nil {
				klog.Errorf("%s: ERROR while stopping node %q: %v", k, node.Name, err)
				return
			}

			// Listening for interruptions on stopCh or wait for the simulated downtime
			sleepInterrupt(time.Duration(k.config.SimulatedDowntime), stopCh)

			klog.Infof("%s: Rebooting %q to repair the node", k, node.Name)
			k.addRebootEvent(node.Name)
			err = util.SSH("sudo reboot", &node, nil)
			if err != nil {
				klog.Errorf("%s: Error while rebooting node %q: %v", k, node.Name, err)
				return
			}
		}()
	}
	wg.Wait()
}

func (k *NodeKiller) addStopServicesEvent(nodeName string) {
	k.recorder.record(stopServices, nodeName)
}

func (k *NodeKiller) addRebootEvent(nodeName string) {
	k.recorder.record(rebootNode, nodeName)
}

// Summary logs NodeKiller execution
func (k *NodeKiller) Summary() string {
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("%s: Recorded following events\n", k))
	for _, e := range k.recorder.events {
		sb.WriteString(fmt.Sprintf("%s: At %v %v happend for node %s\n", k, e.time.Format(time.UnixDate), e.action, e.nodeName))
	}
	return sb.String()
}

func (k *NodeKiller) String() string {
	return "NodeKiller"
}

// Utility method to put the current routine to sleep or break the sleep if stopCh closes
// Note of warning: if stopCh is already closed the process may not sleep at all.
func sleepInterrupt(duration time.Duration, stopCh <-chan struct{}) {
	select {
	case <-stopCh:
		break
	case <-time.After(duration):
		break
	}
}
<|endoftext|>


<commit_before>package weed_server

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net/http"
	"strconv"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (ms *MasterServer) collectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	collection, ok := ms.Topo.FindCollection(r.FormValue("collection"))
	if !ok {
		writeJsonError(w, r, http.StatusBadRequest, fmt.Errorf("collection %s does not exist", r.FormValue("collection")))
		return
	}
	for _, server := range collection.ListVolumeServers() {
		err := operation.WithVolumeServerClient(server.Url(), ms.grpcDialOpiton, func(client volume_server_pb.VolumeServerClient) error {
			_, deleteErr := client.DeleteCollection(context.Background(), &volume_server_pb.DeleteCollectionRequest{
				Collection: collection.Name,
			})
			return deleteErr
		})
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
			return
		}
	}
	ms.Topo.DeleteCollection(r.FormValue("collection"))
}

func (ms *MasterServer) dirStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Topology"] = ms.Topo.ToMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) volumeVacuumHandler(w http.ResponseWriter, r *http.Request) {
	gcString := r.FormValue("garbageThreshold")
	gcThreshold := ms.option.GarbageThreshold
	if gcString != "" {
		var err error
		gcThreshold, err = strconv.ParseFloat(gcString, 32)
		if err != nil {
			glog.V(0).Infof("garbageThreshold %s is not a valid float number: %v", gcString, err)
			return
		}
	}
	glog.Infoln("garbageThreshold =", gcThreshold)
	ms.Topo.Vacuum(ms.grpcDialOpiton, gcThreshold, ms.preallocateSize)
	ms.dirStatusHandler(w, r)
}

func (ms *MasterServer) volumeGrowHandler(w http.ResponseWriter, r *http.Request) {
	count := 0
	option, err := ms.getVolumeGrowOption(r)
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
		return
	}
	if err == nil {
		if count, err = strconv.Atoi(r.FormValue("count")); err == nil {
			if ms.Topo.FreeSpace() < int64(count*option.ReplicaPlacement.GetCopyCount()) {
				err = fmt.Errorf("only %d volumes left, not enough for %d", ms.Topo.FreeSpace(), count*option.ReplicaPlacement.GetCopyCount())
			} else {
				count, err = ms.vg.GrowByCountAndType(ms.grpcDialOpiton, count, option, ms.Topo)
			}
		} else {
			err = errors.New("parameter count is not found")
		}
	}
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
	} else {
		writeJsonQuiet(w, r, http.StatusOK, map[string]interface{}{"count": count})
	}
}

func (ms *MasterServer) volumeStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Volumes"] = ms.Topo.ToVolumeMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) redirectHandler(w http.ResponseWriter, r *http.Request) {
	vid, _, _, _, _ := parseURLPath(r.URL.Path)
	collection := r.FormValue("collection")
	location := ms.findVolumeLocation(collection, vid)
	if location.Error == "" {
		loc := location.Locations[rand.Intn(len(location.Locations))]
		var url string
		if r.URL.RawQuery != "" {
			url = util.NormalizeUrl(loc.PublicUrl) + r.URL.Path + "?" + r.URL.RawQuery
		} else {
			url = util.NormalizeUrl(loc.PublicUrl) + r.URL.Path
		}
		http.Redirect(w, r, url, http.StatusMovedPermanently)
	} else {
		writeJsonError(w, r, http.StatusNotFound, fmt.Errorf("volume id %s not found: %s", vid, location.Error))
	}
}

func (ms *MasterServer) selfUrl(r *http.Request) string {
	if r.Host != "" {
		return r.Host
	}
	return "localhost:" + strconv.Itoa(ms.option.Port)
}
func (ms *MasterServer) submitFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		submitForClientHandler(w, r, ms.selfUrl(r), ms.grpcDialOpiton)
	} else {
		masterUrl, err := ms.Topo.Leader()
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
		} else {
			submitForClientHandler(w, r, masterUrl, ms.grpcDialOpiton)
		}
	}
}

func (ms *MasterServer) HasWritableVolume(option *topology.VolumeGrowOption) bool {
	vl := ms.Topo.GetVolumeLayout(option.Collection, option.ReplicaPlacement, option.Ttl)
	return vl.GetActiveVolumeCount(option) > 0
}

func (ms *MasterServer) getVolumeGrowOption(r *http.Request) (*topology.VolumeGrowOption, error) {
	replicationString := r.FormValue("replication")
	if replicationString == "" {
		replicationString = ms.option.DefaultReplicaPlacement
	}
	replicaPlacement, err := storage.NewReplicaPlacementFromString(replicationString)
	if err != nil {
		return nil, err
	}
	ttl, err := needle.ReadTTL(r.FormValue("ttl"))
	if err != nil {
		return nil, err
	}
	memoryMapMaxSizeMB, err := needle.ReadMemoryMapMaxSizeMB(r.FormValue("memorymapmaxsizemb"))
	if err != nil {
		return nil, err
	}

	preallocate := ms.preallocateSize
	if r.FormValue("preallocate") != "" {
		preallocate, err = strconv.ParseInt(r.FormValue("preallocate"), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse int64 preallocate = %s: %v", r.FormValue("preallocate"), err)
		}
	}
	volumeGrowOption := &topology.VolumeGrowOption{
		Collection:         r.FormValue("collection"),
		ReplicaPlacement:   replicaPlacement,
		Ttl:                ttl,
		Prealloacte:        preallocate,
		DataCenter:         r.FormValue("dataCenter"),
		Rack:               r.FormValue("rack"),
		DataNode:           r.FormValue("dataNode"),
		MemoryMapMaxSizeMB: memoryMapMaxSizeMB,
	}
	return volumeGrowOption, nil
}
<commit_msg>fix unreachable code<commit_after>package weed_server

import (
	"context"
	"fmt"
	"math/rand"
	"net/http"
	"strconv"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (ms *MasterServer) collectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	collection, ok := ms.Topo.FindCollection(r.FormValue("collection"))
	if !ok {
		writeJsonError(w, r, http.StatusBadRequest, fmt.Errorf("collection %s does not exist", r.FormValue("collection")))
		return
	}
	for _, server := range collection.ListVolumeServers() {
		err := operation.WithVolumeServerClient(server.Url(), ms.grpcDialOpiton, func(client volume_server_pb.VolumeServerClient) error {
			_, deleteErr := client.DeleteCollection(context.Background(), &volume_server_pb.DeleteCollectionRequest{
				Collection: collection.Name,
			})
			return deleteErr
		})
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
			return
		}
	}
	ms.Topo.DeleteCollection(r.FormValue("collection"))
}

func (ms *MasterServer) dirStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Topology"] = ms.Topo.ToMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) volumeVacuumHandler(w http.ResponseWriter, r *http.Request) {
	gcString := r.FormValue("garbageThreshold")
	gcThreshold := ms.option.GarbageThreshold
	if gcString != "" {
		var err error
		gcThreshold, err = strconv.ParseFloat(gcString, 32)
		if err != nil {
			glog.V(0).Infof("garbageThreshold %s is not a valid float number: %v", gcString, err)
			return
		}
	}
	glog.Infoln("garbageThreshold =", gcThreshold)
	ms.Topo.Vacuum(ms.grpcDialOpiton, gcThreshold, ms.preallocateSize)
	ms.dirStatusHandler(w, r)
}

func (ms *MasterServer) volumeGrowHandler(w http.ResponseWriter, r *http.Request) {
	count := 0
	option, err := ms.getVolumeGrowOption(r)
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
		return
	}

	if count, err = strconv.Atoi(r.FormValue("count")); err == nil {
		if ms.Topo.FreeSpace() < int64(count*option.ReplicaPlacement.GetCopyCount()) {
			err = fmt.Errorf("only %d volumes left, not enough for %d", ms.Topo.FreeSpace(), count*option.ReplicaPlacement.GetCopyCount())
		} else {
			count, err = ms.vg.GrowByCountAndType(ms.grpcDialOpiton, count, option, ms.Topo)
		}
	} else {
		err = fmt.Errorf("can not parse parameter count %s", r.FormValue("count"))
	}

	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
	} else {
		writeJsonQuiet(w, r, http.StatusOK, map[string]interface{}{"count": count})
	}
}

func (ms *MasterServer) volumeStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Volumes"] = ms.Topo.ToVolumeMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) redirectHandler(w http.ResponseWriter, r *http.Request) {
	vid, _, _, _, _ := parseURLPath(r.URL.Path)
	collection := r.FormValue("collection")
	location := ms.findVolumeLocation(collection, vid)
	if location.Error == "" {
		loc := location.Locations[rand.Intn(len(location.Locations))]
		var url string
		if r.URL.RawQuery != "" {
			url = util.NormalizeUrl(loc.PublicUrl) + r.URL.Path + "?" + r.URL.RawQuery
		} else {
			url = util.NormalizeUrl(loc.PublicUrl) + r.URL.Path
		}
		http.Redirect(w, r, url, http.StatusMovedPermanently)
	} else {
		writeJsonError(w, r, http.StatusNotFound, fmt.Errorf("volume id %s not found: %s", vid, location.Error))
	}
}

func (ms *MasterServer) selfUrl(r *http.Request) string {
	if r.Host != "" {
		return r.Host
	}
	return "localhost:" + strconv.Itoa(ms.option.Port)
}
func (ms *MasterServer) submitFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		submitForClientHandler(w, r, ms.selfUrl(r), ms.grpcDialOpiton)
	} else {
		masterUrl, err := ms.Topo.Leader()
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
		} else {
			submitForClientHandler(w, r, masterUrl, ms.grpcDialOpiton)
		}
	}
}

func (ms *MasterServer) HasWritableVolume(option *topology.VolumeGrowOption) bool {
	vl := ms.Topo.GetVolumeLayout(option.Collection, option.ReplicaPlacement, option.Ttl)
	return vl.GetActiveVolumeCount(option) > 0
}

func (ms *MasterServer) getVolumeGrowOption(r *http.Request) (*topology.VolumeGrowOption, error) {
	replicationString := r.FormValue("replication")
	if replicationString == "" {
		replicationString = ms.option.DefaultReplicaPlacement
	}
	replicaPlacement, err := storage.NewReplicaPlacementFromString(replicationString)
	if err != nil {
		return nil, err
	}
	ttl, err := needle.ReadTTL(r.FormValue("ttl"))
	if err != nil {
		return nil, err
	}
	memoryMapMaxSizeMB, err := needle.ReadMemoryMapMaxSizeMB(r.FormValue("memorymapmaxsizemb"))
	if err != nil {
		return nil, err
	}

	preallocate := ms.preallocateSize
	if r.FormValue("preallocate") != "" {
		preallocate, err = strconv.ParseInt(r.FormValue("preallocate"), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse int64 preallocate = %s: %v", r.FormValue("preallocate"), err)
		}
	}
	volumeGrowOption := &topology.VolumeGrowOption{
		Collection:         r.FormValue("collection"),
		ReplicaPlacement:   replicaPlacement,
		Ttl:                ttl,
		Prealloacte:        preallocate,
		DataCenter:         r.FormValue("dataCenter"),
		Rack:               r.FormValue("rack"),
		DataNode:           r.FormValue("dataNode"),
		MemoryMapMaxSizeMB: memoryMapMaxSizeMB,
	}
	return volumeGrowOption, nil
}
<|endoftext|>


<commit_before>package server

import (
	"admin"
	"api/http"
	"cluster"
	log "code.google.com/p/log4go"
	"configuration"
	"coordinator"
	"datastore"
	"wal"
)

type Server struct {
	RaftServer     *coordinator.RaftServer
	ProtobufServer *coordinator.ProtobufServer
	ClusterConfig  *cluster.ClusterConfiguration
	HttpApi        *http.HttpServer
	AdminServer    *admin.HttpServer
	Coordinator    coordinator.Coordinator
	Config         *configuration.Configuration
	RequestHandler *coordinator.ProtobufRequestHandler
	stopped        bool
	writeLog       *wal.WAL
	shardStore     *datastore.LevelDbShardDatastore
}

func NewServer(config *configuration.Configuration) (*Server, error) {
	log.Info("Opening database at %s", config.DataDir)
	shardDb, err := datastore.NewLevelDbShardDatastore(config)
	if err != nil {
		return nil, err
	}

	newClient := func(connectString string) cluster.ServerConnection {
		return coordinator.NewProtobufClient(connectString, config.ProtobufTimeout.Duration)
	}
	writeLog, err := wal.NewWAL(config)
	if err != nil {
		return nil, err
	}

	clusterConfig := cluster.NewClusterConfiguration(config, writeLog, shardDb, newClient)
	raftServer := coordinator.NewRaftServer(config, clusterConfig)
	clusterConfig.LocalRaftName = raftServer.GetRaftName()
	clusterConfig.SetShardCreator(raftServer)
	clusterConfig.CreateFutureShardsAutomaticallyBeforeTimeComes()

	coord := coordinator.NewCoordinatorImpl(config, raftServer, clusterConfig)
	requestHandler := coordinator.NewProtobufRequestHandler(coord, clusterConfig)
	protobufServer := coordinator.NewProtobufServer(config.ProtobufPortString(), requestHandler)

	raftServer.AssignCoordinator(coord)
	httpApi := http.NewHttpServer(config.ApiHttpPortString(), config.AdminAssetsDir, coord, coord, clusterConfig, raftServer)
	httpApi.EnableSsl(config.ApiHttpSslPortString(), config.ApiHttpCertPath)
	adminServer := admin.NewHttpServer(config.AdminAssetsDir, config.AdminHttpPortString())

	return &Server{
		RaftServer:     raftServer,
		ProtobufServer: protobufServer,
		ClusterConfig:  clusterConfig,
		HttpApi:        httpApi,
		Coordinator:    coord,
		AdminServer:    adminServer,
		Config:         config,
		RequestHandler: requestHandler,
		writeLog:       writeLog,
		shardStore:     shardDb}, nil
}

func (self *Server) ListenAndServe() error {
	go self.ProtobufServer.ListenAndServe()

	err := self.RaftServer.ListenAndServe()
	if err != nil {
		return err
	}

	log.Info("Waiting for local server to be added")
	self.ClusterConfig.WaitForLocalServerLoaded()
	self.writeLog.SetServerId(self.ClusterConfig.ServerId())

	log.Info("Recovering from log...")
	err = self.ClusterConfig.RecoverFromWAL()
	if err != nil {
		return err
	}
	log.Info("recovered")

	err = self.Coordinator.(*coordinator.CoordinatorImpl).ConnectToProtobufServers(self.Config.ProtobufConnectionString())
	if err != nil {
		return err
	}
	log.Info("Starting admin interface on port %d", self.Config.AdminHttpPort)
	go self.AdminServer.ListenAndServe()
	log.Info("Starting Http Api server on port %d", self.Config.ApiHttpPort)
	self.HttpApi.ListenAndServe()
	return nil
}

func (self *Server) Stop() {
	if self.stopped {
		return
	}
	self.stopped = true
	self.RaftServer.Close()
	self.HttpApi.Close()
	self.ProtobufServer.Close()
	self.AdminServer.Close()
	self.writeLog.Close()
	self.shardStore.Close()
	// TODO: close admin server and protobuf client connections
	log.Info("Stopping server")
}
<commit_msg>graceful shutdown<commit_after>package server

import (
	"admin"
	"api/http"
	"cluster"
	log "code.google.com/p/log4go"
	"configuration"
	"coordinator"
	"datastore"
	"os"
	"os/signal"
	"syscall"
	"time"
	"wal"
)

type Server struct {
	RaftServer     *coordinator.RaftServer
	ProtobufServer *coordinator.ProtobufServer
	ClusterConfig  *cluster.ClusterConfiguration
	HttpApi        *http.HttpServer
	AdminServer    *admin.HttpServer
	Coordinator    coordinator.Coordinator
	Config         *configuration.Configuration
	RequestHandler *coordinator.ProtobufRequestHandler
	stopped        bool
	writeLog       *wal.WAL
	shardStore     *datastore.LevelDbShardDatastore
}

func NewServer(config *configuration.Configuration) (*Server, error) {
	log.Info("Opening database at %s", config.DataDir)
	shardDb, err := datastore.NewLevelDbShardDatastore(config)
	if err != nil {
		return nil, err
	}

	newClient := func(connectString string) cluster.ServerConnection {
		return coordinator.NewProtobufClient(connectString, config.ProtobufTimeout.Duration)
	}
	writeLog, err := wal.NewWAL(config)
	if err != nil {
		return nil, err
	}

	clusterConfig := cluster.NewClusterConfiguration(config, writeLog, shardDb, newClient)
	raftServer := coordinator.NewRaftServer(config, clusterConfig)
	clusterConfig.LocalRaftName = raftServer.GetRaftName()
	clusterConfig.SetShardCreator(raftServer)
	clusterConfig.CreateFutureShardsAutomaticallyBeforeTimeComes()

	coord := coordinator.NewCoordinatorImpl(config, raftServer, clusterConfig)
	requestHandler := coordinator.NewProtobufRequestHandler(coord, clusterConfig)
	protobufServer := coordinator.NewProtobufServer(config.ProtobufPortString(), requestHandler)

	raftServer.AssignCoordinator(coord)
	httpApi := http.NewHttpServer(config.ApiHttpPortString(), config.AdminAssetsDir, coord, coord, clusterConfig, raftServer)
	httpApi.EnableSsl(config.ApiHttpSslPortString(), config.ApiHttpCertPath)
	adminServer := admin.NewHttpServer(config.AdminAssetsDir, config.AdminHttpPortString())

	return &Server{
		RaftServer:     raftServer,
		ProtobufServer: protobufServer,
		ClusterConfig:  clusterConfig,
		HttpApi:        httpApi,
		Coordinator:    coord,
		AdminServer:    adminServer,
		Config:         config,
		RequestHandler: requestHandler,
		writeLog:       writeLog,
		shardStore:     shardDb}, nil
}

func (self *Server) ListenAndServe() error {
	go self.ProtobufServer.ListenAndServe()

	err := self.RaftServer.ListenAndServe()
	if err != nil {
		return err
	}

	log.Info("Waiting for local server to be added")
	self.ClusterConfig.WaitForLocalServerLoaded()
	self.writeLog.SetServerId(self.ClusterConfig.ServerId())

	log.Info("Recovering from log...")
	err = self.ClusterConfig.RecoverFromWAL()
	if err != nil {
		return err
	}
	log.Info("recovered")

	err = self.Coordinator.(*coordinator.CoordinatorImpl).ConnectToProtobufServers(self.Config.ProtobufConnectionString())
	if err != nil {
		return err
	}
	go self.ListenForSignals()
	log.Info("Starting admin interface on port %d", self.Config.AdminHttpPort)
	go self.AdminServer.ListenAndServe()
	log.Info("Starting Http Api server on port %d", self.Config.ApiHttpPort)
	self.HttpApi.ListenAndServe()
	return nil
}

func (self *Server) Stop() {
	if self.stopped {
		return
	}
	log.Info("Stopping server")
	self.stopped = true
	self.RaftServer.Close()
	self.HttpApi.Close()
	self.ProtobufServer.Close()
	self.AdminServer.Close()
	self.writeLog.Close()
	self.shardStore.Close()
}

func (self *Server) ListenForSignals() {
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT)
	for {
		sig := <-ch
		log.Info("Received signal: %s\n", sig.String())
		switch sig {
		case syscall.SIGINT, syscall.SIGTERM:
			self.Stop()
			time.Sleep(time.Second)
			os.Exit(0)
		}
	}
}
<|endoftext|>


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

package command

import (
	"context"
	"fmt"
	"net/url"
	"strings"
	"time"

	"cloud.google.com/go/datastore"

	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/go/rtx"

	"github.com/spf13/cobra"
)

// createCmd represents the create command
var createCmd = &cobra.Command{
	Use:   "create",
	Short: "Adds a new ePoxy Host record to Datastore",
	Long: `
USAGE:
	**ONLY FOR TESTING**

    Creates a new datastore Host record for ePoxy server. Calling "create" on
    an existing host will overwrite the original.

EXAMPLE:

    # Use the default boot and update stage URLs:
    epoxy_admin create --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35
`,
	Run: runCreate,
}

// fmtURL formats (if needed) and parses the given string as a URL. If the
// resulting URL is invalid, fmtURL panics.
func fmtURL(urlStr string) string {
	if strings.Contains(urlStr, "%s") {
		urlStr = fmt.Sprintf(urlStr, fProject)
	}
	_, err := url.Parse(urlStr)
	rtx.Must(err, "Failed to parse URL: %s", urlStr)
	return urlStr
}

// TODO: add unit tests by masking out NewClient & NewDatstoreConfig. Consider
// promoting the fake datastore types from storage/datastore_test.go to an
// internal fake package.
func runCreate(cmd *cobra.Command, args []string) {
	// Setup Datastore client.
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()

	client, err := datastore.NewClient(ctx, fProject)
	rtx.Must(err, "Failed to create new datastore client")

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)

	h := &storage.Host{
		Name:          cfHostname,
		IPv4Addr:      cfAddress,
		UpdateEnabled: cfUpdate,
		Extensions:    []string{cfExtension},
		Boot: storage.Sequence{
			Stage1ChainURL: fmtURL(cfBootStage1),
			Stage2ChainURL: fmtURL(cfBootStage2),
			Stage3ChainURL: fmtURL(cfBootStage3),
		},
		Update: storage.Sequence{
			Stage1ChainURL: fmtURL(cfUpdateStage1),
			Stage2ChainURL: fmtURL(cfUpdateStage2),
			Stage3ChainURL: fmtURL(cfUpdateStage3),
		},
	}

	// Save the host record.
	err = ds.Save(h)
	rtx.Must(err, "Failed to save new host record")

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h, err = ds.Load(h.Name)
	rtx.Must(err, "Failed to save new host record")
	fmt.Println(h.String())
}

func init() {
	rootCmd.AddCommand(createCmd)

	// Required local flags.
	createCmd.Flags().StringVar(&cfHostname, "hostname", "",
		"Hostname of new record.")
	createCmd.Flags().StringVar(&cfAddress, "address", "",
		"IP address of hostname.")
	createCmd.MarkFlagRequired("hostname")
	createCmd.MarkFlagRequired("address")

	// Local flags which will only apply when "create" is called directly.
	createCmd.Flags().StringVar(&cfExtension, "extension", "allocate_k8s_token",
		"Name of an extension to enable for host.")
	createCmd.Flags().BoolVar(&cfUpdate, "update", false,
		"Set Host.UpdateEnabled to true for an existing Host.")
	createCmd.Flags().StringVar(&cfBootStage1, "boot-stage1",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to stage2 boot.")
	createCmd.Flags().StringVar(&cfBootStage2, "boot-stage2",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to stage3 boot.")
	createCmd.Flags().StringVar(&cfBootStage3, "boot-stage3",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage3post.json",
		"Absolute URL to an action definition to run after running stage3 boot.")
	createCmd.Flags().StringVar(&cfUpdateStage1, "update-stage1",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to stage2 update.")
	createCmd.Flags().StringVar(&cfUpdateStage2, "update-stage2",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to stage3 update.")
	createCmd.Flags().StringVar(&cfUpdateStage3, "update-stage3",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage3post.json",
		"Absolute URL to an action definition to run after running stage3 update.")
}
<commit_msg>Create non-nil CollectedInformation when creating a new Host<commit_after>// Copyright 2018 ePoxy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"context"
	"fmt"
	"net/url"
	"strings"
	"time"

	"cloud.google.com/go/datastore"

	"github.com/m-lab/epoxy/datastorex"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/go/rtx"

	"github.com/spf13/cobra"
)

// createCmd represents the create command
var createCmd = &cobra.Command{
	Use:   "create",
	Short: "Adds a new ePoxy Host record to Datastore",
	Long: `
USAGE:
	**ONLY FOR TESTING**

    Creates a new datastore Host record for ePoxy server. Calling "create" on
    an existing host will overwrite the original.

EXAMPLE:

    # Use the default boot and update stage URLs:
    epoxy_admin create --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35
`,
	Run: runCreate,
}

// fmtURL formats (if needed) and parses the given string as a URL. If the
// resulting URL is invalid, fmtURL panics.
func fmtURL(urlStr string) string {
	if strings.Contains(urlStr, "%s") {
		urlStr = fmt.Sprintf(urlStr, fProject)
	}
	_, err := url.Parse(urlStr)
	rtx.Must(err, "Failed to parse URL: %s", urlStr)
	return urlStr
}

// TODO: add unit tests by masking out NewClient & NewDatstoreConfig. Consider
// promoting the fake datastore types from storage/datastore_test.go to an
// internal fake package.
func runCreate(cmd *cobra.Command, args []string) {
	// Setup Datastore client.
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()

	client, err := datastore.NewClient(ctx, fProject)
	rtx.Must(err, "Failed to create new datastore client")

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)

	h := &storage.Host{
		Name:          cfHostname,
		IPv4Addr:      cfAddress,
		UpdateEnabled: cfUpdate,
		Extensions:    []string{cfExtension},
		Boot: storage.Sequence{
			Stage1ChainURL: fmtURL(cfBootStage1),
			Stage2ChainURL: fmtURL(cfBootStage2),
			Stage3ChainURL: fmtURL(cfBootStage3),
		},
		Update: storage.Sequence{
			Stage1ChainURL: fmtURL(cfUpdateStage1),
			Stage2ChainURL: fmtURL(cfUpdateStage2),
			Stage3ChainURL: fmtURL(cfUpdateStage3),
		},
		CollectedInformation: datastorex.Map{},
	}

	// Save the host record.
	err = ds.Save(h)
	rtx.Must(err, "Failed to save new host record")

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h, err = ds.Load(h.Name)
	rtx.Must(err, "Failed to save new host record")
	fmt.Println(h.String())
}

func init() {
	rootCmd.AddCommand(createCmd)

	// Required local flags.
	createCmd.Flags().StringVar(&cfHostname, "hostname", "",
		"Hostname of new record.")
	createCmd.Flags().StringVar(&cfAddress, "address", "",
		"IP address of hostname.")
	createCmd.MarkFlagRequired("hostname")
	createCmd.MarkFlagRequired("address")

	// Local flags which will only apply when "create" is called directly.
	createCmd.Flags().StringVar(&cfExtension, "extension", "allocate_k8s_token",
		"Name of an extension to enable for host.")
	createCmd.Flags().BoolVar(&cfUpdate, "update", false,
		"Set Host.UpdateEnabled to true for an existing Host.")
	createCmd.Flags().StringVar(&cfBootStage1, "boot-stage1",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to stage2 boot.")
	createCmd.Flags().StringVar(&cfBootStage2, "boot-stage2",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to stage3 boot.")
	createCmd.Flags().StringVar(&cfBootStage3, "boot-stage3",
		"https://storage.googleapis.com/epoxy-%s/stage3_coreos/stage3post.json",
		"Absolute URL to an action definition to run after running stage3 boot.")
	createCmd.Flags().StringVar(&cfUpdateStage1, "update-stage1",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to stage2 update.")
	createCmd.Flags().StringVar(&cfUpdateStage2, "update-stage2",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to stage3 update.")
	createCmd.Flags().StringVar(&cfUpdateStage3, "update-stage3",
		"https://storage.googleapis.com/epoxy-%s/stage3_mlxupdate/stage3post.json",
		"Absolute URL to an action definition to run after running stage3 update.")
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/policy"
)

type workloadListOpts struct {
	*rootOpts
	namespace     string
	allNamespaces bool
}

func newWorkloadList(parent *rootOpts) *workloadListOpts {
	return &workloadListOpts{rootOpts: parent}
}

func (opts *workloadListOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:     "list-workloads",
		Aliases: []string{"list-controllers"}, // Transient backwards compatibility after replacing controller by workload
		Short:   "List workloads currently running in the cluster.",
		Example: makeExample("fluxctl list-workloads"),
		RunE:    opts.RunE,
	}
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Confine query to namespace")
	cmd.Flags().BoolVarP(&opts.allNamespaces, "all-namespaces", "a", false, "Query across all namespaces")
	return cmd
}

func (opts *workloadListOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if opts.allNamespaces {
		opts.namespace = ""
	}

	ctx := context.Background()

	workloads, err := opts.API.ListServices(ctx, opts.namespace)
	if err != nil {
		return err
	}

	sort.Sort(workloadStatusByName(workloads))

	w := newTabwriter()
	fmt.Fprintf(w, "WORKLOAD\tCONTAINER\tIMAGE\tRELEASE\tPOLICY\n")
	for _, workload := range workloads {
		if len(workload.Containers) > 0 {
			c := workload.Containers[0]
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", workload.ID, c.Name, c.Current.ID, workload.Status, policies(workload))
			for _, c := range workload.Containers[1:] {
				fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
			}
		} else {
			fmt.Fprintf(w, "%s\t\t\t\t\n", workload.ID)
		}
	}
	w.Flush()
	return nil
}

type workloadStatusByName []v6.ControllerStatus

func (s workloadStatusByName) Len() int {
	return len(s)
}

func (s workloadStatusByName) Less(a, b int) bool {
	return s[a].ID.String() < s[b].ID.String()
}

func (s workloadStatusByName) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func policies(s v6.ControllerStatus) string {
	var ps []string
	if s.Automated {
		ps = append(ps, string(policy.Automated))
	}
	if s.Locked {
		ps = append(ps, string(policy.Locked))
	}
	if s.Ignore {
		ps = append(ps, string(policy.Ignore))
	}
	sort.Strings(ps)
	return strings.Join(ps, ",")
}
<commit_msg>Always list the status of a workload in fluxctl<commit_after>package main

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/policy"
)

type workloadListOpts struct {
	*rootOpts
	namespace     string
	allNamespaces bool
}

func newWorkloadList(parent *rootOpts) *workloadListOpts {
	return &workloadListOpts{rootOpts: parent}
}

func (opts *workloadListOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:     "list-workloads",
		Aliases: []string{"list-controllers"}, // Transient backwards compatibility after replacing controller by workload
		Short:   "List workloads currently running in the cluster.",
		Example: makeExample("fluxctl list-workloads"),
		RunE:    opts.RunE,
	}
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Confine query to namespace")
	cmd.Flags().BoolVarP(&opts.allNamespaces, "all-namespaces", "a", false, "Query across all namespaces")
	return cmd
}

func (opts *workloadListOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if opts.allNamespaces {
		opts.namespace = ""
	}

	ctx := context.Background()

	workloads, err := opts.API.ListServices(ctx, opts.namespace)
	if err != nil {
		return err
	}

	sort.Sort(workloadStatusByName(workloads))

	w := newTabwriter()
	fmt.Fprintf(w, "WORKLOAD\tCONTAINER\tIMAGE\tRELEASE\tPOLICY\n")
	for _, workload := range workloads {
		if len(workload.Containers) > 0 {
			c := workload.Containers[0]
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", workload.ID, c.Name, c.Current.ID, workload.Status, policies(workload))
			for _, c := range workload.Containers[1:] {
				fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
			}
		} else {
			fmt.Fprintf(w, "%s\t\t\t%s\t%s\n", workload.ID, workload.Status, policies(workload))
		}
	}
	w.Flush()
	return nil
}

type workloadStatusByName []v6.ControllerStatus

func (s workloadStatusByName) Len() int {
	return len(s)
}

func (s workloadStatusByName) Less(a, b int) bool {
	return s[a].ID.String() < s[b].ID.String()
}

func (s workloadStatusByName) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func policies(s v6.ControllerStatus) string {
	var ps []string
	if s.Automated {
		ps = append(ps, string(policy.Automated))
	}
	if s.Locked {
		ps = append(ps, string(policy.Locked))
	}
	if s.Ignore {
		ps = append(ps, string(policy.Ignore))
	}
	sort.Strings(ps)
	return strings.Join(ps, ",")
}
<|endoftext|>