<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"runtime"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/rgbkrk/juno"

	fsnotify "gopkg.in/fsnotify.v1"

	"github.com/gorilla/websocket"
)

// WatchRuntimes watches for new running kernels. This is a placeholder.
func WatchRuntimes() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer watcher.Close()

	done := make(chan bool)

	go func() {
		for {
			select {
			case event := <-watcher.Events:
				log.Println("event:", event)
				if event.Op&fsnotify.Write == fsnotify.Write {
					log.Println("modified file:", event.Name)
				}
			case err := <-watcher.Errors:
				log.Println("error:", err)
			}
		}
	}()

	runtimeDir, err := RuntimeDir()
	if err != nil {
		return err
	}

	err = watcher.Add(runtimeDir)
	if err != nil {
		return err
	}
	<-done

	return err
}

// HomeDir is the home directory for the user. Light wrapper on the homedir lib.
func HomeDir() (string, error) {
	homish, err := homedir.Dir()
	if err != nil {
		return "", fmt.Errorf("Unable to acquire home directory: %v", err)
	}

	home, err := homedir.Expand(homish)
	if err != nil {
		return "", fmt.Errorf("Unable to expand home directory: %v", err)
	}

	return home, nil
}

// ConfigDir gets the Jupyter config directory for this platform and user.
//
// Returns JUPYTER_CONFIG_DIR if defined, else ~/.jupyter
func ConfigDir() (string, error) {
	home, err := HomeDir()
	if err != nil {
		return "", err
	}

	if os.Getenv("JUPYTER_CONFIG_DIR") != "" {
		return os.Getenv("JUPYTER_CONFIG_DIR"), nil
	}

	return path.Join(home, ".jupyter"), nil
}

// DataDir gets the config directory for Jupyter data files. These are
// non-transient, non-configuration files.
//
// Returns JUPYTER_DATA_DIR if defined, else a platform-appropriate path.
func DataDir() (string, error) {
	home, err := HomeDir()
	if err != nil {
		return "", err
	}

	if runtime.GOOS == "darwin" {
		return path.Join(home, "Library", "Jupyter"), nil
	} else if runtime.GOOS == "windows" {
		// Modern Windows
		appdata := os.Getenv("APPDATA")
		if appdata != "" {
			return path.Join(appdata, "jupyter"), nil
		}

		// TODO: jupyter_config_dir() from python
		configDir, err := ConfigDir()
		if err != nil {
			return "", err
		}
		return path.Join(configDir, "data"), nil

	} else { // *nix
		xdg := os.Getenv("XDG_DATA_HOME")
		if xdg != "" {
			xdg = path.Join(home, ".local", "share")
		}
		return path.Join(xdg, "jupyter"), nil
	}

}

// IsNixButNotDarwin returns true if on *nix (but not on OS X)
func IsNixButNotDarwin() bool {
	return runtime.GOOS == "linux" ||
		runtime.GOOS == "dragonfly" || // So many BSDs
		runtime.GOOS == "freebsd" ||
		runtime.GOOS == "netbsd" ||
		runtime.GOOS == "openbsd" ||
		runtime.GOOS == "solaris"
}

// RuntimeDir is the directory of running kernels
func RuntimeDir() (string, error) {
	if os.Getenv("JUPYTER_RUNTIME_DIR") != "" {
		return os.Getenv("JUPYTER_RUNTIME_DIR"), nil
	}

	if IsNixButNotDarwin() && os.Getenv("XDG_RUNTIME_DIR") != "" {
		return path.Join(os.Getenv("XDG_RUNTIME_DIR"), "jupyter"), nil
	}

	dataDir, err := DataDir()
	if err != nil {
		return "", err
	}

	return path.Join(dataDir, "runtime"), nil

}

func main() {
	connFile := flag.String("existing", "", "Path to connection file")
	ioloHub := flag.String("hub", "http://127.0.0.1:8080", "IOLO Hub Base URL")
	flag.Parse()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Println()
	}

	if *connFile == "" {
		flag.Usage()
		log.Fatalln("Need a connection file.")
	}

	// Expects a runtime kernel-*.json
	connInfo, err := juno.OpenConnectionFile(*connFile)

	if err != nil {
		log.Fatalf("%v\n", err)
	}

	iopub, err := juno.NewIOPubSocket(connInfo, "")

	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't start the iopub socket: %v", err)
	}

	defer iopub.Close()

	u, err := url.Parse(*ioloHub)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to parse URL: %v", err)
		os.Exit(3)
	}

	u.Path = "/ws"

	rawConn, err := net.Dial("tcp", u.Host)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to connect to host: %v", err)
		os.Exit(4)
	}

	wsHeaders := http.Header{
		"Origin":                   {u.Scheme + "://" + u.Host},
		"Sec-WebSocket-Extensions": {"permessage-deflate; client_max_window_bits, x-webkit-deflate-frame"},
	}

	writeBufferSize := 1024 * 1024 * 2

	wsConn, resp, err := websocket.NewClient(rawConn, u, wsHeaders, 0, writeBufferSize)
	if err != nil {
		fmt.Fprintf(os.Stderr, "websocket.NewClient Error: %s\nResp:%+v", err, resp)
		os.Exit(5)
	}
	defer wsConn.Close()

	for {
		message, err := iopub.ReadMessage()

		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			continue
		}

		err = wsConn.WriteJSON(message)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Problem writing JSON %v.\nOutput too big?", err)
			break
		}
	}

}
<commit_msg>Set default to iolo.fict.io<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"runtime"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/rgbkrk/juno"

	fsnotify "gopkg.in/fsnotify.v1"

	"github.com/gorilla/websocket"
)

// WatchRuntimes watches for new running kernels. This is a placeholder.
func WatchRuntimes() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer watcher.Close()

	done := make(chan bool)

	go func() {
		for {
			select {
			case event := <-watcher.Events:
				log.Println("event:", event)
				if event.Op&fsnotify.Write == fsnotify.Write {
					log.Println("modified file:", event.Name)
				}
			case err := <-watcher.Errors:
				log.Println("error:", err)
			}
		}
	}()

	runtimeDir, err := RuntimeDir()
	if err != nil {
		return err
	}

	err = watcher.Add(runtimeDir)
	if err != nil {
		return err
	}
	<-done

	return err
}

// HomeDir is the home directory for the user. Light wrapper on the homedir lib.
func HomeDir() (string, error) {
	homish, err := homedir.Dir()
	if err != nil {
		return "", fmt.Errorf("Unable to acquire home directory: %v", err)
	}

	home, err := homedir.Expand(homish)
	if err != nil {
		return "", fmt.Errorf("Unable to expand home directory: %v", err)
	}

	return home, nil
}

// ConfigDir gets the Jupyter config directory for this platform and user.
//
// Returns JUPYTER_CONFIG_DIR if defined, else ~/.jupyter
func ConfigDir() (string, error) {
	home, err := HomeDir()
	if err != nil {
		return "", err
	}

	if os.Getenv("JUPYTER_CONFIG_DIR") != "" {
		return os.Getenv("JUPYTER_CONFIG_DIR"), nil
	}

	return path.Join(home, ".jupyter"), nil
}

// DataDir gets the config directory for Jupyter data files. These are
// non-transient, non-configuration files.
//
// Returns JUPYTER_DATA_DIR if defined, else a platform-appropriate path.
func DataDir() (string, error) {
	home, err := HomeDir()
	if err != nil {
		return "", err
	}

	if runtime.GOOS == "darwin" {
		return path.Join(home, "Library", "Jupyter"), nil
	} else if runtime.GOOS == "windows" {
		// Modern Windows
		appdata := os.Getenv("APPDATA")
		if appdata != "" {
			return path.Join(appdata, "jupyter"), nil
		}

		// TODO: jupyter_config_dir() from python
		configDir, err := ConfigDir()
		if err != nil {
			return "", err
		}
		return path.Join(configDir, "data"), nil

	} else { // *nix
		xdg := os.Getenv("XDG_DATA_HOME")
		if xdg != "" {
			xdg = path.Join(home, ".local", "share")
		}
		return path.Join(xdg, "jupyter"), nil
	}

}

// IsNixButNotDarwin returns true if on *nix (but not on OS X)
func IsNixButNotDarwin() bool {
	return runtime.GOOS == "linux" ||
		runtime.GOOS == "dragonfly" || // So many BSDs
		runtime.GOOS == "freebsd" ||
		runtime.GOOS == "netbsd" ||
		runtime.GOOS == "openbsd" ||
		runtime.GOOS == "solaris"
}

// RuntimeDir is the directory of running kernels
func RuntimeDir() (string, error) {
	if os.Getenv("JUPYTER_RUNTIME_DIR") != "" {
		return os.Getenv("JUPYTER_RUNTIME_DIR"), nil
	}

	if IsNixButNotDarwin() && os.Getenv("XDG_RUNTIME_DIR") != "" {
		return path.Join(os.Getenv("XDG_RUNTIME_DIR"), "jupyter"), nil
	}

	dataDir, err := DataDir()
	if err != nil {
		return "", err
	}

	return path.Join(dataDir, "runtime"), nil

}

func main() {
	connFile := flag.String("existing", "", "Path to connection file")
	ioloHub := flag.String("hub", "http://iolo.fict.io:80", "IOLO Hub Base URL")
	flag.Parse()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Println()
	}

	if *connFile == "" {
		flag.Usage()
		log.Fatalln("Need a connection file.")
	}

	// Expects a runtime kernel-*.json
	connInfo, err := juno.OpenConnectionFile(*connFile)

	if err != nil {
		log.Fatalf("%v\n", err)
	}

	iopub, err := juno.NewIOPubSocket(connInfo, "")

	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't start the iopub socket: %v", err)
	}

	defer iopub.Close()

	u, err := url.Parse(*ioloHub)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to parse URL: %v", err)
		os.Exit(3)
	}

	u.Path = "/ws"

	rawConn, err := net.Dial("tcp", u.Host)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to connect to host: %v", err)
		os.Exit(4)
	}

	wsHeaders := http.Header{
		"Origin":                   {u.Scheme + "://" + u.Host},
		"Sec-WebSocket-Extensions": {"permessage-deflate; client_max_window_bits, x-webkit-deflate-frame"},
	}

	writeBufferSize := 1024 * 1024 * 2

	wsConn, resp, err := websocket.NewClient(rawConn, u, wsHeaders, 0, writeBufferSize)
	if err != nil {
		fmt.Fprintf(os.Stderr, "websocket.NewClient Error: %s\nResp:%+v", err, resp)
		os.Exit(5)
	}
	defer wsConn.Close()

	for {
		message, err := iopub.ReadMessage()

		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			continue
		}

		err = wsConn.WriteJSON(message)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Problem writing JSON %v.\nOutput too big?", err)
			break
		}
	}

}
<|endoftext|>


<commit_before>package server

