<commit_before>package views

/*
Holds the insert task related view handlers, includes the one for file upload
*/
import (
	"crypto/md5"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/thewhitetulip/Tasks/db"
	"github.com/thewhitetulip/Tasks/sessions"
	"github.com/thewhitetulip/Tasks/utils"
)

// UploadedFileHandler is used to handle the uploaded file related requests
func UploadedFileHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	token := r.URL.Path[len("/files/"):]

	//file, err := db.GetFileName(token)
	//if err != nil {
	log.Println("serving file ./files/" + token)
	http.ServeFile(w, r, "./files/"+token)
	//}
}

//AddTaskFunc is used to handle the addition of new task, "/add" URL
func AddTaskFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { // Will work only for POST requests, will redirect to home
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	var filelink string // will store the html when we have files to be uploaded, appened to the note content
	r.ParseForm()
	file, handler, err := r.FormFile("uploadfile")
	if err != nil && handler != nil {
		//Case executed when file is uploaded and yet an error occurs
		log.Println(err)
		message = "Error uploading file"
		http.Redirect(w, r, "/", http.StatusInternalServerError)
	}

	taskPriority, priorityErr := strconv.Atoi(r.FormValue("priority"))

	if priorityErr != nil {
		log.Print(priorityErr)
		message = "Bad task priority"
	}
	priorityList := []int{1, 2, 3}
	found := false
	for _, priority := range priorityList {
		if taskPriority == priority {
			found = true
		}
	}
	//If someone gives us incorrect priority number, we give the priority
	//to that task as 1 i.e. Low
	if !found {
		taskPriority = 1
	}
	var hidden int
	hideTimeline := r.FormValue("hide")
	if hideTimeline != "" {
		hidden = 1
	} else {
		hidden = 0
	}
	// dueDate := r.FormValue("dueDate")
	category := r.FormValue("category")
	title := template.HTMLEscapeString(r.Form.Get("title"))
	content := template.HTMLEscapeString(r.Form.Get("content"))
	formToken := template.HTMLEscapeString(r.Form.Get("CSRFToken"))

	cookie, _ := r.Cookie("csrftoken")
	if formToken == cookie.Value {
		username := sessions.GetCurrentUserName(r)
		if handler != nil {
			// this will be executed whenever a file is uploaded
			r.ParseMultipartForm(32 << 20) //defined maximum size of file
			defer file.Close()
			randomFileName := md5.New()
			io.WriteString(randomFileName, strconv.FormatInt(time.Now().Unix(), 10))
			io.WriteString(randomFileName, handler.Filename)
			token := fmt.Sprintf("%x", randomFileName.Sum(nil))
			f, err := os.OpenFile("./files/"+token, os.O_WRONLY|os.O_CREATE, 0666)
			if err != nil {
				log.Println(err)
				return
			}
			defer f.Close()
			io.Copy(f, file)

			if strings.HasSuffix(handler.Filename, ".png") || strings.HasSuffix(handler.Filename, ".jpg") {
				filelink = "
 "
			} else {
				filelink = "
 " + handler.Filename + ""
			}
			content = content + filelink

			fileTruth := db.AddFile(handler.Filename, token, username)
			if fileTruth != nil {
				message = "Error adding filename in db"
				log.Println("error adding task to db")
			}
		}
		//taskTruth := db.AddTask(title, content, category, taskPriority, username, dueDate)
		taskTruth := db.AddTask(title, content, category, taskPriority, username, hidden)
		if taskTruth != nil {
			message = "Error adding task"
			log.Println("error adding task to db")
			http.Redirect(w, r, "/", http.StatusInternalServerError)
		} else {
			message = "Task added"
			log.Println("added task to db")
			http.Redirect(w, r, "/", http.StatusFound)
		}
	} else {
		log.Println("CSRF mismatch")
		message = "Server Error"
		http.Redirect(w, r, "/", http.StatusInternalServerError)
	}

}

//AddCategoryFunc used to add new categories to the database
func AddCategoryFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { // We respond only to POST requests, redirect to home for others
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	r.ParseForm()
	category := r.Form.Get("category")
	if strings.Trim(category, " ") != "" {
		username := sessions.GetCurrentUserName(r)
		log.Println("adding category")
		err := db.AddCategory(username, category)
		if err != nil {
			message = "Error adding category"
			http.Redirect(w, r, "/", http.StatusBadRequest)
		} else {
			message = "Added category"
			http.Redirect(w, r, "/", http.StatusFound)
		}
	}
}

//EditTaskFunc is used to edit tasks, handles "/edit/" URL
func EditTaskFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	id, err := strconv.Atoi(r.URL.Path[len("/edit/"):])
	if err != nil {
		log.Println(err)
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	redirectURL := utils.GetRedirectUrl(r.Referer())
	username := sessions.GetCurrentUserName(r)
	task, err := db.GetTaskByID(username, id)
	categories := db.GetCategories(username)
	task.Categories = categories
	task.Referer = redirectURL

	if err != nil {
		task.Message = "Error fetching Tasks"
	}
	editTemplate.Execute(w, task)

}

//AddCommentFunc will be used
func AddCommentFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		log.Println(err)
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	r.ParseForm()
	text := r.Form.Get("commentText")
	id := r.Form.Get("taskID")

	idInt, err := strconv.Atoi(id)

	if (err != nil) || (text == "") {
		log.Println("unable to convert into integer")
		message = "Error adding comment"
	} else {
		username := sessions.GetCurrentUserName(r)
		err = db.AddComments(username, idInt, text)

		if err != nil {
			log.Println("unable to insert into db")
			message = "Comment not added"
		} else {
			message = "Comment added"
		}
	}

	http.Redirect(w, r, "/", http.StatusFound)

}
<commit_msg>stores file object with original filename and update references<commit_after>package views

/*
Holds the insert task related view handlers, includes the one for file upload
*/
import (
	"crypto/md5"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/thewhitetulip/Tasks/db"
	"github.com/thewhitetulip/Tasks/sessions"
	"github.com/thewhitetulip/Tasks/utils"
)

// UploadedFileHandler is used to handle the uploaded file related requests
func UploadedFileHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	token := r.URL.Path[len("/files/"):]

	//file, err := db.GetFileName(token)
	//if err != nil {
	log.Println("serving file ./files/" + token)
	http.ServeFile(w, r, "./files/"+token)
	//}
}

//AddTaskFunc is used to handle the addition of new task, "/add" URL
func AddTaskFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { // Will work only for POST requests, will redirect to home
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	var filelink string // will store the html when we have files to be uploaded, appened to the note content
	r.ParseForm()
	file, handler, err := r.FormFile("uploadfile")
	if err != nil && handler != nil {
		//Case executed when file is uploaded and yet an error occurs
		log.Println(err)
		message = "Error uploading file"
		http.Redirect(w, r, "/", http.StatusInternalServerError)
	}

	taskPriority, priorityErr := strconv.Atoi(r.FormValue("priority"))

	if priorityErr != nil {
		log.Print(priorityErr)
		message = "Bad task priority"
	}
	priorityList := []int{1, 2, 3}
	found := false
	for _, priority := range priorityList {
		if taskPriority == priority {
			found = true
		}
	}
	//If someone gives us incorrect priority number, we give the priority
	//to that task as 1 i.e. Low
	if !found {
		taskPriority = 1
	}
	var hidden int
	hideTimeline := r.FormValue("hide")
	if hideTimeline != "" {
		hidden = 1
	} else {
		hidden = 0
	}
	// dueDate := r.FormValue("dueDate")
	category := r.FormValue("category")
	title := template.HTMLEscapeString(r.Form.Get("title"))
	content := template.HTMLEscapeString(r.Form.Get("content"))
	formToken := template.HTMLEscapeString(r.Form.Get("CSRFToken"))

	cookie, _ := r.Cookie("csrftoken")
	if formToken == cookie.Value {
		username := sessions.GetCurrentUserName(r)
		if handler != nil {
			// this will be executed whenever a file is uploaded
			r.ParseMultipartForm(32 << 20) //defined maximum size of file
			defer file.Close()
			htmlFilename := strings.Replace(handler.Filename, " ", "-", -1)
			randomFileName := md5.New()
			io.WriteString(randomFileName, strconv.FormatInt(time.Now().Unix(), 10))
			io.WriteString(randomFileName, htmlFilename)
			token := fmt.Sprintf("%x", randomFileName.Sum(nil))
			f, err := os.OpenFile("./files/"+htmlFilename, os.O_WRONLY|os.O_CREATE, 0666)
			if err != nil {
				log.Println(err)
				return
			}
			defer f.Close()
			io.Copy(f, file)

			if strings.HasSuffix(htmlFilename, ".png") || strings.HasSuffix(htmlFilename, ".jpg") {
				filelink = "
 "
			} else {
				filelink = "
 " + htmlFilename + ""
			}
			content = content + filelink

			fileTruth := db.AddFile(htmlFilename, token, username)
			if fileTruth != nil {
				message = "Error adding filename in db"
				log.Println("error adding task to db")
			}
		}
		//taskTruth := db.AddTask(title, content, category, taskPriority, username, dueDate)
		taskTruth := db.AddTask(title, content, category, taskPriority, username, hidden)
		if taskTruth != nil {
			message = "Error adding task"
			log.Println("error adding task to db")
			http.Redirect(w, r, "/", http.StatusInternalServerError)
		} else {
			message = "Task added"
			log.Println("added task to db")
			http.Redirect(w, r, "/", http.StatusFound)
		}
	} else {
		log.Println("CSRF mismatch")
		message = "Server Error"
		http.Redirect(w, r, "/", http.StatusInternalServerError)
	}

}

//AddCategoryFunc used to add new categories to the database
func AddCategoryFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" { // We respond only to POST requests, redirect to home for others
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	r.ParseForm()
	category := r.Form.Get("category")
	if strings.Trim(category, " ") != "" {
		username := sessions.GetCurrentUserName(r)
		log.Println("adding category")
		err := db.AddCategory(username, category)
		if err != nil {
			message = "Error adding category"
			http.Redirect(w, r, "/", http.StatusBadRequest)
		} else {
			message = "Added category"
			http.Redirect(w, r, "/", http.StatusFound)
		}
	}
}

//EditTaskFunc is used to edit tasks, handles "/edit/" URL
func EditTaskFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}

	id, err := strconv.Atoi(r.URL.Path[len("/edit/"):])
	if err != nil {
		log.Println(err)
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	redirectURL := utils.GetRedirectUrl(r.Referer())
	username := sessions.GetCurrentUserName(r)
	task, err := db.GetTaskByID(username, id)
	categories := db.GetCategories(username)
	task.Categories = categories
	task.Referer = redirectURL

	if err != nil {
		task.Message = "Error fetching Tasks"
	}
	editTemplate.Execute(w, task)

}

//AddCommentFunc will be used
func AddCommentFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		log.Println(err)
		http.Redirect(w, r, "/", http.StatusBadRequest)
		return
	}
	r.ParseForm()
	text := r.Form.Get("commentText")
	id := r.Form.Get("taskID")

	idInt, err := strconv.Atoi(id)

	if (err != nil) || (text == "") {
		log.Println("unable to convert into integer")
		message = "Error adding comment"
	} else {
		username := sessions.GetCurrentUserName(r)
		err = db.AddComments(username, idInt, text)

		if err != nil {
			log.Println("unable to insert into db")
			message = "Comment not added"
		} else {
			message = "Comment added"
		}
	}

	http.Redirect(w, r, "/", http.StatusFound)

}
<|endoftext|>


<commit_before>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom_test

import (
	. "github.com/FactomProject/factom"
	"testing"
)

