<commit_before>// Package grace allows for gracefully waiting for a listener to
// finish serving it's active requests.
package grace

import (
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
)

var (
	// This error is returned by Inherits() when we're not inheriting any fds.
	ErrNotInheriting = errors.New("no inherited listeners")

	// This error is returned by Listener.Accept() when Close is in progress.
	ErrAlreadyClosed = errors.New("already closed")
)

const (
	// Used to indicate a graceful restart in the new process.
	envCountKey = "LISTEN_FDS"

	// The error returned by the standard library when the socket is closed.
	errClosed = "use of closed network connection"

	// Used for the counter chan.
	inc = true
	dec = false
)

// A Listener providing a graceful Close process and can be sent
// across processes using the underlying File descriptor.
type Listener interface {
	net.Listener

	// Will indicate that a Close is requested preventing further Accept. It will
	// also wait for the active connections to be terminated before returning.
	// Note, this won't actually do the close, and is provided as part of the
	// public API for cases where the socket must not be closed (such as systemd
	// activation).
	CloseRequest()

	// Will return the underlying file representing this Listener.
	File() (f *os.File, err error)
}

// A FileListener is a file backed net.Listener.
type FileListener interface {
	net.Listener

	// Will return the underlying file representing this Listener.
	File() (f *os.File, err error)
}

// A goroutine based counter that provides graceful Close for listeners.
type listener struct {
	FileListener
	closed       bool      // Indicates we're already closed.
	closeRequest chan bool // Send a bool here to indicate we want to Close.
	allClosed    chan bool // Receive from here will indicate a clean Close.
	counter      chan bool // Use the inc/dec counters.
}

// Allows for us to notice when the connection is closed.
type conn struct {
	net.Conn
	counter chan bool
}

func (c conn) Close() error {
	c.counter <- dec
	return c.Conn.Close()
}

// Wraps an existing File listener to provide a graceful Close() process.
func NewListener(l FileListener) Listener {
	i := &listener{
		FileListener: l,
		closeRequest: make(chan bool),
		allClosed:    make(chan bool),
		counter:      make(chan bool),
	}
	go i.enabler()
	return i
}

func (l *listener) enabler() {
	var counter uint64
	var change bool
	for {
		select {
		case <-l.closeRequest:
			l.closed = true
		case change = <-l.counter:
			if change == inc {
				counter++
			} else {
				counter--
			}
		}
		if l.closed && counter == 0 {
			l.allClosed <- true
			break
		}
	}
}

func (l *listener) CloseRequest() {
	if l.closed == true {
		return
	}
	l.closeRequest <- true
	<-l.allClosed
}

func (l *listener) Close() error {
	l.CloseRequest()
	return l.FileListener.Close()
}

func (l *listener) Accept() (net.Conn, error) {
	if l.closed == true {
		return nil, ErrAlreadyClosed
	}
	c, err := l.FileListener.Accept()
	if err != nil {
		if strings.HasSuffix(err.Error(), errClosed) {
			return nil, ErrAlreadyClosed
		}
		return nil, err
	}
	l.counter <- inc
	return conn{
		Conn:    c,
		counter: l.counter,
	}, nil
}

// Wait for signals to gracefully terminate or restart the process.
func Wait(listeners []Listener) (err error) {
	ch := make(chan os.Signal, 2)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		switch sig {
		case syscall.SIGTERM:
			var wg sync.WaitGroup
			wg.Add(len(listeners))
			for _, l := range listeners {
				go func(l Listener) {
					if os.Getppid() == 1 { // init provided sockets dont actually close
						l.CloseRequest()
					} else {
						cErr := l.Close()
						if cErr != nil {
							err = cErr
						}
					}
					wg.Done()
				}(l)
			}
			wg.Wait()
			return
		case syscall.SIGUSR2:
			rErr := Restart(listeners)
			if rErr != nil {
				return rErr
			}
		}
	}
	panic("not reached")
}

// Try to inherit listeners from the parent process.
func Inherit() (listeners []Listener, err error) {
	countStr := os.Getenv(envCountKey)
	if countStr == "" {
		return nil, ErrNotInheriting
	}
	count, err := strconv.Atoi(countStr)
	if err != nil {
		return nil, err
	}
	// If we are inheriting, the listeners will begin at fd 3
	for i := 3; i < 3+count; i++ {
		file := os.NewFile(uintptr(i), "listener")
		tmp, err := net.FileListener(file)
		file.Close()
		if err != nil {
			return nil, err
		}
		l := tmp.(*net.TCPListener)
		listeners = append(listeners, NewListener(l))
	}
	return
}

// Start the Close process in the parent. This does not wait for the
// parent to close and simply sends it the TERM signal.
func CloseParent() error {
	ppid := os.Getppid()
	if ppid == 1 { // init provided sockets, for example systemd
		return nil
	}
	return syscall.Kill(ppid, syscall.SIGTERM)
}

// Restart the process passing the given listeners to the new process.
func Restart(listeners []Listener) (err error) {
	if len(listeners) == 0 {
		return errors.New("restart must be given listeners.")
	}
	files := make([]*os.File, len(listeners))
	for i, l := range listeners {
		files[i], err = l.File()
		if err != nil {
			return err
		}
		defer files[i].Close()
		syscall.CloseOnExec(int(files[i].Fd()))
	}
	argv0, err := exec.LookPath(os.Args[0])
	if err != nil {
		return err
	}
	wd, err := os.Getwd()
	if err != nil {
		return err
	}
	allFiles := append([]*os.File{os.Stdin, os.Stdout, os.Stderr}, files...)
	allFiles = append(allFiles, nil)
	_, err = os.StartProcess(argv0, os.Args, &os.ProcAttr{
		Dir:   wd,
		Env:   append(os.Environ(), fmt.Sprintf("%s=%d", envCountKey, len(files))),
		Files: allFiles,
	})
	return err
}
<commit_msg>remove some duplication<commit_after>// Package grace allows for gracefully waiting for a listener to
// finish serving it's active requests.
package grace

import (
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
)

var (
	// This error is returned by Inherits() when we're not inheriting any fds.
	ErrNotInheriting = errors.New("no inherited listeners")

	// This error is returned by Listener.Accept() when Close is in progress.
	ErrAlreadyClosed = errors.New("already closed")
)

const (
	// Used to indicate a graceful restart in the new process.
	envCountKey = "LISTEN_FDS"

	// The error returned by the standard library when the socket is closed.
	errClosed = "use of closed network connection"

	// Used for the counter chan.
	inc = true
	dec = false
)

// A FileListener is a file backed net.Listener.
type FileListener interface {
	net.Listener

	// Will return the underlying file representing this Listener.
	File() (f *os.File, err error)
}

// A Listener providing a graceful Close process and can be sent
// across processes using the underlying File descriptor.
type Listener interface {
	FileListener

	// Will indicate that a Close is requested preventing further Accept. It will
	// also wait for the active connections to be terminated before returning.
	// Note, this won't actually do the close, and is provided as part of the
	// public API for cases where the socket must not be closed (such as systemd
	// activation).
	CloseRequest()
}

// A goroutine based counter that provides graceful Close for listeners.
type listener struct {
	FileListener
	closed       bool      // Indicates we're already closed.
	closeRequest chan bool // Send a bool here to indicate we want to Close.
	allClosed    chan bool // Receive from here will indicate a clean Close.
	counter      chan bool // Use the inc/dec counters.
}

// Allows for us to notice when the connection is closed.
type conn struct {
	net.Conn
	counter chan bool
}

func (c conn) Close() error {
	c.counter <- dec
	return c.Conn.Close()
}

// Wraps an existing File listener to provide a graceful Close() process.
func NewListener(l FileListener) Listener {
	i := &listener{
		FileListener: l,
		closeRequest: make(chan bool),
		allClosed:    make(chan bool),
		counter:      make(chan bool),
	}
	go i.enabler()
	return i
}

func (l *listener) enabler() {
	var counter uint64
	var change bool
	for {
		select {
		case <-l.closeRequest:
			l.closed = true
		case change = <-l.counter:
			if change == inc {
				counter++
			} else {
				counter--
			}
		}
		if l.closed && counter == 0 {
			l.allClosed <- true
			break
		}
	}
}

func (l *listener) CloseRequest() {
	if l.closed == true {
		return
	}
	l.closeRequest <- true
	<-l.allClosed
}

func (l *listener) Close() error {
	l.CloseRequest()
	return l.FileListener.Close()
}

func (l *listener) Accept() (net.Conn, error) {
	if l.closed == true {
		return nil, ErrAlreadyClosed
	}
	c, err := l.FileListener.Accept()
	if err != nil {
		if strings.HasSuffix(err.Error(), errClosed) {
			return nil, ErrAlreadyClosed
		}
		return nil, err
	}
	l.counter <- inc
	return conn{
		Conn:    c,
		counter: l.counter,
	}, nil
}

// Wait for signals to gracefully terminate or restart the process.
func Wait(listeners []Listener) (err error) {
	ch := make(chan os.Signal, 2)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		switch sig {
		case syscall.SIGTERM:
			var wg sync.WaitGroup
			wg.Add(len(listeners))
			for _, l := range listeners {
				go func(l Listener) {
					if os.Getppid() == 1 { // init provided sockets dont actually close
						l.CloseRequest()
					} else {
						cErr := l.Close()
						if cErr != nil {
							err = cErr
						}
					}
					wg.Done()
				}(l)
			}
			wg.Wait()
			return
		case syscall.SIGUSR2:
			rErr := Restart(listeners)
			if rErr != nil {
				return rErr
			}
		}
	}
	panic("not reached")
}

// Try to inherit listeners from the parent process.
func Inherit() (listeners []Listener, err error) {
	countStr := os.Getenv(envCountKey)
	if countStr == "" {
		return nil, ErrNotInheriting
	}
	count, err := strconv.Atoi(countStr)
	if err != nil {
		return nil, err
	}
	// If we are inheriting, the listeners will begin at fd 3
	for i := 3; i < 3+count; i++ {
		file := os.NewFile(uintptr(i), "listener")
		tmp, err := net.FileListener(file)
		file.Close()
		if err != nil {
			return nil, err
		}
		l := tmp.(*net.TCPListener)
		listeners = append(listeners, NewListener(l))
	}
	return
}

// Start the Close process in the parent. This does not wait for the
// parent to close and simply sends it the TERM signal.
func CloseParent() error {
	ppid := os.Getppid()
	if ppid == 1 { // init provided sockets, for example systemd
		return nil
	}
	return syscall.Kill(ppid, syscall.SIGTERM)
}

// Restart the process passing the given listeners to the new process.
func Restart(listeners []Listener) (err error) {
	if len(listeners) == 0 {
		return errors.New("restart must be given listeners.")
	}
	files := make([]*os.File, len(listeners))
	for i, l := range listeners {
		files[i], err = l.File()
		if err != nil {
			return err
		}
		defer files[i].Close()
		syscall.CloseOnExec(int(files[i].Fd()))
	}
	argv0, err := exec.LookPath(os.Args[0])
	if err != nil {
		return err
	}
	wd, err := os.Getwd()
	if err != nil {
		return err
	}
	allFiles := append([]*os.File{os.Stdin, os.Stdout, os.Stderr}, files...)
	allFiles = append(allFiles, nil)
	_, err = os.StartProcess(argv0, os.Args, &os.ProcAttr{
		Dir:   wd,
		Env:   append(os.Environ(), fmt.Sprintf("%s=%d", envCountKey, len(files))),
		Files: allFiles,
	})
	return err
}
<|endoftext|>