import (
	"../bitset"
	"../hashtree"
	"../network"
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

var ERROR_NOT_LOCAL = errors.New("file is not locally available")
var ERROR_LEVEL_LOW = errors.New("the inner hash level is lower than cached")
var ERROR_INDEX_OFF = errors.New("the index of inner hashes is out of range for file of this size")
var ERROR_ALREADY_EXIST = errors.New("the thing was already added or started")

type FileState int

const (
	FILE_UNKNOW FileState = iota
	FILE_PART
	FILE_COMPLETE
)

type Database interface {
	LowestInnerHashes() hashtree.Level
	ImportFromReader(r io.Reader) network.StaticId
	GetState(id network.StaticId) FileState
	GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error)
	GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error)
	StartPart(id network.StaticId) error
	PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error
	PutInnerHashes(id network.StaticId, set network.InnerHashes) (has hashtree.Nodes, complete bool, err error)
	Remove(id network.StaticId)
	Close()
}

func ImportLocalFile(d Database, location string) (id network.StaticId) {
	f, _ := os.Open(location)
	r := bufio.NewReader(f)
	id = d.ImportFromReader(r)
	f.Close()
	return
}

type simpleDatabase struct {
	datafolder        *os.File
	dirname           string
	lowestInnerHashes hashtree.Level
}

func OpenSimpleDatabase(dirname string, lowestInnerHashes hashtree.Level) Database {
	os.MkdirAll(dirname, 0777)
	dir, err := os.Open(dirname)
	if err != nil {
		panic(err)
	}
	d := &simpleDatabase{
		datafolder:        dir,
		dirname:           dirname,
		lowestInnerHashes: lowestInnerHashes,
	}

	return d
}

func (d *simpleDatabase) Close() {
	err := d.datafolder.Close()
	if err != nil {
		panic(err)
	}
	d.datafolder = nil
}

func (d *simpleDatabase) LowestInnerHashes() hashtree.Level {
	return d.lowestInnerHashes
}

var refHash = hashtree.NewFile()

func (d *simpleDatabase) hashNumber(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum + n)
}
func (d *simpleDatabase) hashTopNumber(leafs hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	l := hashtree.Levels(leafs)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum)
}

func (d *simpleDatabase) hashPosition(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	return d.hashNumber(leafs, l, n) * int64(refHash.Size())
}

func (d *simpleDatabase) innerHashListenerFile(hasher hashtree.HashTree, len hashtree.Bytes) *os.File {
	leafs := hasher.Nodes(len)
	top := hasher.Levels(leafs) - 1
	file, err := ioutil.TempFile(d.dirname, "listener-")
	if err != nil {
		panic(err)
	}
	listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
		//TODO: don't save levels lower than needed
		if l == top {
			return //don't need the root here
		}
		b := h.ToBytes()
		off := d.hashPosition(leafs, l, i)
		if _, err := file.WriteAt(b, off); err != nil {
			panic(err)
		}
	}
	hasher.SetInnerHashListener(listener)
	return file
}

func (d *simpleDatabase) ImportFromReader(r io.Reader) network.StaticId {
	f, err := ioutil.TempFile(d.dirname, "import-")
	if err != nil {
		panic(err)
	}
	len, err2 := io.Copy(f, r)
	if err2 != nil {
		panic(err2)
	}
	hasher := hashtree.NewFile()
	hashFile := d.innerHashListenerFile(hasher, hashtree.Bytes(len))

	f.Seek(0, os.SEEK_SET)
	io.Copy(hasher, f)
	id := network.StaticId{
		Hash:   hasher.Sum(nil),
		Length: &len,
	}
	f.Close()
	hashFile.Close()

	err = os.Rename(f.Name(), d.fileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(f.Name())
		} else {
			panic(err)
		}
	}
	err = os.Rename(hashFile.Name(), d.hashFileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(hashFile.Name())
		} else {
			panic(err)
		}
	}
	return id
}

func (d *simpleDatabase) GetState(id network.StaticId) FileState {
	_, err := os.Stat(d.fileNameForId(id))
	if os.IsNotExist(err) {
		_, err = os.Stat(d.hashFileNameForId(id))
		if os.IsNotExist(err) {
			return FILE_UNKNOW
		} else {
			return FILE_PART
		}
	} else {
		return FILE_COMPLETE
	}
}

func (d *simpleDatabase) GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error) {
	f, err := os.Open(d.fileNameForId(id))
	if err != nil {
		return 0, ERROR_NOT_LOCAL
	}
	defer f.Close()
	return f.ReadAt(b, int64(off))
}

func (d *simpleDatabase) GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error) {
	leafs := refHash.Nodes(hashtree.Bytes(id.GetLength()))
	level := hashtree.Level(req.GetHeight())
	from := hashtree.Nodes(req.GetFrom())
	nodes := hashtree.Nodes(req.GetLength())
	if nodes == 0 {
		return req, nil //nothing was requested
	} else if level < d.lowestInnerHashes {
		return req, ERROR_LEVEL_LOW
	} else if level >= refHash.Levels(leafs)-1 || from+nodes > refHash.LevelWidth(leafs, level) {
		return req, ERROR_INDEX_OFF
	}
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return req, ERROR_NOT_LOCAL
	}
	defer f.Close()
	off := d.hashPosition(leafs, level, from)
	b := make([]byte, refHash.Size()*int(nodes))
	if _, err := f.ReadAt(b, off); err != nil {
		panic(err)
	}
	req.Hashes = b
	return req, nil
}

func (d *simpleDatabase) StartPart(id network.StaticId) error {
	_, err := os.Stat(d.hashFileNameForId(id))
	if os.IsNotExist(err) {
		pf, err := os.Create(d.partFileNameForId(id))
		if err != nil {
			return err
		}
		defer pf.Close()
		hf, err2 := os.Create(d.hashFileNameForId(id))
		if err2 != nil {
			return err2
		}
		defer hf.Close()
		return nil
	} else {
		return ERROR_ALREADY_EXIST
	}

}
func (d *simpleDatabase) PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error {
	return nil
}

func (d *simpleDatabase) PutInnerHashes(id network.StaticId, set network.InnerHashes) (has hashtree.Nodes, complete bool, err error) {
	leafs := id.Blocks()
	bits := bitset.OpenCountingFileBacked(d.haveHashNameForId(id), int(d.hashTopNumber(leafs)-1))
	defer bits.Close()
	f, err := os.OpenFile(d.hashFileNameForId(id), os.O_RDWR, 0666)
	if err != nil {
		return 0, false, ERROR_NOT_LOCAL
	}
	defer f.Close()

	writeHash := func(realL hashtree.Level, realN hashtree.Nodes, b []byte) {
		off := d.hashPosition(leafs, realL, realN)
		if _, err := f.WriteAt(b, off); err != nil {
			panic(err)
		}
		n := int(d.hashNumber(leafs, realL, realN))
		bits.Set(n)
	}

	hashBuffer := make([]byte, refHash.Size())
	splited := set.SplitLocalSummable(&id)
	for _, hashes := range splited {
		rootL, rootN := hashes.LocalRoot(leafs)
		key := int(d.hashNumber(leafs, rootL, rootN))
		if key == bits.Capacity() {
			//this is root
		} else if !bits.Get(key) {
			continue // this part of hashes can not be verified, skiped
		}
		sum := hashes.LocalSum()
		if key == bits.Capacity() {
			copy(hashBuffer, id.GetHash())
		} else {
			off := d.hashPosition(leafs, rootL, rootN)
			if _, err := f.ReadAt(hashBuffer, off); err != nil {
				panic(err)
			}
		}
		if bytes.Equal(sum, hashBuffer) {
			//verified, now save
			listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
				realL := set.GetHeightL() + l
				realN := i + rootN<<uint32(rootL-l)
				if realL == rootL {
					return //don't need the root here, already verified
				}
				b := h.ToBytes()
				writeHash(realL, realN, b)

				//propagate down nodes with single branch
				for realL > d.LowestInnerHashes() &&
					realN+1 == hashtree.LevelWidth(leafs, realL) &&
					hashtree.LevelWidth(leafs, realL-1)%2 == 1 {

					realL--
					realN = hashtree.LevelWidth(leafs, realL) - 1
					writeHash(realL, realN, b)
				}
			}
			hasher := hashtree.NewTree()
			hasher.SetInnerHashListener(listener)
			hasher.Write(hashes.GetHashes())
			hasher.Sum(nil)
		}
	}
	return hashtree.Nodes(bits.Count()), bits.Full(), nil
}

func remove(filename string) {
	err := os.Remove(filename)
	if err != nil && !os.IsNotExist(err) {
		panic(err)
	}
}

func (d *simpleDatabase) Remove(id network.StaticId) {
	remove(d.havePartNameForId(id))
	remove(d.haveHashNameForId(id))
	remove(d.partFileNameForId(id))
	remove(d.fileNameForId(id))
	remove(d.hashFileNameForId(id))
}

func (d *simpleDatabase) fileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/F-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) hashFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/H-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) partFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/P-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) haveHashNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hH-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) havePartNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hP-%s", d.datafolder.Name(), id.CompactId())
}
<commit_msg>partial implement PutAt, also needs testing<commit_after>package server

import (
	"../bitset"
	"../hashtree"
	"../network"
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

var ERROR_NOT_LOCAL = errors.New("file is not locally available")
var ERROR_NOT_PART = errors.New("file is not downloading")
var ERROR_LEVEL_LOW = errors.New("the inner hash level is lower than cached")
var ERROR_INDEX_OFF = errors.New("the index of inner hashes is out of range for file of this size")
var ERROR_ALREADY_EXIST = errors.New("the thing was already added or started")

type FileState int

const (
	FILE_UNKNOW FileState = iota
	FILE_PART
	FILE_COMPLETE
)

type Database interface {
	LowestInnerHashes() hashtree.Level
	ImportFromReader(r io.Reader) network.StaticId
	GetState(id network.StaticId) FileState
	GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error)
	GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error)
	StartPart(id network.StaticId) error
	PutAt(b []byte, id network.StaticId, off hashtree.Bytes) (has hashtree.Nodes, complete bool, err error)
	PutInnerHashes(id network.StaticId, set network.InnerHashes) (has hashtree.Nodes, complete bool, err error)
	Remove(id network.StaticId)
	Close()
}

func ImportLocalFile(d Database, location string) (id network.StaticId) {
	f, _ := os.Open(location)
	r := bufio.NewReader(f)
	id = d.ImportFromReader(r)
	f.Close()
	return
}

type simpleDatabase struct {
	datafolder        *os.File
	dirname           string
	lowestInnerHashes hashtree.Level
}

func OpenSimpleDatabase(dirname string, lowestInnerHashes hashtree.Level) Database {
	os.MkdirAll(dirname, 0777)
	dir, err := os.Open(dirname)
	if err != nil {
		panic(err)
	}
	d := &simpleDatabase{
		datafolder:        dir,
		dirname:           dirname,
		lowestInnerHashes: lowestInnerHashes,
	}

	return d
}

func (d *simpleDatabase) Close() {
	err := d.datafolder.Close()
	if err != nil {
		panic(err)
	}
	d.datafolder = nil
}

func (d *simpleDatabase) LowestInnerHashes() hashtree.Level {
	return d.lowestInnerHashes
}

var refHash = hashtree.NewFile()

func (d *simpleDatabase) hashNumber(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum + n)
}
func (d *simpleDatabase) hashTopNumber(leafs hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	l := hashtree.Levels(leafs)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum)
}