func TestFactoidToFactoshi(t *testing.T) {
	v1 := "1.0001"
	v2 := "100000000.00000001"
	v3 := ".01"

	t.Log(FactoidToFactoshi(v1))
	t.Log(FactoidToFactoshi(v2))
	t.Log(FactoidToFactoshi(v3))
}
<commit_msg>updated test for factoid conversion<commit_after>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom_test

import (
	. "github.com/FactomProject/factom"
	"testing"
)

func TestFactoidToFactoshi(t *testing.T) {
	v1 := "1.0001"
	v2 := "100000000.00000001"
	v3 := ".01"

	e1 := uint64(100010000)
	e2 := uint64(10000000000000001)
	e3 := uint64(1000000)

	if r1 := FactoidToFactoshi(v1); r1 != e1 {
		t.Errorf("r1=%d expecting %d", r1, e1)
	}

	if r2 := FactoidToFactoshi(v2); r2 != e2 {
		t.Errorf("r2=%d expecting %d", r2, e2)
	}

	if r3 := FactoidToFactoshi(v3); r3 != e3 {
		t.Errorf("r3=%d expecting %d", r3, e3)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go-NetCDF Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package netcdf

// #include <netcdf.h>
import "C"

// FileMode represents a file's mode.
type FileMode C.int

// File modes for Open or Create
const (
	SHARE FileMode = C.NC_SHARE // share updates, limit cacheing
)

// File modes for Open
const (
	NOWRITE FileMode = C.NC_NOWRITE // set read-only access
	WRITE   FileMode = C.NC_WRITE   // set read-write access
)

// File modes for Create
const (
	CLOBBER       FileMode = C.NC_CLOBBER       // destroy existing file
	NOCLOBBER     FileMode = C.NC_NOCLOBBER     // don't destroy existing file
	CLASSIC_MODEL FileMode = C.NC_CLASSIC_MODEL // enforce classic model
	NETCDF4       FileMode = C.NC_NETCDF4       // use netCDF-4/HDF5 format
	OFFSET_64BIT  FileMode = C.NC_64BIT_OFFSET  // use large (64-bit) file offsets
)

// Type is a netCDF external data type.
type Type C.nc_type

const (
	BYTE   Type = C.NC_BYTE   // signed 1 byte integer
	CHAR   Type = C.NC_CHAR   // ISO/ASCII character
	SHORT  Type = C.NC_SHORT  // signed 2 byte integer
	INT    Type = C.NC_INT    // signed 4 byte integer
	LONG   Type = C.NC_LONG   // deprecated, but required for backward compatibility.
	FLOAT  Type = C.NC_FLOAT  // single precision floating point number
	DOUBLE Type = C.NC_DOUBLE // double precision floating point number
	UBYTE  Type = C.NC_UBYTE  // unsigned 1 byte int
	USHORT Type = C.NC_USHORT // unsigned 2-byte int
	UINT   Type = C.NC_UINT   // unsigned 4-byte int
	INT64  Type = C.NC_INT64  // signed 8-byte int
	UINT64 Type = C.NC_UINT64 // unsigned 8-byte int
	STRING Type = C.NC_STRING // string
)

var typeNames = map[Type]string{
	BYTE:   "BYTE",
	CHAR:   "CHAR",
	SHORT:  "SHORT",
	INT:    "INT",
	FLOAT:  "FLOAT",
	DOUBLE: "DOUBLE",
	UBYTE:  "UBYTE",
	USHORT: "USHORT",
	UINT:   "UINT",
	INT64:  "INT64",
	UINT64: "UINT64",
	STRING: "STRING",
}

// String converts a Type to its string representation.
func (t Type) String() string {
	return typeNames[t]
}
<commit_msg>Add commentary about types<commit_after>// Copyright 2014 The Go-NetCDF Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package netcdf

// #include <netcdf.h>
import "C"

// FileMode represents a file's mode.
type FileMode C.int

// File modes for Open or Create
const (
	SHARE FileMode = C.NC_SHARE // share updates, limit cacheing
)

// File modes for Open
const (
	NOWRITE FileMode = C.NC_NOWRITE // set read-only access
	WRITE   FileMode = C.NC_WRITE   // set read-write access
)

// File modes for Create
const (
	CLOBBER       FileMode = C.NC_CLOBBER       // destroy existing file
	NOCLOBBER     FileMode = C.NC_NOCLOBBER     // don't destroy existing file
	CLASSIC_MODEL FileMode = C.NC_CLASSIC_MODEL // enforce classic model
	NETCDF4       FileMode = C.NC_NETCDF4       // use netCDF-4/HDF5 format
	OFFSET_64BIT  FileMode = C.NC_64BIT_OFFSET  // use large (64-bit) file offsets
)

// Type is a netCDF external data type.
type Type C.nc_type

// Type declarations according to C standards
const (
	BYTE   Type = C.NC_BYTE   // signed 1 byte integer
	CHAR   Type = C.NC_CHAR   // ISO/ASCII character
	SHORT  Type = C.NC_SHORT  // signed 2 byte integer
	INT    Type = C.NC_INT    // signed 4 byte integer
	LONG   Type = C.NC_LONG   // deprecated, but required for backward compatibility.
	FLOAT  Type = C.NC_FLOAT  // single precision floating point number
	DOUBLE Type = C.NC_DOUBLE // double precision floating point number
	UBYTE  Type = C.NC_UBYTE  // unsigned 1 byte int
	USHORT Type = C.NC_USHORT // unsigned 2-byte int
	UINT   Type = C.NC_UINT   // unsigned 4-byte int
	INT64  Type = C.NC_INT64  // signed 8-byte int
	UINT64 Type = C.NC_UINT64 // unsigned 8-byte int
	STRING Type = C.NC_STRING // string
)

var typeNames = map[Type]string{
	BYTE:   "BYTE",
	CHAR:   "CHAR",
	SHORT:  "SHORT",
	INT:    "INT",
	FLOAT:  "FLOAT",
	DOUBLE: "DOUBLE",
	UBYTE:  "UBYTE",
	USHORT: "USHORT",
	UINT:   "UINT",
	INT64:  "INT64",
	UINT64: "UINT64",
	STRING: "STRING",
}

// String converts a Type to its string representation.
func (t Type) String() string {
	return typeNames[t]
}
<|endoftext|>


<commit_before>/*
Package network implements collectd's binary network protocol.
*/
package network // import "collectd.org/network"

// Well-known addresses and port.
const (
	DefaultIPv4Address = "239.192.74.66"
	DefaultIPv6Address = "ff18::efc0:4a42"
	DefaultService     = "25826"
)

// Default size of "Buffer". This is based on the maximum bytes that fit into
// an Ethernet frame without fragmentation:
//   <Ethernet frame> - (<IPv6 header> + <UDP header>) = 1500 - (40 + 8) = 1452
const DefaultBufferSize = 1452

// Numeric data source type identifiers.
const (
	dsTypeCounter = 0
	dsTypeGauge   = 1
	dsTypeDerive  = 2
)

// IDs of the various "parts", i.e. subcomponents of a packet.
const (
	typeHost           = 0x0000
	typeTime           = 0x0001
	typeTimeHR         = 0x0008
	typePlugin         = 0x0002
	typePluginInstance = 0x0003
	typeType           = 0x0004
	typeTypeInstance   = 0x0005
	typeValues         = 0x0006
	typeInterval       = 0x0007
	typeIntervalHR     = 0x0009
	typeSignSHA256     = 0x0200
	typeEncryptAES256  = 0x0210
)

// SecurityLevel determines whether data is signed, encrypted or used without
// any protection.
type SecurityLevel int

// Predefined security levels. "None" is used for plain text.
const (
	None SecurityLevel = iota
	Sign
	Encrypt
)
<commit_msg>Package network: Add the DefaultPort constant.<commit_after>/*
Package network implements collectd's binary network protocol.
*/
package network // import "collectd.org/network"

// Well-known addresses and port.
const (
	DefaultIPv4Address = "239.192.74.66"
	DefaultIPv6Address = "ff18::efc0:4a42"
	DefaultService     = "25826"
	DefaultPort        = 25826
)

// Default size of "Buffer". This is based on the maximum bytes that fit into
// an Ethernet frame without fragmentation:
//   <Ethernet frame> - (<IPv6 header> + <UDP header>) = 1500 - (40 + 8) = 1452
const DefaultBufferSize = 1452

// Numeric data source type identifiers.
const (
	dsTypeCounter = 0
	dsTypeGauge   = 1
	dsTypeDerive  = 2
)

// IDs of the various "parts", i.e. subcomponents of a packet.
const (
	typeHost           = 0x0000
	typeTime           = 0x0001
	typeTimeHR         = 0x0008
	typePlugin         = 0x0002
	typePluginInstance = 0x0003
	typeType           = 0x0004
	typeTypeInstance   = 0x0005
	typeValues         = 0x0006
	typeInterval       = 0x0007
	typeIntervalHR     = 0x0009
	typeSignSHA256     = 0x0200
	typeEncryptAES256  = 0x0210
)

// SecurityLevel determines whether data is signed, encrypted or used without
// any protection.
type SecurityLevel int

// Predefined security levels. "None" is used for plain text.
const (
	None SecurityLevel = iota
	Sign
	Encrypt
)
<|endoftext|>


<commit_before>package aws

import (
	"encoding/json"
	"log"
	"os/exec"
	"strings"

	"github.com/cockroachdb/roachprod/vm"
	"github.com/pkg/errors"
)

// Both M5 and I3 machines expose their EBS or local SSD volumes as NVMe block devices, but
// the actual device numbers vary a bit between the two types.
// This user-data script will create a filesystem, mount the data volume, and chmod 777.
// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/nvme-ebs-volumes.html
const awsStartupScript = `#!/usr/bin/env bash
set -x
disknum=0
for d in $(ls /dev/nvme?n1); do
  if ! mount | grep ${d}; then
    let "disknum++"
    echo "Disk ${d} not mounted, creating..."
    mountpoint="/mnt/data${disknum}"
    mkdir -p "${mountpoint}"
    mkfs.ext4 ${d}
    mount -o discard,defaults ${d} ${mountpoint}
    chmod 777 ${mountpoint}
    echo "${d} ${mountpoint} ext4 discard,defaults 1 1" | tee -a /etc/fstab
  else
    echo "Disk ${disknum}: ${d} already mounted, skipping..."
  fi
done
if [ "${disknum}" -eq "0" ]; then
  echo "No disks mounted, creating /mnt/data1"
  mkdir -p /mnt/data1
  chmod 777 /mnt/data1
fi
`

// runCommand is used to invoke an AWS command for which no output is expected.
func runCommand(args []string) error {
	cmd := exec.Command("aws", args...)

	_, err := cmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			log.Println(string(exitErr.Stderr))
		}
		return errors.Wrapf(err, "failed to run: aws %s", strings.Join(args, " "))
	}
	return nil
}

// runJSONCommand invokes an aws command and parses the json output.
func runJSONCommand(args []string, parsed interface{}) error {
	cmd := exec.Command("aws", args...)

	rawJSON, err := cmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			log.Println(string(exitErr.Stderr))
		}
		return errors.Wrapf(err, "failed to run: aws %s", strings.Join(args, " "))
	}

	if err := json.Unmarshal(rawJSON, &parsed); err != nil {
		return errors.Wrapf(err, "failed to parse json %s", rawJSON)
	}

	return nil
}

// splitMap splits a list of `key:value` pairs into a map.
func splitMap(data []string) (map[string]string, error) {
	ret := make(map[string]string, len(data))
	for _, part := range data {
		parts := strings.Split(part, ":")
		if len(parts) != 2 {
			return nil, errors.Errorf("Could not split Region:AMI: %s", part)
		}
		ret[parts[0]] = parts[1]
	}
	return ret, nil
}