<commit_before>package continuous

// MorphologicalComputationW [...]
func MorphologicalComputationW(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false)
}

// MorphologicalComputationA [...]
func MorphologicalComputationA(w2a1w1 [][]float64, w2Indices, a1Indices, w1Indices []int, k int) float64 {
	return FrenzelPompe(w2a1w1, w2Indices, a1Indices, w1Indices, k, false)
}

// MorphologicalComputationCW1 [...]
func MorphologicalComputationCW(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger1(w2w1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger1(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationCW2 [...]
func MorphologicalComputationCW(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger2(w2w1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger2(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWA1 = I(W;{W,A}) - I(W';A)
func MorphologicalComputationWA(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false) - KraskovStoegbauerGrassberger1(w2a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWA2 = I(W;{W,A}) - I(W';A)
func MorphologicalComputationWA(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false) - KraskovStoegbauerGrassberger2(w2a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWS = I(W;{W,S}) - I(W';S)
func MorphologicalComputationWS(w2w1s1 [][]float64, w2Indices, w1Indices, s1Idices []int, k int) float64 {
	return FrenzelPompe(w2w1s1, w2Indices, w1Indices, s1Indices, k, false) - KraskovStoegbauerGrassberger(w2s1, w2Indices, s1Indices, k, false)
}

// MorphologicalComputationWS1 = I(W;{W,S}) - I(W';S)
func MorphologicalComputationWS(w2w1s1 [][]float64, w2Indices, w1Indices, s1Idices []int, k int) float64 {
	return FrenzelPompe(w2w1s1, w2Indices, w1Indices, s1Indices, k, false) - KraskovStoegbauerGrassberger1(w2s1, w2Indices, s1Indices, k, false)
}

// MorphologicalComputationWS2 = I(W;{W,S}) - I(W';S)
func MorphologicalComputationWS(w2w1s1 [][]float64, w2Indices, w1Indices, s1Idices []int, k int) float64 {
	return FrenzelPompe(w2w1s1, w2Indices, w1Indices, s1Indices, k, false) - KraskovStoegbauerGrassberger2(w2s1, w2Indices, s1Indices, k, false)
}

// MorphologicalComputationMI1 [...]
func MorphologicalComputationMI(w2w1s1a1 [][]float64, w2Indices, w1Indices, s1Idices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger1(w2w1s1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger1(a1s1, a1Indices, s1Idices, k, false)
}

// MorphologicalComputationMI2 [...]
func MorphologicalComputationMI(w2w1s1a1 [][]float64, w2Indices, w1Indices, s1Idices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger2(w2w1s1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger2(a1s1, a1Indices, s1Idices, k, false)
}
<commit_msg>Fixed test cases<commit_after>package continuous

// MorphologicalComputationW [...]
func MorphologicalComputationW(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false)
}

// MorphologicalComputationA [...]
func MorphologicalComputationA(w2a1w1 [][]float64, w2Indices, a1Indices, w1Indices []int, k int) float64 {
	return FrenzelPompe(w2a1w1, w2Indices, a1Indices, w1Indices, k, false)
}

// MorphologicalComputationCW1 [...]
func MorphologicalComputationCW1(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger1(w2w1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger1(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationCW2 [...]
func MorphologicalComputationCW2(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger2(w2w1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger2(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWA1 = I(W;{W,A}) - I(W';A)
func MorphologicalComputationWA1(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false) - KraskovStoegbauerGrassberger1(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWA2 = I(W;{W,A}) - I(W';A)
func MorphologicalComputationWA(w2w1a1 [][]float64, w2Indices, w1Indices, a1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1a1, w2Indices, w1Indices, a1Indices, k, false) - KraskovStoegbauerGrassberger2(w2w1a1, w2Indices, a1Indices, k, false)
}

// MorphologicalComputationWS1 = I(W;{W,S}) - I(W';S)
func MorphologicalComputationWS1(w2w1s1 [][]float64, w2Indices, w1Indices, s1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1s1, w2Indices, w1Indices, s1Indices, k, false) - KraskovStoegbauerGrassberger1(w2w1s1, w2Indices, s1Indices, k, false)
}

// MorphologicalComputationWS2 = I(W;{W,S}) - I(W';S)
func MorphologicalComputationWS2(w2w1s1 [][]float64, w2Indices, w1Indices, s1Indices []int, k int) float64 {
	return FrenzelPompe(w2w1s1, w2Indices, w1Indices, s1Indices, k, false) - KraskovStoegbauerGrassberger2(w2w1s1, w2Indices, s1Indices, k, false)
}

// MorphologicalComputationMI1 [...]
func MorphologicalComputationMI1(w2w1s1a1 [][]float64, w2Indices, w1Indices, s1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger1(w2w1s1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger1(w2w1s1a1, a1Indices, s1Indices, k, false)
}

// MorphologicalComputationMI2 [...]
func MorphologicalComputationMI2(w2w1s1a1 [][]float64, w2Indices, w1Indices, s1Indices, a1Indices []int, k int) float64 {
	return KraskovStoegbauerGrassberger2(w2w1s1a1, w2Indices, w1Indices, k, false) - KraskovStoegbauerGrassberger2(w2w1s1a1, a1Indices, s1Indices, k, false)
}
<|endoftext|>


<commit_before>package docker

import (
	"fmt"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

// A Graph is a store for versioned filesystem images and the relationship between them.
type Graph struct {
	Root    string
	idIndex *utils.TruncIndex
}

// NewGraph instantiates a new graph at the given root path in the filesystem.
// `root` will be created if it doesn't exist.
func NewGraph(root string) (*Graph, error) {
	abspath, err := filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	// Create the root directory if it doesn't exists
	if err := os.MkdirAll(root, 0700); err != nil && !os.IsExist(err) {
		return nil, err
	}
	graph := &Graph{
		Root:    abspath,
		idIndex: utils.NewTruncIndex(),
	}
	if err := graph.restore(); err != nil {
		return nil, err
	}
	return graph, nil
}

func (graph *Graph) restore() error {
	dir, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, v := range dir {
		id := v.Name()
		graph.idIndex.Add(id)
	}
	return nil
}

// FIXME: Implement error subclass instead of looking at the error text
// Note: This is the way golang implements os.IsNotExists on Plan9
func (graph *Graph) IsNotExist(err error) bool {
	return err != nil && (strings.Contains(err.Error(), "does not exist") || strings.Contains(err.Error(), "No such"))
}

// Exists returns true if an image is registered at the given id.
// If the image doesn't exist or if an error is encountered, false is returned.
func (graph *Graph) Exists(id string) bool {
	if _, err := graph.Get(id); err != nil {
		return false
	}
	return true
}

// Get returns the image with the given id, or an error if the image doesn't exist.
func (graph *Graph) Get(name string) (*Image, error) {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return nil, err
	}
	// FIXME: return nil when the image doesn't exist, instead of an error
	img, err := LoadImage(graph.imageRoot(id))
	if err != nil {
		return nil, err
	}
	if img.ID != id {
		return nil, fmt.Errorf("Image stored at '%s' has wrong id '%s'", id, img.ID)
	}
	img.graph = graph
	if img.Size == 0 {
		root, err := img.root()
		if err != nil {
			return nil, err
		}
		if err := StoreSize(img, root); err != nil {
			return nil, err
		}
	}
	return img, nil
}

// Create creates a new image and registers it in the graph.
func (graph *Graph) Create(layerData Archive, container *Container, comment, author string, config *Config) (*Image, error) {
	img := &Image{
		ID:            GenerateID(),
		Comment:       comment,
		Created:       time.Now(),
		DockerVersion: VERSION,
		Author:        author,
		Config:        config,
		Architecture:  "x86_64",
	}
	if container != nil {
		img.Parent = container.Image
		img.Container = container.ID
		img.ContainerConfig = *container.Config
	}
	if err := graph.Register(nil, layerData, img); err != nil {
		return nil, err
	}
	return img, nil
}

// Register imports a pre-existing image into the graph.
// FIXME: pass img as first argument
func (graph *Graph) Register(jsonData []byte, layerData Archive, img *Image) error {
	if err := ValidateID(img.ID); err != nil {
		return err
	}
	// (This is a convenience to save time. Race conditions are taken care of by os.Rename)
	if graph.Exists(img.ID) {
		return fmt.Errorf("Image %s already exists", img.ID)
	}
	tmp, err := graph.Mktemp("")
	defer os.RemoveAll(tmp)
	if err != nil {
		return fmt.Errorf("Mktemp failed: %s", err)
	}
	if err := StoreImage(img, jsonData, layerData, tmp); err != nil {
		return err
	}
	// Commit
	if err := os.Rename(tmp, graph.imageRoot(img.ID)); err != nil {
		return err
	}
	img.graph = graph
	graph.idIndex.Add(img.ID)
	return nil
}

// TempLayerArchive creates a temporary archive of the given image's filesystem layer.
//   The archive is stored on disk and will be automatically deleted as soon as has been read.
//   If output is not nil, a human-readable progress bar will be written to it.
//   FIXME: does this belong in Graph? How about MktempFile, let the caller use it for archives?
func (graph *Graph) TempLayerArchive(id string, compression Compression, sf *utils.StreamFormatter, output io.Writer) (*TempArchive, error) {
	image, err := graph.Get(id)
	if err != nil {
		return nil, err
	}
	tmp, err := graph.tmp()
	if err != nil {
		return nil, err
	}
	archive, err := image.TarLayer(compression)
	if err != nil {
		return nil, err
	}
	return NewTempArchive(utils.ProgressReader(ioutil.NopCloser(archive), 0, output, sf.FormatProgress("", "Buffering to disk", "%v/%v (%v)"), sf, true), tmp.Root)
}

// Mktemp creates a temporary sub-directory inside the graph's filesystem.
func (graph *Graph) Mktemp(id string) (string, error) {
	if id == "" {
		id = GenerateID()
	}
	tmp, err := graph.tmp()
	if err != nil {
		return "", fmt.Errorf("Couldn't create temp: %s", err)
	}
	if tmp.Exists(id) {
		return "", fmt.Errorf("Image %s already exists", id)
	}
	return tmp.imageRoot(id), nil
}

// getDockerInitLayer returns the path of a layer containing a mountpoint suitable
// for bind-mounting dockerinit into the container. The mountpoint is simply an
// empty file at /.dockerinit
//
// This extra layer is used by all containers as the top-most ro layer. It protects
// the container from unwanted side-effects on the rw layer.
func (graph *Graph) getDockerInitLayer() (string, error) {
	tmp, err := graph.tmp()
	if err != nil {
		return "", err
	}
	initLayer := tmp.imageRoot("_dockerinit")
	if err := os.Mkdir(initLayer, 0755); err != nil && !os.IsExist(err) {
		// If directory already existed, keep going.
		// For all other errors, abort.
		return "", err
	}

	for pth, typ := range map[string]string{
		"/dev/pts":         "dir",
		"/dev/shm":         "dir",
		"/proc":            "dir",
		"/sys":             "dir",
		"/.dockerinit":     "file",
		"/etc/resolv.conf": "file",
		// "var/run": "dir",
		// "var/lock": "dir",
	} {
		if _, err := os.Stat(path.Join(initLayer, pth)); err != nil {
			if os.IsNotExist(err) {
				switch typ {
				case "dir":
					if err := os.MkdirAll(path.Join(initLayer, pth), 0755); err != nil {
						return "", err
					}
				case "file":
					if err := os.MkdirAll(path.Join(initLayer, path.Dir(pth)), 0755); err != nil {
						return "", err
					}

					if f, err := os.OpenFile(path.Join(initLayer, pth), os.O_CREATE, 0755); err != nil {
						return "", err
					} else {
						f.Close()
					}
				}
			} else {
				return "", err
			}
		}
	}

	// Layer is ready to use, if it wasn't before.
	return initLayer, nil
}

func (graph *Graph) tmp() (*Graph, error) {
	// Changed to _tmp from :tmp:, because it messed with ":" separators in aufs branch syntax...
	return NewGraph(path.Join(graph.Root, "_tmp"))
}

// Check if given error is "not empty".
// Note: this is the way golang does it internally with os.IsNotExists.
func isNotEmpty(err error) bool {
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		err = pe.Err
	case *os.LinkError:
		err = pe.Err
	}
	return strings.Contains(err.Error(), " not empty")
}

// Delete atomically removes an image from the graph.
func (graph *Graph) Delete(name string) error {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return err
	}
	tmp, err := graph.Mktemp("")
	if err != nil {
		return err
	}
	graph.idIndex.Delete(id)
	err = os.Rename(graph.imageRoot(id), tmp)
	if err != nil {
		return err
	}
	return os.RemoveAll(tmp)
}

// Map returns a list of all images in the graph, addressable by ID.
func (graph *Graph) Map() (map[string]*Image, error) {
	// FIXME: this should replace All()
	all, err := graph.All()
	if err != nil {
		return nil, err
	}
	images := make(map[string]*Image, len(all))
	for _, image := range all {
		images[image.ID] = image
	}
	return images, nil
}

// All returns a list of all images in the graph.
func (graph *Graph) All() ([]*Image, error) {
	var images []*Image
	err := graph.WalkAll(func(image *Image) {
		images = append(images, image)
	})
	return images, err
}

// WalkAll iterates over each image in the graph, and passes it to a handler.
// The walking order is undetermined.
func (graph *Graph) WalkAll(handler func(*Image)) error {
	files, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, st := range files {
		if img, err := graph.Get(st.Name()); err != nil {
			// Skip image
			continue
		} else if handler != nil {
			handler(img)
		}
	}
	return nil
}

// ByParent returns a lookup table of images by their parent.
// If an image of id ID has 3 children images, then the value for key ID
// will be a list of 3 images.
// If an image has no children, it will not have an entry in the table.
func (graph *Graph) ByParent() (map[string][]*Image, error) {
	byParent := make(map[string][]*Image)
	err := graph.WalkAll(func(image *Image) {
		parent, err := graph.Get(image.Parent)
		if err != nil {
			return
		}
		if children, exists := byParent[parent.ID]; exists {
			byParent[parent.ID] = []*Image{image}
		} else {
			byParent[parent.ID] = append(children, image)
		}
	})
	return byParent, err
}

// Heads returns all heads in the graph, keyed by id.
// A head is an image which is not the parent of another image in the graph.
func (graph *Graph) Heads() (map[string]*Image, error) {
	heads := make(map[string]*Image)
	byParent, err := graph.ByParent()
	if err != nil {
		return nil, err
	}
	err = graph.WalkAll(func(image *Image) {
		// If it's not in the byParent lookup table, then
		// it's not a parent -> so it's a head!
		if _, exists := byParent[image.ID]; !exists {
			heads[image.ID] = image
		}
	})
	return heads, err
}

func (graph *Graph) imageRoot(id string) string {
	return path.Join(graph.Root, id)
}
<commit_msg>Fix Graph ByParent() to generate list of child images per parent image.<commit_after>package docker

import (
	"fmt"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

// A Graph is a store for versioned filesystem images and the relationship between them.
type Graph struct {
	Root    string
	idIndex *utils.TruncIndex
}

// NewGraph instantiates a new graph at the given root path in the filesystem.
// `root` will be created if it doesn't exist.
func NewGraph(root string) (*Graph, error) {
	abspath, err := filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	// Create the root directory if it doesn't exists
	if err := os.MkdirAll(root, 0700); err != nil && !os.IsExist(err) {
		return nil, err
	}
	graph := &Graph{
		Root:    abspath,
		idIndex: utils.NewTruncIndex(),
	}
	if err := graph.restore(); err != nil {
		return nil, err
	}
	return graph, nil
}

func (graph *Graph) restore() error {
	dir, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, v := range dir {
		id := v.Name()
		graph.idIndex.Add(id)
	}
	return nil
}

// FIXME: Implement error subclass instead of looking at the error text
// Note: This is the way golang implements os.IsNotExists on Plan9
func (graph *Graph) IsNotExist(err error) bool {
	return err != nil && (strings.Contains(err.Error(), "does not exist") || strings.Contains(err.Error(), "No such"))
}

// Exists returns true if an image is registered at the given id.
// If the image doesn't exist or if an error is encountered, false is returned.
func (graph *Graph) Exists(id string) bool {
	if _, err := graph.Get(id); err != nil {
		return false
	}
	return true
}

// Get returns the image with the given id, or an error if the image doesn't exist.
func (graph *Graph) Get(name string) (*Image, error) {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return nil, err
	}
	// FIXME: return nil when the image doesn't exist, instead of an error
	img, err := LoadImage(graph.imageRoot(id))
	if err != nil {
		return nil, err
	}
	if img.ID != id {
		return nil, fmt.Errorf("Image stored at '%s' has wrong id '%s'", id, img.ID)
	}
	img.graph = graph
	if img.Size == 0 {
		root, err := img.root()
		if err != nil {
			return nil, err
		}
		if err := StoreSize(img, root); err != nil {
			return nil, err
		}
	}
	return img, nil
}

// Create creates a new image and registers it in the graph.
func (graph *Graph) Create(layerData Archive, container *Container, comment, author string, config *Config) (*Image, error) {
	img := &Image{
		ID:            GenerateID(),
		Comment:       comment,
		Created:       time.Now(),
		DockerVersion: VERSION,
		Author:        author,
		Config:        config,
		Architecture:  "x86_64",
	}
	if container != nil {
		img.Parent = container.Image
		img.Container = container.ID
		img.ContainerConfig = *container.Config
	}
	if err := graph.Register(nil, layerData, img); err != nil {
		return nil, err
	}
	return img, nil
}

// Register imports a pre-existing image into the graph.
// FIXME: pass img as first argument
func (graph *Graph) Register(jsonData []byte, layerData Archive, img *Image) error {
	if err := ValidateID(img.ID); err != nil {
		return err
	}
	// (This is a convenience to save time. Race conditions are taken care of by os.Rename)
	if graph.Exists(img.ID) {
		return fmt.Errorf("Image %s already exists", img.ID)
	}
	tmp, err := graph.Mktemp("")
	defer os.RemoveAll(tmp)
	if err != nil {
		return fmt.Errorf("Mktemp failed: %s", err)
	}
	if err := StoreImage(img, jsonData, layerData, tmp); err != nil {
		return err
	}
	// Commit
	if err := os.Rename(tmp, graph.imageRoot(img.ID)); err != nil {
		return err
	}
	img.graph = graph
	graph.idIndex.Add(img.ID)
	return nil
}

// TempLayerArchive creates a temporary archive of the given image's filesystem layer.
//   The archive is stored on disk and will be automatically deleted as soon as has been read.
//   If output is not nil, a human-readable progress bar will be written to it.
//   FIXME: does this belong in Graph? How about MktempFile, let the caller use it for archives?
func (graph *Graph) TempLayerArchive(id string, compression Compression, sf *utils.StreamFormatter, output io.Writer) (*TempArchive, error) {
	image, err := graph.Get(id)
	if err != nil {
		return nil, err
	}
	tmp, err := graph.tmp()
	if err != nil {
		return nil, err
	}
	archive, err := image.TarLayer(compression)
	if err != nil {
		return nil, err
	}
	return NewTempArchive(utils.ProgressReader(ioutil.NopCloser(archive), 0, output, sf.FormatProgress("", "Buffering to disk", "%v/%v (%v)"), sf, true), tmp.Root)
}

// Mktemp creates a temporary sub-directory inside the graph's filesystem.
func (graph *Graph) Mktemp(id string) (string, error) {
	if id == "" {
		id = GenerateID()
	}
	tmp, err := graph.tmp()
	if err != nil {
		return "", fmt.Errorf("Couldn't create temp: %s", err)
	}
	if tmp.Exists(id) {
		return "", fmt.Errorf("Image %s already exists", id)
	}
	return tmp.imageRoot(id), nil
}

// getDockerInitLayer returns the path of a layer containing a mountpoint suitable
// for bind-mounting dockerinit into the container. The mountpoint is simply an
// empty file at /.dockerinit
//
// This extra layer is used by all containers as the top-most ro layer. It protects
// the container from unwanted side-effects on the rw layer.
func (graph *Graph) getDockerInitLayer() (string, error) {
	tmp, err := graph.tmp()
	if err != nil {
		return "", err
	}
	initLayer := tmp.imageRoot("_dockerinit")
	if err := os.Mkdir(initLayer, 0755); err != nil && !os.IsExist(err) {
		// If directory already existed, keep going.
		// For all other errors, abort.
		return "", err
	}

	for pth, typ := range map[string]string{
		"/dev/pts":         "dir",
		"/dev/shm":         "dir",
		"/proc":            "dir",
		"/sys":             "dir",
		"/.dockerinit":     "file",
		"/etc/resolv.conf": "file",
		// "var/run": "dir",
		// "var/lock": "dir",
	} {
		if _, err := os.Stat(path.Join(initLayer, pth)); err != nil {
			if os.IsNotExist(err) {
				switch typ {
				case "dir":
					if err := os.MkdirAll(path.Join(initLayer, pth), 0755); err != nil {
						return "", err
					}
				case "file":
					if err := os.MkdirAll(path.Join(initLayer, path.Dir(pth)), 0755); err != nil {
						return "", err
					}

					if f, err := os.OpenFile(path.Join(initLayer, pth), os.O_CREATE, 0755); err != nil {
						return "", err
					} else {
						f.Close()
					}
				}
			} else {
				return "", err
			}
		}
	}

	// Layer is ready to use, if it wasn't before.
	return initLayer, nil
}

func (graph *Graph) tmp() (*Graph, error) {
	// Changed to _tmp from :tmp:, because it messed with ":" separators in aufs branch syntax...
	return NewGraph(path.Join(graph.Root, "_tmp"))
}

// Check if given error is "not empty".
// Note: this is the way golang does it internally with os.IsNotExists.
func isNotEmpty(err error) bool {
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		err = pe.Err
	case *os.LinkError:
		err = pe.Err
	}
	return strings.Contains(err.Error(), " not empty")
}

// Delete atomically removes an image from the graph.
func (graph *Graph) Delete(name string) error {
	id, err := graph.idIndex.Get(name)
	if err != nil {
		return err
	}
	tmp, err := graph.Mktemp("")
	if err != nil {
		return err
	}
	graph.idIndex.Delete(id)
	err = os.Rename(graph.imageRoot(id), tmp)
	if err != nil {
		return err
	}
	return os.RemoveAll(tmp)
}

// Map returns a list of all images in the graph, addressable by ID.
func (graph *Graph) Map() (map[string]*Image, error) {
	// FIXME: this should replace All()
	all, err := graph.All()
	if err != nil {
		return nil, err
	}
	images := make(map[string]*Image, len(all))
	for _, image := range all {
		images[image.ID] = image
	}
	return images, nil
}

// All returns a list of all images in the graph.
func (graph *Graph) All() ([]*Image, error) {
	var images []*Image
	err := graph.WalkAll(func(image *Image) {
		images = append(images, image)
	})
	return images, err
}

// WalkAll iterates over each image in the graph, and passes it to a handler.
// The walking order is undetermined.
func (graph *Graph) WalkAll(handler func(*Image)) error {
	files, err := ioutil.ReadDir(graph.Root)
	if err != nil {
		return err
	}
	for _, st := range files {
		if img, err := graph.Get(st.Name()); err != nil {
			// Skip image
			continue
		} else if handler != nil {
			handler(img)
		}
	}
	return nil
}

// ByParent returns a lookup table of images by their parent.
// If an image of id ID has 3 children images, then the value for key ID
// will be a list of 3 images.
// If an image has no children, it will not have an entry in the table.
func (graph *Graph) ByParent() (map[string][]*Image, error) {
	byParent := make(map[string][]*Image)
	err := graph.WalkAll(func(image *Image) {
		parent, err := graph.Get(image.Parent)
		if err != nil {
			return
		}
		if children, exists := byParent[parent.ID]; exists {
			byParent[parent.ID] = append(children, image)
		} else {
			byParent[parent.ID] = []*Image{image}
		}
	})
	return byParent, err
}

// Heads returns all heads in the graph, keyed by id.
// A head is an image which is not the parent of another image in the graph.
func (graph *Graph) Heads() (map[string]*Image, error) {
	heads := make(map[string]*Image)
	byParent, err := graph.ByParent()
	if err != nil {
		return nil, err
	}
	err = graph.WalkAll(func(image *Image) {
		// If it's not in the byParent lookup table, then
		// it's not a parent -> so it's a head!
		if _, exists := byParent[image.ID]; !exists {
			heads[image.ID] = image
		}
	})
	return heads, err
}

func (graph *Graph) imageRoot(id string) string {
	return path.Join(graph.Root, id)
}
<|endoftext|>


<commit_before>// gogl provides a framework for representing and working with graphs.
package gogl

// Constants defining graph capabilities and behaviors.
const (
	E_DIRECTED, EM_DIRECTED = 1 << iota, 1<<iota - 1
	E_UNDIRECTED, EM_UNDIRECTED
	E_WEIGHTED, EM_WEIGHTED
	E_TYPED, EM_TYPED
	E_SIGNED, EM_SIGNED
	E_LOOPS, EM_LOOPS
	E_MULTIGRAPH, EM_MULTIGRAPH
)

type Vertex interface{}

type Edge struct {
	Tail, Head Vertex
}

type Graph interface {
	EachVertex(f func(vertex Vertex))
	EachEdge(f func(edge Edge))
	EachAdjacent(vertex Vertex, f func(adjacent Vertex))
	HasVertex(vertex Vertex) bool
	Order() uint
	Size() uint
	GetSubgraph([]Vertex) Graph
}

type MutableGraph interface {
	Graph
	AddVertex(v interface{}) bool
	RemoveVertex(v interface{}) bool
}

type DirectedGraph interface {
	Graph
	Transpose() DirectedGraph
	IsAcyclic() bool
	GetCycles() [][]interface{}
}

type MutableDirectedGraph interface {
	MutableGraph
	DirectedGraph
	addDirectedEdge(source interface{}, target interface{}) bool
	removeDirectedEdge(source interface{}, target interface{}) bool
}
<commit_msg>Fix interfaces so they actually work.<commit_after>// gogl provides a framework for representing and working with graphs.
package gogl

// Constants defining graph capabilities and behaviors.
const (
	E_DIRECTED, EM_DIRECTED = 1 << iota, 1<<iota - 1
	E_UNDIRECTED, EM_UNDIRECTED
	E_WEIGHTED, EM_WEIGHTED
	E_TYPED, EM_TYPED
	E_SIGNED, EM_SIGNED
	E_LOOPS, EM_LOOPS
	E_MULTIGRAPH, EM_MULTIGRAPH
)

type Vertex interface{}

type Edge struct {
	Tail, Head Vertex
}

type Graph interface {
	EachVertex(f func(vertex Vertex))
	EachEdge(f func(edge Edge))
	EachAdjacent(vertex Vertex, f func(adjacent Vertex))
	HasVertex(vertex Vertex) bool
	Order() uint
	Size() uint
	AddVertex(v interface{}) bool
	RemoveVertex(v interface{}) bool
}

type DirectedGraph interface {
	Graph
	Transpose() DirectedGraph
	IsAcyclic() bool
	GetCycles() [][]interface{}
	addDirectedEdge(source interface{}, target interface{}) bool
	removeDirectedEdge(source interface{}, target interface{}) bool
}
<|endoftext|>


<commit_before>package gogl

import (
	"fmt"
)

/* Vertex structures */

// As a rule, gogl tries to place as low a requirement on its vertices as
// possible. This is because, from a purely graph theoretic perspective,
// vertices are inert. Boring, even. Graphs are more about the topology, the
// characteristics of the edges connecting the points than the points
// themselves. Your use case cares about the content of your vertices, but gogl
// does not.  Consequently, anything can act as a vertex.
type Vertex interface{}

/* Atomic graph interfaces */

// A VertexEnumerator iteratively enumerates vertices.
type VertexEnumerator interface {
	EachVertex(f func(Vertex))
}

// An EdgeEnumerator iteratively enumerates edges.
type EdgeEnumerator interface {
	EachEdge(f func(Edge))
}

// An AdjacencyEnumerator iteratively enumerates a given vertex's adjacent vertices.
type AdjacencyEnumerator interface {
	EachAdjacent(start Vertex, f func(adjacent Vertex))
}

// A VertexMembershipChecker can indicate the presence of a vertex.
type VertexMembershipChecker interface {
	HasVertex(Vertex) bool // Whether or not the vertex is present in the set
}

// An InOutDegreeChecker reports the number of edges incident to a given vertex.
// TODO use this
type DegreeChecker interface {
	Degree(Vertex) (degree int, exists bool) // Number of incident edges; if vertex is present
}

// An InOutDegreeChecker reports the number of in or out-edges a given vertex has.
type InOutDegreeChecker interface {
	InDegree(Vertex) (degree int, exists bool)  // Number of in-edges; if vertex is present
	OutDegree(Vertex) (degree int, exists bool) // Number of out-edges; if vertex is present
}

// An EdgeMembershipChecker can indicate the presence of an edge.
type EdgeMembershipChecker interface {
	HasEdge(Edge) bool
}

// A VertexSetMutator allows the addition and removal of vertices from a set.
type VertexSetMutator interface {
	EnsureVertex(...Vertex)
	RemoveVertex(...Vertex)
}

// An EdgeSetMutator allows the addition and removal of edges from a set.
type EdgeSetMutator interface {
	AddEdges(edges ...Edge)
	RemoveEdges(edges ...Edge)
}

// A WeightedEdgeSetMutator allows the addition and removal of weighted edges from a set.
type WeightedEdgeSetMutator interface {
	AddEdges(edges ...WeightedEdge)
	RemoveEdges(edges ...WeightedEdge)
}

// A LabeledEdgeSetMutator allows the addition and removal of labeled edges from a set.
type LabeledEdgeSetMutator interface {
	AddEdges(edges ...LabeledEdge)
	RemoveEdges(edges ...LabeledEdge)
}

// A PropertyEdgeSetMutator allows the addition and removal of data edges from a set.
type PropertyEdgeSetMutator interface {
	AddEdges(edges ...PropertyEdge)
	RemoveEdges(edges ...PropertyEdge)
}

// A Transposer produces a transposed version of a DirectedGraph.
type Transposer interface {
	Transpose() DirectedGraph
}

/* Aggregate graph interfaces */

// Graph is gogl's most basic interface: it contains only the methods that
// *every* type of graph implements.
//
// Graph is intentionally underspecified: both directed and undirected graphs
// implement it; simple graphs, multigraphs, weighted, labeled, or any
// combination thereof.
//
// The semantics of some of these methods vary slightly from one graph type
// to another, but in general, the basic Graph methods are supplemented, not
// superceded, by the methods in more specific interfaces.
//
// Graph is a purely read oriented interface; the various Mutable*Graph
// interfaces contain the methods for writing.
type Graph interface {
	VertexEnumerator        // Allows enumerated traversal of vertices
	EdgeEnumerator          // Allows enumerated traversal of edges
	AdjacencyEnumerator     // Allows enumerated traversal of a vertex's adjacent vertices
	VertexMembershipChecker // Allows inspection of contained vertices
	EdgeMembershipChecker   // Allows inspection of contained edges
	InOutDegreeChecker      // Reports in- and out-degree of vertices
	Order() int             // Total number of vertices in the graph
	Size() int              // Total number of edges in the graph
}

// DirectedGraph describes a Graph all of whose edges are directed.
//
// Implementing DirectedGraph is the only unambiguous signal gogl provides
// that a graph's edges are directed.
type DirectedGraph interface {
	Graph
	Transposer // DirectedGraphs can produce a transpose of themselves
}

// MutableGraph describes a graph with basic edges (no weighting, labeling, etc.)
// that can be modified freely by adding or removing vertices or edges.
type MutableGraph interface {
	Graph
	VertexSetMutator
	EdgeSetMutator
}

// A simple graph is in opposition to a multigraph: it disallows loops and
// parallel edges.
type SimpleGraph interface {
	Graph
	Density() float64
}

// A weighted graph is a graph subtype where the edges have a numeric weight;
// as described by the WeightedEdge interface, this weight is a signed int.
//
// WeightedGraphs have both the HasEdge() and HasWeightedEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its weight.
//
// HasWeightedEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge weights are the same.
type WeightedGraph interface {
	Graph
	HasWeightedEdge(e WeightedEdge) bool
	EachWeightedEdge(f func(edge WeightedEdge))
}

// MutableWeightedGraph is the mutable version of a weighted graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only weighted edges can be present in the graph.
type MutableWeightedGraph interface {
	WeightedGraph
	VertexSetMutator
	WeightedEdgeSetMutator
}

// A labeled graph is a graph subtype where the edges have an identifier;
// as described by the LabeledEdge interface, this identifier is a string.
//
// LabeledGraphs have both the HasEdge() and HasLabeledEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its label.
//
// HasLabeledEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge labels are the same.
type LabeledGraph interface {
	Graph
	HasLabeledEdge(e LabeledEdge) bool
	EachLabeledEdge(f func(edge LabeledEdge))
}

// LabeledWeightedGraph is the mutable version of a labeled graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only labeled edges can be present in the graph.
type MutableLabeledGraph interface {
	LabeledGraph
	VertexSetMutator
	LabeledEdgeSetMutator
}

// A data graph is a graph subtype where the edges carry arbitrary Go data;
// as described by the PropertyEdge interface, this identifier is an interface{}.
//
// PropertyGraphs have both the HasEdge() and HasPropertyEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its label.
//
// HasPropertyEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge data is the same. Simple comparison will typically be used
// to establish data equality, which means that using noncomparables (a slice,
// map, or non-pointer struct containing a slice or a map) for the data will
// cause a panic.
type PropertyGraph interface {
	Graph
	HasPropertyEdge(e PropertyEdge) bool
	EachPropertyEdge(f func(edge PropertyEdge))
}

// MutablePropertyGraph is the mutable version of a propety graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only property edges can be present in the graph.
type MutablePropertyGraph interface {
	PropertyGraph
	VertexSetMutator
	PropertyEdgeSetMutator
}

/* Graph creation */

type GraphFactory func() interface{}

var Graphs = make(map[string]GraphFactory, 0)

// Creates a new graph instance.
//
// You will need to type assert the returned graph to the interface appropriate
// for your use case: Graph, DirectedGraph, MutableGraph, WeightedGraph, etc.
func New(name string) (graph interface{}, err error) {
	if _, exists := Graphs[name]; !exists {
		return nil, fmt.Errorf("No graph is registered with the name %q", name)
	}

	return Graphs[name](), nil
}

func RegisterGraph(name string, factory GraphFactory) error {
	if _, exists := Graphs[name]; exists {
		return fmt.Errorf("A graph is already registered with the name %q", name)
	}

	g := factory()

	if _, ok := g.(Graph); ok {
		return nil
	} else if _, ok := g.(WeightedGraph); ok {
		return nil
	}

	return fmt.Errorf("Value returned from factory does not implement a known Graph interface")
}

func init() {
	RegisterGraph("basic.directed", func() interface{} {
		return NewDirected()
	})
	RegisterGraph("basic.undirected", func() interface{} {
		return NewUndirected()
	})
	RegisterGraph("weighted.directed", func() interface{} {
		return NewWeightedDirected()
	})
	RegisterGraph("weighted.undirected", func() interface{} {
		return NewWeightedUndirected()
	})
}
<commit_msg>Add incidence edge/arc enumerator interfaces.<commit_after>package gogl

import (
	"fmt"
)

/* Vertex structures */

// As a rule, gogl tries to place as low a requirement on its vertices as
// possible. This is because, from a purely graph theoretic perspective,
// vertices are inert. Boring, even. Graphs are more about the topology, the
// characteristics of the edges connecting the points than the points
// themselves. Your use case cares about the content of your vertices, but gogl
// does not.  Consequently, anything can act as a vertex.
type Vertex interface{}

/* Atomic graph interfaces */

// A VertexEnumerator iteratively enumerates vertices.
type VertexEnumerator interface {
	EachVertex(f func(Vertex))
}

// An EdgeEnumerator iteratively enumerates edges.
type EdgeEnumerator interface {
	EachEdge(f func(Edge))
}

// An IncidentEdgeEnumerator iteratively enumerates a given vertex's incident edges.
type IncidentEdgeEnumerator interface {
	EachEdgeIncidentTo(Vertex, f func(Edge)) bool
}

// An IncidentArcEnumerator iteratively enumerates a given vertex's incident arcs (directed edges).
// One enumerator provides inbound edges, the other outbound edges.
type IncidentArcEnumerator interface {
	EachArcFrom(Vertex, f func(Edge)) bool
	EachArcTo(Vertex, f func(Edge)) bool
}

// An AdjacencyEnumerator iteratively enumerates a given vertex's adjacent vertices.
type AdjacencyEnumerator interface {
	EachAdjacent(start Vertex, f func(adjacent Vertex))
}

// A VertexMembershipChecker can indicate the presence of a vertex.
type VertexMembershipChecker interface {
	HasVertex(Vertex) bool // Whether or not the vertex is present in the set
}

// An InOutDegreeChecker reports the number of edges incident to a given vertex.
// TODO use this
type DegreeChecker interface {
	DegreeOf(Vertex) (degree int, exists bool) // Number of incident edges; if vertex is present
}

// An InOutDegreeChecker reports the number of in or out-edges a given vertex has.
type InOutDegreeChecker interface {
	InDegree(Vertex) (degree int, exists bool)  // Number of in-edges; if vertex is present
	OutDegree(Vertex) (degree int, exists bool) // Number of out-edges; if vertex is present
}

// An EdgeMembershipChecker can indicate the presence of an edge.
type EdgeMembershipChecker interface {
	HasEdge(Edge) bool
}

// A VertexSetMutator allows the addition and removal of vertices from a set.
type VertexSetMutator interface {
	EnsureVertex(...Vertex)
	RemoveVertex(...Vertex)
}

// An EdgeSetMutator allows the addition and removal of edges from a set.
type EdgeSetMutator interface {
	AddEdges(edges ...Edge)
	RemoveEdges(edges ...Edge)
}

// A WeightedEdgeSetMutator allows the addition and removal of weighted edges from a set.
type WeightedEdgeSetMutator interface {
	AddEdges(edges ...WeightedEdge)
	RemoveEdges(edges ...WeightedEdge)
}

// A LabeledEdgeSetMutator allows the addition and removal of labeled edges from a set.
type LabeledEdgeSetMutator interface {
	AddEdges(edges ...LabeledEdge)
	RemoveEdges(edges ...LabeledEdge)
}

// A PropertyEdgeSetMutator allows the addition and removal of data edges from a set.
type PropertyEdgeSetMutator interface {
	AddEdges(edges ...PropertyEdge)
	RemoveEdges(edges ...PropertyEdge)
}

// A Transposer produces a transposed version of a DirectedGraph.
type Transposer interface {
	Transpose() DirectedGraph
}

/* Aggregate graph interfaces */

// Graph is gogl's most basic interface: it contains only the methods that
// *every* type of graph implements.
//
// Graph is intentionally underspecified: both directed and undirected graphs
// implement it; simple graphs, multigraphs, weighted, labeled, or any
// combination thereof.
//
// The semantics of some of these methods vary slightly from one graph type
// to another, but in general, the basic Graph methods are supplemented, not
// superceded, by the methods in more specific interfaces.
//
// Graph is a purely read oriented interface; the various Mutable*Graph
// interfaces contain the methods for writing.
type Graph interface {
	VertexEnumerator        // Allows enumerated traversal of vertices
	EdgeEnumerator          // Allows enumerated traversal of edges
	AdjacencyEnumerator     // Allows enumerated traversal of a vertex's adjacent vertices
	VertexMembershipChecker // Allows inspection of contained vertices
	EdgeMembershipChecker   // Allows inspection of contained edges
	InOutDegreeChecker      // Reports in- and out-degree of vertices
	Order() int             // Total number of vertices in the graph
	Size() int              // Total number of edges in the graph
}

// DirectedGraph describes a Graph all of whose edges are directed.
//
// Implementing DirectedGraph is the only unambiguous signal gogl provides
// that a graph's edges are directed.
type DirectedGraph interface {
	Graph
	Transposer // DirectedGraphs can produce a transpose of themselves
}

// MutableGraph describes a graph with basic edges (no weighting, labeling, etc.)
// that can be modified freely by adding or removing vertices or edges.
type MutableGraph interface {
	Graph
	VertexSetMutator
	EdgeSetMutator
}

// A simple graph is in opposition to a multigraph: it disallows loops and
// parallel edges.
type SimpleGraph interface {
	Graph
	Density() float64
}

// A weighted graph is a graph subtype where the edges have a numeric weight;
// as described by the WeightedEdge interface, this weight is a signed int.
//
// WeightedGraphs have both the HasEdge() and HasWeightedEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its weight.
//
// HasWeightedEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge weights are the same.
type WeightedGraph interface {
	Graph
	HasWeightedEdge(e WeightedEdge) bool
	EachWeightedEdge(f func(edge WeightedEdge))
}

// MutableWeightedGraph is the mutable version of a weighted graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only weighted edges can be present in the graph.
type MutableWeightedGraph interface {
	WeightedGraph
	VertexSetMutator
	WeightedEdgeSetMutator
}

// A labeled graph is a graph subtype where the edges have an identifier;
// as described by the LabeledEdge interface, this identifier is a string.
//
// LabeledGraphs have both the HasEdge() and HasLabeledEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its label.
//
// HasLabeledEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge labels are the same.
type LabeledGraph interface {
	Graph
	HasLabeledEdge(e LabeledEdge) bool
	EachLabeledEdge(f func(edge LabeledEdge))
}

// LabeledWeightedGraph is the mutable version of a labeled graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only labeled edges can be present in the graph.
type MutableLabeledGraph interface {
	LabeledGraph
	VertexSetMutator
	LabeledEdgeSetMutator
}

// A data graph is a graph subtype where the edges carry arbitrary Go data;
// as described by the PropertyEdge interface, this identifier is an interface{}.
//
// PropertyGraphs have both the HasEdge() and HasPropertyEdge() methods.
// Correct implementations should treat the difference as a matter of strictness:
//
// HasEdge() should return true as long as an edge exists
// connecting the two given vertices (respecting directed or undirected as
// appropriate), regardless of its label.
//
// HasPropertyEdge() should return true iff an edge exists connecting the
// two given vertices (respecting directed or undirected as appropriate),
// AND if the edge data is the same. Simple comparison will typically be used
// to establish data equality, which means that using noncomparables (a slice,
// map, or non-pointer struct containing a slice or a map) for the data will
// cause a panic.
type PropertyGraph interface {
	Graph
	HasPropertyEdge(e PropertyEdge) bool
	EachPropertyEdge(f func(edge PropertyEdge))
}

// MutablePropertyGraph is the mutable version of a propety graph. Its
// AddEdges() method is incompatible with MutableGraph, guaranteeing
// only property edges can be present in the graph.
type MutablePropertyGraph interface {
	PropertyGraph
	VertexSetMutator
	PropertyEdgeSetMutator
}

/* Graph creation */

type GraphFactory func() interface{}

var Graphs = make(map[string]GraphFactory, 0)

// Creates a new graph instance.
//
// You will need to type assert the returned graph to the interface appropriate
// for your use case: Graph, DirectedGraph, MutableGraph, WeightedGraph, etc.
func New(name string) (graph interface{}, err error) {
	if _, exists := Graphs[name]; !exists {
		return nil, fmt.Errorf("No graph is registered with the name %q", name)
	}

	return Graphs[name](), nil
}

func RegisterGraph(name string, factory GraphFactory) error {
	if _, exists := Graphs[name]; exists {
		return fmt.Errorf("A graph is already registered with the name %q", name)
	}

	g := factory()

	if _, ok := g.(Graph); ok {
		return nil
	} else if _, ok := g.(WeightedGraph); ok {
		return nil
	}

	return fmt.Errorf("Value returned from factory does not implement a known Graph interface")
}

func init() {
	RegisterGraph("basic.directed", func() interface{} {
		return NewDirected()
	})
	RegisterGraph("basic.undirected", func() interface{} {
		return NewUndirected()
	})
	RegisterGraph("weighted.directed", func() interface{} {
		return NewWeightedDirected()
	})
	RegisterGraph("weighted.undirected", func() interface{} {
		return NewWeightedUndirected()
	})
}
<|endoftext|>


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

package gcscaching_test

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestIntegration(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type IntegrationTest struct {
	cache   gcscaching.StatCache
	clock   timeutil.SimulatedClock
	wrapped gcs.Bucket

	bucket gcs.Bucket
}

func init() { RegisterTestSuite(&IntegrationTest{}) }

func (t *IntegrationTest) SetUp(ti *TestInfo) {
	// Set up a fixed, non-zero time.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// Set up dependencies.
	const cacheCapacity = 100
	t.cache = gcscaching.NewStatCache(cacheCapacity)
	t.wrapped = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	t.bucket = gcscaching.NewFastStatBucket(
		ttl,
		t.cache,
		&t.clock,
		t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *IntegrationTest) StatUnknownTwice() {
	var err error
	req := &gcs.StatObjectRequest{
		Name: "taco",
	}

	// First
	_, err = t.bucket.StatObject(context.Background(), req)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Second
	_, err = t.bucket.StatObject(context.Background(), req)
	ExpectThat(err, HasSameTypeAs(&gcs.NotFoundError{}))
}

func (t *IntegrationTest) CreateThenStat() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) ListThenStat() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) CreateThenUpdateThenStat() {
	AssertFalse(true, "TODO")
}

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

package gcscaching_test

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestIntegration(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type IntegrationTest struct {
	cache   gcscaching.StatCache
	clock   timeutil.SimulatedClock
	wrapped gcs.Bucket

	bucket gcs.Bucket
}

func init() { RegisterTestSuite(&IntegrationTest{}) }

func (t *IntegrationTest) SetUp(ti *TestInfo) {
	// Set up a fixed, non-zero time.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// Set up dependencies.
	const cacheCapacity = 100
	t.cache = gcscaching.NewStatCache(cacheCapacity)
	t.wrapped = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	t.bucket = gcscaching.NewFastStatBucket(
		ttl,
		t.cache,
		&t.clock,
		t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *IntegrationTest) StatDoesntCacheNotFoundErrors() {
	const name = "taco"
	var err error

	// Stat an unknown object.
	req := &gcs.StatObjectRequest{
		Name: name,
	}

	_, err = t.bucket.StatObject(context.Background(), req)
	AssertThat(err, HasSameTypeAs(&gcs.NotFoundError{}))

	// Create the object through the back door.
	_, err = gcsutil.CreateObject(context.Background(), t.wrapped, name, "")
	AssertEq(nil, err)

	// Stat again. We should now see the object.
	o, err := t.bucket.StatObject(context.Background(), req)
	AssertEq(nil, err)
	ExpectNe(nil, o)
}

func (t *IntegrationTest) CreateThenStat() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) ListThenStat() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) CreateThenUpdateThenStat() {
	AssertFalse(true, "TODO")
}

func (t *IntegrationTest) DeleteThenStat() {
	AssertFalse(true, "TODO")
}
<|endoftext|>


<commit_before>package errors

import (
	"errors"
	"fmt"
	"strings"
)

// New returns an error that formats as the given text.
func New(text string) error {
	return errors.New(text)
}

// wrapperError satisfies the error interface.
type wrapperError struct {
	msg    string
	detail []string
	data   map[string]interface{}
	stack  []StackFrame
	root   error
}

// It satisfies the error interface.
func (e wrapperError) Error() string {
	return e.msg
}

// Root returns the original error that was wrapped by one or more
// calls to Wrap. If e does not wrap other errors, it will be returned
// as-is.
func Root(e error) error {
	if wErr, ok := e.(wrapperError); ok {
		return wErr.root
	}
	return e
}

// wrap adds a context message and stack trace to err and returns a new error
// containing the new context. This function is meant to be composed within
// other exported functions, such as Wrap and WithDetail.
// The argument stackSkip is the number of stack frames to ascend when
// generating stack straces, where 0 is the caller of wrap.
func wrap(err error, msg string, stackSkip int) error {
	if err == nil {
		return nil
	}

	werr, ok := err.(wrapperError)
	if !ok {
		werr.root = err
		werr.msg = err.Error()
		werr.stack = getStack(stackSkip+2, stackTraceSize)
	}
	if msg != "" {
		werr.msg = msg + ": " + werr.msg
	}

	return werr
}

// Wrap adds a context message and stack trace to err and returns a new error
// with the new context. Arguments are handled as in fmt.Print.
// Use Root to recover the original error wrapped by one or more calls to Wrap.
// Use Stack to recover the stack trace.
// Wrap returns nil if err is nil.
func Wrap(err error, a ...interface{}) error {
	return wrap(err, fmt.Sprint(a...), 1)
}

// Wrapf is like Wrap, but arguments are handled as in fmt.Printf.
func Wrapf(err error, format string, a ...interface{}) error {
	return wrap(err, fmt.Sprintf(format, a...), 1)
}

// WithDetail returns a new error that wraps
// err as a chain error messsage containing text
// as its additional context.
// Function Detail will return the given text
// when called on the new error value.
func WithDetail(err error, text string) error {
	if err == nil {
		return nil
	}
	if text == "" {
		return err
	}
	e1 := wrap(err, text, 1).(wrapperError)
	e1.detail = append(e1.detail, text)
	return e1
}

// WithDetailf is like WithDetail, except it formats
// the detail message as in fmt.Printf.
// Function Detail will return the formatted text
// when called on the new error value.
func WithDetailf(err error, format string, v ...interface{}) error {
	if err == nil {
		return nil
	}
	text := fmt.Sprintf(format, v...)
	e1 := wrap(err, text, 1).(wrapperError)
	e1.detail = append(e1.detail, text)
	return e1
}

// Detail returns the detail message contained in err, if any.
// An error has a detail message if it was made by WithDetail
// or WithDetailf.
func Detail(err error) string {
	wrapper, _ := err.(wrapperError)
	return strings.Join(wrapper.detail, "; ")
}

// withData returns a new error that wraps err
// as a chain error message containing v as
// an extra data item.
// Calling Data on the returned error yields v.
// Note that if err already has a data item,
// it will not be accessible via the returned error value.
func withData(err error, v map[string]interface{}) error {
	if err == nil {
		return nil
	}
	e1 := wrap(err, "", 1).(wrapperError)
	e1.data = v
	return e1
}

// WithData returns a new error that wraps err
// as a chain error message containing a value of type
// map[string]interface{} as an extra data item.
// The map contains the values in the map in err,
// if any, plus the items in keyval.
// Keyval takes the form
//   k1, v1, k2, v2, ...
// Values kN must be strings.
// Calling Data on the returned error yields the map.
// Note that if err already has a data item of any other type,
// it will not be accessible via the returned error value.
func WithData(err error, keyval ...interface{}) error {
	// TODO(kr): add vet check for odd-length keyval and non-string keys
	newkv := make(map[string]interface{})
	for k, v := range Data(err) {
		newkv[k] = v
	}
	for i := 0; i < len(keyval); i += 2 {
		newkv[keyval[i].(string)] = keyval[i+1]
	}
	return withData(err, newkv)
}

// Data returns the data item in err, if any.
func Data(err error) map[string]interface{} {
	wrapper, _ := err.(wrapperError)
	return wrapper.data
}

// Sub returns an error with all the associated data from old,
// including stack trace, detail, message, and data, but using
// new as the root error. It also wraps the newly-created error
// with the formatted error string from old.
//
// Sub returns nil when either new or old is nil.
//
// Use this when you need to substitute a new root error in place
// of an existing error that may already hold a stack trace
// or other metadata.
func Sub(new, old error) error {
	if wrapper, ok := old.(wrapperError); ok && new != nil {
		wrapper.root = Root(new)
		wrapper.msg = new.Error()
		new = wrapper
	}
	if old == nil {
		return nil
	}
	return Wrap(new, old.Error())
}
<commit_msg>errors: clarify documentation<commit_after>package errors

import (
	"errors"
	"fmt"
	"strings"
)

// New returns an error that formats as the given text.
func New(text string) error {
	return errors.New(text)
}

// wrapperError satisfies the error interface.
type wrapperError struct {
	msg    string
	detail []string
	data   map[string]interface{}
	stack  []StackFrame
	root   error
}

// It satisfies the error interface.
func (e wrapperError) Error() string {
	return e.msg
}

// Root returns the original error that was wrapped by one or more
// calls to Wrap. If e does not wrap other errors, it will be returned
// as-is.
func Root(e error) error {
	if wErr, ok := e.(wrapperError); ok {
		return wErr.root
	}
	return e
}

// wrap adds a context message and stack trace to err and returns a new error
// containing the new context. This function is meant to be composed within
// other exported functions, such as Wrap and WithDetail.
// The argument stackSkip is the number of stack frames to ascend when
// generating stack straces, where 0 is the caller of wrap.
func wrap(err error, msg string, stackSkip int) error {
	if err == nil {
		return nil
	}

	werr, ok := err.(wrapperError)
	if !ok {
		werr.root = err
		werr.msg = err.Error()
		werr.stack = getStack(stackSkip+2, stackTraceSize)
	}
	if msg != "" {
		werr.msg = msg + ": " + werr.msg
	}

	return werr
}

// Wrap adds a context message and stack trace to err and returns a new error
// with the new context. Arguments are handled as in fmt.Print.
// Use Root to recover the original error wrapped by one or more calls to Wrap.
// Use Stack to recover the stack trace.
// Wrap returns nil if err is nil.
func Wrap(err error, a ...interface{}) error {
	return wrap(err, fmt.Sprint(a...), 1)
}

// Wrapf is like Wrap, but arguments are handled as in fmt.Printf.
func Wrapf(err error, format string, a ...interface{}) error {
	return wrap(err, fmt.Sprintf(format, a...), 1)
}

// WithDetail returns a new error that wraps
// err as a chain error messsage containing text
// as its additional context.
// Function Detail will return the given text
// when called on the new error value.
func WithDetail(err error, text string) error {
	if err == nil {
		return nil
	}
	if text == "" {
		return err
	}
	e1 := wrap(err, text, 1).(wrapperError)
	e1.detail = append(e1.detail, text)
	return e1
}

// WithDetailf is like WithDetail, except it formats
// the detail message as in fmt.Printf.
// Function Detail will return the formatted text
// when called on the new error value.
func WithDetailf(err error, format string, v ...interface{}) error {
	if err == nil {
		return nil
	}
	text := fmt.Sprintf(format, v...)
	e1 := wrap(err, text, 1).(wrapperError)
	e1.detail = append(e1.detail, text)
	return e1
}

// Detail returns the detail message contained in err, if any.
// An error has a detail message if it was made by WithDetail
// or WithDetailf.
func Detail(err error) string {
	wrapper, _ := err.(wrapperError)
	return strings.Join(wrapper.detail, "; ")
}

// withData returns a new error that wraps err
// as a chain error message containing v as
// an extra data item.
// Calling Data on the returned error yields v.
// Note that if err already has a data item,
// it will not be accessible via the returned error value.
func withData(err error, v map[string]interface{}) error {
	if err == nil {
		return nil
	}
	e1 := wrap(err, "", 1).(wrapperError)
	e1.data = v
	return e1
}

// WithData returns a new error that wraps err
// as a chain error message containing a value of type
// map[string]interface{} as an extra data item.
// The map contains the values in the map in err,
// if any, plus the items in keyval.
// Keyval takes the form
//   k1, v1, k2, v2, ...
// Values kN must be strings.
// Calling Data on the returned error yields the map.
// Note that if err already has a data item of any other type,
// it will not be accessible via the returned error value.
func WithData(err error, keyval ...interface{}) error {
	// TODO(kr): add vet check for odd-length keyval and non-string keys
	newkv := make(map[string]interface{})
	for k, v := range Data(err) {
		newkv[k] = v
	}
	for i := 0; i < len(keyval); i += 2 {
		newkv[keyval[i].(string)] = keyval[i+1]
	}
	return withData(err, newkv)
}

// Data returns the data item in err, if any.
func Data(err error) map[string]interface{} {
	wrapper, _ := err.(wrapperError)
	return wrapper.data
}

// Sub returns an error containing root as its root and
// taking all other metadata (stack trace, detail, message,
// and data items) from err.
//
// Sub returns nil when either root or err is nil.
//
// Use this when you need to substitute a new root error in place
// of an existing error that may already hold a stack trace
// or other metadata.
func Sub(root, err error) error {
	if wrapper, ok := err.(wrapperError); ok && root != nil {
		wrapper.root = Root(root)
		wrapper.msg = root.Error()
		root = wrapper
	}
	if err == nil {
		return nil
	}
	return Wrap(root, err.Error())
}
<|endoftext|>


<commit_before>package rpc

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"sync"
	"time"
)

// MuxConn is able to multiplex multiple streams on top of any
// io.ReadWriteCloser. These streams act like TCP connections (Dial, Accept,
// Close, full duplex, etc.).
//
// The underlying io.ReadWriteCloser is expected to guarantee delivery
// and ordering, such as TCP. Congestion control and such aren't implemented
// by the streams, so that is also up to the underlying connection.
//
// MuxConn works using a fairly dumb multiplexing technique of simply
// framing every piece of data sent into a prefix + data format. Streams
// are established using a subset of the TCP protocol. Only a subset is
// necessary since we assume ordering on the underlying RWC.
type MuxConn struct {
	curId   uint32
	rwc     io.ReadWriteCloser
	streams map[uint32]*Stream
	mu      sync.RWMutex
	wlock   sync.Mutex
}

type muxPacketType byte

const (
	muxPacketSyn muxPacketType = iota
	muxPacketSynAck
	muxPacketAck
	muxPacketFin
	muxPacketData
)

// Create a new MuxConn around any io.ReadWriteCloser.
func NewMuxConn(rwc io.ReadWriteCloser) *MuxConn {
	m := &MuxConn{
		rwc:     rwc,
		streams: make(map[uint32]*Stream),
	}

	go m.loop()

	return m
}

// Close closes the underlying io.ReadWriteCloser. This will also close
// all streams that are open.
func (m *MuxConn) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Close all the streams
	for _, w := range m.streams {
		w.Close()
	}
	m.streams = make(map[uint32]*Stream)

	// Close the actual connection. This will also force the loop
	// to end since it'll read EOF or closed connection.
	return m.rwc.Close()
}

// Accept accepts a multiplexed connection with the given ID. This
// will block until a request is made to connect.
func (m *MuxConn) Accept(id uint32) (io.ReadWriteCloser, error) {
	stream, err := m.openStream(id)
	if err != nil {
		return nil, err
	}

	// If the stream isn't closed, then it is already open somehow
	stream.mu.Lock()
	defer stream.mu.Unlock()
	if stream.state != streamStateSynRecv && stream.state != streamStateClosed {
		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
	}

	if stream.state == streamStateClosed {
		// Go into the listening state and wait for a syn
		stream.setState(streamStateListen)
		if err := stream.waitState(streamStateSynRecv); err != nil {
			return nil, err
		}
	}

	if stream.state == streamStateSynRecv {
		// Send a syn-ack
		if _, err := m.write(stream.id, muxPacketSynAck, nil); err != nil {
			return nil, err
		}
	}

	if err := stream.waitState(streamStateEstablished); err != nil {
		return nil, err
	}

	return stream, nil
}

// Dial opens a connection to the remote end using the given stream ID.
// An Accept on the remote end will only work with if the IDs match.
func (m *MuxConn) Dial(id uint32) (io.ReadWriteCloser, error) {
	stream, err := m.openStream(id)
	if err != nil {
		return nil, err
	}

	// If the stream isn't closed, then it is already open somehow
	stream.mu.Lock()
	defer stream.mu.Unlock()
	if stream.state != streamStateClosed {
		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
	}

	// Open a connection
	if _, err := m.write(stream.id, muxPacketSyn, nil); err != nil {
		return nil, err
	}
	stream.setState(streamStateSynSent)

	if err := stream.waitState(streamStateEstablished); err != nil {
		return nil, err
	}

	m.write(id, muxPacketAck, nil)
	return stream, nil
}

// NextId returns the next available stream ID that isn't currently
// taken.
func (m *MuxConn) NextId() uint32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	for {
		result := m.curId
		m.curId++
		if _, ok := m.streams[result]; !ok {
			return result
		}
	}
}

func (m *MuxConn) openStream(id uint32) (*Stream, error) {
	// First grab a read-lock if we have the stream already we can
	// cheaply return it.
	m.mu.RLock()
	if stream, ok := m.streams[id]; ok {
		m.mu.RUnlock()
		return stream, nil
	}

	// Now acquire a full blown write lock so we can create the stream
	m.mu.RUnlock()
	m.mu.Lock()
	defer m.mu.Unlock()

	// Make sure we attempt to use the next biggest stream ID
	if id >= m.curId {
		m.curId = id + 1
	}

	// We have to check this again because there is a time period
	// above where we couldn't lost this lock.
	if stream, ok := m.streams[id]; ok {
		return stream, nil
	}

	// Create the stream object and channel where data will be sent to
	dataR, dataW := io.Pipe()
	writeCh := make(chan []byte, 256)

	// Set the data channel so we can write to it.
	stream := &Stream{
		id:          id,
		mux:         m,
		reader:      dataR,
		writeCh:     writeCh,
		stateChange: make(map[chan<- streamState]struct{}),
	}
	stream.setState(streamStateClosed)

	// Start the goroutine that will read from the queue and write
	// data out.
	go func() {
		defer dataW.Close()

		for {
			data := <-writeCh
			if data == nil {
				// A nil is a tombstone letting us know we're done
				// accepting data.
				return
			}

			if _, err := dataW.Write(data); err != nil {
				return
			}
		}
	}()

	m.streams[id] = stream
	return m.streams[id], nil
}

func (m *MuxConn) loop() {
	// Force close every stream that we know about when we exit so
	// that they all read EOF and don't block forever.
	defer func() {
		log.Printf("[INFO] Mux connection loop exiting")
		m.mu.Lock()
		defer m.mu.Unlock()
		for _, w := range m.streams {
			w.mu.Lock()
			w.remoteClose()
			w.mu.Unlock()
		}
	}()

	var id uint32
	var packetType muxPacketType
	var length int32
	for {
		if err := binary.Read(m.rwc, binary.BigEndian, &id); err != nil {
			log.Printf("[ERR] Error reading stream ID: %s", err)
			return
		}
		if err := binary.Read(m.rwc, binary.BigEndian, &packetType); err != nil {
			log.Printf("[ERR] Error reading packet type: %s", err)
			return
		}
		if err := binary.Read(m.rwc, binary.BigEndian, &length); err != nil {
			log.Printf("[ERR] Error reading length: %s", err)
			return
		}

		// TODO(mitchellh): probably would be better to re-use a buffer...
		data := make([]byte, length)
		if length > 0 {
			if _, err := m.rwc.Read(data); err != nil {
				log.Printf("[ERR] Error reading data: %s", err)
				return
			}
		}

		stream, err := m.openStream(id)
		if err != nil {
			log.Printf("[ERR] Error opening stream %d: %s", id, err)
			return
		}

		log.Printf("[TRACE] Stream %d received packet %d", id, packetType)
		switch packetType {
		case muxPacketSyn:
			stream.mu.Lock()
			switch stream.state {
			case streamStateClosed:
				fallthrough
			case streamStateListen:
				stream.setState(streamStateSynRecv)
			default:
				log.Printf("[ERR] Syn received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketAck:
			stream.mu.Lock()
			switch stream.state {
			case streamStateSynRecv:
				stream.setState(streamStateEstablished)
			case streamStateFinWait1:
				stream.setState(streamStateFinWait2)
			default:
				log.Printf("[ERR] Ack received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketSynAck:
			stream.mu.Lock()
			switch stream.state {
			case streamStateSynSent:
				stream.setState(streamStateEstablished)
			default:
				log.Printf("[ERR] SynAck received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketFin:
			stream.mu.Lock()
			switch stream.state {
			case streamStateEstablished:
				stream.setState(streamStateCloseWait)
				m.write(id, muxPacketAck, nil)

				// Close the writer on our end since we won't receive any
				// more data.
				stream.writeCh <- nil
			case streamStateFinWait1:
				fallthrough
			case streamStateFinWait2:
				stream.remoteClose()

				m.mu.Lock()
				delete(m.streams, stream.id)
				m.mu.Unlock()
			default:
				log.Printf("[ERR] Fin received for stream %d in state: %d", id, stream.state)
			}
			stream.mu.Unlock()

		case muxPacketData:
			stream.mu.Lock()
			switch stream.state {
			case streamStateFinWait1:
				fallthrough
			case streamStateFinWait2:
				fallthrough
			case streamStateEstablished:
				select {
				case stream.writeCh <- data:
				default:
					panic(fmt.Sprintf("Failed to write data, buffer full for stream %d", id))
				}
			default:
				log.Printf("[ERR] Data received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		}
	}
}

func (m *MuxConn) write(id uint32, dataType muxPacketType, p []byte) (int, error) {
	m.wlock.Lock()
	defer m.wlock.Unlock()

	if err := binary.Write(m.rwc, binary.BigEndian, id); err != nil {
		return 0, err
	}
	if err := binary.Write(m.rwc, binary.BigEndian, byte(dataType)); err != nil {
		return 0, err
	}
	if err := binary.Write(m.rwc, binary.BigEndian, int32(len(p))); err != nil {
		return 0, err
	}
	if len(p) == 0 {
		return 0, nil
	}
	return m.rwc.Write(p)
}

// Stream is a single stream of data and implements io.ReadWriteCloser.
// A Stream is full-duplex so you can write data as well as read data.
type Stream struct {
	id           uint32
	mux          *MuxConn
	reader       io.Reader
	state        streamState
	stateChange  map[chan<- streamState]struct{}
	stateUpdated time.Time
	mu           sync.Mutex
	writeCh      chan<- []byte
}

type streamState byte

const (
	streamStateClosed streamState = iota
	streamStateListen
	streamStateSynRecv
	streamStateSynSent
	streamStateEstablished
	streamStateFinWait1
	streamStateFinWait2
	streamStateCloseWait
)

func (s *Stream) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.state != streamStateEstablished && s.state != streamStateCloseWait {
		return fmt.Errorf("Stream in bad state: %d", s.state)
	}

	if s.state == streamStateEstablished {
		s.setState(streamStateFinWait1)
	} else {
		s.remoteClose()
	}

	s.mux.write(s.id, muxPacketFin, nil)
	return nil
}

func (s *Stream) Read(p []byte) (int, error) {
	return s.reader.Read(p)
}

func (s *Stream) Write(p []byte) (int, error) {
	s.mu.Lock()
	state := s.state
	s.mu.Unlock()

	if state != streamStateEstablished && state != streamStateCloseWait {
		return 0, fmt.Errorf("Stream %d in bad state to send: %d", s.id, state)
	}

	return s.mux.write(s.id, muxPacketData, p)
}

func (s *Stream) remoteClose() {
	s.setState(streamStateClosed)
	s.writeCh <- nil
}

func (s *Stream) setState(state streamState) {
	log.Printf("[TRACE] Stream %d went to state %d", s.id, state)
	s.state = state
	s.stateUpdated = time.Now().UTC()
	for ch, _ := range s.stateChange {
		select {
		case ch <- state:
		default:
		}
	}
}

func (s *Stream) waitState(target streamState) error {
	// Register a state change listener to wait for changes
	stateCh := make(chan streamState, 10)
	s.stateChange[stateCh] = struct{}{}
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		delete(s.stateChange, stateCh)
	}()

	state := <-stateCh
	if state == target {
		return nil
	} else {
		return fmt.Errorf("Stream %d went to bad state: %d", s.id, state)
	}
}
<commit_msg>packer/rpc: Clean up old streams [GH-708]<commit_after>package rpc

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"sync"
	"time"
)

// MuxConn is able to multiplex multiple streams on top of any
// io.ReadWriteCloser. These streams act like TCP connections (Dial, Accept,
// Close, full duplex, etc.).
//
// The underlying io.ReadWriteCloser is expected to guarantee delivery
// and ordering, such as TCP. Congestion control and such aren't implemented
// by the streams, so that is also up to the underlying connection.
//
// MuxConn works using a fairly dumb multiplexing technique of simply
// framing every piece of data sent into a prefix + data format. Streams
// are established using a subset of the TCP protocol. Only a subset is
// necessary since we assume ordering on the underlying RWC.
type MuxConn struct {
	curId   uint32
	rwc     io.ReadWriteCloser
	streams map[uint32]*Stream
	mu      sync.RWMutex
	wlock   sync.Mutex
	doneCh  chan struct{}
}

type muxPacketType byte

const (
	muxPacketSyn muxPacketType = iota
	muxPacketSynAck
	muxPacketAck
	muxPacketFin
	muxPacketData
)

// Create a new MuxConn around any io.ReadWriteCloser.
func NewMuxConn(rwc io.ReadWriteCloser) *MuxConn {
	m := &MuxConn{
		rwc:     rwc,
		streams: make(map[uint32]*Stream),
		doneCh:  make(chan struct{}),
	}

	go m.cleaner()
	go m.loop()

	return m
}

// Close closes the underlying io.ReadWriteCloser. This will also close
// all streams that are open.
func (m *MuxConn) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Close all the streams
	for _, w := range m.streams {
		w.Close()
	}
	m.streams = make(map[uint32]*Stream)

	// Close the actual connection. This will also force the loop
	// to end since it'll read EOF or closed connection.
	return m.rwc.Close()
}

// Accept accepts a multiplexed connection with the given ID. This
// will block until a request is made to connect.
func (m *MuxConn) Accept(id uint32) (io.ReadWriteCloser, error) {
	stream, err := m.openStream(id)
	if err != nil {
		return nil, err
	}

	// If the stream isn't closed, then it is already open somehow
	stream.mu.Lock()
	defer stream.mu.Unlock()
	if stream.state != streamStateSynRecv && stream.state != streamStateClosed {
		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
	}

	if stream.state == streamStateClosed {
		// Go into the listening state and wait for a syn
		stream.setState(streamStateListen)
		if err := stream.waitState(streamStateSynRecv); err != nil {
			return nil, err
		}
	}

	if stream.state == streamStateSynRecv {
		// Send a syn-ack
		if _, err := m.write(stream.id, muxPacketSynAck, nil); err != nil {
			return nil, err
		}
	}

	if err := stream.waitState(streamStateEstablished); err != nil {
		return nil, err
	}

	return stream, nil
}

// Dial opens a connection to the remote end using the given stream ID.
// An Accept on the remote end will only work with if the IDs match.
func (m *MuxConn) Dial(id uint32) (io.ReadWriteCloser, error) {
	stream, err := m.openStream(id)
	if err != nil {
		return nil, err
	}

	// If the stream isn't closed, then it is already open somehow
	stream.mu.Lock()
	defer stream.mu.Unlock()
	if stream.state != streamStateClosed {
		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
	}

	// Open a connection
	if _, err := m.write(stream.id, muxPacketSyn, nil); err != nil {
		return nil, err
	}
	stream.setState(streamStateSynSent)

	if err := stream.waitState(streamStateEstablished); err != nil {
		return nil, err
	}

	m.write(id, muxPacketAck, nil)
	return stream, nil
}

// NextId returns the next available stream ID that isn't currently
// taken.
func (m *MuxConn) NextId() uint32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	for {
		result := m.curId
		m.curId++
		if _, ok := m.streams[result]; !ok {
			return result
		}
	}
}

func (m *MuxConn) openStream(id uint32) (*Stream, error) {
	// First grab a read-lock if we have the stream already we can
	// cheaply return it.
	m.mu.RLock()
	if stream, ok := m.streams[id]; ok {
		m.mu.RUnlock()
		return stream, nil
	}

	// Now acquire a full blown write lock so we can create the stream
	m.mu.RUnlock()
	m.mu.Lock()
	defer m.mu.Unlock()

	// Make sure we attempt to use the next biggest stream ID
	if id >= m.curId {
		m.curId = id + 1
	}

	// We have to check this again because there is a time period
	// above where we couldn't lost this lock.
	if stream, ok := m.streams[id]; ok {
		return stream, nil
	}

	// Create the stream object and channel where data will be sent to
	dataR, dataW := io.Pipe()
	writeCh := make(chan []byte, 256)

	// Set the data channel so we can write to it.
	stream := &Stream{
		id:          id,
		mux:         m,
		reader:      dataR,
		writeCh:     writeCh,
		stateChange: make(map[chan<- streamState]struct{}),
	}
	stream.setState(streamStateClosed)

	// Start the goroutine that will read from the queue and write
	// data out.
	go func() {
		defer dataW.Close()

		for {
			data := <-writeCh
			if data == nil {
				// A nil is a tombstone letting us know we're done
				// accepting data.
				return
			}

			if _, err := dataW.Write(data); err != nil {
				return
			}
		}
	}()

	m.streams[id] = stream
	return m.streams[id], nil
}

func (m *MuxConn) cleaner() {
	for {
		done := false
		select {
		case <-time.After(500 * time.Millisecond):
		case <-m.doneCh:
			done = true
		}

		m.mu.Lock()
		for id, s := range m.streams {
			s.mu.Lock()
			if s.state == streamStateClosed {
				delete(m.streams, id)
			}
			s.mu.Unlock()
		}

		if done {
			for _, s := range m.streams {
				s.mu.Lock()
				s.closeWriter()
				s.mu.Unlock()
			}
		}
		m.mu.Unlock()

		if done {
			return
		}
	}
}

func (m *MuxConn) loop() {
	// Force close every stream that we know about when we exit so
	// that they all read EOF and don't block forever.
	defer func() {
		log.Printf("[INFO] Mux connection loop exiting")
		close(m.doneCh)
	}()

	var id uint32
	var packetType muxPacketType
	var length int32
	for {
		if err := binary.Read(m.rwc, binary.BigEndian, &id); err != nil {
			log.Printf("[ERR] Error reading stream ID: %s", err)
			return
		}
		if err := binary.Read(m.rwc, binary.BigEndian, &packetType); err != nil {
			log.Printf("[ERR] Error reading packet type: %s", err)
			return
		}
		if err := binary.Read(m.rwc, binary.BigEndian, &length); err != nil {
			log.Printf("[ERR] Error reading length: %s", err)
			return
		}

		// TODO(mitchellh): probably would be better to re-use a buffer...
		data := make([]byte, length)
		if length > 0 {
			if _, err := m.rwc.Read(data); err != nil {
				log.Printf("[ERR] Error reading data: %s", err)
				return
			}
		}

		stream, err := m.openStream(id)
		if err != nil {
			log.Printf("[ERR] Error opening stream %d: %s", id, err)
			return
		}

		log.Printf("[TRACE] Stream %d received packet %d", id, packetType)
		switch packetType {
		case muxPacketSyn:
			stream.mu.Lock()
			switch stream.state {
			case streamStateClosed:
				fallthrough
			case streamStateListen:
				stream.setState(streamStateSynRecv)
			default:
				log.Printf("[ERR] Syn received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketAck:
			stream.mu.Lock()
			switch stream.state {
			case streamStateSynRecv:
				stream.setState(streamStateEstablished)
			case streamStateFinWait1:
				stream.setState(streamStateFinWait2)
			case streamStateLastAck:
				stream.closeWriter()
				fallthrough
			case streamStateClosing:
				stream.setState(streamStateClosed)
			default:
				log.Printf("[ERR] Ack received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketSynAck:
			stream.mu.Lock()
			switch stream.state {
			case streamStateSynSent:
				stream.setState(streamStateEstablished)
			default:
				log.Printf("[ERR] SynAck received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		case muxPacketFin:
			stream.mu.Lock()
			switch stream.state {
			case streamStateEstablished:
				stream.closeWriter()
				stream.setState(streamStateCloseWait)
				m.write(id, muxPacketAck, nil)
			case streamStateFinWait2:
				stream.closeWriter()
				stream.setState(streamStateClosed)
				m.write(id, muxPacketAck, nil)
			case streamStateFinWait1:
				stream.closeWriter()
				stream.setState(streamStateClosing)
				m.write(id, muxPacketAck, nil)
			default:
				log.Printf("[ERR] Fin received for stream %d in state: %d", id, stream.state)
			}
			stream.mu.Unlock()

		case muxPacketData:
			stream.mu.Lock()
			switch stream.state {
			case streamStateFinWait1:
				fallthrough
			case streamStateFinWait2:
				fallthrough
			case streamStateEstablished:
				select {
				case stream.writeCh <- data:
				default:
					panic(fmt.Sprintf("Failed to write data, buffer full for stream %d", id))
				}
			default:
				log.Printf("[ERR] Data received for stream in state: %d", stream.state)
			}
			stream.mu.Unlock()
		}
	}
}

func (m *MuxConn) write(id uint32, dataType muxPacketType, p []byte) (int, error) {
	m.wlock.Lock()
	defer m.wlock.Unlock()

	if err := binary.Write(m.rwc, binary.BigEndian, id); err != nil {
		return 0, err
	}
	if err := binary.Write(m.rwc, binary.BigEndian, byte(dataType)); err != nil {
		return 0, err
	}
	if err := binary.Write(m.rwc, binary.BigEndian, int32(len(p))); err != nil {
		return 0, err
	}
	if len(p) == 0 {
		return 0, nil
	}
	return m.rwc.Write(p)
}

// Stream is a single stream of data and implements io.ReadWriteCloser.
// A Stream is full-duplex so you can write data as well as read data.
type Stream struct {
	id           uint32
	mux          *MuxConn
	reader       io.Reader
	state        streamState
	stateChange  map[chan<- streamState]struct{}
	stateUpdated time.Time
	mu           sync.Mutex
	writeCh      chan<- []byte
}

type streamState byte

const (
	streamStateClosed streamState = iota
	streamStateListen
	streamStateSynRecv
	streamStateSynSent
	streamStateEstablished
	streamStateFinWait1
	streamStateFinWait2
	streamStateCloseWait
	streamStateClosing
	streamStateLastAck
)

func (s *Stream) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.state != streamStateEstablished && s.state != streamStateCloseWait {
		return fmt.Errorf("Stream in bad state: %d", s.state)
	}

	if s.state == streamStateEstablished {
		s.setState(streamStateFinWait1)
	} else {
		s.setState(streamStateLastAck)
	}

	s.mux.write(s.id, muxPacketFin, nil)
	return nil
}

func (s *Stream) Read(p []byte) (int, error) {
	return s.reader.Read(p)
}

func (s *Stream) Write(p []byte) (int, error) {
	s.mu.Lock()
	state := s.state
	s.mu.Unlock()

	if state != streamStateEstablished && state != streamStateCloseWait {
		return 0, fmt.Errorf("Stream %d in bad state to send: %d", s.id, state)
	}

	return s.mux.write(s.id, muxPacketData, p)
}

func (s *Stream) closeWriter() {
	s.writeCh <- nil
}

func (s *Stream) setState(state streamState) {
	log.Printf("[TRACE] Stream %d went to state %d", s.id, state)
	s.state = state
	s.stateUpdated = time.Now().UTC()
	for ch, _ := range s.stateChange {
		select {
		case ch <- state:
		default:
		}
	}
}

func (s *Stream) waitState(target streamState) error {
	// Register a state change listener to wait for changes
	stateCh := make(chan streamState, 10)
	s.stateChange[stateCh] = struct{}{}
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		delete(s.stateChange, stateCh)
	}()

	state := <-stateCh
	if state == target {
		return nil
	} else {
		return fmt.Errorf("Stream %d went to bad state: %d", s.id, state)
	}
}
<|endoftext|>