func (d *simpleDatabase) hashPosition(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	return d.hashNumber(leafs, l, n) * int64(refHash.Size())
}

func (d *simpleDatabase) innerHashListenerFile(hasher hashtree.HashTree, len hashtree.Bytes) *os.File {
	leafs := hasher.Nodes(len)
	top := hasher.Levels(leafs) - 1
	file, err := ioutil.TempFile(d.dirname, "listener-")
	if err != nil {
		panic(err)
	}
	listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
		//TODO: don't save levels lower than needed
		if l == top {
			return //don't need the root here
		}
		b := h.ToBytes()
		off := d.hashPosition(leafs, l, i)
		if _, err := file.WriteAt(b, off); err != nil {
			panic(err)
		}
	}
	hasher.SetInnerHashListener(listener)
	return file
}

func (d *simpleDatabase) ImportFromReader(r io.Reader) network.StaticId {
	f, err := ioutil.TempFile(d.dirname, "import-")
	if err != nil {
		panic(err)
	}
	len, err2 := io.Copy(f, r)
	if err2 != nil {
		panic(err2)
	}
	hasher := hashtree.NewFile()
	hashFile := d.innerHashListenerFile(hasher, hashtree.Bytes(len))

	f.Seek(0, os.SEEK_SET)
	io.Copy(hasher, f)
	id := network.StaticId{
		Hash:   hasher.Sum(nil),
		Length: &len,
	}
	f.Close()
	hashFile.Close()

	err = os.Rename(f.Name(), d.fileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(f.Name())
		} else {
			panic(err)
		}
	}
	err = os.Rename(hashFile.Name(), d.hashFileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(hashFile.Name())
		} else {
			panic(err)
		}
	}
	return id
}

func (d *simpleDatabase) GetState(id network.StaticId) FileState {
	_, err := os.Stat(d.fileNameForId(id))
	if os.IsNotExist(err) {
		_, err = os.Stat(d.hashFileNameForId(id))
		if os.IsNotExist(err) {
			return FILE_UNKNOW
		} else {
			return FILE_PART
		}
	} else {
		return FILE_COMPLETE
	}
}

func (d *simpleDatabase) GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error) {
	f, err := os.Open(d.fileNameForId(id))
	if err != nil {
		return 0, ERROR_NOT_LOCAL
	}
	defer f.Close()
	return f.ReadAt(b, int64(off))
}

func (d *simpleDatabase) GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error) {
	leafs := refHash.Nodes(hashtree.Bytes(id.GetLength()))
	level := hashtree.Level(req.GetHeight())
	from := hashtree.Nodes(req.GetFrom())
	nodes := hashtree.Nodes(req.GetLength())
	if nodes == 0 {
		return req, nil //nothing was requested
	} else if level < d.lowestInnerHashes {
		return req, ERROR_LEVEL_LOW
	} else if level >= refHash.Levels(leafs)-1 || from+nodes > refHash.LevelWidth(leafs, level) {
		return req, ERROR_INDEX_OFF
	}
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return req, ERROR_NOT_LOCAL
	}
	defer f.Close()
	off := d.hashPosition(leafs, level, from)
	b := make([]byte, refHash.Size()*int(nodes))
	if _, err := f.ReadAt(b, off); err != nil {
		panic(err)
	}
	req.Hashes = b
	return req, nil
}

func (d *simpleDatabase) StartPart(id network.StaticId) error {
	_, err := os.Stat(d.hashFileNameForId(id))
	if os.IsNotExist(err) {
		pf, err := os.Create(d.partFileNameForId(id))
		if err != nil {
			return err
		}
		defer pf.Close()
		hf, err2 := os.Create(d.hashFileNameForId(id))
		if err2 != nil {
			return err2
		}
		defer hf.Close()
		return nil
	} else {
		return ERROR_ALREADY_EXIST
	}

}
func (d *simpleDatabase) PutAt(b []byte, id network.StaticId, off hashtree.Bytes) (has hashtree.Nodes, complete bool, err error) {
	blocks := id.Blocks()
	f, err := os.OpenFile(d.partFileNameForId(id), os.O_RDWR, 0666)
	if err != nil {
		return 0, false, ERROR_NOT_PART
	}
	defer f.Close()
	bits := bitset.OpenCountingFileBacked(d.havePartNameForId(id), int(blocks))
	defer bits.Close()

	if off%hashtree.FILE_BLOCK_SIZE != 0 {
		panic("offset must start from a block")
	}
	startingNode := hashtree.Nodes(off / hashtree.FILE_BLOCK_SIZE)

	if off >= id.Bytes() {
		b = nil
	} else if off+hashtree.Bytes(len(b)) >= id.Bytes() {
		b = b[:id.Bytes()-off]
	}
	bLen := hashtree.Bytes(len(b))

	hash := hashtree.NewFile()

	for i := hashtree.Bytes(0); i < bLen; i += hashtree.FILE_BLOCK_SIZE {
		end := i + hashtree.FILE_BLOCK_SIZE
		if end > bLen {
			end = bLen
		}
		partData := b[i:end]
		hash.Write(partData)
		sum := hash.Sum(nil)
		hash.Reset()
		nthBlock := startingNode + hashtree.Nodes(i/hashtree.FILE_BLOCK_SIZE)
		exp, err := d.GetInnerHashes(id, network.NewInnerHashes(0, nthBlock, 1, nil))
		if err != nil {
			panic(err)
		}
		if bytes.Equal(exp.GetHashes(), sum) {
			_, err := f.WriteAt(partData, int64(nthBlock*hashtree.FILE_BLOCK_SIZE))
			if err != nil {
				panic(err)
			}
			bits.Set(int(nthBlock))
		}
	}
	if bits.Full() {
		//todo: move file to completed
	}
	return hashtree.Nodes(bits.Count()), bits.Full(), nil
}

func (d *simpleDatabase) PutInnerHashes(id network.StaticId, set network.InnerHashes) (has hashtree.Nodes, complete bool, err error) {
	leafs := id.Blocks()
	f, err := os.OpenFile(d.hashFileNameForId(id), os.O_RDWR, 0666)
	if err != nil {
		return 0, false, ERROR_NOT_PART
	}
	defer f.Close()
	bits := bitset.OpenCountingFileBacked(d.haveHashNameForId(id), int(d.hashTopNumber(leafs)-1))
	defer bits.Close()

	writeHash := func(realL hashtree.Level, realN hashtree.Nodes, b []byte) {
		off := d.hashPosition(leafs, realL, realN)
		if _, err := f.WriteAt(b, off); err != nil {
			panic(err)
		}
		n := int(d.hashNumber(leafs, realL, realN))
		bits.Set(n)
	}

	hashBuffer := make([]byte, refHash.Size())
	splited := set.SplitLocalSummable(&id)
	for _, hashes := range splited {
		rootL, rootN := hashes.LocalRoot(leafs)
		key := int(d.hashNumber(leafs, rootL, rootN))
		if key == bits.Capacity() {
			//this is root
		} else if !bits.Get(key) {
			continue // this part of hashes can not be verified, skiped
		}
		sum := hashes.LocalSum()
		if key == bits.Capacity() {
			copy(hashBuffer, id.GetHash())
		} else {
			off := d.hashPosition(leafs, rootL, rootN)
			if _, err := f.ReadAt(hashBuffer, off); err != nil {
				panic(err)
			}
		}
		if bytes.Equal(sum, hashBuffer) {
			//verified, now save
			listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
				realL := set.GetHeightL() + l
				realN := i + rootN<<uint32(rootL-l)
				if realL == rootL {
					return //don't need the root here, already verified
				}
				b := h.ToBytes()
				writeHash(realL, realN, b)

				//propagate down nodes with single branch
				for realL > d.LowestInnerHashes() &&
					realN+1 == hashtree.LevelWidth(leafs, realL) &&
					hashtree.LevelWidth(leafs, realL-1)%2 == 1 {

					realL--
					realN = hashtree.LevelWidth(leafs, realL) - 1
					writeHash(realL, realN, b)
				}
			}
			hasher := hashtree.NewTree()
			hasher.SetInnerHashListener(listener)
			hasher.Write(hashes.GetHashes())
			hasher.Sum(nil)
		}
	}
	return hashtree.Nodes(bits.Count()), bits.Full(), nil
}

func remove(filename string) {
	err := os.Remove(filename)
	if err != nil && !os.IsNotExist(err) {
		panic(err)
	}
}

func (d *simpleDatabase) Remove(id network.StaticId) {
	remove(d.havePartNameForId(id))
	remove(d.haveHashNameForId(id))
	remove(d.partFileNameForId(id))
	remove(d.fileNameForId(id))
	remove(d.hashFileNameForId(id))
}

func (d *simpleDatabase) fileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/F-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) hashFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/H-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) partFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/P-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) haveHashNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hH-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) havePartNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hP-%s", d.datafolder.Name(), id.CompactId())
}
<|endoftext|>


<commit_before>package shardkv

import (
  "container/list"
  "sync"
  "time"
  "shardmaster"
  "strconv"
  "labix.org/v2/mgo"
  "labix.org/v2/mgo/bson"
  "fmt"
)

//
// A simple LRU cache implementation which supports
// basic put and get operations.
//
type ShardCache struct {
  mu sync.Mutex

  // recent usage list and lookup table
  list *list.List // tracks recent usage
  table map[int]map[string]*list.Element // table for fast lookup

  // approximation of current size
  size uint64

  // cache size limitation
  capacity uint64
}

type entry struct {
  key string
  value string
  time_accessed time.Time
}

func MakeCache(capacity uint64) *ShardCache {
  cache := &ShardCache{
    list:     list.New(),
    table:    make(map[int](map[string]*list.Element)),
    capacity: capacity,
  }
  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
  return cache
}

func (cache *ShardCache) Size() uint64 {
  return cache.size
}

func (cache *ShardCache) Get(key string) (v string, ok bool) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    cache.add(key, value)
    return "", false
  }
  cache.touch(element)
  return element.Value.(*entry).value, true
}

func (cache *ShardCache) Put(key string, value string) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element != nil {
    cache.update(element, value)
  } else {
    cache.add(key, value)
  }
}

func (cache *ShardCache) Delete(key string) bool {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return false
  }

  cache.list.Remove(element)
  delete(cache.table[shard], key)
  cache.size -= uint64(len(element.Value.(*entry).value) + len(element.Value.(*entry).key))
  return true
}

func (cache *ShardCache) Clear() {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  // replace old state with initial state
  // old references are garbage collected
  cache.list.Init()
  cache.table = make(map[int]map[string]*list.Element)
  cache.size = 0

  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
}

func (cache *ShardCache) ClearShard(shard int) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  cache.table[shard] = make(map[string]*list.Element)
}

func (cache *ShardCache) Keys() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  keys := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    keys = append(keys, e.Value.(*entry).key)
  }
  return keys
}

func (cache *ShardCache) Values() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  values := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    values = append(values, e.Value.(*entry).value)
  }
  return values
}