// regionMap collates VM instances by their region.
func regionMap(vms vm.List) (map[string]vm.List, error) {
	// Fan out the work by region
	byRegion := make(map[string]vm.List)
	for _, m := range vms {
		region, err := zoneToRegion(m.Zone)
		if err != nil {
			return nil, err
		}
		byRegion[region] = append(byRegion[region], m)
	}
	return byRegion, nil
}

// zoneToRegion converts an availability zone like us-east-2a to the zone name us-east-2
func zoneToRegion(zone string) (string, error) {
	return zone[0 : len(zone)-1], nil
}
<commit_msg>Force aws output to json<commit_after>package aws

import (
	"encoding/json"
	"log"
	"os/exec"
	"strings"

	"github.com/cockroachdb/roachprod/vm"
	"github.com/pkg/errors"
)

// Both M5 and I3 machines expose their EBS or local SSD volumes as NVMe block devices, but
// the actual device numbers vary a bit between the two types.
// This user-data script will create a filesystem, mount the data volume, and chmod 777.
// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/nvme-ebs-volumes.html
const awsStartupScript = `#!/usr/bin/env bash
set -x
disknum=0
for d in $(ls /dev/nvme?n1); do
  if ! mount | grep ${d}; then
    let "disknum++"
    echo "Disk ${d} not mounted, creating..."
    mountpoint="/mnt/data${disknum}"
    mkdir -p "${mountpoint}"
    mkfs.ext4 ${d}
    mount -o discard,defaults ${d} ${mountpoint}
    chmod 777 ${mountpoint}
    echo "${d} ${mountpoint} ext4 discard,defaults 1 1" | tee -a /etc/fstab
  else
    echo "Disk ${disknum}: ${d} already mounted, skipping..."
  fi
done
if [ "${disknum}" -eq "0" ]; then
  echo "No disks mounted, creating /mnt/data1"
  mkdir -p /mnt/data1
  chmod 777 /mnt/data1
fi
`

// runCommand is used to invoke an AWS command for which no output is expected.
func runCommand(args []string) error {
	cmd := exec.Command("aws", args...)

	_, err := cmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			log.Println(string(exitErr.Stderr))
		}
		return errors.Wrapf(err, "failed to run: aws %s", strings.Join(args, " "))
	}
	return nil
}

// runJSONCommand invokes an aws command and parses the json output.
func runJSONCommand(args []string, parsed interface{}) error {
	// force json output in case the user has overridden the default behavior
	args = append(args[:len(args):len(args)], "--output", "json")
	cmd := exec.Command("aws", args...)

	rawJSON, err := cmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			log.Println(string(exitErr.Stderr))
		}
		return errors.Wrapf(err, "failed to run: aws %s", strings.Join(args, " "))
	}

	if err := json.Unmarshal(rawJSON, &parsed); err != nil {
		return errors.Wrapf(err, "failed to parse json %s", rawJSON)
	}

	return nil
}

// splitMap splits a list of `key:value` pairs into a map.
func splitMap(data []string) (map[string]string, error) {
	ret := make(map[string]string, len(data))
	for _, part := range data {
		parts := strings.Split(part, ":")
		if len(parts) != 2 {
			return nil, errors.Errorf("Could not split Region:AMI: %s", part)
		}
		ret[parts[0]] = parts[1]
	}
	return ret, nil
}

// regionMap collates VM instances by their region.
func regionMap(vms vm.List) (map[string]vm.List, error) {
	// Fan out the work by region
	byRegion := make(map[string]vm.List)
	for _, m := range vms {
		region, err := zoneToRegion(m.Zone)
		if err != nil {
			return nil, err
		}
		byRegion[region] = append(byRegion[region], m)
	}
	return byRegion, nil
}

// zoneToRegion converts an availability zone like us-east-2a to the zone name us-east-2
func zoneToRegion(zone string) (string, error) {
	return zone[0 : len(zone)-1], nil
}
<|endoftext|>


<commit_before>package iradix

import (
	"bytes"

	"github.com/hashicorp/golang-lru/simplelru"
)

const (
	// defaultModifiedCache is the default size of the modified node
	// cache used per transaction. This is used to cache the updates
	// to the nodes near the root, while the leaves do not need to be
	// cached. This is important for very large transactions to prevent
	// the modified cache from growing to be enormous.
	defaultModifiedCache = 8192
)

// Tree implements an immutable radix tree. This can be treated as a
// Dictionary abstract data type. The main advantage over a standard
// hash map is prefix-based lookups and ordered iteration. The immutability
// means that it is safe to concurrently read from a Tree without any
// coordination.
type Tree struct {
	root *Node
	size int
}

// New returns an empty Tree
func New() *Tree {
	t := &Tree{root: &Node{}}
	return t
}

// Len is used to return the number of elements in the tree
func (t *Tree) Len() int {
	return t.size
}

// Txn is a transaction on the tree. This transaction is applied
// atomically and returns a new tree when committed. A transaction
// is not thread safe, and should only be used by a single goroutine.
type Txn struct {
	root     *Node
	size     int
	modified *simplelru.LRU

	mutatedNode  map[*Node]struct{}
	mutatedLeaf  map[*leafNode]struct{}
	notifyMutate bool
}

// Txn starts a new transaction that can be used to mutate the tree
func (t *Tree) Txn() *Txn {
	txn := &Txn{
		root: t.root,
		size: t.size,
	}
	return txn
}

// NotifyMutate can be used to toggle if mutation cause a notification
// to the affected nodes. This must be enabled before any modifications are
// made within the transaction.
func (t *Txn) NotifyMutate(notify bool) error {
	if len(t.modified) > 0 && notify {
		return fmt.Errorf("transaction already in progress")
	}
	t.notifyMutate = notify
	return nil
}

// writeNode returns a node to be modified, if the current
// node as already been modified during the course of
// the transaction, it is used in-place.
func (t *Txn) writeNode(n *Node) *Node {
	// Ensure the modified set exists
	if t.modified == nil {
		lru, err := simplelru.NewLRU(defaultModifiedCache, nil)
		if err != nil {
			panic(err)
		}
		t.modified = lru
	}

	// If this node has already been modified, we can
	// continue to use it during this transaction.
	if _, ok := t.modified.Get(n); ok {
		return n
	}

	// Mark this node as being mutated
	if t.notifyMutate {
		if t.mutatedNode == nil {
			t.mutatedNode = make(map[*Node]struct{})
		}
		t.mutatedNode[n] = struct{}{}
	}

	// Copy the existing node
	nc := new(Node)
	nc.mutateCh = make(chan struct{})
	if n.prefix != nil {
		nc.prefix = make([]byte, len(n.prefix))
		copy(nc.prefix, n.prefix)
	}
	if n.leaf != nil {
		nc.leaf = n.leaf
	}
	if len(n.edges) != 0 {
		nc.edges = make([]edge, len(n.edges))
		copy(nc.edges, n.edges)
	}

	// Mark this node as modified
	t.modified.Add(nc, nil)
	return nc
}

// mutateLeaf is used to mark a leaf as being mutated
// for the purposes of notification.
func (t *Txn) mutateLeaf(leaf *leafNode) {
	if !t.notifyMutate {
		return
	}
	if t.mutatedLeaf == nil {
		t.mutatedLeaf = make(map[*leafNode]struct{})
	}
	t.mutatedLeaf[leaf] = struct{}{}
}

// insert does a recursive insertion
func (t *Txn) insert(n *Node, k, search []byte, v interface{}) (*Node, interface{}, bool) {
	// Handle key exhaution
	if len(search) == 0 {
		nc := t.writeNode(n)
		nc.leaf = &leafNode{
			mutateCh: make(chan struct{}),
			key:      k,
			val:      v,
		}
		if n.isLeaf() {
			t.mutateLeaf(n.leaf)
			return nc, n.leaf.val, true
		} else {
			return nc, nil, false
		}
	}

	// Look for the edge
	idx, child := n.getEdge(search[0])

	// No edge, create one
	if child == nil {
		e := edge{
			label: search[0],
			node: &Node{
				mutateCh: make(chan struct{}),
				leaf: &leafNode{
					mutateCh: make(chan struct{}),
					key:      k,
					val:      v,
				},
				prefix: search,
			},
		}
		nc := t.writeNode(n)
		nc.addEdge(e)
		return nc, nil, false
	}

	// Determine longest prefix of the search key on match
	commonPrefix := longestPrefix(search, child.prefix)
	if commonPrefix == len(child.prefix) {
		search = search[commonPrefix:]
		newChild, oldVal, didUpdate := t.insert(child, k, search, v)
		if newChild != nil {
			nc := t.writeNode(n)
			nc.edges[idx].node = newChild
			return nc, oldVal, didUpdate
		}
		return nil, oldVal, didUpdate
	}

	// Split the node
	nc := t.writeNode(n)
	splitNode := &Node{
		mutateCh: make(chan struct{}),
		prefix:   search[:commonPrefix],
	}
	nc.replaceEdge(edge{
		label: search[0],
		node:  splitNode,
	})

	// Restore the existing child node
	modChild := t.writeNode(child)
	splitNode.addEdge(edge{
		label: modChild.prefix[commonPrefix],
		node:  modChild,
	})
	modChild.prefix = modChild.prefix[commonPrefix:]

	// Create a new leaf node
	leaf := &leafNode{
		mutateCh: make(chan struct{}),
		key:      k,
		val:      v,
	}

	// If the new key is a subset, add to to this node
	search = search[commonPrefix:]
	if len(search) == 0 {
		splitNode.leaf = leaf
		return nc, nil, false
	}

	// Create a new edge for the node
	splitNode.addEdge(edge{
		label: search[0],
		node: &Node{
			mutateCh: make(chan struct{}),
			leaf:     leaf,
			prefix:   search,
		},
	})
	return nc, nil, false
}

// delete does a recursive deletion
func (t *Txn) delete(parent, n *Node, search []byte) (*Node, *leafNode) {
	// Check for key exhaution
	if len(search) == 0 {
		if !n.isLeaf() {
			return nil, nil
		}

		// Remove the leaf node
		nc := t.writeNode(n)
		nc.leaf = nil
		t.mutateLeaf(n.leaf)

		// Check if this node should be merged
		if n != t.root && len(nc.edges) == 1 {
			nc.mergeChild()
		}
		return nc, n.leaf
	}

	// Look for an edge
	label := search[0]
	idx, child := n.getEdge(label)
	if child == nil || !bytes.HasPrefix(search, child.prefix) {
		return nil, nil
	}

	// Consume the search prefix
	search = search[len(child.prefix):]
	newChild, leaf := t.delete(n, child, search)
	if newChild == nil {
		return nil, nil
	}

	// Copy this node
	nc := t.writeNode(n)

	// Delete the edge if the node has no edges
	if newChild.leaf == nil && len(newChild.edges) == 0 {
		nc.delEdge(label)
		if n != t.root && len(nc.edges) == 1 && !nc.isLeaf() {
			nc.mergeChild()
		}
	} else {
		nc.edges[idx].node = newChild
	}
	return nc, leaf
}

// Insert is used to add or update a given key. The return provides
// the previous value and a bool indicating if any was set.
func (t *Txn) Insert(k []byte, v interface{}) (interface{}, bool) {
	newRoot, oldVal, didUpdate := t.insert(t.root, k, k, v)
	if newRoot != nil {
		t.root = newRoot
	}
	if !didUpdate {
		t.size++
	}
	return oldVal, didUpdate
}