func (cache *ShardCache) KVPairs() []KVPair {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  kvpairs := make([]KVPair, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    kvpairs = append(kvpairs, KVPair{key2shard(e.Value.(*entry).key), e.Value.(*entry).key, e.Value.(*entry).value})
  }
  return kvpairs
}

func (cache *ShardCache) update(element *list.Element, value string) {
  valueSize := len(value)
  sizeDiff := valueSize - len(element.Value.(*entry).value)
  element.Value.(*entry).value = value
  cache.size += uint64(sizeDiff)
  cache.touch(element)
  cache.maintainSize()
}

func (cache *ShardCache) touch(element *list.Element) {
  cache.list.MoveToFront(element)
  element.Value.(*entry).time_accessed = time.Now()
}

func (cache *ShardCache) add(key string, value string) {
  shard := key2shard(key)
  newEntry := &entry{key, value, time.Now()}
  element := cache.list.PushFront(newEntry)
  cache.table[shard][key] = element
  cache.size += uint64(len(newEntry.value) + len(newEntry.key))
  cache.maintainSize()
}

func (cache *ShardCache) maintainSize() {
  //deleted := list.New()
  for cache.size > cache.capacity {
    delElem := cache.list.Back()
    delEntry := delElem.Value.(*entry)
    shard := key2shard(delEntry.key)
    delete(cache.table[shard], delEntry.key)
    cache.list.Remove(delElem)
    cache.size -= uint64(len(delEntry.value) + len(delEntry.key))
  }
}

type Storage struct {
  mu sync.Mutex

  cache *ShardCache

  // disk storage state
  dbSession *mgo.Session
  db *mgo.Collection
  snapshots *mgo.Collection
  dedupsnaps *mgo.Collection

  // logging for background writes
  writeLog map[int]WriteOp
  applied int
}

type WriteOp struct {
  shard int
  key string
  value string
  dbok bool
}

func (st *Storage) makeCache(capacity uint64) {
  st.cache = MakeCache(capacity)
}

func (st *Storage) connectToDiskDB(url string) {
  var err error
  st.dbSession, err = mgo.Dial(url)
  if err != nil {
    panic(err)
  }
  st.db = st.dbSession.DB("db").C("kvstore")
  st.snapshots = st.dbSession.DB("db").C("snapshots")
}

func (st *Storage) DBClear() {
  st.db.RemoveAll(bson.M{})
  check, _ := st.snapshots.Find(bson.M{}).Count()
  if (check > 0) {
    st.snapshots.RemoveAll(bson.M{})
    st.dedupsnaps.RemoveAll(bson.M{})
  }
}

func (st *Storage) CacheClear() {
  st.cache.Clear()
}

func (st *Storage) Clear() {
  st.DBClear()
  st.CacheClear()
}

func MakeStorage(capacity uint64, dbURL string) *Storage {
  st := new(Storage)
  st.makeCache(capacity)
  st.connectToDiskDB(dbURL)
  st.writeLog = make(map[int]WriteOp)

  fmt.Printf("Making storage...\n")

  go st.writeInBackground()
  return st
}

func (st *Storage) Get(key string, shardNum int) string {

  value, ok := st.cache.Get(key)
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      ok = false
    } else {
      value = result.Value
    }
  }
  return value
}

func (st *Storage) Put(key string, value string, doHash bool, shardNum int) string {
  prev, ok := st.cache.Get(key)
  var dbok bool
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      dbok = false
    } else {
      prev = result.Value
      dbok = true
    }
  }

  //var deleted *list.List // list of cache values removed to send to disk
  if doHash {
    toBeHashed := prev + value
    hash := strconv.Itoa(int(hash(toBeHashed)))
    st.cache.Put(key, hash)
    //deleted = st.cache.Put(key, hash)
  } else {
    st.cache.Put(key, value)
    //deleted = st.cache.Put(key, value)
  }

  st.writeLog[st.applied] = WriteOp{shardNum, key, value, dbok}
  st.applied++

  // insert removed cache entries one at a time into DB, possibly faster if done together?
  /*for e := st.cache.list.Front(); e != nil; e = e.Next() {
    entry := e.(*entry)
    shard := key2shard(entry.key)
    err := st.db.Insert(&KVPair{shard, entry.key, entry.value})
    if err != nil {
      panic(err)
    }
  }*/
  return prev
}

func (st *Storage) CreateSnapshot(confignum int, dedup map[string]ClientReply) {
  cachedata := st.cache.KVPairs()
  results := []KVPair{}
  index := 0
  for len(results) >= 100 {
    st.db.Find(bson.M{}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
    for i := 0; i < len(results); i++ {
      st.db.Insert(&SnapshotKV{confignum, results[i].Shard, results[i].Key, results[i].Value, false})
    }
    index++
  }
  for i := 0; i < len(cachedata); i++ {
    st.db.Insert(&SnapshotKV{confignum, cachedata[i].Shard, cachedata[i].Key, cachedata[i].Value, true})
  }
  for key, value := range dedup {
    st.dedupsnaps.Insert(&SnapshotDedup{confignum, key, value.Value, value.Err, value.Counter})
  }
}

func (st *Storage) ReadSnapshotDB(confignum int, shardnum int, index int, cache bool) (p map[string]string, fin bool) {
  piece := make(map[string]string)
  results := []SnapshotKV{}
  if cache {
    st.snapshots.Find(bson.M{"cache": "true", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  } else {
    st.snapshots.Find(bson.M{"cache": "false", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  }
  for i := 0; i < len(results); i++ {
    piece[results[i].Key] = results[i].Value
  }
  if len(results) < GrabSize {
    return piece, true
  }
  return piece, false
}

func (st *Storage) ReadSnapshotDedup(confignum int) map[string]ClientReply {
  dedup := make(map[string]ClientReply)
  results := []SnapshotDedup{}
  st.dedupsnaps.Find(bson.M{"config": confignum}).All(&results)
  for i := 0; i < len(results); i++ {
    dedup[results[i].Key] = ClientReply{results[i].Value, results[i].Err, results[i].Counter}
  }
  return dedup
}

func (st *Storage) writeInBackground() {
  current := 0
  for {
    if st.applied > current {
      currentWrite := st.writeLog[current]
      var err error
      if !currentWrite.dbok {
        err = st.db.Insert(&KVPair{currentWrite.shard, currentWrite.key, currentWrite.value})
      } else {
        err = st.db.Update(bson.M{"shard": currentWrite.shard, "key": currentWrite.key}, bson.M{"$set": bson.M{"value": currentWrite.value}})
      }
      if err != nil {
        panic(err)
      }
      delete(st.writeLog, current)
      current++
    }
    time.Sleep(25 * time.Millisecond)
  }
}

func (st *Storage) closeDBConnection() {
  st.dbSession.Close()
}
<commit_msg>Fixed storage creation problems<commit_after>package shardkv

import (
  "container/list"
  "sync"
  "time"
  "shardmaster"
  "strconv"
  "labix.org/v2/mgo"
  "labix.org/v2/mgo/bson"
  "fmt"
)

//
// A simple LRU cache implementation which supports
// basic put and get operations.
//
type ShardCache struct {
  mu sync.Mutex

  // recent usage list and lookup table
  list *list.List // tracks recent usage
  table map[int]map[string]*list.Element // table for fast lookup

  // approximation of current size
  size uint64

  // cache size limitation
  capacity uint64
}

type entry struct {
  key string
  value string
  time_accessed time.Time
}

func MakeCache(capacity uint64) *ShardCache {
  cache := &ShardCache{
    list:     list.New(),
    table:    make(map[int](map[string]*list.Element)),
    capacity: capacity,
  }
  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
  return cache
}

func (cache *ShardCache) Size() uint64 {
  return cache.size
}

func (cache *ShardCache) Get(key string) (v string, ok bool) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return "", false
  }
  cache.touch(element)
  return element.Value.(*entry).value, true
}

func (cache *ShardCache) Put(key string, value string) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element != nil {
    cache.update(element, value)
  } else {
    cache.add(key, value)
  }
}

func (cache *ShardCache) Delete(key string) bool {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  shard := key2shard(key)
  element := cache.table[shard][key]
  if element == nil {
    return false
  }

  cache.list.Remove(element)
  delete(cache.table[shard], key)
  cache.size -= uint64(len(element.Value.(*entry).value) + len(element.Value.(*entry).key))
  return true
}

func (cache *ShardCache) Clear() {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  // replace old state with initial state
  // old references are garbage collected
  cache.list.Init()
  cache.table = make(map[int]map[string]*list.Element)
  cache.size = 0

  for i := 0; i < shardmaster.NShards; i++ {
    cache.table[i] = make(map[string]*list.Element)
  }
}

func (cache *ShardCache) ClearShard(shard int) {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  cache.table[shard] = make(map[string]*list.Element)
}

func (cache *ShardCache) Keys() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  keys := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    keys = append(keys, e.Value.(*entry).key)
  }
  return keys
}

func (cache *ShardCache) Values() []string {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  values := make([]string, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    values = append(values, e.Value.(*entry).value)
  }
  return values
}

func (cache *ShardCache) KVPairs() []KVPair {
  cache.mu.Lock()
  defer cache.mu.Unlock()

  kvpairs := make([]KVPair, 0, cache.list.Len())
  for e := cache.list.Front(); e != nil; e = e.Next() {
    kvpairs = append(kvpairs, KVPair{key2shard(e.Value.(*entry).key), e.Value.(*entry).key, e.Value.(*entry).value})
  }
  return kvpairs
}

func (cache *ShardCache) update(element *list.Element, value string) {
  valueSize := len(value)
  sizeDiff := valueSize - len(element.Value.(*entry).value)
  element.Value.(*entry).value = value
  cache.size += uint64(sizeDiff)
  cache.touch(element)
  cache.maintainSize()
}

func (cache *ShardCache) touch(element *list.Element) {
  cache.list.MoveToFront(element)
  element.Value.(*entry).time_accessed = time.Now()
}

func (cache *ShardCache) add(key string, value string) {
  shard := key2shard(key)
  newEntry := &entry{key, value, time.Now()}
  element := cache.list.PushFront(newEntry)
  cache.table[shard][key] = element
  cache.size += uint64(len(newEntry.value) + len(newEntry.key))
  cache.maintainSize()
}

func (cache *ShardCache) maintainSize() {
  //deleted := list.New()
  for cache.size > cache.capacity {
    delElem := cache.list.Back()
    delEntry := delElem.Value.(*entry)
    shard := key2shard(delEntry.key)
    delete(cache.table[shard], delEntry.key)
    cache.list.Remove(delElem)
    cache.size -= uint64(len(delEntry.value) + len(delEntry.key))
  }
}

type Storage struct {
  mu sync.Mutex

  cache *ShardCache

  // disk storage state
  dbSession *mgo.Session
  db *mgo.Collection
  snapshots *mgo.Collection
  dedupsnaps *mgo.Collection

  // logging for background writes
  writeLog map[int]WriteOp
  applied int
}

type WriteOp struct {
  shard int
  key string
  value string
  dbok bool
}

func (st *Storage) makeCache(capacity uint64) {
  st.cache = MakeCache(capacity)
}

func (st *Storage) connectToDiskDB(url string) {
  var err error
  st.dbSession, err = mgo.Dial(url)
  if err != nil {
    panic(err)
  }
  st.db = st.dbSession.DB("db").C("kvstore")
  st.snapshots = st.dbSession.DB("db").C("snapshots")
  st.dedupsnaps = st.dbSession.DB("db").C("dedupsnaps")
}

func (st *Storage) DBClear() {
  st.db.RemoveAll(bson.M{})
  check, _ := st.snapshots.Find(bson.M{}).Count()
  if (check > 0) {
    st.snapshots.RemoveAll(bson.M{})
    st.dedupsnaps.RemoveAll(bson.M{})
  }
}

func (st *Storage) CacheClear() {
  st.cache.Clear()
}

func (st *Storage) Clear() {
  st.DBClear()
  st.CacheClear()
}

func MakeStorage(capacity uint64, dbURL string) *Storage {
  st := new(Storage)
  st.makeCache(capacity)
  st.connectToDiskDB(dbURL)
  st.writeLog = make(map[int]WriteOp)

  fmt.Printf("Making storage...\n")

  go st.writeInBackground()
  return st
}

func (st *Storage) Get(key string, shardNum int) string {

  value, ok := st.cache.Get(key)
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      ok = false
    } else {
      st.cache.Put(key, result.Value)
      value = result.Value
    }
  }
  return value
}

func (st *Storage) Put(key string, value string, doHash bool, shardNum int) string {
  prev, ok := st.cache.Get(key)
  var dbok bool
  if !ok {
    result := KVPair{}
    err := st.db.Find(bson.M{"shard": shardNum, "key": key}).One(&result)
    if err != nil {
      dbok = false
    } else {
      prev = result.Value
      dbok = true
    }
  }

  //var deleted *list.List // list of cache values removed to send to disk
  if doHash {
    toBeHashed := prev + value
    hash := strconv.Itoa(int(hash(toBeHashed)))
    st.cache.Put(key, hash)
    //deleted = st.cache.Put(key, hash)
  } else {
    st.cache.Put(key, value)
    //deleted = st.cache.Put(key, value)
  }

  st.writeLog[st.applied] = WriteOp{shardNum, key, value, dbok}
  st.applied++

  // insert removed cache entries one at a time into DB, possibly faster if done together?
  /*for e := st.cache.list.Front(); e != nil; e = e.Next() {
    entry := e.(*entry)
    shard := key2shard(entry.key)
    err := st.db.Insert(&KVPair{shard, entry.key, entry.value})
    if err != nil {
      panic(err)
    }
  }*/
  return prev
}

func (st *Storage) CreateSnapshot(confignum int, dedup map[string]ClientReply) {
  cachedata := st.cache.KVPairs()
  results := []KVPair{}
  index := 0
  for len(results) >= 100 {
    st.db.Find(bson.M{}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
    for i := 0; i < len(results); i++ {
      st.db.Insert(&SnapshotKV{confignum, results[i].Shard, results[i].Key, results[i].Value, false})
    }
    index++
  }
  for i := 0; i < len(cachedata); i++ {
    st.db.Insert(&SnapshotKV{confignum, cachedata[i].Shard, cachedata[i].Key, cachedata[i].Value, true})
  }
  for key, value := range dedup {
    st.dedupsnaps.Insert(&SnapshotDedup{confignum, key, value.Value, value.Err, value.Counter})
  }
}

func (st *Storage) ReadSnapshotDB(confignum int, shardnum int, index int, cache bool) (p map[string]string, fin bool) {
  piece := make(map[string]string)
  results := []SnapshotKV{}
  if cache {
    st.snapshots.Find(bson.M{"cache": "true", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  } else {
    st.snapshots.Find(bson.M{"cache": "false", "config": confignum}).Skip(index * GrabSize).Limit(GrabSize).All(&results)
  }
  for i := 0; i < len(results); i++ {
    piece[results[i].Key] = results[i].Value
  }
  if len(results) < GrabSize {
    return piece, true
  }
  return piece, false
}

func (st *Storage) ReadSnapshotDedup(confignum int) map[string]ClientReply {
  dedup := make(map[string]ClientReply)
  results := []SnapshotDedup{}
  st.dedupsnaps.Find(bson.M{"config": confignum}).All(&results)
  for i := 0; i < len(results); i++ {
    dedup[results[i].Key] = ClientReply{results[i].Value, results[i].Err, results[i].Counter}
  }
  return dedup
}

func (st *Storage) writeInBackground() {
  current := 0
  for {
    if st.applied > current {
      currentWrite := st.writeLog[current]
      var err error
      if !currentWrite.dbok {
        err = st.db.Insert(&KVPair{currentWrite.shard, currentWrite.key, currentWrite.value})
      } else {
        err = st.db.Update(bson.M{"shard": currentWrite.shard, "key": currentWrite.key}, bson.M{"$set": bson.M{"value": currentWrite.value}})
      }
      if err != nil {
        panic(err)
      }
      delete(st.writeLog, current)
      current++
    }
    time.Sleep(25 * time.Millisecond)
  }
}

func (st *Storage) closeDBConnection() {
  st.dbSession.Close()
}
<|endoftext|>


<commit_before>// Package terminal provides a tty emulation and session handlings that is
// supported via Screen
package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"sync"
	"syscall"
	"time"
	"unicode/utf8"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/terminal/pty"
)

type Terminal struct {
	InputHook    func()
	Log          kite.Logger
	screenrcPath string

	Users      map[string]*User
	sync.Mutex // protects Users
}

func New(log kite.Logger, screenPath string) *Terminal {
	return &Terminal{
		Users: make(map[string]*User),
		Log:   log,
	}
}

func (t *Terminal) HasLimit(username string) bool {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		// not available yet so it's good to go
		return false
	}

	return user.HasLimit()
}

// AddUserSession adds the given username and session
func (t *Terminal) AddUserSession(username, session string, server *Server) {
	t.Lock()
	defer t.Unlock()

	// check if it's exists, if not go and lazy initialize a new user instance
	var user *User
	var ok bool
	user, ok = t.Users[username]
	if !ok {
		user = NewUser(username)
	}

	user.AddSession(session, server)
	t.Users[username] = user
}

// DeleteUserSession deletes the given session from the Users map
func (t *Terminal) DeleteUserSession(username, session string) {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		return // nothing to do
	}

	user.DeleteSession(session)

	if len(user.Sessions) == 0 {
		delete(t.Users, username)
	}
}

// RenameUserSession renames the given users session in the Users map
func (t *Terminal) RenameUserSession(username, oldName, newName string) error {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		return errors.New("username not available")
	}

	user.RenameSession(oldName, newName)
	t.Users[username] = user

	return nil
}

// KillSession kills the given screen session
func (t *Terminal) KillSession(r *kite.Request) (interface{}, error) {
	var params struct {
		Session string
	}

	if r.Args.One().Unmarshal(&params) != nil {
		return nil, errors.New("{ session: [string] }")
	}

	if params.Session == "" {
		return nil, errors.New("session is empty")
	}

	if err := killSession(params.Session); err != nil {
		return nil, err
	}

	t.DeleteUserSession(r.Username, params.Session)

	return true, nil
}

// KillSessions kills all available screen sessions
func (t *Terminal) KillSessions(r *kite.Request) (interface{}, error) {
	if err := killSessions(r.Username); err != nil {
		return nil, err
	}

	return true, nil
}

// RenameSession renames the given session to the new session name
func (t *Terminal) RenameSession(r *kite.Request) (interface{}, error) {
	var params struct {
		OldName string `json:"oldName"`
		NewName string `json:"newName"`
	}

	if r.Args.One().Unmarshal(&params) != nil {
		return nil, errors.New("{ oldName: [string] newName: [string] }")
	}

	if params.OldName == "" {
		return nil, errors.New("session name empty")
	}

	if params.NewName == "" {
		return nil, errors.New("session name to be renamed is empty")
	}

	if err := renameSession(params.OldName, params.NewName); err != nil {
		return nil, err
	}

	t.RenameUserSession(r.Username, params.OldName, params.NewName)

	return true, nil
}

// GetSessions return a list of curren active screen sessions
func (t *Terminal) GetSessions(r *kite.Request) (interface{}, error) {
	user, err := user.Current()
	if err != nil {
		return nil, fmt.Errorf("Could not get home dir: %s", err)
	}

	sessions := screenSessions(user.Username)
	if len(sessions) == 0 {
		return nil, errors.New("no sessions available")
	}

	return sessions, nil
}

// Connect creates and open a new TTY instance. It returns a *Server instance
// so every caller can send and receive from the connected TTY end.
func (t *Terminal) Connect(r *kite.Request) (interface{}, error) {
	var params struct {
		Remote       Remote
		Session      string
		SizeX, SizeY int
		Mode         string
	}

	if err := r.Args.One().Unmarshal(&params); err != nil {
		return nil, fmt.Errorf("{ remote: [object], session: %s, noScreen: [bool] }, err: %s",
			params.Session, err)
	}

	if params.SizeX <= 0 || params.SizeY <= 0 {
		return nil, fmt.Errorf("{ sizeX: %d, sizeY: %d } { raw JSON : %v }", params.SizeX, params.SizeY, r.Args.One())
	}

	user, err := user.Current()
	if err != nil {
		return nil, fmt.Errorf("Could not get home dir: %s", err)
	}

	if params.Mode == "create" && t.HasLimit(r.Username) {
		return nil, errors.New("session limit has reached")
	}

	command, err := newCommand(params.Mode, params.Session, user.Username)
	if err != nil {
		return nil, err
	}

	// get pty and tty descriptors
	p, err := pty.NewPTY()
	if err != nil {
		return nil, err
	}

	// We will return this object to the client.
	server := &Server{
		Session:   command.Session,
		remote:    params.Remote,
		pty:       p,
		inputHook: t.InputHook,
	}
	server.setSize(float64(params.SizeX), float64(params.SizeY))

	t.AddUserSession(r.Username, command.Session, server)

	// wrap the command with sudo -i for initiation login shell. This is needed
	// in order to have Environments and other to be initialized correctly.
	// check also if klient was started in root mode or not.
	var args []string
	if os.Geteuid() == 0 {
		args = []string{"-i", command.Name}
	} else {
		args = []string{"-i", "-u", "#" + user.Uid, "--", command.Name}
	}

	// check if we have custom screenrc path and there is a file for it. If yes
	// use it for screen binary otherwise it'll just start without any screenrc.
	if t.screenrcPath != "" {
		if _, err := os.Stat(t.screenrcPath); err == nil {
			args = append(args, "-c", t.screenrcPath)
		}
	}

	args = append(args, command.Args...)
	cmd := exec.Command("/usr/bin/sudo", args...)

	// For test use this, sudo is not going to work
	// cmd := exec.Command(command.Name, command.Args...)

	cmd.Env = []string{"TERM=xterm-256color", "HOME=" + user.HomeDir}
	cmd.Stdin = server.pty.Slave
	cmd.Stdout = server.pty.Slave
	cmd.Dir = user.HomeDir
	// cmd.Stderr = server.pty.Slave

	// Open in background, this is needed otherwise the process will be killed
	// if you hit close on the client side.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setctty: true, Setsid: true}
	err = cmd.Start()
	if err != nil {
		fmt.Println("could not start", err)
	}

	// Wait until the shell process is closed and notify the client.
	go func() {
		err := cmd.Wait()
		if err != nil {
			fmt.Println("cmd.wait err", err)
		}

		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded.Call()

		t.DeleteUserSession(r.Username, command.Session)
	}()

	// Read the STDOUT from shell process and send to the connected client.
	go func() {
		buf := make([]byte, (4096)-utf8.UTFMax, 4096)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			// Rate limiting...
			if server.throttling {
				s := time.Now().Unix()
				if server.currentSecond != s {
					server.currentSecond = s
					server.messageCounter = 0
					server.byteCounter = 0
					server.lineFeeedCounter = 0
				}
				server.messageCounter += 1
				server.byteCounter += n
				server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
				if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
					time.Sleep(time.Second)
				}
			}

			server.remote.Output.Call(string(filterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	return server, nil
}

// CloseSessions closes all active session for the given username and deletes
// it from the internal user map
func (t *Terminal) CloseSessions(username string) {
	t.Log.Info("Closing terminal sessions of user '%s'", username)

	t.Lock()
	user, ok := t.Users[username]
	t.Unlock()

	if !ok {
		return
	}

	user.CloseSessions()

	t.Lock()
	delete(t.Users, username)
	t.Unlock()
}

func filterInvalidUTF8(buf []byte) []byte {
	i := 0
	j := 0
	for {
		r, l := utf8.DecodeRune(buf[i:])
		if l == 0 {
			break
		}
		if r < 0xD800 {
			if i != j {
				copy(buf[j:], buf[i:i+l])
			}
			j += l
		}
		i += l
	}
	return buf[:j]
}
<commit_msg>terminal: pass screenpath<commit_after>// Package terminal provides a tty emulation and session handlings that is
// supported via Screen
package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"sync"
	"syscall"
	"time"
	"unicode/utf8"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/terminal/pty"
)

type Terminal struct {
	InputHook    func()
	Log          kite.Logger
	screenrcPath string

	Users      map[string]*User
	sync.Mutex // protects Users
}

func New(log kite.Logger, screenPath string) *Terminal {
	return &Terminal{
		Users:        make(map[string]*User),
		screenrcPath: screenPath,
		Log:          log,
	}
}

func (t *Terminal) HasLimit(username string) bool {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		// not available yet so it's good to go
		return false
	}

	return user.HasLimit()
}