// Delete is used to delete a given key. Returns the old value if any,
// and a bool indicating if the key was set.
func (t *Txn) Delete(k []byte) (interface{}, bool) {
	newRoot, leaf := t.delete(nil, t.root, k)
	if newRoot != nil {
		t.root = newRoot
	}
	if leaf != nil {
		t.size--
		return leaf.val, true
	}
	return nil, false
}

// Root returns the current root of the radix tree within this
// transaction. The root is not safe across insert and delete operations,
// but can be used to read the current state during a transaction.
func (t *Txn) Root() *Node {
	return t.root
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Txn) Get(k []byte) (interface{}, bool) {
	return t.root.Get(k)
}

// GetWatch is used to lookup a specific key, returning
// the watch channel, value and if it was found
func (t *Txn) GetWatch(k []byte) (<-chan struct{}, interface{}, bool) {
	return t.root.GetWatch(k)
}

// Commit is used to finalize the transaction and return a new tree
func (t *Txn) Commit() *Tree {
	nt := &Tree{t.root, t.size}
	if t.notifyMutate {
		for leaf := range t.mutatedLeaf {
			close(leaf.mutateCh)
		}
		for node := range t.mutatedNode {
			close(node.mutateCh)
		}
	}
	t.modified = nil
	t.mutatedNode = nil
	t.mutatedLeaf = nil
	return nt
}

// Insert is used to add or update a given key. The return provides
// the new tree, previous value and a bool indicating if any was set.
func (t *Tree) Insert(k []byte, v interface{}) (*Tree, interface{}, bool) {
	txn := t.Txn()
	old, ok := txn.Insert(k, v)
	return txn.Commit(), old, ok
}

// Delete is used to delete a given key. Returns the new tree,
// old value if any, and a bool indicating if the key was set.
func (t *Tree) Delete(k []byte) (*Tree, interface{}, bool) {
	txn := t.Txn()
	old, ok := txn.Delete(k)
	return txn.Commit(), old, ok
}

// Root returns the root node of the tree which can be used for richer
// query operations.
func (t *Tree) Root() *Node {
	return t.root
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Tree) Get(k []byte) (interface{}, bool) {
	return t.root.Get(k)
}

// longestPrefix finds the length of the shared prefix
// of two strings
func longestPrefix(k1, k2 []byte) int {
	max := len(k1)
	if l := len(k2); l < max {
		max = l
	}
	var i int
	for i = 0; i < max; i++ {
		if k1[i] != k2[i] {
			break
		}
	}
	return i
}

// concat two byte slices, returning a third new copy
func concat(a, b []byte) []byte {
	c := make([]byte, len(a)+len(b))
	copy(c, a)
	copy(c[len(a):], b)
	return c
}
<commit_msg>Create mutateCh on root node<commit_after>package iradix

import (
	"bytes"

	"github.com/hashicorp/golang-lru/simplelru"
)

const (
	// defaultModifiedCache is the default size of the modified node
	// cache used per transaction. This is used to cache the updates
	// to the nodes near the root, while the leaves do not need to be
	// cached. This is important for very large transactions to prevent
	// the modified cache from growing to be enormous.
	defaultModifiedCache = 8192
)

// Tree implements an immutable radix tree. This can be treated as a
// Dictionary abstract data type. The main advantage over a standard
// hash map is prefix-based lookups and ordered iteration. The immutability
// means that it is safe to concurrently read from a Tree without any
// coordination.
type Tree struct {
	root *Node
	size int
}

// New returns an empty Tree
func New() *Tree {
	t := &Tree{
		root: &Node{
			mutateCh: make(chan struct{}),
		},
	}
	return t
}

// Len is used to return the number of elements in the tree
func (t *Tree) Len() int {
	return t.size
}

// Txn is a transaction on the tree. This transaction is applied
// atomically and returns a new tree when committed. A transaction
// is not thread safe, and should only be used by a single goroutine.
type Txn struct {
	root     *Node
	size     int
	modified *simplelru.LRU

	mutatedNode  map[*Node]struct{}
	mutatedLeaf  map[*leafNode]struct{}
	notifyMutate bool
}

// Txn starts a new transaction that can be used to mutate the tree
func (t *Tree) Txn() *Txn {
	txn := &Txn{
		root: t.root,
		size: t.size,
	}
	return txn
}

// NotifyMutate can be used to toggle if mutation cause a notification
// to the affected nodes. This must be enabled before any modifications are
// made within the transaction.
func (t *Txn) NotifyMutate(notify bool) error {
	if len(t.modified) > 0 && notify {
		return fmt.Errorf("transaction already in progress")
	}
	t.notifyMutate = notify
	return nil
}

// writeNode returns a node to be modified, if the current
// node as already been modified during the course of
// the transaction, it is used in-place.
func (t *Txn) writeNode(n *Node) *Node {
	// Ensure the modified set exists
	if t.modified == nil {
		lru, err := simplelru.NewLRU(defaultModifiedCache, nil)
		if err != nil {
			panic(err)
		}
		t.modified = lru
	}

	// If this node has already been modified, we can
	// continue to use it during this transaction.
	if _, ok := t.modified.Get(n); ok {
		return n
	}

	// Mark this node as being mutated
	if t.notifyMutate {
		if t.mutatedNode == nil {
			t.mutatedNode = make(map[*Node]struct{})
		}
		t.mutatedNode[n] = struct{}{}
	}

	// Copy the existing node
	nc := new(Node)
	nc.mutateCh = make(chan struct{})
	if n.prefix != nil {
		nc.prefix = make([]byte, len(n.prefix))
		copy(nc.prefix, n.prefix)
	}
	if n.leaf != nil {
		nc.leaf = n.leaf
	}
	if len(n.edges) != 0 {
		nc.edges = make([]edge, len(n.edges))
		copy(nc.edges, n.edges)
	}

	// Mark this node as modified
	t.modified.Add(nc, nil)
	return nc
}

// mutateLeaf is used to mark a leaf as being mutated
// for the purposes of notification.
func (t *Txn) mutateLeaf(leaf *leafNode) {
	if !t.notifyMutate {
		return
	}
	if t.mutatedLeaf == nil {
		t.mutatedLeaf = make(map[*leafNode]struct{})
	}
	t.mutatedLeaf[leaf] = struct{}{}
}

// insert does a recursive insertion
func (t *Txn) insert(n *Node, k, search []byte, v interface{}) (*Node, interface{}, bool) {
	// Handle key exhaution
	if len(search) == 0 {
		nc := t.writeNode(n)
		nc.leaf = &leafNode{
			mutateCh: make(chan struct{}),
			key:      k,
			val:      v,
		}
		if n.isLeaf() {
			t.mutateLeaf(n.leaf)
			return nc, n.leaf.val, true
		} else {
			return nc, nil, false
		}
	}

	// Look for the edge
	idx, child := n.getEdge(search[0])

	// No edge, create one
	if child == nil {
		e := edge{
			label: search[0],
			node: &Node{
				mutateCh: make(chan struct{}),
				leaf: &leafNode{
					mutateCh: make(chan struct{}),
					key:      k,
					val:      v,
				},
				prefix: search,
			},
		}
		nc := t.writeNode(n)
		nc.addEdge(e)
		return nc, nil, false
	}

	// Determine longest prefix of the search key on match
	commonPrefix := longestPrefix(search, child.prefix)
	if commonPrefix == len(child.prefix) {
		search = search[commonPrefix:]
		newChild, oldVal, didUpdate := t.insert(child, k, search, v)
		if newChild != nil {
			nc := t.writeNode(n)
			nc.edges[idx].node = newChild
			return nc, oldVal, didUpdate
		}
		return nil, oldVal, didUpdate
	}

	// Split the node
	nc := t.writeNode(n)
	splitNode := &Node{
		mutateCh: make(chan struct{}),
		prefix:   search[:commonPrefix],
	}
	nc.replaceEdge(edge{
		label: search[0],
		node:  splitNode,
	})

	// Restore the existing child node
	modChild := t.writeNode(child)
	splitNode.addEdge(edge{
		label: modChild.prefix[commonPrefix],
		node:  modChild,
	})
	modChild.prefix = modChild.prefix[commonPrefix:]

	// Create a new leaf node
	leaf := &leafNode{
		mutateCh: make(chan struct{}),
		key:      k,
		val:      v,
	}

	// If the new key is a subset, add to to this node
	search = search[commonPrefix:]
	if len(search) == 0 {
		splitNode.leaf = leaf
		return nc, nil, false
	}

	// Create a new edge for the node
	splitNode.addEdge(edge{
		label: search[0],
		node: &Node{
			mutateCh: make(chan struct{}),
			leaf:     leaf,
			prefix:   search,
		},
	})
	return nc, nil, false
}

// delete does a recursive deletion
func (t *Txn) delete(parent, n *Node, search []byte) (*Node, *leafNode) {
	// Check for key exhaution
	if len(search) == 0 {
		if !n.isLeaf() {
			return nil, nil
		}

		// Remove the leaf node
		nc := t.writeNode(n)
		nc.leaf = nil
		t.mutateLeaf(n.leaf)

		// Check if this node should be merged
		if n != t.root && len(nc.edges) == 1 {
			nc.mergeChild()
		}
		return nc, n.leaf
	}

	// Look for an edge
	label := search[0]
	idx, child := n.getEdge(label)
	if child == nil || !bytes.HasPrefix(search, child.prefix) {
		return nil, nil
	}

	// Consume the search prefix
	search = search[len(child.prefix):]
	newChild, leaf := t.delete(n, child, search)
	if newChild == nil {
		return nil, nil
	}

	// Copy this node
	nc := t.writeNode(n)

	// Delete the edge if the node has no edges
	if newChild.leaf == nil && len(newChild.edges) == 0 {
		nc.delEdge(label)
		if n != t.root && len(nc.edges) == 1 && !nc.isLeaf() {
			nc.mergeChild()
		}
	} else {
		nc.edges[idx].node = newChild
	}
	return nc, leaf
}

// Insert is used to add or update a given key. The return provides
// the previous value and a bool indicating if any was set.
func (t *Txn) Insert(k []byte, v interface{}) (interface{}, bool) {
	newRoot, oldVal, didUpdate := t.insert(t.root, k, k, v)
	if newRoot != nil {
		t.root = newRoot
	}
	if !didUpdate {
		t.size++
	}
	return oldVal, didUpdate
}

// Delete is used to delete a given key. Returns the old value if any,
// and a bool indicating if the key was set.
func (t *Txn) Delete(k []byte) (interface{}, bool) {
	newRoot, leaf := t.delete(nil, t.root, k)
	if newRoot != nil {
		t.root = newRoot
	}
	if leaf != nil {
		t.size--
		return leaf.val, true
	}
	return nil, false
}

// Root returns the current root of the radix tree within this
// transaction. The root is not safe across insert and delete operations,
// but can be used to read the current state during a transaction.
func (t *Txn) Root() *Node {
	return t.root
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Txn) Get(k []byte) (interface{}, bool) {
	return t.root.Get(k)
}

// GetWatch is used to lookup a specific key, returning
// the watch channel, value and if it was found
func (t *Txn) GetWatch(k []byte) (<-chan struct{}, interface{}, bool) {
	return t.root.GetWatch(k)
}

// Commit is used to finalize the transaction and return a new tree
func (t *Txn) Commit() *Tree {
	nt := &Tree{t.root, t.size}
	if t.notifyMutate {
		for leaf := range t.mutatedLeaf {
			close(leaf.mutateCh)
		}
		for node := range t.mutatedNode {
			close(node.mutateCh)
		}
	}
	t.modified = nil
	t.mutatedNode = nil
	t.mutatedLeaf = nil
	return nt
}

// Insert is used to add or update a given key. The return provides
// the new tree, previous value and a bool indicating if any was set.
func (t *Tree) Insert(k []byte, v interface{}) (*Tree, interface{}, bool) {
	txn := t.Txn()
	old, ok := txn.Insert(k, v)
	return txn.Commit(), old, ok
}

// Delete is used to delete a given key. Returns the new tree,
// old value if any, and a bool indicating if the key was set.
func (t *Tree) Delete(k []byte) (*Tree, interface{}, bool) {
	txn := t.Txn()
	old, ok := txn.Delete(k)
	return txn.Commit(), old, ok
}

// Root returns the root node of the tree which can be used for richer
// query operations.
func (t *Tree) Root() *Node {
	return t.root
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Tree) Get(k []byte) (interface{}, bool) {
	return t.root.Get(k)
}

// longestPrefix finds the length of the shared prefix
// of two strings
func longestPrefix(k1, k2 []byte) int {
	max := len(k1)
	if l := len(k2); l < max {
		max = l
	}
	var i int
	for i = 0; i < max; i++ {
		if k1[i] != k2[i] {
			break
		}
	}
	return i
}

// concat two byte slices, returning a third new copy
func concat(a, b []byte) []byte {
	c := make([]byte, len(a)+len(b))
	copy(c, a)
	copy(c[len(a):], b)
	return c
}
<|endoftext|>


<commit_before>package main

type SimpleLibraryResponse struct {
    book Book
    registeredCopyCount int
    availableCopyCount int
    err error
}

func (r *SimpleLibraryResponse) GetBook(fmt.Stringer, error) {
    return r.book, r.err
}

func (r *SimpleLibraryResponse) GetAvailability(int, int) {
    return availableCopyCount, registeredCopyCount
}
<commit_msg>fix: forgot to import "fmt"<commit_after>package main

import (
    "fmt"
)

type SimpleLibraryResponse struct {
    book Book
    registeredCopyCount int
    availableCopyCount int
    err error
}

func (r *SimpleLibraryResponse) GetBook(fmt.Stringer, error) {
    return r.book, r.err
}

func (r *SimpleLibraryResponse) GetAvailability(int, int) {
    return availableCopyCount, registeredCopyCount
}
<|endoftext|>


<commit_before>package prometheus

import (
	"fmt"
	"runtime"
	"runtime/debug"
	"time"
)

type goCollector struct {
	goroutines Gauge
	gcDesc     *Desc

	// metrics to describe and collect
	metrics memStatsMetrics
}

// NewGoCollector returns a collector which exports metrics about the current
// go process.
func NewGoCollector() *goCollector {
	return &goCollector{
		goroutines: NewGauge(GaugeOpts{
			Namespace: "go",
			Name:      "goroutines",
			Help:      "Number of goroutines that currently exist.",
		}),
		gcDesc: NewDesc(
			"go_gc_duration_seconds",
			"A summary of the GC invocation durations.",
			nil, nil),
		metrics: memStatsMetrics{
			{
				desc: NewDesc(
					memstatNamespace("alloc_bytes"),
					"Number of bytes allocated and still in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Alloc) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("alloc_bytes_total"),
					"Total number of bytes allocated, even if freed.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.TotalAlloc) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("sys_bytes"),
					"Number of bytes obtained by system. Sum of all system allocations.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Sys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("lookups_total"),
					"Total number of pointer lookups.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Lookups) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mallocs_total"),
					"Total number of mallocs.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Mallocs) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("frees_total"),
					"Total number of frees.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Frees) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_alloc_bytes"),
					"Number of heap bytes allocated and still in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapAlloc) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_sys_bytes"),
					"Number of heap bytes obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_idle_bytes"),
					"Number of heap bytes waiting to be used.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapIdle) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_inuse_bytes"),
					"Number of heap bytes that are in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_released_bytes"),
					"Number of bytes in heap released to OS.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapReleased) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_objects"),
					"Number of allocated objects.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapObjects) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("stack_inuse_bytes"),
					"Number of bytes in use by the stack allocator.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.StackInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("stack_sys_bytes"),
					"Number of bytes obtained from system for stack allocator.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.StackSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mspan_inuse_bytes"),
					"Number of bytes in use by mspan structures.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MSpanInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mspan_sys_bytes"),
					"Number of bytes used for mspan structures obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MSpanSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mcache_inuse_bytes"),
					"Number of bytes in use by mcache structures.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MCacheInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mcache_sys_bytes"),
					"Number of bytes used for mcache structures obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MCacheSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("buck_hash_sys_bytes"),
					"Number of bytes used by the profiling bucket hash table.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.BuckHashSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("gc_sys_bytes"),
					"Number of bytes used for garbage collection system metadata.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.GCSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("other_sys_bytes"),
					"Number of bytes used for other system allocations.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.OtherSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("next_gc_bytes"),
					"Number of heap bytes when next garbage collection will take place.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.NextGC) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("last_gc_time_seconds"),
					"Number of seconds since 1970 of last garbage collection.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.LastGC*10 ^ 9) },
				valType: GaugeValue,
			},
		},
	}
}

func memstatNamespace(s string) string {
	return fmt.Sprintf("go_memstats_%s", s)
}

// Describe returns all descriptions of the collector.
func (c *goCollector) Describe(ch chan<- *Desc) {
	ch <- c.goroutines.Desc()
	ch <- c.gcDesc

	for _, i := range c.metrics {
		ch <- i.desc
	}
}

// Collect returns the current state of all metrics of the collector.
func (c *goCollector) Collect(ch chan<- Metric) {
	c.goroutines.Set(float64(runtime.NumGoroutine()))
	ch <- c.goroutines

	var stats debug.GCStats
	stats.PauseQuantiles = make([]time.Duration, 5)
	debug.ReadGCStats(&stats)

	quantiles := make(map[float64]float64)
	for idx, pq := range stats.PauseQuantiles[1:] {
		quantiles[float64(idx+1)/float64(len(stats.PauseQuantiles)-1)] = pq.Seconds()
	}
	quantiles[0.0] = stats.PauseQuantiles[0].Seconds()
	ch <- MustNewConstSummary(c.gcDesc, uint64(stats.NumGC), float64(stats.PauseTotal.Seconds()), quantiles)

	ms := &runtime.MemStats{}
	runtime.ReadMemStats(ms)
	for _, i := range c.metrics {
		ch <- MustNewConstMetric(i.desc, i.valType, i.eval(ms))
	}
}

// memStatsMetrics provide description, value, and value type for memstat metrics.
type memStatsMetrics []struct {
	desc    *Desc
	eval    func(*runtime.MemStats) float64
	valType ValueType
}
<commit_msg>heap bytes released total<commit_after>package prometheus

import (
	"fmt"
	"runtime"
	"runtime/debug"
	"time"
)

type goCollector struct {
	goroutines Gauge
	gcDesc     *Desc

	// metrics to describe and collect
	metrics memStatsMetrics
}

// NewGoCollector returns a collector which exports metrics about the current
// go process.
func NewGoCollector() *goCollector {
	return &goCollector{
		goroutines: NewGauge(GaugeOpts{
			Namespace: "go",
			Name:      "goroutines",
			Help:      "Number of goroutines that currently exist.",
		}),
		gcDesc: NewDesc(
			"go_gc_duration_seconds",
			"A summary of the GC invocation durations.",
			nil, nil),
		metrics: memStatsMetrics{
			{
				desc: NewDesc(
					memstatNamespace("alloc_bytes"),
					"Number of bytes allocated and still in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Alloc) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("alloc_bytes_total"),
					"Total number of bytes allocated, even if freed.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.TotalAlloc) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("sys_bytes"),
					"Number of bytes obtained by system. Sum of all system allocations.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Sys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("lookups_total"),
					"Total number of pointer lookups.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Lookups) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mallocs_total"),
					"Total number of mallocs.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Mallocs) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("frees_total"),
					"Total number of frees.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.Frees) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_alloc_bytes"),
					"Number of heap bytes allocated and still in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapAlloc) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_sys_bytes"),
					"Number of heap bytes obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_idle_bytes"),
					"Number of heap bytes waiting to be used.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapIdle) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_inuse_bytes"),
					"Number of heap bytes that are in use.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_released_bytes_total"),
					"Total number of heap bytes released to OS.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapReleased) },
				valType: CounterValue,
			}, {
				desc: NewDesc(
					memstatNamespace("heap_objects"),
					"Number of allocated objects.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.HeapObjects) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("stack_inuse_bytes"),
					"Number of bytes in use by the stack allocator.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.StackInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("stack_sys_bytes"),
					"Number of bytes obtained from system for stack allocator.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.StackSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mspan_inuse_bytes"),
					"Number of bytes in use by mspan structures.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MSpanInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mspan_sys_bytes"),
					"Number of bytes used for mspan structures obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MSpanSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mcache_inuse_bytes"),
					"Number of bytes in use by mcache structures.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MCacheInuse) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("mcache_sys_bytes"),
					"Number of bytes used for mcache structures obtained from system.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.MCacheSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("buck_hash_sys_bytes"),
					"Number of bytes used by the profiling bucket hash table.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.BuckHashSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("gc_sys_bytes"),
					"Number of bytes used for garbage collection system metadata.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.GCSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("other_sys_bytes"),
					"Number of bytes used for other system allocations.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.OtherSys) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("next_gc_bytes"),
					"Number of heap bytes when next garbage collection will take place.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.NextGC) },
				valType: GaugeValue,
			}, {
				desc: NewDesc(
					memstatNamespace("last_gc_time_seconds"),
					"Number of seconds since 1970 of last garbage collection.",
					nil, nil,
				),
				eval:    func(ms *runtime.MemStats) float64 { return float64(ms.LastGC*10 ^ 9) },
				valType: GaugeValue,
			},
		},
	}
}

func memstatNamespace(s string) string {
	return fmt.Sprintf("go_memstats_%s", s)
}

// Describe returns all descriptions of the collector.
func (c *goCollector) Describe(ch chan<- *Desc) {
	ch <- c.goroutines.Desc()
	ch <- c.gcDesc

	for _, i := range c.metrics {
		ch <- i.desc
	}
}

// Collect returns the current state of all metrics of the collector.
func (c *goCollector) Collect(ch chan<- Metric) {
	c.goroutines.Set(float64(runtime.NumGoroutine()))
	ch <- c.goroutines

	var stats debug.GCStats
	stats.PauseQuantiles = make([]time.Duration, 5)
	debug.ReadGCStats(&stats)

	quantiles := make(map[float64]float64)
	for idx, pq := range stats.PauseQuantiles[1:] {
		quantiles[float64(idx+1)/float64(len(stats.PauseQuantiles)-1)] = pq.Seconds()
	}
	quantiles[0.0] = stats.PauseQuantiles[0].Seconds()
	ch <- MustNewConstSummary(c.gcDesc, uint64(stats.NumGC), float64(stats.PauseTotal.Seconds()), quantiles)

	ms := &runtime.MemStats{}
	runtime.ReadMemStats(ms)
	for _, i := range c.metrics {
		ch <- MustNewConstMetric(i.desc, i.valType, i.eval(ms))
	}
}

// memStatsMetrics provide description, value, and value type for memstat metrics.
type memStatsMetrics []struct {
	desc    *Desc
	eval    func(*runtime.MemStats) float64
	valType ValueType
}
<|endoftext|>


<commit_before>package node

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"github.com/anycable/anycable-go/utils"
	"github.com/apex/log"
	"github.com/gorilla/websocket"
	nanoid "github.com/matoous/go-nanoid"
)

const (
	// DefaultCloseStatus is what it states)
	DefaultCloseStatus = 3000

	writeWait      = 10 * time.Second
	maxMessageSize = 512
	pingInterval   = 3 * time.Second
)