// AddUserSession adds the given username and session
func (t *Terminal) AddUserSession(username, session string, server *Server) {
	t.Lock()
	defer t.Unlock()

	// check if it's exists, if not go and lazy initialize a new user instance
	var user *User
	var ok bool
	user, ok = t.Users[username]
	if !ok {
		user = NewUser(username)
	}

	user.AddSession(session, server)
	t.Users[username] = user
}

// DeleteUserSession deletes the given session from the Users map
func (t *Terminal) DeleteUserSession(username, session string) {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		return // nothing to do
	}

	user.DeleteSession(session)

	if len(user.Sessions) == 0 {
		delete(t.Users, username)
	}
}

// RenameUserSession renames the given users session in the Users map
func (t *Terminal) RenameUserSession(username, oldName, newName string) error {
	t.Lock()
	defer t.Unlock()

	user, ok := t.Users[username]
	if !ok {
		return errors.New("username not available")
	}

	user.RenameSession(oldName, newName)
	t.Users[username] = user

	return nil
}

// KillSession kills the given screen session
func (t *Terminal) KillSession(r *kite.Request) (interface{}, error) {
	var params struct {
		Session string
	}

	if r.Args.One().Unmarshal(&params) != nil {
		return nil, errors.New("{ session: [string] }")
	}

	if params.Session == "" {
		return nil, errors.New("session is empty")
	}

	if err := killSession(params.Session); err != nil {
		return nil, err
	}

	t.DeleteUserSession(r.Username, params.Session)

	return true, nil
}

// KillSessions kills all available screen sessions
func (t *Terminal) KillSessions(r *kite.Request) (interface{}, error) {
	if err := killSessions(r.Username); err != nil {
		return nil, err
	}

	return true, nil
}

// RenameSession renames the given session to the new session name
func (t *Terminal) RenameSession(r *kite.Request) (interface{}, error) {
	var params struct {
		OldName string `json:"oldName"`
		NewName string `json:"newName"`
	}

	if r.Args.One().Unmarshal(&params) != nil {
		return nil, errors.New("{ oldName: [string] newName: [string] }")
	}

	if params.OldName == "" {
		return nil, errors.New("session name empty")
	}

	if params.NewName == "" {
		return nil, errors.New("session name to be renamed is empty")
	}

	if err := renameSession(params.OldName, params.NewName); err != nil {
		return nil, err
	}

	t.RenameUserSession(r.Username, params.OldName, params.NewName)

	return true, nil
}

// GetSessions return a list of curren active screen sessions
func (t *Terminal) GetSessions(r *kite.Request) (interface{}, error) {
	user, err := user.Current()
	if err != nil {
		return nil, fmt.Errorf("Could not get home dir: %s", err)
	}

	sessions := screenSessions(user.Username)
	if len(sessions) == 0 {
		return nil, errors.New("no sessions available")
	}

	return sessions, nil
}

// Connect creates and open a new TTY instance. It returns a *Server instance
// so every caller can send and receive from the connected TTY end.
func (t *Terminal) Connect(r *kite.Request) (interface{}, error) {
	var params struct {
		Remote       Remote
		Session      string
		SizeX, SizeY int
		Mode         string
	}

	if err := r.Args.One().Unmarshal(&params); err != nil {
		return nil, fmt.Errorf("{ remote: [object], session: %s, noScreen: [bool] }, err: %s",
			params.Session, err)
	}

	if params.SizeX <= 0 || params.SizeY <= 0 {
		return nil, fmt.Errorf("{ sizeX: %d, sizeY: %d } { raw JSON : %v }", params.SizeX, params.SizeY, r.Args.One())
	}

	user, err := user.Current()
	if err != nil {
		return nil, fmt.Errorf("Could not get home dir: %s", err)
	}

	if params.Mode == "create" && t.HasLimit(r.Username) {
		return nil, errors.New("session limit has reached")
	}

	command, err := newCommand(params.Mode, params.Session, user.Username)
	if err != nil {
		return nil, err
	}

	// get pty and tty descriptors
	p, err := pty.NewPTY()
	if err != nil {
		return nil, err
	}

	// We will return this object to the client.
	server := &Server{
		Session:   command.Session,
		remote:    params.Remote,
		pty:       p,
		inputHook: t.InputHook,
	}
	server.setSize(float64(params.SizeX), float64(params.SizeY))

	t.AddUserSession(r.Username, command.Session, server)

	// wrap the command with sudo -i for initiation login shell. This is needed
	// in order to have Environments and other to be initialized correctly.
	// check also if klient was started in root mode or not.
	var args []string
	if os.Geteuid() == 0 {
		args = []string{"-i", command.Name}
	} else {
		args = []string{"-i", "-u", "#" + user.Uid, "--", command.Name}
	}

	// check if we have custom screenrc path and there is a file for it. If yes
	// use it for screen binary otherwise it'll just start without any screenrc.
	if t.screenrcPath != "" {
		if _, err := os.Stat(t.screenrcPath); err == nil {
			args = append(args, "-c", t.screenrcPath)
		}
	}

	args = append(args, command.Args...)
	cmd := exec.Command("/usr/bin/sudo", args...)

	// For test use this, sudo is not going to work
	// cmd := exec.Command(command.Name, command.Args...)

	cmd.Env = []string{"TERM=xterm-256color", "HOME=" + user.HomeDir}
	cmd.Stdin = server.pty.Slave
	cmd.Stdout = server.pty.Slave
	cmd.Dir = user.HomeDir
	// cmd.Stderr = server.pty.Slave

	// Open in background, this is needed otherwise the process will be killed
	// if you hit close on the client side.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setctty: true, Setsid: true}
	err = cmd.Start()
	if err != nil {
		fmt.Println("could not start", err)
	}

	// Wait until the shell process is closed and notify the client.
	go func() {
		err := cmd.Wait()
		if err != nil {
			fmt.Println("cmd.wait err", err)
		}

		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded.Call()

		t.DeleteUserSession(r.Username, command.Session)
	}()

	// Read the STDOUT from shell process and send to the connected client.
	go func() {
		buf := make([]byte, (4096)-utf8.UTFMax, 4096)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			// Rate limiting...
			if server.throttling {
				s := time.Now().Unix()
				if server.currentSecond != s {
					server.currentSecond = s
					server.messageCounter = 0
					server.byteCounter = 0
					server.lineFeeedCounter = 0
				}
				server.messageCounter += 1
				server.byteCounter += n
				server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
				if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
					time.Sleep(time.Second)
				}
			}

			server.remote.Output.Call(string(filterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	return server, nil
}

// CloseSessions closes all active session for the given username and deletes
// it from the internal user map
func (t *Terminal) CloseSessions(username string) {
	t.Log.Info("Closing terminal sessions of user '%s'", username)

	t.Lock()
	user, ok := t.Users[username]
	t.Unlock()

	if !ok {
		return
	}

	user.CloseSessions()

	t.Lock()
	delete(t.Users, username)
	t.Unlock()
}

func filterInvalidUTF8(buf []byte) []byte {
	i := 0
	j := 0
	for {
		r, l := utf8.DecodeRune(buf[i:])
		if l == 0 {
			break
		}
		if r < 0xD800 {
			if i != j {
				copy(buf[j:], buf[i:i+l])
			}
			j += l
		}
		i += l
	}
	return buf[:j]
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/config/module"
)

// InputMode defines what sort of input will be asked for when Input
// is called on Context.
type InputMode byte

const (
	// InputModeVar asks for all variables
	InputModeVar InputMode = 1 << iota

	// InputModeVarUnset asks for variables which are not set yet
	InputModeVarUnset

	// InputModeProvider asks for provider variables
	InputModeProvider

	// InputModeStd is the standard operating mode and asks for both variables
	// and providers.
	InputModeStd = InputModeVar | InputModeProvider
)

// ContextOpts are the user-configurable options to create a context with
// NewContext.
type ContextOpts struct {
	Destroy      bool
	Diff         *Diff
	Hooks        []Hook
	Module       *module.Tree
	Parallelism  int
	State        *State
	Providers    map[string]ResourceProviderFactory
	Provisioners map[string]ResourceProvisionerFactory
	Targets      []string
	Variables    map[string]string

	UIInput UIInput
}

// Context represents all the context that Terraform needs in order to
// perform operations on infrastructure. This structure is built using
// NewContext. See the documentation for that.
type Context struct {
	destroy      bool
	diff         *Diff
	diffLock     sync.RWMutex
	hooks        []Hook
	module       *module.Tree
	providers    map[string]ResourceProviderFactory
	provisioners map[string]ResourceProvisionerFactory
	sh           *stopHook
	state        *State
	stateLock    sync.RWMutex
	targets      []string
	uiInput      UIInput
	variables    map[string]string

	l                   sync.Mutex // Lock acquired during any task
	parallelSem         Semaphore
	providerInputConfig map[string]map[string]interface{}
	runCh               <-chan struct{}
}

// NewContext creates a new Context structure.
//
// Once a Context is creator, the pointer values within ContextOpts
// should not be mutated in any way, since the pointers are copied, not
// the values themselves.
func NewContext(opts *ContextOpts) *Context {
	// Copy all the hooks and add our stop hook. We don't append directly
	// to the Config so that we're not modifying that in-place.
	sh := new(stopHook)
	hooks := make([]Hook, len(opts.Hooks)+1)
	copy(hooks, opts.Hooks)
	hooks[len(opts.Hooks)] = sh

	state := opts.State
	if state == nil {
		state = new(State)
		state.init()
	}

	// Determine parallelism, default to 10. We do this both to limit
	// CPU pressure but also to have an extra guard against rate throttling
	// from providers.
	par := opts.Parallelism
	if par == 0 {
		par = 10
	}

	// Setup the variables. We first take the variables given to us.
	// We then merge in the variables set in the environment.
	variables := make(map[string]string)
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, VarEnvPrefix) {
			continue
		}

		// Strip off the prefix and get the value after the first "="
		idx := strings.Index(v, "=")
		k := v[len(VarEnvPrefix):idx]
		v = v[idx+1:]

		// Override the command-line set variable
		variables[k] = v
	}
	for k, v := range opts.Variables {
		variables[k] = v
	}

	return &Context{
		destroy:      opts.Destroy,
		diff:         opts.Diff,
		hooks:        hooks,
		module:       opts.Module,
		providers:    opts.Providers,
		provisioners: opts.Provisioners,
		state:        state,
		targets:      opts.Targets,
		uiInput:      opts.UIInput,
		variables:    variables,

		parallelSem:         NewSemaphore(par),
		providerInputConfig: make(map[string]map[string]interface{}),
		sh:                  sh,
	}
}