// Session represents active client
type Session struct {
	node          *Node
	ws            *websocket.Conn
	path          string
	headers       map[string]string
	subscriptions map[string]bool
	send          chan []byte
	closed        bool
	connected     bool
	mu            sync.Mutex
	pingTimer     *time.Timer

	UID         string
	Identifiers string
	Log         *log.Entry
}

type pingMessage struct {
	Type    string      `json:"type"`
	Message interface{} `json:"message"`
}

func (p *pingMessage) toJSON() []byte {
	jsonStr, err := json.Marshal(&p)
	if err != nil {
		panic("Failed to build ping JSON 😲")
	}
	return jsonStr
}

// NewSession build a new Session struct from ws connetion and http request
func NewSession(node *Node, ws *websocket.Conn, request *http.Request) (*Session, error) {
	path := request.URL.String()
	headers := utils.FetchHeaders(request, node.Config.Headers)

	session := &Session{
		node:          node,
		ws:            ws,
		path:          path,
		headers:       headers,
		subscriptions: make(map[string]bool),
		send:          make(chan []byte, 256),
		closed:        false,
		connected:     false,
	}

	uid, err := nanoid.Nanoid()

	if err != nil {
		defer session.Close("Nanoid Error")
		return nil, err
	}

	session.UID = uid

	ctx := log.WithFields(log.Fields{
		"sid": session.UID,
	})

	session.Log = ctx

	err = node.Authenticate(session, path, &headers)

	if err != nil {
		defer session.Close("Auth Error")
		return nil, err
	}

	go session.SendMessages()

	session.addPing()

	return session, nil
}

// SendMessages waits for incoming messages and send them to the client connection
func (s *Session) SendMessages() {
	defer s.Disconnect("Write Failed")
	for {
		select {
		case message, ok := <-s.send:
			if !ok {
				return
			}

			err := s.write(message, time.Now().Add(writeWait))

			if err != nil {
				return
			}
		}
	}
}

func (s *Session) write(message []byte, deadline time.Time) error {
	s.ws.SetWriteDeadline(deadline)

	w, err := s.ws.NextWriter(websocket.TextMessage)

	if err != nil {
		return err
	}

	w.Write(message)

	return w.Close()
}

// Send data to client connection
func (s *Session) Send(msg []byte) {
	select {
	case s.send <- msg:
	default:
		s.mu.Lock()

		if s.send != nil {
			close(s.send)
			defer s.Disconnect("Write failed")
		}

		defer s.mu.Unlock()
		s.send = nil
	}
}

// ReadMessages reads messages from ws connection and send them to node
func (s *Session) ReadMessages() {
	// s.ws.SetReadLimit(MaxMessageSize)

	defer s.Disconnect("")

	for {
		_, message, err := s.ws.ReadMessage()

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				s.Log.Debugf("Websocket read error: %v", err)
			}
			break
		}

		s.node.HandleCommand(s, message)
	}
}

// Disconnect enqueues RPC disconnect request and closes the connection
func (s *Session) Disconnect(reason string) {
	s.mu.Lock()
	if !s.connected {
		s.node.Disconnect(s)
	}
	s.connected = false
	s.mu.Unlock()

	s.Close(reason)
}

// Close websocket connection with the specified reason
func (s *Session) Close(reason string) {
	s.mu.Lock()
	if s.closed {
		return
	}
	s.closed = true
	s.mu.Unlock()

	if s.pingTimer != nil {
		s.pingTimer.Stop()
	}

	// TODO: make deadline and status code configurable
	deadline := time.Now().Add(time.Second)
	msg := websocket.FormatCloseMessage(DefaultCloseStatus, reason)
	s.ws.WriteControl(websocket.CloseMessage, msg, deadline)
	s.ws.Close()
}

func (s *Session) sendPing() {
	deadline := time.Now().Add(pingInterval / 2)
	err := s.write(newPingMessage(), deadline)

	if err == nil {
		s.addPing()
	}
}

func (s *Session) addPing() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return
	}

	s.pingTimer = time.AfterFunc(pingInterval, s.sendPing)
}

func newPingMessage() []byte {
	return (&pingMessage{Type: "ping", Message: time.Now().Unix()}).toJSON()
}
<commit_msg>Fix concurrent writes<commit_after>package node

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"github.com/anycable/anycable-go/utils"
	"github.com/apex/log"
	"github.com/gorilla/websocket"
	nanoid "github.com/matoous/go-nanoid"
)

const (
	// DefaultCloseStatus is what it states)
	DefaultCloseStatus = 3000

	writeWait      = 10 * time.Second
	maxMessageSize = 512
	pingInterval   = 3 * time.Second
)

// Session represents active client
type Session struct {
	node          *Node
	ws            *websocket.Conn
	path          string
	headers       map[string]string
	subscriptions map[string]bool
	send          chan []byte
	closed        bool
	connected     bool
	mu            sync.Mutex
	pingTimer     *time.Timer

	UID         string
	Identifiers string
	Log         *log.Entry
}

type pingMessage struct {
	Type    string      `json:"type"`
	Message interface{} `json:"message"`
}

func (p *pingMessage) toJSON() []byte {
	jsonStr, err := json.Marshal(&p)
	if err != nil {
		panic("Failed to build ping JSON 😲")
	}
	return jsonStr
}

// NewSession build a new Session struct from ws connetion and http request
func NewSession(node *Node, ws *websocket.Conn, request *http.Request) (*Session, error) {
	path := request.URL.String()
	headers := utils.FetchHeaders(request, node.Config.Headers)

	session := &Session{
		node:          node,
		ws:            ws,
		path:          path,
		headers:       headers,
		subscriptions: make(map[string]bool),
		send:          make(chan []byte, 256),
		closed:        false,
		connected:     false,
	}

	uid, err := nanoid.Nanoid()

	if err != nil {
		defer session.Close("Nanoid Error")
		return nil, err
	}

	session.UID = uid

	ctx := log.WithFields(log.Fields{
		"sid": session.UID,
	})

	session.Log = ctx

	err = node.Authenticate(session, path, &headers)

	if err != nil {
		defer session.Close("Auth Error")
		return nil, err
	}

	go session.SendMessages()

	session.addPing()

	return session, nil
}

// SendMessages waits for incoming messages and send them to the client connection
func (s *Session) SendMessages() {
	defer s.Disconnect("Write Failed")
	for {
		select {
		case message, ok := <-s.send:
			if !ok {
				return
			}

			err := s.write(message, time.Now().Add(writeWait))

			if err != nil {
				return
			}
		}
	}
}

func (s *Session) write(message []byte, deadline time.Time) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.ws.SetWriteDeadline(deadline)

	w, err := s.ws.NextWriter(websocket.TextMessage)

	if err != nil {
		return err
	}

	w.Write(message)

	return w.Close()
}

// Send data to client connection
func (s *Session) Send(msg []byte) {
	select {
	case s.send <- msg:
	default:
		s.mu.Lock()

		if s.send != nil {
			close(s.send)
			defer s.Disconnect("Write failed")
		}

		defer s.mu.Unlock()
		s.send = nil
	}
}

// ReadMessages reads messages from ws connection and send them to node
func (s *Session) ReadMessages() {
	// s.ws.SetReadLimit(MaxMessageSize)

	defer s.Disconnect("")

	for {
		_, message, err := s.ws.ReadMessage()

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				s.Log.Debugf("Websocket read error: %v", err)
			}
			break
		}

		s.node.HandleCommand(s, message)
	}
}

// Disconnect enqueues RPC disconnect request and closes the connection
func (s *Session) Disconnect(reason string) {
	s.mu.Lock()
	if !s.connected {
		s.node.Disconnect(s)
	}
	s.connected = false
	s.mu.Unlock()

	s.Close(reason)
}

// Close websocket connection with the specified reason
func (s *Session) Close(reason string) {
	s.mu.Lock()
	if s.closed {
		return
	}
	s.closed = true
	s.mu.Unlock()

	if s.pingTimer != nil {
		s.pingTimer.Stop()
	}

	// TODO: make deadline and status code configurable
	deadline := time.Now().Add(time.Second)
	msg := websocket.FormatCloseMessage(DefaultCloseStatus, reason)
	s.ws.WriteControl(websocket.CloseMessage, msg, deadline)
	s.ws.Close()
}

func (s *Session) sendPing() {
	deadline := time.Now().Add(pingInterval / 2)
	err := s.write(newPingMessage(), deadline)

	if err == nil {
		s.addPing()
	}
}

func (s *Session) addPing() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return
	}

	s.pingTimer = time.AfterFunc(pingInterval, s.sendPing)
}

func newPingMessage() []byte {
	return (&pingMessage{Type: "ping", Message: time.Now().Unix()}).toJSON()
}
<|endoftext|>


<commit_before>package restic

import (
	"os"
	"syscall"
)

func (node *Node) OpenForReading() (*os.File, error) {
	file, err := os.OpenFile(node.path, os.O_RDONLY, 0)
	if os.IsPermission(err) {
		return os.OpenFile(node.path, os.O_RDONLY, 0)
	}
	return file, err
}

// mknod() creates a filesystem node (file, device
// special file, or named pipe) named pathname, with attributes
// specified by mode and dev.
var mknod = func(path string, mode uint32, dev int) (err error) {
	panic("mknod not implemented")
}

// Windows doesn't need lchown
var lchown = func(path string, uid int, gid int) (err error) {
	return nil
}

func (node Node) restoreSymlinkTimestamps(path string, utimes [2]syscall.Timespec) error {
	return nil
}

type statWin syscall.Win32FileAttributeData

func toStatT(i interface{}) (statT, bool) {
	if i == nil {
		return nil, false
	}
	s, ok := i.(*syscall.Win32FileAttributeData)
	if ok && s != nil {
		return statWin(*s), true
	}
	return nil, false
}

func (s statWin) dev() uint64   { return 0 }
func (s statWin) ino() uint64   { return 0 }
func (s statWin) nlink() uint64 { return 0 }
func (s statWin) uid() uint32   { return 0 }
func (s statWin) gid() uint32   { return 0 }
func (s statWin) rdev() uint64  { return 0 }

func (s statWin) size() int64 {
	return int64(s.FileSizeLow) | (int64(s.FileSizeHigh) << 32)
}

func (s statWin) atim() syscall.Timespec {
	return syscall.NsecToTimespec(s.LastAccessTime.Nanoseconds())
}

func (s statWin) mtim() syscall.Timespec {
	return syscall.NsecToTimespec(s.LastWriteTime.Nanoseconds())
}

func (s statWin) ctim() syscall.Timespec {
	return syscall.NsecToTimespec(s.CreationTime.Nanoseconds())
}
<commit_msg>Remove redundant code.<commit_after>package restic

import (
	"os"
	"syscall"
)

func (node *Node) OpenForReading() (*os.File, error) {
	return os.OpenFile(node.path, os.O_RDONLY, 0)
}

// mknod() creates a filesystem node (file, device
// special file, or named pipe) named pathname, with attributes
// specified by mode and dev.
var mknod = func(path string, mode uint32, dev int) (err error) {
	panic("mknod not implemented")
}

// Windows doesn't need lchown
var lchown = func(path string, uid int, gid int) (err error) {
	return nil
}

func (node Node) restoreSymlinkTimestamps(path string, utimes [2]syscall.Timespec) error {
	return nil
}

type statWin syscall.Win32FileAttributeData

func toStatT(i interface{}) (statT, bool) {
	if i == nil {
		return nil, false
	}
	s, ok := i.(*syscall.Win32FileAttributeData)
	if ok && s != nil {
		return statWin(*s), true
	}
	return nil, false
}