type ContextGraphOpts struct {
	Validate bool
	Verbose  bool
}

// Graph returns the graph for this config.
func (c *Context) Graph(g *ContextGraphOpts) (*Graph, error) {
	return c.graphBuilder(g).Build(RootModulePath)
}

// GraphBuilder returns the GraphBuilder that will be used to create
// the graphs for this context.
func (c *Context) graphBuilder(g *ContextGraphOpts) GraphBuilder {
	// TODO test
	providers := make([]string, 0, len(c.providers))
	for k, _ := range c.providers {
		providers = append(providers, k)
	}

	provisioners := make([]string, 0, len(c.provisioners))
	for k, _ := range c.provisioners {
		provisioners = append(provisioners, k)
	}

	return &BuiltinGraphBuilder{
		Root:         c.module,
		Diff:         c.diff,
		Providers:    providers,
		Provisioners: provisioners,
		State:        c.state,
		Targets:      c.targets,
		Destroy:      c.destroy,
		Validate:     g.Validate,
		Verbose:      g.Verbose,
	}
}

// Input asks for input to fill variables and provider configurations.
// This modifies the configuration in-place, so asking for Input twice
// may result in different UI output showing different current values.
func (c *Context) Input(mode InputMode) error {
	v := c.acquireRun()
	defer c.releaseRun(v)

	if mode&InputModeVar != 0 {
		// Walk the variables first for the root module. We walk them in
		// alphabetical order for UX reasons.
		rootConf := c.module.Config()
		names := make([]string, len(rootConf.Variables))
		m := make(map[string]*config.Variable)
		for i, v := range rootConf.Variables {
			names[i] = v.Name
			m[v.Name] = v
		}
		sort.Strings(names)
		for _, n := range names {
			// If we only care about unset variables, then if the variabel
			// is set, continue on.
			if mode&InputModeVarUnset != 0 {
				if _, ok := c.variables[n]; ok {
					continue
				}
			}

			v := m[n]
			switch v.Type() {
			case config.VariableTypeUnknown:
				continue
			case config.VariableTypeMap:
				continue
			case config.VariableTypeString:
				// Good!
			default:
				panic(fmt.Sprintf("Unknown variable type: %#v", v.Type()))
			}

			var defaultString string
			if v.Default != nil {
				defaultString = v.Default.(string)
			}

			// Ask the user for a value for this variable
			var value string
			for {
				var err error
				value, err = c.uiInput.Input(&InputOpts{
					Id:          fmt.Sprintf("var.%s", n),
					Query:       fmt.Sprintf("var.%s", n),
					Default:     defaultString,
					Description: v.Description,
				})
				if err != nil {
					return fmt.Errorf(
						"Error asking for %s: %s", n, err)
				}

				if value == "" && v.Required() {
					// Redo if it is required.
					continue
				}

				if value == "" {
					// No value, just exit the loop. With no value, we just
					// use whatever is currently set in variables.
					break
				}

				break
			}

			if value != "" {
				c.variables[n] = value
			}
		}
	}

	if mode&InputModeProvider != 0 {
		// Build the graph
		graph, err := c.Graph(&ContextGraphOpts{Validate: true})
		if err != nil {
			return err
		}

		// Do the walk
		if _, err := c.walk(graph, walkInput); err != nil {
			return err
		}
	}

	return nil
}

// Apply applies the changes represented by this context and returns
// the resulting state.
//
// In addition to returning the resulting state, this context is updated
// with the latest state.
func (c *Context) Apply() (*State, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	// Copy our own state
	c.state = c.state.DeepCopy()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	_, err = c.walk(graph, walkApply)

	// Clean out any unused things
	c.state.prune()

	return c.state, err
}

// Plan generates an execution plan for the given context.
//
// The execution plan encapsulates the context and can be stored
// in order to reinstantiate a context later for Apply.
//
// Plan also updates the diff of this context to be the diff generated
// by the plan, so Apply can be called after.
func (c *Context) Plan() (*Plan, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	p := &Plan{
		Module: c.module,
		Vars:   c.variables,
		State:  c.state,
	}

	var operation walkOperation
	if c.destroy {
		operation = walkPlanDestroy
	} else {
		// Set our state to be something temporary. We do this so that
		// the plan can update a fake state so that variables work, then
		// we replace it back with our old state.
		old := c.state
		if old == nil {
			c.state = &State{}
			c.state.init()
		} else {
			c.state = old.DeepCopy()
		}
		defer func() {
			c.state = old
		}()

		operation = walkPlan
	}

	// Setup our diff
	c.diffLock.Lock()
	c.diff = new(Diff)
	c.diff.init()
	c.diffLock.Unlock()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	if _, err := c.walk(graph, operation); err != nil {
		return nil, err
	}
	p.Diff = c.diff

	// Now that we have a diff, we can build the exact graph that Apply will use
	// and catch any possible cycles during the Plan phase.
	if _, err := c.Graph(&ContextGraphOpts{Validate: true}); err != nil {
		return nil, err
	}

	return p, nil
}

// Refresh goes through all the resources in the state and refreshes them
// to their latest state. This will update the state that this context
// works with, along with returning it.
//
// Even in the case an error is returned, the state will be returned and
// will potentially be partially updated.
func (c *Context) Refresh() (*State, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	// Copy our own state
	c.state = c.state.DeepCopy()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	if _, err := c.walk(graph, walkRefresh); err != nil {
		return nil, err
	}

	// Clean out any unused things
	c.state.prune()

	return c.state, nil
}

// Stop stops the running task.
//
// Stop will block until the task completes.
func (c *Context) Stop() {
	c.l.Lock()
	ch := c.runCh

	// If we aren't running, then just return
	if ch == nil {
		c.l.Unlock()
		return
	}

	// Tell the hook we want to stop
	c.sh.Stop()

	// Wait for us to stop
	c.l.Unlock()
	<-ch
}

// Validate validates the configuration and returns any warnings or errors.
func (c *Context) Validate() ([]string, []error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	var errs error

	// Validate the configuration itself
	if err := c.module.Validate(); err != nil {
		errs = multierror.Append(errs, err)
	}

	// This only needs to be done for the root module, since inter-module
	// variables are validated in the module tree.
	if config := c.module.Config(); config != nil {
		// Validate the user variables
		if err := smcUserVariables(config, c.variables); len(err) > 0 {
			errs = multierror.Append(errs, err...)
		}
	}

	// Build the graph so we can walk it and run Validate on nodes.
	// We also validate the graph generated here, but this graph doesn't
	// necessarily match the graph that Plan will generate, so we'll validate the
	// graph again later after Planning.
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, []error{err}
	}

	// Walk
	walker, err := c.walk(graph, walkValidate)
	if err != nil {
		return nil, multierror.Append(errs, err).Errors
	}

	// Return the result
	rerrs := multierror.Append(errs, walker.ValidationErrors...)
	return walker.ValidationWarnings, rerrs.Errors
}

// Module returns the module tree associated with this context.
func (c *Context) Module() *module.Tree {
	return c.module
}

// Variables will return the mapping of variables that were defined
// for this Context. If Input was called, this mapping may be different
// than what was given.
func (c *Context) Variables() map[string]string {
	return c.variables
}

// SetVariable sets a variable after a context has already been built.
func (c *Context) SetVariable(k, v string) {
	c.variables[k] = v
}

func (c *Context) acquireRun() chan<- struct{} {
	c.l.Lock()
	defer c.l.Unlock()

	// Wait for no channel to exist
	for c.runCh != nil {
		c.l.Unlock()
		ch := c.runCh
		<-ch
		c.l.Lock()
	}

	ch := make(chan struct{})
	c.runCh = ch
	return ch
}

func (c *Context) releaseRun(ch chan<- struct{}) {
	c.l.Lock()
	defer c.l.Unlock()

	close(ch)
	c.runCh = nil
	c.sh.Reset()
}

func (c *Context) walk(
	graph *Graph, operation walkOperation) (*ContextGraphWalker, error) {
	// Walk the graph
	log.Printf("[INFO] Starting graph walk: %s", operation.String())
	walker := &ContextGraphWalker{Context: c, Operation: operation}
	return walker, graph.Walk(walker)
}
<commit_msg>terraform: fail early if module validation fails<commit_after>package terraform

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/config/module"
)

// InputMode defines what sort of input will be asked for when Input
// is called on Context.
type InputMode byte

const (
	// InputModeVar asks for all variables
	InputModeVar InputMode = 1 << iota

	// InputModeVarUnset asks for variables which are not set yet
	InputModeVarUnset

	// InputModeProvider asks for provider variables
	InputModeProvider

	// InputModeStd is the standard operating mode and asks for both variables
	// and providers.
	InputModeStd = InputModeVar | InputModeProvider
)

// ContextOpts are the user-configurable options to create a context with
// NewContext.
type ContextOpts struct {
	Destroy      bool
	Diff         *Diff
	Hooks        []Hook
	Module       *module.Tree
	Parallelism  int
	State        *State
	Providers    map[string]ResourceProviderFactory
	Provisioners map[string]ResourceProvisionerFactory
	Targets      []string
	Variables    map[string]string

	UIInput UIInput
}