func (s statWin) dev() uint64   { return 0 }
func (s statWin) ino() uint64   { return 0 }
func (s statWin) nlink() uint64 { return 0 }
func (s statWin) uid() uint32   { return 0 }
func (s statWin) gid() uint32   { return 0 }
func (s statWin) rdev() uint64  { return 0 }

func (s statWin) size() int64 {
	return int64(s.FileSizeLow) | (int64(s.FileSizeHigh) << 32)
}

func (s statWin) atim() syscall.Timespec {
	return syscall.NsecToTimespec(s.LastAccessTime.Nanoseconds())
}

func (s statWin) mtim() syscall.Timespec {
	return syscall.NsecToTimespec(s.LastWriteTime.Nanoseconds())
}

func (s statWin) ctim() syscall.Timespec {
	return syscall.NsecToTimespec(s.CreationTime.Nanoseconds())
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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

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

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

package main

import (
	"strings"
	"text/template"

	"github.com/spf13/pflag"
	"github.com/nodetemple/cobra"
	"github.com/nodetemple/nodetemple/version"
)

var (
	cmdHelp = &cobra.Command{
		Use:   "help [command]",
		Short: "Print usage information",
		Long:  "Print usage information for any command",
		Run:   cmdNodectl.HelpFunc(),
		PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
		PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	}

	commandUsageTemplate *template.Template
	templFuncs           = template.FuncMap{
		"descToLines": func(s string) []string {
			return strings.Split(strings.Trim(s, "\n\t "), "\n")
		},
		"cmdName": func(cmd *cobra.Command, startCmd *cobra.Command) string {
			parts := []string{cmd.Name()}
			for cmd.HasParent() && cmd.Parent().Name() != startCmd.Name() {
				cmd = cmd.Parent()
				parts = append([]string{cmd.Name()}, parts...)
			}
			return strings.Join(parts, " ")
		},
	}
)

func init() {
	commandUsage := `
{{ $cmd := .Cmd }}\
{{ $cmdname := cmdName .Cmd .Cmd.Root }}\
NAME:
{{ if not .Cmd.HasParent }}\
{{printf "\t%s - %s" .Cmd.Name .Cmd.Short}}
{{else}}\
{{printf "\t%s - %s" $cmdname .Cmd.Short}}
{{end}}\

USAGE:
{{printf "\t%s" .Cmd.UseLine}}
{{ if not .Cmd.HasParent }}\

VERSION:
{{printf "\t%s" .Version}}
{{end}}\
{{if .Cmd.HasSubCommands}}\

COMMANDS:
{{range .SubCommands}}\
{{ $cmdname := cmdName . $cmd }}\
{{ if .Runnable }}\
{{printf "\t%s\t%s" $cmdname .Short}}
{{end}}\
{{end}}\
{{end}}\
{{ if .Cmd.Long }}\

DESCRIPTION:
{{range $line := descToLines .Cmd.Long}}{{printf "\t%s" $line}}
{{end}}\
{{end}}\
{{if .Cmd.HasLocalFlags}}\

{{ if not .Cmd.HasParent }}GLOBAL {{end}}FLAGS:
{{.Cmd.LocalFlags.FlagUsages}}\
{{ if not .Cmd.HasParent }}
Global flags can also be configured via upper-case environment variables prefixed with "{{.EnvFlag}}_"
For example: "--some-flag" => "{{.EnvFlag}}_SOME_FLAG"
{{end}}\
{{end}}\
{{if .Cmd.HasInheritedFlags}}\

GLOBAL FLAGS:
{{.Cmd.InheritedFlags.FlagUsages}}
Global flags can also be configured via upper-case environment variables prefixed with "{{.EnvFlag}}_"
For example: "--some-flag" => "{{.EnvFlag}}_SOME_FLAG"
{{end}}\
{{ if .Cmd.HasSubCommands }}
Run "{{.Cmd.CommandPath}} help [command]" for more information about a specific command usage
{{end}}\
{{ if .Cmd.HasParent }}
Run "{{.Executable}} help" for more information about a general usage
{{end}}`[1:]

	commandUsageTemplate = template.Must(template.New("command_usage").Funcs(templFuncs).Parse(strings.Replace(commandUsage, "\\\n", "", -1)))
}

func getSubCommands(cmd *cobra.Command) []*cobra.Command {
	subCommands := []*cobra.Command{}
	for _, subCmd := range cmd.Commands() {
		subCommands = append(subCommands, subCmd)
		subCommands = append(subCommands, getSubCommands(subCmd)...)
	}
	return subCommands
}

func usageFunc(cmd *cobra.Command) error {
	subCommands := getSubCommands(cmd)
	err := commandUsageTemplate.Execute(tabOut, struct {
		Executable  string
		Cmd         *cobra.Command
		CmdFlags    *pflag.FlagSet
		SubCommands []*cobra.Command
		EnvFlag     string
		Version     string
	}{
		cliName,
		cmd,
		cmd.Flags(),
		subCommands,
		strings.ToUpper(cliName),
		version.Version,
	})
	tabOut.Flush()
	return err
}
<commit_msg>Detailed usage texts<commit_after>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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

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

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

package main

import (
	"strings"
	"text/template"

	"github.com/spf13/pflag"
	"github.com/nodetemple/cobra"
	"github.com/nodetemple/nodetemple/version"
)

var (
	cmdHelp = &cobra.Command{
		Use:   "help [command]",
		Short: "Print usage information",
		Long:  "Print usage information for any command",
		Run:   cmdNodectl.HelpFunc(),
		PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
		PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	}

	commandUsageTemplate *template.Template
	templFuncs           = template.FuncMap{
		"descToLines": func(s string) []string {
			return strings.Split(strings.Trim(s, "\n\t "), "\n")
		},
		"cmdName": func(cmd *cobra.Command, startCmd *cobra.Command) string {
			parts := []string{cmd.Name()}
			for cmd.HasParent() && cmd.Parent().Name() != startCmd.Name() {
				cmd = cmd.Parent()
				parts = append([]string{cmd.Name()}, parts...)
			}
			return strings.Join(parts, " ")
		},
	}
)

func init() {
	commandUsage := `
{{ $cmd := .Cmd }}\
{{ $cmdname := cmdName .Cmd .Cmd.Root }}\
NAME:
{{ if not .Cmd.HasParent }}\
{{printf "\t%s - %s" .Cmd.Name .Cmd.Short}}
{{else}}\
{{printf "\t%s - %s" $cmdname .Cmd.Short}}
{{end}}\

USAGE:
{{printf "\t%s" .Cmd.UseLine}}
{{ if not .Cmd.HasParent }}\

VERSION:
{{printf "\t%s" .Version}}
{{end}}\
{{if .Cmd.HasSubCommands}}\

COMMANDS:
{{range .SubCommands}}\
{{ $cmdname := cmdName . $cmd }}\
{{ if .Runnable }}\
{{printf "\t%s\t%s" $cmdname .Short}}
{{end}}\
{{end}}\
{{end}}\
{{ if .Cmd.Long }}\

DESCRIPTION:
{{range $line := descToLines .Cmd.Long}}{{printf "\t%s" $line}}
{{end}}\
{{end}}\
{{if .Cmd.HasLocalFlags}}\

{{ if not .Cmd.HasParent }}GLOBAL {{end}}FLAGS:
{{.Cmd.LocalFlags.FlagUsages}}\
{{ if not .Cmd.HasParent }}
Global flags can also be configured via upper-case environment variables prefixed with '{{.EnvFlag}}_'
For example: '--some-flag' => '{{.EnvFlag}}_SOME_FLAG'
{{end}}\
{{end}}\
{{if .Cmd.HasInheritedFlags}}\

GLOBAL FLAGS:
{{.Cmd.InheritedFlags.FlagUsages}}
Global flags can also be configured via upper-case environment variables prefixed with '{{.EnvFlag}}_'
For example: '--some-flag' => '{{.EnvFlag}}_SOME_FLAG'
{{end}}\
{{ if .Cmd.HasSubCommands }}
Run '{{.Cmd.CommandPath}} help [command]' for more information about a specific command usage
{{end}}\
{{ if .Cmd.HasParent }}
Run '{{.Executable}} help' for more information about a general usage
{{end}}`[1:]

	commandUsageTemplate = template.Must(template.New("command_usage").Funcs(templFuncs).Parse(strings.Replace(commandUsage, "\\\n", "", -1)))
}

func getSubCommands(cmd *cobra.Command) []*cobra.Command {
	subCommands := []*cobra.Command{}
	for _, subCmd := range cmd.Commands() {
		subCommands = append(subCommands, subCmd)
		subCommands = append(subCommands, getSubCommands(subCmd)...)
	}
	return subCommands
}

func usageFunc(cmd *cobra.Command) error {
	subCommands := getSubCommands(cmd)
	err := commandUsageTemplate.Execute(tabOut, struct {
		Executable  string
		Cmd         *cobra.Command
		CmdFlags    *pflag.FlagSet
		SubCommands []*cobra.Command
		EnvFlag     string
		Version     string
	}{
		cliName,
		cmd,
		cmd.Flags(),
		subCommands,
		strings.ToUpper(cliName),
		version.Version,
	})
	tabOut.Flush()
	return err
}
<|endoftext|>


<commit_before>package util

import "reflect"

func CheckValue(values ...string) bool {
	for _, value := range values {
		if value == "" {
			return false
		}
	}

	return true
}

// passing struct and check all the attribute
// false if there is empty value
func CheckAttribute(input interface{}) bool {
	object := reflect.ValueOf(input)

	for index := 0; index < object.NumField(); index++ {

		if IsZeroOfUnderlyingType(object.Field(index).Interface()) {
			return false
		}
	}

	return true
}

func IsZeroOfUnderlyingType(objectValue interface{}) bool {
	return reflect.DeepEqual(objectValue, reflect.Zero(reflect.TypeOf(objectValue)).Interface())
}
<commit_msg>add some comment<commit_after>package util

import "reflect"

// Check if all input values in not nil.
func CheckValue(values ...string) bool {
	for _, value := range values {
		if value == "" {
			return false
		}
	}

	return true
}

// passing struct and check all the attribute,
// false if there is empty value.
func CheckAttribute(input interface{}) bool {
	object := reflect.ValueOf(input)

	for index := 0; index < object.NumField(); index++ {

		if IsZeroOfUnderlyingType(object.Field(index).Interface()) {
			return false
		}
	}

	return true
}

// compare the object value to non assign value.
func IsZeroOfUnderlyingType(objectValue interface{}) bool {
	return reflect.DeepEqual(objectValue, reflect.Zero(reflect.TypeOf(objectValue)).Interface())
}
<|endoftext|>


<commit_before>package state

import (
	"errors"
	"regexp"
	"strings"

	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/utils"
)

//go:generate go tool yacc -o ./constraints_gen.go  ./constraints_parser.y

var uniqWhat = []string{
	"hostname",
}

var likeWhat = []string{
	"hostname",
	"agentid",
}

var equalWhat = []string{
	"hostname",
	"agentid",
}

type ConstraintParamHolder struct {
	Slot  *Slot
	Offer *mesos.Offer
}

type Statement interface {
	Eval() bool
	Valid() error
	SetContext(ctx *ConstraintParamHolder)
}

// not (unique hostname)
type NotStatement struct {
	Op1 Statement
}

func (ns *NotStatement) Eval() bool {
	return !ns.Op1.Eval()
}

func (ns *NotStatement) Valid() error {
	err := ns.Op1.Valid()
	if err != nil {
		return err
	}

	return nil
}

func (ns *NotStatement) SetContext(ctx *ConstraintParamHolder) {
	ns.Op1.SetContext(ctx)
}

// and (unique hostname) (unique ip)
// and (not (unique hostname)) (unique ip)
type AndStatement struct {
	Op1 Statement
	Op2 Statement
}

func (as *AndStatement) Eval() bool {
	return as.Op2.Eval() && as.Op1.Eval()
}

func (as *AndStatement) Valid() error {
	err := as.Op1.Valid()
	if err != nil {
		return err
	}

	err1 := as.Op2.Valid()
	if err1 != nil {
		return err1
	}

	return nil
}

func (as *AndStatement) SetContext(ctx *ConstraintParamHolder) {
	as.Op1.SetContext(ctx)
	as.Op2.SetContext(ctx)
}

// or (like ip foobar) (unique hostname)
type OrStatement struct {
	ConstraintParamHolder
	Op1 Statement
	Op2 Statement
}

func (os *OrStatement) Eval() bool {
	return os.Op2.Eval() || os.Op1.Eval()
}

func (os *OrStatement) Valid() error {
	err := os.Op1.Valid()
	if err != nil {
		return err
	}

	err1 := os.Op2.Valid()
	if err1 != nil {
		return err1
	}

	return nil
}

func (os *OrStatement) SetContext(ctx *ConstraintParamHolder) {
	os.Op1.SetContext(ctx)
	os.Op2.SetContext(ctx)
}

// (unique hostname)
type UniqueStatment struct {
	ConstraintParamHolder
	What string
}

func (us *UniqueStatment) Eval() bool {
	if us.What == "hostname" {
		slotsOnHost := OfferAllocatorInstance().SlotsByHostname(us.Offer.GetHostname())
		for _, slotOnHost := range slotsOnHost { // slots belongs to same app on same host
			if strings.SplitN(slotOnHost, "-", 2)[1] == strings.SplitN(us.Slot.ID, "-", 2)[1] {
				return false
			}
		}
	}

	if us.What == "agentid" {
		slotsOnAgent := OfferAllocatorInstance().SlotsByAgentID(*us.Offer.GetAgentId().Value)
		for _, slotOnAgent := range slotsOnAgent { // slots belongs to same app on same agentID
			if strings.SplitN(slotOnAgent, "-", 2)[1] == strings.SplitN(us.Slot.ID, "-", 2)[1] {
				return false
			}
		}
	}

	return true
}

func (us *UniqueStatment) Valid() error {
	if utils.SliceContains(uniqWhat, us.What) {
		return nil
	} else {
		return errors.New("only hostname is supported for the time being")
	}
}

func (us *UniqueStatment) SetContext(ctx *ConstraintParamHolder) {
	us.Offer = ctx.Offer
	us.Slot = ctx.Slot
}

// like hostname foobar*
type LikeStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (ls *LikeStatement) Eval() bool {
	r := regexp.MustCompile(ls.Regex)
	if ls.What == "hostname" {
		return r.MatchString(ls.Offer.GetHostname())
	}

	if ls.What == "agentid" {
		return r.MatchString(*ls.Offer.GetAgentId().Value)
	}

	// user defined attributes match
	for _, attr := range ls.Offer.Attributes {
		if attr.GetName() == ls.What && attr.GetType() == mesos.Value_TEXT {
			return r.MatchString(*attr.GetText().Value)
		}
	}

	return false
}

func (ls *LikeStatement) Valid() error {
	if utils.SliceContains(likeWhat, ls.What) {
		return nil
	} else {
		return errors.New("only hostname, agentid are supported")
	}
}

func (ls *LikeStatement) SetContext(ctx *ConstraintParamHolder) {
	ls.Offer = ctx.Offer
	ls.Slot = ctx.Slot
}

// equal hostname xxxx
type EqualStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (ls *EqualStatement) Eval() bool {
	if ls.What == "hostname" {
		return ls.Offer.GetHostname() == ls.Regex
	}

	if ls.What == "agentid" {
		return *ls.Offer.GetAgentId().Value == ls.Regex
	}

	// user defined attributes match
	for _, attr := range ls.Offer.Attributes {
		if attr.GetName() == ls.What && attr.GetType() == mesos.Value_TEXT {
			return *attr.GetText().Value == ls.Regex
		}
	}

	return false
}

func (ls *EqualStatement) Valid() error {
	if ls.What != "" && ls.Regex != "" {
		return nil
	} else {
		return errors.New("Equal statement must contain two operand")
	}
}

func (ls *EqualStatement) SetContext(ctx *ConstraintParamHolder) {
	ls.Offer = ctx.Offer
	ls.Slot = ctx.Slot
}

// contains hostname barfoo
type ContainsStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (cs *ContainsStatement) Eval() bool {
	if cs.What == "hostname" {
		return strings.Contains(cs.Offer.GetHostname(), cs.Regex)
	}

	if cs.What == "agentid" {
		return strings.Contains(*cs.Offer.GetAgentId().Value, cs.Regex)
	}

	// user defined attributes match
	for _, attr := range cs.Offer.Attributes {
		if attr.GetName() == cs.What && attr.GetType() == mesos.Value_TEXT {
			return strings.Contains(*attr.GetText().Value, cs.Regex)
		}
	}

	return false
}

func (cs *ContainsStatement) Valid() error {
	if utils.SliceContains(likeWhat, cs.What) {
		return nil
	} else {
		return errors.New("only hostname, ip, agentid are supported")
	}
}

func (cs *ContainsStatement) SetContext(ctx *ConstraintParamHolder) {
	cs.Offer = ctx.Offer
	cs.Slot = ctx.Slot
}
<commit_msg>Remove valid check for like/contains operator<commit_after>package state

import (
	"errors"
	"regexp"
	"strings"

	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/utils"
)

//go:generate go tool yacc -o ./constraints_gen.go  ./constraints_parser.y

var uniqWhat = []string{
	"hostname",
}

type ConstraintParamHolder struct {
	Slot  *Slot
	Offer *mesos.Offer
}

type Statement interface {
	Eval() bool
	Valid() error
	SetContext(ctx *ConstraintParamHolder)
}

// not (unique hostname)
type NotStatement struct {
	Op1 Statement
}

func (ns *NotStatement) Eval() bool {
	return !ns.Op1.Eval()
}

func (ns *NotStatement) Valid() error {
	err := ns.Op1.Valid()
	if err != nil {
		return err
	}

	return nil
}

func (ns *NotStatement) SetContext(ctx *ConstraintParamHolder) {
	ns.Op1.SetContext(ctx)
}

// and (unique hostname) (unique ip)
// and (not (unique hostname)) (unique ip)
type AndStatement struct {
	Op1 Statement
	Op2 Statement
}

func (as *AndStatement) Eval() bool {
	return as.Op2.Eval() && as.Op1.Eval()
}

func (as *AndStatement) Valid() error {
	err := as.Op1.Valid()
	if err != nil {
		return err
	}

	err1 := as.Op2.Valid()
	if err1 != nil {
		return err1
	}

	return nil
}

func (as *AndStatement) SetContext(ctx *ConstraintParamHolder) {
	as.Op1.SetContext(ctx)
	as.Op2.SetContext(ctx)
}

// or (like ip foobar) (unique hostname)
type OrStatement struct {
	ConstraintParamHolder
	Op1 Statement
	Op2 Statement
}

func (os *OrStatement) Eval() bool {
	return os.Op2.Eval() || os.Op1.Eval()
}

func (os *OrStatement) Valid() error {
	err := os.Op1.Valid()
	if err != nil {
		return err
	}

	err1 := os.Op2.Valid()
	if err1 != nil {
		return err1
	}

	return nil
}

func (os *OrStatement) SetContext(ctx *ConstraintParamHolder) {
	os.Op1.SetContext(ctx)
	os.Op2.SetContext(ctx)
}

// (unique hostname)
type UniqueStatment struct {
	ConstraintParamHolder
	What string
}

func (us *UniqueStatment) Eval() bool {
	if us.What == "hostname" {
		slotsOnHost := OfferAllocatorInstance().SlotsByHostname(us.Offer.GetHostname())
		for _, slotOnHost := range slotsOnHost { // slots belongs to same app on same host
			if strings.SplitN(slotOnHost, "-", 2)[1] == strings.SplitN(us.Slot.ID, "-", 2)[1] {
				return false
			}
		}
	}

	if us.What == "agentid" {
		slotsOnAgent := OfferAllocatorInstance().SlotsByAgentID(*us.Offer.GetAgentId().Value)
		for _, slotOnAgent := range slotsOnAgent { // slots belongs to same app on same agentID
			if strings.SplitN(slotOnAgent, "-", 2)[1] == strings.SplitN(us.Slot.ID, "-", 2)[1] {
				return false
			}
		}
	}

	return true
}

func (us *UniqueStatment) Valid() error {
	if utils.SliceContains(uniqWhat, us.What) {
		return nil
	} else {
		return errors.New("only hostname is supported for the time being")
	}
}

func (us *UniqueStatment) SetContext(ctx *ConstraintParamHolder) {
	us.Offer = ctx.Offer
	us.Slot = ctx.Slot
}

// like hostname foobar*
type LikeStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (ls *LikeStatement) Eval() bool {
	r := regexp.MustCompile(ls.Regex)
	if ls.What == "hostname" {
		return r.MatchString(ls.Offer.GetHostname())
	}

	if ls.What == "agentid" {
		return r.MatchString(*ls.Offer.GetAgentId().Value)
	}

	// user defined attributes match
	for _, attr := range ls.Offer.Attributes {
		if attr.GetName() == ls.What && attr.GetType() == mesos.Value_TEXT {
			return r.MatchString(*attr.GetText().Value)
		}
	}

	return false
}

func (ls *LikeStatement) Valid() error {
	return nil
}

func (ls *LikeStatement) SetContext(ctx *ConstraintParamHolder) {
	ls.Offer = ctx.Offer
	ls.Slot = ctx.Slot
}

// equal hostname xxxx
type EqualStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (ls *EqualStatement) Eval() bool {
	if ls.What == "hostname" {
		return ls.Offer.GetHostname() == ls.Regex
	}

	if ls.What == "agentid" {
		return *ls.Offer.GetAgentId().Value == ls.Regex
	}

	// user defined attributes match
	for _, attr := range ls.Offer.Attributes {
		if attr.GetName() == ls.What && attr.GetType() == mesos.Value_TEXT {
			return *attr.GetText().Value == ls.Regex
		}
	}

	return false
}

func (ls *EqualStatement) Valid() error {
	if ls.What != "" && ls.Regex != "" {
		return nil
	} else {
		return errors.New("Equal statement must contain two operand")
	}
}

func (ls *EqualStatement) SetContext(ctx *ConstraintParamHolder) {
	ls.Offer = ctx.Offer
	ls.Slot = ctx.Slot
}

// contains hostname barfoo
type ContainsStatement struct {
	ConstraintParamHolder
	What  string
	Regex string
}

func (cs *ContainsStatement) Eval() bool {
	if cs.What == "hostname" {
		return strings.Contains(cs.Offer.GetHostname(), cs.Regex)
	}

	if cs.What == "agentid" {
		return strings.Contains(*cs.Offer.GetAgentId().Value, cs.Regex)
	}

	// user defined attributes match
	for _, attr := range cs.Offer.Attributes {
		if attr.GetName() == cs.What && attr.GetType() == mesos.Value_TEXT {
			return strings.Contains(*attr.GetText().Value, cs.Regex)
		}
	}

	return false
}

func (cs *ContainsStatement) Valid() error {
	return nil
}

func (cs *ContainsStatement) SetContext(ctx *ConstraintParamHolder) {
	cs.Offer = ctx.Offer
	cs.Slot = ctx.Slot
}
<|endoftext|>