// Context represents all the context that Terraform needs in order to
// perform operations on infrastructure. This structure is built using
// NewContext. See the documentation for that.
type Context struct {
	destroy      bool
	diff         *Diff
	diffLock     sync.RWMutex
	hooks        []Hook
	module       *module.Tree
	providers    map[string]ResourceProviderFactory
	provisioners map[string]ResourceProvisionerFactory
	sh           *stopHook
	state        *State
	stateLock    sync.RWMutex
	targets      []string
	uiInput      UIInput
	variables    map[string]string

	l                   sync.Mutex // Lock acquired during any task
	parallelSem         Semaphore
	providerInputConfig map[string]map[string]interface{}
	runCh               <-chan struct{}
}

// NewContext creates a new Context structure.
//
// Once a Context is creator, the pointer values within ContextOpts
// should not be mutated in any way, since the pointers are copied, not
// the values themselves.
func NewContext(opts *ContextOpts) *Context {
	// Copy all the hooks and add our stop hook. We don't append directly
	// to the Config so that we're not modifying that in-place.
	sh := new(stopHook)
	hooks := make([]Hook, len(opts.Hooks)+1)
	copy(hooks, opts.Hooks)
	hooks[len(opts.Hooks)] = sh

	state := opts.State
	if state == nil {
		state = new(State)
		state.init()
	}

	// Determine parallelism, default to 10. We do this both to limit
	// CPU pressure but also to have an extra guard against rate throttling
	// from providers.
	par := opts.Parallelism
	if par == 0 {
		par = 10
	}

	// Setup the variables. We first take the variables given to us.
	// We then merge in the variables set in the environment.
	variables := make(map[string]string)
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, VarEnvPrefix) {
			continue
		}

		// Strip off the prefix and get the value after the first "="
		idx := strings.Index(v, "=")
		k := v[len(VarEnvPrefix):idx]
		v = v[idx+1:]

		// Override the command-line set variable
		variables[k] = v
	}
	for k, v := range opts.Variables {
		variables[k] = v
	}

	return &Context{
		destroy:      opts.Destroy,
		diff:         opts.Diff,
		hooks:        hooks,
		module:       opts.Module,
		providers:    opts.Providers,
		provisioners: opts.Provisioners,
		state:        state,
		targets:      opts.Targets,
		uiInput:      opts.UIInput,
		variables:    variables,

		parallelSem:         NewSemaphore(par),
		providerInputConfig: make(map[string]map[string]interface{}),
		sh:                  sh,
	}
}

type ContextGraphOpts struct {
	Validate bool
	Verbose  bool
}

// Graph returns the graph for this config.
func (c *Context) Graph(g *ContextGraphOpts) (*Graph, error) {
	return c.graphBuilder(g).Build(RootModulePath)
}

// GraphBuilder returns the GraphBuilder that will be used to create
// the graphs for this context.
func (c *Context) graphBuilder(g *ContextGraphOpts) GraphBuilder {
	// TODO test
	providers := make([]string, 0, len(c.providers))
	for k, _ := range c.providers {
		providers = append(providers, k)
	}

	provisioners := make([]string, 0, len(c.provisioners))
	for k, _ := range c.provisioners {
		provisioners = append(provisioners, k)
	}

	return &BuiltinGraphBuilder{
		Root:         c.module,
		Diff:         c.diff,
		Providers:    providers,
		Provisioners: provisioners,
		State:        c.state,
		Targets:      c.targets,
		Destroy:      c.destroy,
		Validate:     g.Validate,
		Verbose:      g.Verbose,
	}
}

// Input asks for input to fill variables and provider configurations.
// This modifies the configuration in-place, so asking for Input twice
// may result in different UI output showing different current values.
func (c *Context) Input(mode InputMode) error {
	v := c.acquireRun()
	defer c.releaseRun(v)

	if mode&InputModeVar != 0 {
		// Walk the variables first for the root module. We walk them in
		// alphabetical order for UX reasons.
		rootConf := c.module.Config()
		names := make([]string, len(rootConf.Variables))
		m := make(map[string]*config.Variable)
		for i, v := range rootConf.Variables {
			names[i] = v.Name
			m[v.Name] = v
		}
		sort.Strings(names)
		for _, n := range names {
			// If we only care about unset variables, then if the variabel
			// is set, continue on.
			if mode&InputModeVarUnset != 0 {
				if _, ok := c.variables[n]; ok {
					continue
				}
			}

			v := m[n]
			switch v.Type() {
			case config.VariableTypeUnknown:
				continue
			case config.VariableTypeMap:
				continue
			case config.VariableTypeString:
				// Good!
			default:
				panic(fmt.Sprintf("Unknown variable type: %#v", v.Type()))
			}

			var defaultString string
			if v.Default != nil {
				defaultString = v.Default.(string)
			}

			// Ask the user for a value for this variable
			var value string
			for {
				var err error
				value, err = c.uiInput.Input(&InputOpts{
					Id:          fmt.Sprintf("var.%s", n),
					Query:       fmt.Sprintf("var.%s", n),
					Default:     defaultString,
					Description: v.Description,
				})
				if err != nil {
					return fmt.Errorf(
						"Error asking for %s: %s", n, err)
				}

				if value == "" && v.Required() {
					// Redo if it is required.
					continue
				}

				if value == "" {
					// No value, just exit the loop. With no value, we just
					// use whatever is currently set in variables.
					break
				}

				break
			}

			if value != "" {
				c.variables[n] = value
			}
		}
	}

	if mode&InputModeProvider != 0 {
		// Build the graph
		graph, err := c.Graph(&ContextGraphOpts{Validate: true})
		if err != nil {
			return err
		}

		// Do the walk
		if _, err := c.walk(graph, walkInput); err != nil {
			return err
		}
	}

	return nil
}

// Apply applies the changes represented by this context and returns
// the resulting state.
//
// In addition to returning the resulting state, this context is updated
// with the latest state.
func (c *Context) Apply() (*State, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	// Copy our own state
	c.state = c.state.DeepCopy()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	_, err = c.walk(graph, walkApply)

	// Clean out any unused things
	c.state.prune()

	return c.state, err
}

// Plan generates an execution plan for the given context.
//
// The execution plan encapsulates the context and can be stored
// in order to reinstantiate a context later for Apply.
//
// Plan also updates the diff of this context to be the diff generated
// by the plan, so Apply can be called after.
func (c *Context) Plan() (*Plan, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	p := &Plan{
		Module: c.module,
		Vars:   c.variables,
		State:  c.state,
	}

	var operation walkOperation
	if c.destroy {
		operation = walkPlanDestroy
	} else {
		// Set our state to be something temporary. We do this so that
		// the plan can update a fake state so that variables work, then
		// we replace it back with our old state.
		old := c.state
		if old == nil {
			c.state = &State{}
			c.state.init()
		} else {
			c.state = old.DeepCopy()
		}
		defer func() {
			c.state = old
		}()

		operation = walkPlan
	}

	// Setup our diff
	c.diffLock.Lock()
	c.diff = new(Diff)
	c.diff.init()
	c.diffLock.Unlock()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	if _, err := c.walk(graph, operation); err != nil {
		return nil, err
	}
	p.Diff = c.diff

	// Now that we have a diff, we can build the exact graph that Apply will use
	// and catch any possible cycles during the Plan phase.
	if _, err := c.Graph(&ContextGraphOpts{Validate: true}); err != nil {
		return nil, err
	}

	return p, nil
}

// Refresh goes through all the resources in the state and refreshes them
// to their latest state. This will update the state that this context
// works with, along with returning it.
//
// Even in the case an error is returned, the state will be returned and
// will potentially be partially updated.
func (c *Context) Refresh() (*State, error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	// Copy our own state
	c.state = c.state.DeepCopy()

	// Build the graph
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, err
	}

	// Do the walk
	if _, err := c.walk(graph, walkRefresh); err != nil {
		return nil, err
	}

	// Clean out any unused things
	c.state.prune()

	return c.state, nil
}

// Stop stops the running task.
//
// Stop will block until the task completes.
func (c *Context) Stop() {
	c.l.Lock()
	ch := c.runCh

	// If we aren't running, then just return
	if ch == nil {
		c.l.Unlock()
		return
	}

	// Tell the hook we want to stop
	c.sh.Stop()

	// Wait for us to stop
	c.l.Unlock()
	<-ch
}

// Validate validates the configuration and returns any warnings or errors.
func (c *Context) Validate() ([]string, []error) {
	v := c.acquireRun()
	defer c.releaseRun(v)

	var errs error

	// Validate the configuration itself
	if err := c.module.Validate(); err != nil {
		errs = multierror.Append(errs, err)
	}

	// This only needs to be done for the root module, since inter-module
	// variables are validated in the module tree.
	if config := c.module.Config(); config != nil {
		// Validate the user variables
		if err := smcUserVariables(config, c.variables); len(err) > 0 {
			errs = multierror.Append(errs, err...)
		}
	}

	// If we have errors at this point, the graphing has no chance,
	// so just bail early.
	if errs != nil {
		return nil, []error{errs}
	}

	// Build the graph so we can walk it and run Validate on nodes.
	// We also validate the graph generated here, but this graph doesn't
	// necessarily match the graph that Plan will generate, so we'll validate the
	// graph again later after Planning.
	graph, err := c.Graph(&ContextGraphOpts{Validate: true})
	if err != nil {
		return nil, []error{err}
	}

	// Walk
	walker, err := c.walk(graph, walkValidate)
	if err != nil {
		return nil, multierror.Append(errs, err).Errors
	}

	// Return the result
	rerrs := multierror.Append(errs, walker.ValidationErrors...)
	return walker.ValidationWarnings, rerrs.Errors
}

// Module returns the module tree associated with this context.
func (c *Context) Module() *module.Tree {
	return c.module
}

// Variables will return the mapping of variables that were defined
// for this Context. If Input was called, this mapping may be different
// than what was given.
func (c *Context) Variables() map[string]string {
	return c.variables
}

// SetVariable sets a variable after a context has already been built.
func (c *Context) SetVariable(k, v string) {
	c.variables[k] = v
}

func (c *Context) acquireRun() chan<- struct{} {
	c.l.Lock()
	defer c.l.Unlock()

	// Wait for no channel to exist
	for c.runCh != nil {
		c.l.Unlock()
		ch := c.runCh
		<-ch
		c.l.Lock()
	}

	ch := make(chan struct{})
	c.runCh = ch
	return ch
}

func (c *Context) releaseRun(ch chan<- struct{}) {
	c.l.Lock()
	defer c.l.Unlock()

	close(ch)
	c.runCh = nil
	c.sh.Reset()
}

func (c *Context) walk(
	graph *Graph, operation walkOperation) (*ContextGraphWalker, error) {
	// Walk the graph
	log.Printf("[INFO] Starting graph walk: %s", operation.String())
	walker := &ContextGraphWalker{Context: c, Operation: operation}
	return walker, graph.Walk(walker)
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.3"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = ""

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<commit_msg>release: clean up after v0.9.3<commit_after>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.4"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = "dev"

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<|endoftext|>