<commit_before>/*
Package hamt32 implements a functional Hash Array Mapped Trie (HAMT).
It is called hamt32 because this package is using 32 nodes for each level of
the Trie. The term functional is used to imply immutable and persistent.

The 30bits of hash are separated into six 5bit values that constitue the hash
path of any Key in this Trie. However, not all six levels of the Trie are used.
As many levels (six or less) are used to find a unique location
for the leaf to be placed within the Trie.

If all six levels of the Trie are used for two or more key/val pairs then a
special collision leaf will be used to store those key/val pairs,  at the sixth
level of the Trie.
*/
package hamt32

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/lleo/go-hamt/key"
	"github.com/pkg/errors"
)

// Nbits constant is the number of bits(5) a 30bit hash value is split into,
// to provied the indexes of a HAMT.
const Nbits uint = 5

// MaxDepth constant is the maximum depth(5) of Nbits values that constitute
// the path in a HAMT, from [0..MaxDepth]for a total of MaxDepth+1(6) levels.
// Nbits*(MaxDepth+1) == HASHBITS (ie 5*(5+1) == 30).
const MaxDepth uint = 5

// TableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 1<<Nbits (ie 2^5 == 32).
const TableCapacity uint = 1 << Nbits

func hashPathMask(depth uint) uint32 {
	return uint32(1<<(depth*Nbits)) - 1
}

// Create a string of the form "/%02d/%02d..." to describe a hashPath of
// a given depth.
//
// If you want hashPathString() to include the current idx, you Must
// add one to depth. You may need to do this because you are creating
// a table to be put at the idx'th slot of the current table.
func hashPathString(hashPath uint32, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func h30ToString(h30 uint32) string {
	return hashPathString(h30, MaxDepth)
}

func StringToH30(s string) uint32 {
	if !strings.HasPrefix(s, "/") {
		panic(errors.New("does not start with '/'"))
	}
	var s0 = s[1:]
	var as = strings.Split(s0, "/")

	var h30 uint32 = 0
	for i, s1 := range as {
		var ui, err = strconv.ParseUint(s1, 10, int(Nbits))
		if err != nil {
			panic(errors.Wrap(err, fmt.Sprintf("strconv.ParseUint(%q, %d, %d) failed", s1, 10, Nbits)))
		}
		h30 |= uint32(ui << (uint(i) * Nbits))
		//fmt.Printf("%d: h30 = %q %2d %#02x %05b\n", i, s1, ui, ui, ui)
	}

	return h30
}

// h30ToBitStr is for printf debugging use
func h30ToBitStr(h30 uint32) string {
	var strs = make([]string, MaxDepth+1)

	for depth := uint(0); depth <= MaxDepth; depth++ {
		var idx = index(h30, depth)
		strs[MaxDepth-depth] = fmt.Sprintf("%05b", idx)
	}

	return strings.Join(strs, " ")
}

//indexMask() generates a Nbits(5-bit) mask for a given depth
func indexMask(depth uint) uint32 {
	return uint32((1<<Nbits)-1) << (depth * Nbits)
}

//index() calculates a Nbits(5-bit) integer based on the hash and depth
func index(h30 uint32, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h30 & idxMask) >> (depth * Nbits))
	return idx
}

//buildHashPath(hashPath, idx, depth)
//hashPath will be depth Nbits long
func buildHashPath(hashPath uint32, idx, depth uint) uint32 {
	var mask uint32 = (1 << ((depth - 1) * Nbits)) - 1
	hashPath = hashPath & mask

	return hashPath | uint32(idx<<((depth-1)*Nbits))
}

type keyVal struct {
	key key.Key
	val interface{}
}

// GradeTables variable controls whether Hamt structures will upgrade/
// downgrade compressed/full tables. This variable and FullTableInit
// should not be changed during the lifetime of any Hamt structure.
// Default: true
var GradeTables = true

// FullTableInit variable controls whether the initial new table type is
// fullTable, else the initial new table type is compressedTable.
// Default: false
var FullTableInit = false

// UpgradeThreshold is a variable that defines when a compressedTable meats
// or exceeds that number of entries, then that table will be upgraded to
// a fullTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/2.
var UpgradeThreshold = TableCapacity * 2 / 3

// DowngradeThreshold is a variable that defines when a fullTable becomes
// lower than that number of entries, then that table will be downgraded to
// a compressedTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/4.
var DowngradeThreshold = TableCapacity / 4

type Hamt struct {
	root     tableI
	nentries uint
}

func (h Hamt) IsEmpty() bool {
	//return h.root == nil
	//return h.nentries == 0
	//return h.root == nil && h.nentries == 0
	return h == Hamt{}
}

func (h Hamt) Root() tableI {
	return h.root
}

func (h Hamt) Nentries() uint {
	return h.nentries
}

func createRootTable(leaf leafI) tableI {
	if FullTableInit {
		return createRootFullTable(leaf)
	}
	return createRootCompressedTable(leaf)
}

//func createTable(depth uint, leaf1 leafI, k key.Key, v interface{}) tableI {
func createTable(depth uint, leaf1 leafI, leaf2 flatLeaf) tableI {
	//var hashPath = k.Hash30() & hashPathMask(depth)
	//var leaf2 = *newFlatLeaf(k, v)

	if FullTableInit {
		return createFullTable(depth, leaf1, leaf2)
	}
	return createCompressedTable(depth, leaf1, leaf2)
}

// copyUp is ONLY called on a fresh copy of the current Hamt. Hence, modifying
// it is allowed.
func (nh *Hamt) persist(oldTable, newTable tableI, path tableStack) {
	if path.isEmpty() {
		nh.root = newTable
		return
	}

	var depth = uint(path.len())
	var parentDepth = depth - 1

	var parentIdx = index(oldTable.Hash30(), parentDepth)

	var oldParent = path.pop()
	var newParent tableI

	if newTable == nil {
		newParent = oldParent.remove(parentIdx)
	} else {
		newParent = oldParent.replace(parentIdx, newTable)
	}

	nh.persist(oldParent, newParent, path) //recurses at most MaxDepth-1 times

	return
}

func (h Hamt) find(k key.Key) (path tableStack, leaf leafI, idx uint) {
	if h.IsEmpty() {
		return nil, nil, 0
	}

	path = newTableStack()
	var curTable = h.root

	var h30 = k.Hash30()
	var depth uint
	var curNode nodeI

	for depth = 0; depth < MaxDepth; depth++ {
		path.push(curTable)
		idx = index(h30, depth)
		curNode = curTable.Get(idx)

		switch n := curNode.(type) {
		case nil:
			return path, nil, idx
		case leafI:
			return path, n, idx
		case tableI:
			curTable = n
			// exit switch then loop for
		default:
			panic(fmt.Sprintf("switch default case: depth=%d; idx=%d; curNode unknown type = %T; value = %v; path=%s", depth, idx, n, n, path))
		}
	}
	if depth == MaxDepth {
		path.push(curTable)
		idx = index(h30, depth)
		curNode = curTable.Get(idx)

		if curNode == nil {
			return path, nil, idx
		} else if leaf, isLeaf := curNode.(leafI); isLeaf {
			return path, leaf, idx
		} else {
			panic(fmt.Sprintf("depth,%d == MaxDepth: %d; idx=%d; unknown type = %T; value = %v; path=%s", depth, MaxDepth, idx, curNode, curNode, path))
		}
	}

	panic("SHOULD NEVER GET HERE!")
}

// Get(k) retrieves the value for a given key from the Hamt. The bool
// represents whether the key was found.
func (h Hamt) Get(k key.Key) (interface{}, bool) {
	var _, leaf, _ = h.find(k)

	//var depth = path.len()
	//var curTable = path.pop()

	if leaf == nil {
		return nil, false
	}

	var val, found = leaf.get(k)
	if !found {
		return nil, false
	}

	return val, true
}

// Put new key/val pair into Hamt, returning a new persistant Hamt and a bool
// indicating if the key/val pair was added(true) or mearly updated(false).
func (h Hamt) Put(k key.Key, v interface{}) (Hamt, bool) {
	var nh Hamt = h //copy by value

	var path, leaf, idx = h.find(k)

	if path == nil { // h.IsEmpty()
		nh.root = createRootTable(newFlatLeaf(k, v))
		nh.nentries++
		return nh, true
	}

	var curTable = path.pop()
	var depth = uint(path.len())

	var newTable tableI
	var added bool

	if leaf == nil {
		newTable = curTable.insert(idx, newFlatLeaf(k, v))
		added = true
	} else {
		if leaf.Hash30() == k.Hash30() {
			var newLeaf leafI
			newLeaf, added = leaf.put(k, v)
			newTable = curTable.replace(idx, newLeaf)
		} else {
			var tmpTable = createTable(depth+1, leaf, *newFlatLeaf(k, v))
			newTable = curTable.replace(idx, tmpTable)
			added = true
		}
	}

	if added {
		nh.nentries++
	}

	nh.persist(curTable, newTable, path)

	return nh, added
}

// Hamt.Del(k) returns a new Hamt, the value deleted, and a boolean that
// specifies whether or not the key was deleted (eg it didn't exist to start
// with). Therefor you must always test deleted before using the new *Hamt
// value.
func (h Hamt) Del(k key.Key) (Hamt, interface{}, bool) {
	var nh Hamt = h // copy by value

	var path, leaf, idx = h.find(k)

	if path == nil { // h.IsEmpty()
		return nh, nil, false
	}

	var curTable = path.pop()
	//var depth = uint(path.len())

	var newTable tableI
	var val interface{}
	var deleted bool

	if leaf == nil {
		//return nh, val, deleted
		return h, nil, false
	} else {
		var newLeaf leafI
		newLeaf, val, deleted = leaf.del(k)

		if !deleted {
			//return nh, val, deleted
			return h, nil, false
		}

		if newLeaf == nil {
			newTable = curTable.remove(idx)
		} else {
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	if deleted {
		nh.nentries--
	}

	nh.persist(curTable, newTable, path)

	return nh, val, deleted
}

func (h Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root)
}

const halfIndent = "  "
const fullIndent = "    "

func (h Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent+fullIndent, true)
		str += indent + "}end\n"
		return str
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<commit_msg>commented out a random printf debugging func<commit_after>/*
Package hamt32 implements a functional Hash Array Mapped Trie (HAMT).
It is called hamt32 because this package is using 32 nodes for each level of
the Trie. The term functional is used to imply immutable and persistent.

The 30bits of hash are separated into six 5bit values that constitue the hash
path of any Key in this Trie. However, not all six levels of the Trie are used.
As many levels (six or less) are used to find a unique location
for the leaf to be placed within the Trie.

If all six levels of the Trie are used for two or more key/val pairs then a
special collision leaf will be used to store those key/val pairs,  at the sixth
level of the Trie.
*/
package hamt32

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/lleo/go-hamt/key"
	"github.com/pkg/errors"
)

// Nbits constant is the number of bits(5) a 30bit hash value is split into,
// to provied the indexes of a HAMT.
const Nbits uint = 5

// MaxDepth constant is the maximum depth(5) of Nbits values that constitute
// the path in a HAMT, from [0..MaxDepth]for a total of MaxDepth+1(6) levels.
// Nbits*(MaxDepth+1) == HASHBITS (ie 5*(5+1) == 30).
const MaxDepth uint = 5

// TableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 1<<Nbits (ie 2^5 == 32).
const TableCapacity uint = 1 << Nbits

func hashPathMask(depth uint) uint32 {
	return uint32(1<<(depth*Nbits)) - 1
}

// Create a string of the form "/%02d/%02d..." to describe a hashPath of
// a given depth.
//
// If you want hashPathString() to include the current idx, you Must
// add one to depth. You may need to do this because you are creating
// a table to be put at the idx'th slot of the current table.
func hashPathString(hashPath uint32, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func h30ToString(h30 uint32) string {
	return hashPathString(h30, MaxDepth)
}

func StringToH30(s string) uint32 {
	if !strings.HasPrefix(s, "/") {
		panic(errors.New("does not start with '/'"))
	}
	var s0 = s[1:]
	var as = strings.Split(s0, "/")

	var h30 uint32 = 0
	for i, s1 := range as {
		var ui, err = strconv.ParseUint(s1, 10, int(Nbits))
		if err != nil {
			panic(errors.Wrap(err, fmt.Sprintf("strconv.ParseUint(%q, %d, %d) failed", s1, 10, Nbits)))
		}
		h30 |= uint32(ui << (uint(i) * Nbits))
		//fmt.Printf("%d: h30 = %q %2d %#02x %05b\n", i, s1, ui, ui, ui)
	}

	return h30
}

//// h30ToBitStr is for printf debugging use
//func h30ToBitStr(h30 uint32) string {
//	var strs = make([]string, MaxDepth+1)
//
//	for depth := uint(0); depth <= MaxDepth; depth++ {
//		var idx = index(h30, depth)
//		strs[MaxDepth-depth] = fmt.Sprintf("%05b", idx)
//	}
//
//	return strings.Join(strs, " ")
//}

//indexMask() generates a Nbits(5-bit) mask for a given depth
func indexMask(depth uint) uint32 {
	return uint32((1<<Nbits)-1) << (depth * Nbits)
}

//index() calculates a Nbits(5-bit) integer based on the hash and depth
func index(h30 uint32, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h30 & idxMask) >> (depth * Nbits))
	return idx
}

//buildHashPath(hashPath, idx, depth)
//hashPath will be depth Nbits long
func buildHashPath(hashPath uint32, idx, depth uint) uint32 {
	var mask uint32 = (1 << ((depth - 1) * Nbits)) - 1
	hashPath = hashPath & mask

	return hashPath | uint32(idx<<((depth-1)*Nbits))
}

type keyVal struct {
	key key.Key
	val interface{}
}

// GradeTables variable controls whether Hamt structures will upgrade/
// downgrade compressed/full tables. This variable and FullTableInit
// should not be changed during the lifetime of any Hamt structure.
// Default: true
var GradeTables = true

// FullTableInit variable controls whether the initial new table type is
// fullTable, else the initial new table type is compressedTable.
// Default: false
var FullTableInit = false

// UpgradeThreshold is a variable that defines when a compressedTable meats
// or exceeds that number of entries, then that table will be upgraded to
// a fullTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/2.
var UpgradeThreshold = TableCapacity * 2 / 3

// DowngradeThreshold is a variable that defines when a fullTable becomes
// lower than that number of entries, then that table will be downgraded to
// a compressedTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/4.
var DowngradeThreshold = TableCapacity / 4

type Hamt struct {
	root     tableI
	nentries uint
}

func (h Hamt) IsEmpty() bool {
	//return h.root == nil
	//return h.nentries == 0
	//return h.root == nil && h.nentries == 0
	return h == Hamt{}
}

func (h Hamt) Root() tableI {
	return h.root
}

func (h Hamt) Nentries() uint {
	return h.nentries
}

func createRootTable(leaf leafI) tableI {
	if FullTableInit {
		return createRootFullTable(leaf)
	}
	return createRootCompressedTable(leaf)
}

//func createTable(depth uint, leaf1 leafI, k key.Key, v interface{}) tableI {
func createTable(depth uint, leaf1 leafI, leaf2 flatLeaf) tableI {
	//var hashPath = k.Hash30() & hashPathMask(depth)
	//var leaf2 = *newFlatLeaf(k, v)

	if FullTableInit {
		return createFullTable(depth, leaf1, leaf2)
	}
	return createCompressedTable(depth, leaf1, leaf2)
}

// copyUp is ONLY called on a fresh copy of the current Hamt. Hence, modifying
// it is allowed.
func (nh *Hamt) persist(oldTable, newTable tableI, path tableStack) {
	if path.isEmpty() {
		nh.root = newTable
		return
	}

	var depth = uint(path.len())
	var parentDepth = depth - 1

	var parentIdx = index(oldTable.Hash30(), parentDepth)

	var oldParent = path.pop()
	var newParent tableI

	if newTable == nil {
		newParent = oldParent.remove(parentIdx)
	} else {
		newParent = oldParent.replace(parentIdx, newTable)
	}

	nh.persist(oldParent, newParent, path) //recurses at most MaxDepth-1 times

	return
}

func (h Hamt) find(k key.Key) (path tableStack, leaf leafI, idx uint) {
	if h.IsEmpty() {
		return nil, nil, 0
	}

	path = newTableStack()
	var curTable = h.root

	var h30 = k.Hash30()
	var depth uint
	var curNode nodeI

	for depth = 0; depth < MaxDepth; depth++ {
		path.push(curTable)
		idx = index(h30, depth)
		curNode = curTable.Get(idx)

		switch n := curNode.(type) {
		case nil:
			return path, nil, idx
		case leafI:
			return path, n, idx
		case tableI:
			curTable = n
			// exit switch then loop for
		default:
			panic(fmt.Sprintf("switch default case: depth=%d; idx=%d; curNode unknown type = %T; value = %v; path=%s", depth, idx, n, n, path))
		}
	}
	if depth == MaxDepth {
		path.push(curTable)
		idx = index(h30, depth)
		curNode = curTable.Get(idx)

		if curNode == nil {
			return path, nil, idx
		} else if leaf, isLeaf := curNode.(leafI); isLeaf {
			return path, leaf, idx
		} else {
			panic(fmt.Sprintf("depth,%d == MaxDepth: %d; idx=%d; unknown type = %T; value = %v; path=%s", depth, MaxDepth, idx, curNode, curNode, path))
		}
	}

	panic("SHOULD NEVER GET HERE!")
}

// Get(k) retrieves the value for a given key from the Hamt. The bool
// represents whether the key was found.
func (h Hamt) Get(k key.Key) (interface{}, bool) {
	var _, leaf, _ = h.find(k)

	//var depth = path.len()
	//var curTable = path.pop()

	if leaf == nil {
		return nil, false
	}

	var val, found = leaf.get(k)
	if !found {
		return nil, false
	}

	return val, true
}

// Put new key/val pair into Hamt, returning a new persistant Hamt and a bool
// indicating if the key/val pair was added(true) or mearly updated(false).
func (h Hamt) Put(k key.Key, v interface{}) (Hamt, bool) {
	var nh Hamt = h //copy by value

	var path, leaf, idx = h.find(k)

	if path == nil { // h.IsEmpty()
		nh.root = createRootTable(newFlatLeaf(k, v))
		nh.nentries++
		return nh, true
	}

	var curTable = path.pop()
	var depth = uint(path.len())

	var newTable tableI
	var added bool

	if leaf == nil {
		newTable = curTable.insert(idx, newFlatLeaf(k, v))
		added = true
	} else {
		if leaf.Hash30() == k.Hash30() {
			var newLeaf leafI
			newLeaf, added = leaf.put(k, v)
			newTable = curTable.replace(idx, newLeaf)
		} else {
			var tmpTable = createTable(depth+1, leaf, *newFlatLeaf(k, v))
			newTable = curTable.replace(idx, tmpTable)
			added = true
		}
	}

	if added {
		nh.nentries++
	}

	nh.persist(curTable, newTable, path)

	return nh, added
}

// Hamt.Del(k) returns a new Hamt, the value deleted, and a boolean that
// specifies whether or not the key was deleted (eg it didn't exist to start
// with). Therefor you must always test deleted before using the new *Hamt
// value.
func (h Hamt) Del(k key.Key) (Hamt, interface{}, bool) {
	var nh Hamt = h // copy by value

	var path, leaf, idx = h.find(k)

	if path == nil { // h.IsEmpty()
		return nh, nil, false
	}

	var curTable = path.pop()
	//var depth = uint(path.len())

	var newTable tableI
	var val interface{}
	var deleted bool

	if leaf == nil {
		//return nh, val, deleted
		return h, nil, false
	} else {
		var newLeaf leafI
		newLeaf, val, deleted = leaf.del(k)

		if !deleted {
			//return nh, val, deleted
			return h, nil, false
		}

		if newLeaf == nil {
			newTable = curTable.remove(idx)
		} else {
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	if deleted {
		nh.nentries--
	}

	nh.persist(curTable, newTable, path)

	return nh, val, deleted
}

func (h Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root)
}

const halfIndent = "  "
const fullIndent = "    "

func (h Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent+fullIndent, true)
		str += indent + "}end\n"
		return str
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/nsf/termbox-go"
)

const (
	pauseAfterKeypress = (1500 * time.Millisecond)
	redrawPause        = 15 * time.Millisecond
)

var (
	global_lastkeypress int64
)

func getRoot() string {
	if len(flag.Args()) == 0 {
		return "."
	} else {
		return flag.Arg(0)
	}
}

// hf --cmd=emacs ~/go/src/github.com/hugows/ happy
var cmd = flag.String("cmd", "vim", "command to run")

// var termkey *TermboxEventWrapper

// strings.Replace(tw.Text, " ", "+", -1)

func main() {
	flag.Parse()

	var rview ResultsView

	root := getRoot()
	fi, err := os.Stat(root)
	if err != nil {
		fmt.Println(err)
		return
	}

	if !fi.IsDir() {
		fmt.Println(root, "is NOT a folder")
		return
	}

	err = termbox.Init()
	if err != nil {
		panic(err)
	}
	termbox.SetInputMode(termbox.InputEsc)

	resultset := new(ResultSet)

	w, h := termbox.Size()
	modeline := NewModeline(0, h-1, w)
	cmdline := new(CommandLine)

	idleTimer := time.NewTimer(1 * time.Hour)

	fileCh := walkFiles(getRoot())
	termboxEventCh := make(chan termbox.Event)

	forceDrawCh := make(chan bool, 100)
	forceSortCh := make(chan bool, 100)

	timeLastUser := time.Now().Add(-1 * time.Hour)
	timeLastFilter := time.Now()

	go func() {
		for {
			ev := termbox.PollEvent()
			if ev.Type == termbox.EventKey {
				timeLastUser = time.Now()
				global_lastkeypress = timeLastUser.UnixNano()
			}
			termboxEventCh <- ev
		}
	}()

	go func() {
		for {
			<-forceSortCh
			filtered := resultset.Filter(global_lastkeypress, modeline.Contents())
			rview.Update(filtered.results)
			cmdline.Update(rview.GetSelected())
			forceDrawCh <- true
		}
	}()

	// Command name is:
	// os.Args[0]

	modeline.Draw(&rview)
	cmdline.Draw(0, h-2, w)
	rview.SetSize(0, 0, w, h-2)
	termbox.Flush()

	for {
		select {
		case <-forceDrawCh:
			/* redraw */

		case <-idleTimer.C:
			idleTimer = time.NewTimer(1 * time.Hour)
			if !modeline.paused {
				modeline.LastFile()
				fileCh = nil
			}

		case filename, ok := <-fileCh:
			if time.Since(timeLastUser) > pauseAfterKeypress {
				modeline.Unpause()
			} else {
				modeline.Pause()
			}

			if ok {
				resultset.Insert(filename)
			}

			if !modeline.paused && time.Since(timeLastFilter) > (15*time.Millisecond) {
				forceSortCh <- true
				timeLastFilter = time.Now()

				if !ok {
					modeline.LastFile()
					fileCh = nil
				}
			} else if !ok {
				idleTimer.Reset(redrawPause)
				fileCh = nil
			}

		case ev := <-termboxEventCh:
			if fileCh != nil {
				idleTimer.Reset(pauseAfterKeypress)
			} else {
				modeline.Unpause()
			}

			switch ev.Type {
			case termbox.EventKey:
				switch ev.Key {
				case termbox.KeyEsc, termbox.KeyCtrlC:
					termbox.Close()
					return
				case termbox.KeyEnter:
					termbox.Close()
					// runCmdWithArgs(rview.FormatSelected())
					return
				case termbox.KeyCtrlT:
					rview.ToggleMarkAll()
				case termbox.KeyArrowUp, termbox.KeyCtrlP:
					cmdline.Update(rview.SelectPrevious())
				case termbox.KeyArrowDown, termbox.KeyCtrlN:
					cmdline.Update(rview.SelectNext())
				case termbox.KeyArrowLeft, termbox.KeyCtrlB:
					modeline.input.MoveCursorOneRuneBackward()
				case termbox.KeyArrowRight, termbox.KeyCtrlF:
					modeline.input.MoveCursorOneRuneForward()
				case termbox.KeyBackspace, termbox.KeyBackspace2:
					modeline.input.DeleteRuneBackward()
					forceSortCh <- true
				case termbox.KeyDelete, termbox.KeyCtrlD:
					modeline.input.DeleteRuneForward()
					forceSortCh <- true
				case termbox.KeySpace:
					rview.ToggleMark()
				case termbox.KeyCtrlK:
					modeline.input.DeleteTheRestOfTheLine()
					forceSortCh <- true
				case termbox.KeyHome, termbox.KeyCtrlA:
					modeline.input.MoveCursorToBeginningOfTheLine()
				case termbox.KeyEnd, termbox.KeyCtrlE:
					modeline.input.MoveCursorToEndOfTheLine()
				default:
					if ev.Ch != 0 {
						modeline.input.InsertRune(ev.Ch)
						forceSortCh <- true
					}
				}
			case termbox.EventError:
				panic(ev.Err)
			}

			// fmt.Println(modeline.Contents())
		}

		modeline.Draw(&rview)
		cmdline.Draw(0, h-2, w)
		rview.Draw()
		termbox.Flush()
	}

}
<commit_msg>missing forceSort<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/nsf/termbox-go"
)

const (
	pauseAfterKeypress = (1500 * time.Millisecond)
	redrawPause        = 30 * time.Millisecond
)

var (
	global_lastkeypress int64
)

func getRoot() string {
	if len(flag.Args()) == 0 {
		return "."
	} else {
		return flag.Arg(0)
	}
}

// hf --cmd=emacs ~/go/src/github.com/hugows/ happy
var cmd = flag.String("cmd", "vim", "command to run")

func main() {
	flag.Parse()

	var rview ResultsView

	root := getRoot()
	fi, err := os.Stat(root)
	if err != nil {
		fmt.Println(err)
		return
	}

	if !fi.IsDir() {
		fmt.Println(root, "is NOT a folder")
		return
	}

	err = termbox.Init()
	if err != nil {
		panic(err)
	}
	termbox.SetInputMode(termbox.InputEsc)

	resultset := new(ResultSet)

	w, h := termbox.Size()
	modeline := NewModeline(0, h-1, w)
	cmdline := new(CommandLine)

	idleTimer := time.NewTimer(1 * time.Hour)

	fileCh := walkFiles(getRoot())
	termboxEventCh := make(chan termbox.Event)

	forceDrawCh := make(chan bool, 100)
	forceSortCh := make(chan bool, 100)

	timeLastUser := time.Now().Add(-1 * time.Hour)
	timeLastFilter := time.Now()

	go func() {
		for {
			ev := termbox.PollEvent()
			if ev.Type == termbox.EventKey {
				timeLastUser = time.Now()
				global_lastkeypress = timeLastUser.UnixNano()
			}
			termboxEventCh <- ev
		}
	}()

	go func() {
		for {
			<-forceSortCh
			filtered := resultset.Filter(global_lastkeypress, modeline.Contents())
			rview.Update(filtered.results)
			cmdline.Update(rview.GetSelected())
			forceDrawCh <- true
		}
	}()

	// Command name is:
	// os.Args[0]

	modeline.Draw(&rview)
	cmdline.Draw(0, h-2, w)
	rview.SetSize(0, 0, w, h-2)
	termbox.Flush()

	for {
		select {
		case <-forceDrawCh:
			/* redraw */

		case <-idleTimer.C:
			idleTimer = time.NewTimer(1 * time.Hour)
			if !modeline.paused {
				modeline.LastFile()
				fileCh = nil
				forceSortCh <- true
			}

		case filename, ok := <-fileCh:
			if time.Since(timeLastUser) > pauseAfterKeypress {
				modeline.Unpause()
			} else {
				modeline.Pause()
			}

			if ok {
				resultset.Insert(filename)
			}

			if !modeline.paused && time.Since(timeLastFilter) > redrawPause {
				forceSortCh <- true
				timeLastFilter = time.Now()

				if !ok {
					modeline.LastFile()
					fileCh = nil
				}
			} else if !ok {
				idleTimer.Reset(redrawPause)
				fileCh = nil
			}

		case ev := <-termboxEventCh:
			if fileCh != nil {
				idleTimer.Reset(pauseAfterKeypress)
			} else {
				modeline.Unpause()
			}

			switch ev.Type {
			case termbox.EventKey:
				switch ev.Key {
				case termbox.KeyEsc, termbox.KeyCtrlC:
					termbox.Close()
					return
				case termbox.KeyEnter:
					termbox.Close()
					// runCmdWithArgs(rview.FormatSelected())
					return
				case termbox.KeyCtrlT:
					rview.ToggleMarkAll()
				case termbox.KeyArrowUp, termbox.KeyCtrlP:
					cmdline.Update(rview.SelectPrevious())
				case termbox.KeyArrowDown, termbox.KeyCtrlN:
					cmdline.Update(rview.SelectNext())
				case termbox.KeyArrowLeft, termbox.KeyCtrlB:
					modeline.input.MoveCursorOneRuneBackward()
				case termbox.KeyArrowRight, termbox.KeyCtrlF:
					modeline.input.MoveCursorOneRuneForward()
				case termbox.KeyBackspace, termbox.KeyBackspace2:
					modeline.input.DeleteRuneBackward()
					forceSortCh <- true
				case termbox.KeyDelete, termbox.KeyCtrlD:
					modeline.input.DeleteRuneForward()
					forceSortCh <- true
				case termbox.KeySpace:
					rview.ToggleMark()
				case termbox.KeyCtrlK:
					modeline.input.DeleteTheRestOfTheLine()
					forceSortCh <- true
				case termbox.KeyHome, termbox.KeyCtrlA:
					modeline.input.MoveCursorToBeginningOfTheLine()
				case termbox.KeyEnd, termbox.KeyCtrlE:
					modeline.input.MoveCursorToEndOfTheLine()
				default:
					if ev.Ch != 0 {
						modeline.input.InsertRune(ev.Ch)
						forceSortCh <- true
					}
				}
			case termbox.EventError:
				panic(ev.Err)
			}

			// fmt.Println(modeline.Contents())
		}

		modeline.Draw(&rview)
		cmdline.Draw(0, h-2, w)
		rview.Draw()
		termbox.Flush()
	}

}
<|endoftext|>


<commit_before>package pop3

import (
	"bds/lib/maildir"
	"bufio"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

type MailDirGetter interface {
	// get a user's maildir
	GetUserMaildir(user string) (maildir.MailDir, error)
}

type mailDirGetter string

func (md mailDirGetter) GetUserMaildir(user string) (m maildir.MailDir, err error) {
	m = maildir.MailDir(md)
	return
}

// a maildir getter that always uses 1 directory
func NewMailDirGetter(path string) MailDirGetter {
	return mailDirGetter(path)
}

// pop3 server
type Server struct {
	// function that obtains a maildir given a user
	GetMailDir MailDirGetter
	// server name
	name string
}

func New() *Server {
	host, _ := os.Hostname()
	return &Server{
		name: host,
	}
}

// get all messages in maildir
func (s *Server) getMessages(user string) (msgs []maildir.Message, err error) {
	// get user's maildir
	var md maildir.MailDir
	if s.GetMailDir == nil {
		err = errors.New("could't find maildir")
		return
	}
	md, err = s.GetMailDir.GetUserMaildir(user)
	if err != nil {
		return
	}
	// move new mail into cur
	var ms []maildir.Message
	ms, err = md.ListNew()
	if err == nil {
		for _, m := range ms {
			err = md.ProcessNew(m)
			if err != nil {
				log.Errorf("error processing maildir: %s", err.Error())
			}
		}
	}
	// get all current files
	msgs, err = md.ListCur()
	return
}

func (s *Server) checkUser(user, passwd string) (allowed bool) {
	// TODO: implement
	allowed = true
	return
}

func (s *Server) checkDigest(user, digest string) (allowed bool) {
	// TODO: implement
	allowed = true
	return
}

// get all messages and octet count
func (s *Server) obtainMessages(user string) (msgs []maildir.Message, o int64, err error) {
	msgs, err = s.getMessages(user)
	if err == nil {
		for _, msg := range msgs {
			var info os.FileInfo
			info, err = os.Stat(msg.Filepath())
			if err == nil && !info.IsDir() {
				o += info.Size()
			}
		}
	}
	return
}

// pop3 session handler
type pop3Session struct {
	// network connection
	c *textproto.Conn
	// parent server
	s *Server
	// are we in transaction state?
	transaction bool
	// current user
	user string
	// messages we have in this transaction
	msgs []maildir.Message
	// how many octets we have for all messages
	octs int64
	// messages to delete
	dels []maildir.Message
}

// run pop3 session mainloop
func (p *pop3Session) Run() {
	// send banner
	err := p.c.PrintfLine("+OK POP3 Server Ready <%d.%d@%s>", os.Getpid(), time.Now().Unix(), p.s.name)
	for err == nil {
		var line string
		line, err = p.c.ReadLine()
		if err == nil {
			if strings.ToUpper(line) == "QUIT" {
				// check for quit command
				err = p.OK("k bai")
				break
			} else if p.transaction {
				err = p.handleTransactionLine(line)
			} else {
				err = p.handleLine(line)
			}
		}
	}
	if err != nil && err != io.EOF {
		log.Errorf("error in pop3 session: %s", err.Error())
	}
	// close connection
	p.c.Close()
	// delete old messages
	for _, msg := range p.dels {
		os.Remove(msg.Filepath())
	}
}

// handle line when in transaction mode
func (p *pop3Session) handleTransactionLine(line string) (err error) {
	var info os.FileInfo
	var idx int
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "DELE":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid, add it to delete
				p.dels = append(p.dels, p.msgs[idx-1])
			} else {
				// invalid
				err = p.Error(err.Error())
			}
		}
	case "RETR":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				msg := p.msgs[idx-1].Filepath()
				var f *os.File
				f, err = os.Open(msg)
				if err == nil {
					r := bufio.NewReader(f)
					info, err = f.Stat()
					if err == nil {
						err = p.c.PrintfLine("+OK %d octets", info.Size())
						for err == nil {
							// send line
							line, err = r.ReadString(10)
							if err == io.EOF {
								err = nil
								break
							} else if err == nil {
								line = strings.Trim(line, "\r")
								line = strings.Trim(line, "\n")
								if line == "." {
									line = " ."
								}
								// send line
								err = p.c.PrintfLine(line)
							} else {
								// error
								break
							}
						}
					}
					f.Close()
					// end
					err = p.c.PrintfLine(".")
				} else {
					err = p.Error(err.Error())
				}
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			err = p.Error("invalid syntax")
		}
		break
	case "UIDL":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				err = p.OK(parts[1] + " " + p.msgs[idx-1].Filename())
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			// all messages
			p.OK("")
			dw := p.c.DotWriter()
			for idx, msg := range p.msgs {
				fmt.Fprintf(dw, "%d %s\r\n", idx, msg.Filename())
			}
			// FLUSH :D
			err = dw.Close()
		}
		// begin
		break
	case "STAT":
		// begin
		_, err = p.c.W.WriteString("+OK ")
		if err == nil {
			// write list
			for idx, _ := range p.msgs {
				_, err = fmt.Fprintf(p.c.W, "%d ", 1+idx)
				if err != nil {
					break
				}
			}
		}
		if err == nil {
			// done writing list
			err = p.c.PrintfLine("")
		} else {
			err = p.Error(err.Error())
		}
		break
	case "LIST":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil {
				if idx > 0 && idx <= len(p.msgs) {
					info, err = os.Stat(p.msgs[idx-1].Filepath())
					if err == nil {
						err = p.c.PrintfLine("+OK %d %d", idx, info.Size())
					}
				} else {
					// no existing message
					err = p.Error("no such message")
				}
			}
		} else {
			// all messages
			err = p.c.PrintfLine("+OK %d messages (%d octets)", len(p.msgs), p.octs)
			if err == nil {
				dw := p.c.DotWriter()
				for i, msg := range p.msgs {
					info, err = os.Stat(msg.Filepath())
					if err == nil {
						// write out entry
						fmt.Fprintf(dw, "%d %d\r\n", i+1, info.Size())
					} else {
						log.Errorf("error in pop3, stat(): %s", err.Error())
					}
				}
				// FLUSH IT !!! :-DDDDD
				err = dw.Close()
			}
		}
		break
	default:
		err = p.Error("bad command")
	}
	return
}

// handle 1 line of input when not in transaction mode
func (p *pop3Session) handleLine(line string) (err error) {
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "NOOP":
		err = p.OK("")
		break
	case "APOP":
		if len(parts) == 3 {
			if p.s.checkDigest(parts[1], parts[2]) {
				// load messages
				p.msgs, p.octs, err = p.s.obtainMessages(p.user)
				if err == nil {
					err = p.OK("maildrop is go for access")
					p.transaction = err == nil
				} else {
					err = p.Error(err.Error())
				}
			} else {
				err = p.Error("permission denied")
			}
		} else {
			err = p.Error("Bad APOP")
		}
		break
	case "PASS":
		if p.s.checkUser(p.user, line[5:]) {
			p.msgs, p.octs, err = p.s.obtainMessages(p.user)
			if err == nil {
				err = p.c.PrintfLine("+OK %s maildrop logged in, you have %d messages (%d octets)", p.user, len(p.msgs), p.octs)
				p.transaction = err == nil
			} else {
				err = p.Error(err.Error())
			}
		} else {
			err = p.Error("bad login")
		}
		break
	case "USER":
		if len(parts) > 1 {
			p.user = line[5:]
		}
		err = p.OK("anonymous access enabled")
		break
	default:
		err = p.Error("bad command")
	}
	return
}

func (p *pop3Session) OK(msg string) (err error) {
	if len(msg) == 0 {
		err = p.c.PrintfLine("+OK")
	} else {
		err = p.c.PrintfLine("+OK %s", msg)
	}
	return
}

// send error
func (p *pop3Session) Error(msg string) (err error) {
	err = p.c.PrintfLine("-ERR %s", msg)
	return
}

// serve sessions with connections accepted from a net.Listener
func (s *Server) Serve(l net.Listener) (err error) {
	for err == nil {
		var c net.Conn
		c, err = l.Accept()
		if err == nil {
			p := &pop3Session{
				c: textproto.NewConn(c),
				s: s,
			}
			go p.Run()
		}
	}
	return
}
<commit_msg>fix DELE command, have it send a reply<commit_after>package pop3

import (
	"bds/lib/maildir"
	"bufio"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

type MailDirGetter interface {
	// get a user's maildir
	GetUserMaildir(user string) (maildir.MailDir, error)
}

type mailDirGetter string

func (md mailDirGetter) GetUserMaildir(user string) (m maildir.MailDir, err error) {
	m = maildir.MailDir(md)
	return
}

// a maildir getter that always uses 1 directory
func NewMailDirGetter(path string) MailDirGetter {
	return mailDirGetter(path)
}

// pop3 server
type Server struct {
	// function that obtains a maildir given a user
	GetMailDir MailDirGetter
	// server name
	name string
}

func New() *Server {
	host, _ := os.Hostname()
	return &Server{
		name: host,
	}
}

// get all messages in maildir
func (s *Server) getMessages(user string) (msgs []maildir.Message, err error) {
	// get user's maildir
	var md maildir.MailDir
	if s.GetMailDir == nil {
		err = errors.New("could't find maildir")
		return
	}
	md, err = s.GetMailDir.GetUserMaildir(user)
	if err != nil {
		return
	}
	// move new mail into cur
	var ms []maildir.Message
	ms, err = md.ListNew()
	if err == nil {
		for _, m := range ms {
			err = md.ProcessNew(m)
			if err != nil {
				log.Errorf("error processing maildir: %s", err.Error())
			}
		}
	}
	// get all current files
	msgs, err = md.ListCur()
	return
}

func (s *Server) checkUser(user, passwd string) (allowed bool) {
	// TODO: implement
	allowed = true
	return
}

func (s *Server) checkDigest(user, digest string) (allowed bool) {
	// TODO: implement
	allowed = true
	return
}

// get all messages and octet count
func (s *Server) obtainMessages(user string) (msgs []maildir.Message, o int64, err error) {
	msgs, err = s.getMessages(user)
	if err == nil {
		for _, msg := range msgs {
			var info os.FileInfo
			info, err = os.Stat(msg.Filepath())
			if err == nil && !info.IsDir() {
				o += info.Size()
			}
		}
	}
	return
}

// pop3 session handler
type pop3Session struct {
	// network connection
	c *textproto.Conn
	// parent server
	s *Server
	// are we in transaction state?
	transaction bool
	// current user
	user string
	// messages we have in this transaction
	msgs []maildir.Message
	// how many octets we have for all messages
	octs int64
	// messages to delete
	dels []maildir.Message
}

// run pop3 session mainloop
func (p *pop3Session) Run() {
	// send banner
	err := p.c.PrintfLine("+OK POP3 Server Ready <%d.%d@%s>", os.Getpid(), time.Now().Unix(), p.s.name)
	for err == nil {
		var line string
		line, err = p.c.ReadLine()
		if err == nil {
			if strings.ToUpper(line) == "QUIT" {
				// check for quit command
				err = p.OK("k bai")
				break
			} else if p.transaction {
				err = p.handleTransactionLine(line)
			} else {
				err = p.handleLine(line)
			}
		}
	}
	if err != nil && err != io.EOF {
		log.Errorf("error in pop3 session: %s", err.Error())
	}
	// close connection
	p.c.Close()
	// delete old messages
	for _, msg := range p.dels {
		os.Remove(msg.Filepath())
	}
}

// handle line when in transaction mode
func (p *pop3Session) handleTransactionLine(line string) (err error) {
	var info os.FileInfo
	var idx int
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "DELE":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid, add it to delete
				p.dels = append(p.dels, p.msgs[idx-1])
				p.OK("")
			} else {
				// invalid
				err = p.Error(err.Error())
			}
		}
	case "RETR":
		if len(parts) == 2 {
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				msg := p.msgs[idx-1].Filepath()
				var f *os.File
				f, err = os.Open(msg)
				if err == nil {
					r := bufio.NewReader(f)
					info, err = f.Stat()
					if err == nil {
						err = p.c.PrintfLine("+OK %d octets", info.Size())
						for err == nil {
							// send line
							line, err = r.ReadString(10)
							if err == io.EOF {
								err = nil
								break
							} else if err == nil {
								line = strings.Trim(line, "\r")
								line = strings.Trim(line, "\n")
								if line == "." {
									line = " ."
								}
								// send line
								err = p.c.PrintfLine(line)
							} else {
								// error
								break
							}
						}
					}
					f.Close()
					// end
					err = p.c.PrintfLine(".")
				} else {
					err = p.Error(err.Error())
				}
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			err = p.Error("invalid syntax")
		}
		break
	case "UIDL":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil && (idx > 0 && idx <= len(p.msgs)) {
				// valid
				err = p.OK(parts[1] + " " + p.msgs[idx-1].Filename())
			} else {
				// invalid
				err = p.Error("bad message")
			}
		} else {
			// all messages
			p.OK("")
			dw := p.c.DotWriter()
			for idx, msg := range p.msgs {
				fmt.Fprintf(dw, "%d %s\r\n", idx, msg.Filename())
			}
			// FLUSH :D
			err = dw.Close()
		}
		// begin
		break
	case "STAT":
		// begin
		_, err = p.c.W.WriteString("+OK ")
		if err == nil {
			// write list
			for idx, _ := range p.msgs {
				_, err = fmt.Fprintf(p.c.W, "%d ", 1+idx)
				if err != nil {
					break
				}
			}
		}
		if err == nil {
			// done writing list
			err = p.c.PrintfLine("")
		} else {
			err = p.Error(err.Error())
		}
		break
	case "LIST":
		if len(parts) == 2 {
			// 1 message
			idx, err = strconv.Atoi(parts[1])
			if err == nil {
				if idx > 0 && idx <= len(p.msgs) {
					info, err = os.Stat(p.msgs[idx-1].Filepath())
					if err == nil {
						err = p.c.PrintfLine("+OK %d %d", idx, info.Size())
					}
				} else {
					// no existing message
					err = p.Error("no such message")
				}
			}
		} else {
			// all messages
			err = p.c.PrintfLine("+OK %d messages (%d octets)", len(p.msgs), p.octs)
			if err == nil {
				dw := p.c.DotWriter()
				for i, msg := range p.msgs {
					info, err = os.Stat(msg.Filepath())
					if err == nil {
						// write out entry
						fmt.Fprintf(dw, "%d %d\r\n", i+1, info.Size())
					} else {
						log.Errorf("error in pop3, stat(): %s", err.Error())
					}
				}
				// FLUSH IT !!! :-DDDDD
				err = dw.Close()
			}
		}
		break
	default:
		err = p.Error("bad command")
	}
	return
}

// handle 1 line of input when not in transaction mode
func (p *pop3Session) handleLine(line string) (err error) {
	parts := strings.Split(line, " ")
	cmd := strings.ToUpper(parts[0])
	switch cmd {
	case "NOOP":
		err = p.OK("")
		break
	case "APOP":
		if len(parts) == 3 {
			if p.s.checkDigest(parts[1], parts[2]) {
				// load messages
				p.msgs, p.octs, err = p.s.obtainMessages(p.user)
				if err == nil {
					err = p.OK("maildrop is go for access")
					p.transaction = err == nil
				} else {
					err = p.Error(err.Error())
				}
			} else {
				err = p.Error("permission denied")
			}
		} else {
			err = p.Error("Bad APOP")
		}
		break
	case "PASS":
		if p.s.checkUser(p.user, line[5:]) {
			p.msgs, p.octs, err = p.s.obtainMessages(p.user)
			if err == nil {
				err = p.c.PrintfLine("+OK %s maildrop logged in, you have %d messages (%d octets)", p.user, len(p.msgs), p.octs)
				p.transaction = err == nil
			} else {
				err = p.Error(err.Error())
			}
		} else {
			err = p.Error("bad login")
		}
		break
	case "USER":
		if len(parts) > 1 {
			p.user = line[5:]
		}
		err = p.OK("anonymous access enabled")
		break
	default:
		err = p.Error("bad command")
	}
	return
}

func (p *pop3Session) OK(msg string) (err error) {
	if len(msg) == 0 {
		err = p.c.PrintfLine("+OK")
	} else {
		err = p.c.PrintfLine("+OK %s", msg)
	}
	return
}

// send error
func (p *pop3Session) Error(msg string) (err error) {
	err = p.c.PrintfLine("-ERR %s", msg)
	return
}

// serve sessions with connections accepted from a net.Listener
func (s *Server) Serve(l net.Listener) (err error) {
	for err == nil {
		var c net.Conn
		c, err = l.Accept()
		if err == nil {
			p := &pop3Session{
				c: textproto.NewConn(c),
				s: s,
			}
			go p.Run()
		}
	}
	return
}
<|endoftext|>


<commit_before>// Package edict implements a parser for the EDICT2 Japanese/English dictionary.
package edict

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// Gloss encodes an English definition for a Japanese word.
type Gloss struct {
	Definition  string   // English translation.
	Information []Detail // Information about this particular definition.
	Xref        []string // Xref to related entries (to the Kanji key), "see also".
}

// Entry encodes a line of edict2 input.
type Entry struct {
	Kanji              []string // Kanji key.
	Kana               []string // Kana transcription of keys.
	Information        []Detail // Information about the word; part of speech, conjugation type, etc.
	Gloss              []Gloss  // The "glosses", English definitions, ordered by frequency.
	Sequence           string   // The entry's unique identifier.
	RecordingAvailable bool     // True if an audio clip of the entry reading is available from the JapanesePod101.com site.
}

// String formats an Entry as a single line; not in the edict2 format, but familiar enough.
func (e Entry) String() string {
	recording := ""
	if e.RecordingAvailable {
		recording = "X"
	}

	return fmt.Sprintf("%v %v /%v %v/%s%s/", e.Kanji, e.Kana, e.Information, e.Gloss, e.Sequence, recording)
}

// These lines contain the record separator as part of the entry, making the whole thing
// signficantly more difficult to parse.  We'll skip these for now, and I'll patch the dictionary to
// not do this :)
var blacklist = []int{31179, 104168, 104171}

func Parse(in io.Reader) ([]Entry, error) {
	result := []Entry{}
	scanner := bufio.NewScanner(in)
	line := 0
lines:
	for scanner.Scan() {
		line++
		entry, err := parseLine(scanner.Text())
		if err != nil {
			for _, knownBadLine := range blacklist {
				if knownBadLine == line {
					continue lines
				}
			}
			return result, fmt.Errorf("parse: line %d: %s", line, err)
		}
		result = append(result, entry)
	}

	if err := scanner.Err(); err != nil {
		return result, fmt.Errorf("parse: past EOF (line %d): %s", line, err)
	}

	return result, nil
}

// Regular expressions for parsing entry lines.
var (
	// TODO(jrockway): kanji field is optional, according to the docs.  Key part looks like
	// "key1;key2;... [reading1;reading2;...] " (note the space at the end)
	parseKeys = regexp.MustCompile(`^([^[:space:]]+) \[([^\]]+)\] `)
)

type parseGlossState int

const (
	start parseGlossState = iota
	capture
	closed
	definition
)

func parseIdentifier(s string) (*Detail, *string, *string) {
	if _, err := strconv.Atoi(s); err == nil {
		return nil, nil, nil
	} else if strings.HasPrefix(s, "See ") {
		word := strings.TrimPrefix(s, "See ")
		return nil, &word, nil
	} else if detail, ok := DetailFor[s]; ok {
		return &detail, nil, nil
	} else {
		return nil, nil, &s
	}
}

func parseGloss(gloss string) (def string, details []Detail, xrefs []string, err error) {
	gloss = strings.TrimSpace(gloss)

	// This is the state machine for parsing the gloss.  We start in the start state, looking
	// for a ( starting an identifier, or the start of a definition (anything other than an
	// opening paren).  Upon seeing a ( we transition to capture, capturing everything that's
	// not a ).  Upon reaching the ), we then transition to closed.  In the closed state, we
	// look for a space, and finding it, transition to start.  At the end of the loop, we must
	// be in the definition-capture state.  If not, we raise an error.
	state := start
	captured := []rune{}
	defcapture := []rune{}

	for idx, c := range gloss {
		switch state {
		case start:
			if c == '(' {
				state = capture
			} else {
				state = definition
				defcapture = append(defcapture, c)
			}
		case definition:
			defcapture = append(defcapture, c)
		case capture:
			if c == ')' {
				state = closed
				d, x, u := parseIdentifier(string(captured))

				if d != nil {
					details = append(details, *d)
				} else if x != nil {
					xrefs = append(xrefs, *x)
				} else if u != nil {
					defcapture = append(defcapture, '(')
					for _, c := range captured {
						defcapture = append(defcapture, c)
					}
					defcapture = append(defcapture, ')')
					state = definition
				}
				captured = []rune{}
			} else {
				captured = append(captured, c)
			}
		case closed:
			if c == ' ' {
				state = start
			} else {
				err = fmt.Errorf("unexpected '%c' while in closed state (expecting space)", c)
			}
		default:
			err = fmt.Errorf("in unexpected state %v at byte %d '%c'", state, idx, c)
		}
	}

	if state != definition {
		err = fmt.Errorf("not in definition state after parsing:\ndetails=%v, xref=%v, def=%s", details, xrefs, def)
		return
	}

	def = string(defcapture)

	return
}

func parseLine(line string) (Entry, error) {
	result := Entry{}
	parts := strings.Split(line, "/")
	last := parts[len(parts)-1]
	if last != "" {
		return result, fmt.Errorf("parseLine: last component should be blank, but is %s", last)
	}

	// Parse the sequence number part, since having this in the result makes misparsing lines
	// easier to grep for.
	result.Sequence = parts[len(parts)-2]
	if strings.HasSuffix(result.Sequence, "X") {
		result.RecordingAvailable = true
		result.Sequence = strings.TrimSuffix(result.Sequence, "X")
	}
	result.Sequence = seqParts[1]

	// Parse the first part into Kanji/Kana fields.
	key := parts[0]
	keyParts := parseKeys.FindStringSubmatch(key)
	if len(keyParts) == 0 {
		result.Kanji = strings.Split(key, ";")
	} else if keyParts[0] != key || len(keyParts) != 3 {
		return result, fmt.Errorf("incomplete match on key '%s':\n got '%v'", key, keyParts)
	} else {
		result.Kanji = strings.Split(keyParts[1], ";")
		result.Kana = strings.Split(keyParts[2], ";")
	}

	// Next we get some details from the first gloss.
	glosses := []string{parts[1]}

	if len(parts) > 4 {
		// If there's more than one gloss, the entry-wide details come before the (1)
		// marker.
		firstGlossParts := strings.Split(parts[1], "(1)")
		if len(firstGlossParts) == 2 {
			_, detail, xref, err := parseGloss(firstGlossParts[0] + "fake definition")
			if err != nil {
				return result, fmt.Errorf("parsing entry details: %s", err)
			}
			if len(xref) != 0 {
				return result, fmt.Errorf("unexpected xref in global details section")
			}
			result.Information = detail
			glosses[0] = firstGlossParts[1]
		}
	}

	// We already have the first gloss in glosses, add the rest here.
	if len(parts) > 4 {
		for _, gloss := range parts[2 : len(parts)-2] {
			glosses = append(glosses, gloss)
		}
	}

	result.Gloss = []Gloss{}
	for _, gloss := range glosses {
		if gloss == "(P)" { // what a terrible file format
			result.Information = append(result.Information, Common)
			continue
		}

		def, detail, xref, err := parseGloss(gloss)
		if err != nil {
			return result, fmt.Errorf("parsing gloss %s got err %s", gloss, err)
		}
		result.Gloss = append(result.Gloss, Gloss{def, detail, xref})
	}

	// In the event that there's only one gloss, transfer the details to the entry.
	if len(parts) <= 4 {
		result.Information = result.Gloss[0].Information
		result.Gloss[0].Information = []Detail{}
	}

	return result, nil
}
<commit_msg>factor out parseKey, for possible future speedups<commit_after>// Package edict implements a parser for the EDICT2 Japanese/English dictionary.
package edict

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// Gloss encodes an English definition for a Japanese word.
type Gloss struct {
	Definition  string   // English translation.
	Information []Detail // Information about this particular definition.
	Xref        []string // Xref to related entries (to the Kanji key), "see also".
}

// Entry encodes a line of edict2 input.
type Entry struct {
	Kanji              []string // Kanji key.
	Kana               []string // Kana transcription of keys.
	Information        []Detail // Information about the word; part of speech, conjugation type, etc.
	Gloss              []Gloss  // The "glosses", English definitions, ordered by frequency.
	Sequence           string   // The entry's unique identifier.
	RecordingAvailable bool     // True if an audio clip of the entry reading is available from the JapanesePod101.com site.
}

// String formats an Entry as a single line; not in the edict2 format, but familiar enough.
func (e Entry) String() string {
	recording := ""
	if e.RecordingAvailable {
		recording = "X"
	}

	return fmt.Sprintf("%v %v /%v %v/%s%s/", e.Kanji, e.Kana, e.Information, e.Gloss, e.Sequence, recording)
}

// These lines contain the record separator as part of the entry, making the whole thing
// signficantly more difficult to parse.  We'll skip these for now, and I'll patch the dictionary to
// not do this :)
var blacklist = []int{31179, 104168, 104171}

func Parse(in io.Reader) ([]Entry, error) {
	result := []Entry{}
	scanner := bufio.NewScanner(in)
	line := 0
lines:
	for scanner.Scan() {
		line++
		entry, err := parseLine(scanner.Text())
		if err != nil {
			for _, knownBadLine := range blacklist {
				if knownBadLine == line {
					continue lines
				}
			}
			return result, fmt.Errorf("parse: line %d: %s", line, err)
		}
		result = append(result, entry)
	}

	if err := scanner.Err(); err != nil {
		return result, fmt.Errorf("parse: past EOF (line %d): %s", line, err)
	}

	return result, nil
}

// Regular expressions for parsing entry lines.
var (
	// TODO(jrockway): kanji field is optional, according to the docs.  Key part looks like
	// "key1;key2;... [reading1;reading2;...] " (note the space at the end)
	parseKeys = regexp.MustCompile(`^([^[:space:]]+) \[([^\]]+)\] `)
)

type parseGlossState int

const (
	start parseGlossState = iota
	capture
	closed
	definition
)

func parseIdentifier(s string) (*Detail, *string, *string) {
	if _, err := strconv.Atoi(s); err == nil {
		return nil, nil, nil
	} else if strings.HasPrefix(s, "See ") {
		word := strings.TrimPrefix(s, "See ")
		return nil, &word, nil
	} else if detail, ok := DetailFor[s]; ok {
		return &detail, nil, nil
	} else {
		return nil, nil, &s
	}
}

func parseGloss(gloss string) (def string, details []Detail, xrefs []string, err error) {
	gloss = strings.TrimSpace(gloss)

	// This is the state machine for parsing the gloss.  We start in the start state, looking
	// for a ( starting an identifier, or the start of a definition (anything other than an
	// opening paren).  Upon seeing a ( we transition to capture, capturing everything that's
	// not a ).  Upon reaching the ), we then transition to closed.  In the closed state, we
	// look for a space, and finding it, transition to start.  At the end of the loop, we must
	// be in the definition-capture state.  If not, we raise an error.
	state := start
	captured := []rune{}
	defcapture := []rune{}

	for idx, c := range gloss {
		switch state {
		case start:
			if c == '(' {
				state = capture
			} else {
				state = definition
				defcapture = append(defcapture, c)
			}
		case definition:
			defcapture = append(defcapture, c)
		case capture:
			if c == ')' {
				state = closed
				d, x, u := parseIdentifier(string(captured))

				if d != nil {
					details = append(details, *d)
				} else if x != nil {
					xrefs = append(xrefs, *x)
				} else if u != nil {
					defcapture = append(defcapture, '(')
					for _, c := range captured {
						defcapture = append(defcapture, c)
					}
					defcapture = append(defcapture, ')')
					state = definition
				}
				captured = []rune{}
			} else {
				captured = append(captured, c)
			}
		case closed:
			if c == ' ' {
				state = start
			} else {
				err = fmt.Errorf("unexpected '%c' while in closed state (expecting space)", c)
			}
		default:
			err = fmt.Errorf("in unexpected state %v at byte %d '%c'", state, idx, c)
		}
	}

	if state != definition {
		err = fmt.Errorf("not in definition state after parsing:\ndetails=%v, xref=%v, def=%s", details, xrefs, def)
		return
	}

	def = string(defcapture)

	return
}

func parseKey(key string) (kanji []string, kana []string, err error) {
	// Parse the first part into Kanji/Kana fields.
	keyParts := parseKeys.FindStringSubmatch(key)
	if len(keyParts) == 0 {
		kanji = strings.Split(key, ";")
	} else if keyParts[0] != key || len(keyParts) != 3 {
		err = fmt.Errorf("incomplete match on key '%s':\n got '%v'", key, keyParts)
		return
	} else {
		kanji = strings.Split(keyParts[1], ";")
		kana = strings.Split(keyParts[2], ";")
	}
	return
}

func parseLine(line string) (Entry, error) {
	result := Entry{}
	parts := strings.Split(line, "/")
	last := parts[len(parts)-1]
	if last != "" {
		return result, fmt.Errorf("parseLine: last component should be blank, but is %s", last)
	}

	// Parse the sequence number part, since having this in the result makes misparsing lines
	// easier to grep for.
	result.Sequence = parts[len(parts)-2]
	if strings.HasSuffix(result.Sequence, "X") {
		result.RecordingAvailable = true
		result.Sequence = strings.TrimSuffix(result.Sequence, "X")
	}

	var err error;
	result.Kanji, result.Kana, err = parseKey(parts[0])
	if err != nil {
		return result, err
	}

	// Next we get some details from the first gloss.
	glosses := []string{parts[1]}

	if len(parts) > 4 {
		// If there's more than one gloss, the entry-wide details come before the (1)
		// marker.
		firstGlossParts := strings.Split(parts[1], "(1)")
		if len(firstGlossParts) == 2 {
			_, detail, xref, err := parseGloss(firstGlossParts[0] + "fake definition")
			if err != nil {
				return result, fmt.Errorf("parsing entry details: %s", err)
			}
			if len(xref) != 0 {
				return result, fmt.Errorf("unexpected xref in global details section")
			}
			result.Information = detail
			glosses[0] = firstGlossParts[1]
		}
	}

	// We already have the first gloss in glosses, add the rest here.
	if len(parts) > 4 {
		for _, gloss := range parts[2 : len(parts)-2] {
			glosses = append(glosses, gloss)
		}
	}

	result.Gloss = []Gloss{}
	for _, gloss := range glosses {
		if gloss == "(P)" { // what a terrible file format
			result.Information = append(result.Information, Common)
			continue
		}

		def, detail, xref, err := parseGloss(gloss)
		if err != nil {
			return result, fmt.Errorf("parsing gloss %s got err %s", gloss, err)
		}
		result.Gloss = append(result.Gloss, Gloss{def, detail, xref})
	}

	// In the event that there's only one gloss, transfer the details to the entry.
	if len(parts) <= 4 {
		result.Information = result.Gloss[0].Information
		result.Gloss[0].Information = []Detail{}
	}

	return result, nil
}
<|endoftext|>


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

// Hello is a trivial example of a main package.
package main

import (
	"fmt"

	"code.google.com/p/go.example/newmath"
)

func main() {
	fmt.Printf("Hello, world.  Sqrt(2) = %v\n", newmath.Sqrt(2))
}
<commit_msg>Update hello.go<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Hello is a trivial example of a main package.
package main

import (
	"fmt"

	//"code.google.com/p/go.example/newmath"

	"newmath"
)

func main() {
	fmt.Printf("Hello, world.  Sqrt(2) = %v\n", newmath.Sqrt(2))
}
<|endoftext|>


<commit_before>package helpers

import (
	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"

	"reflect"

	"strings"

	"crypto/tls"

	"net"

	"fmt"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/pkg/errors"
)

var (
	mDbSession  *mgo.Session
	mDbDatabase string
)

type mgoLogger struct {
}

func (mgol mgoLogger) Output(calldepth int, s string) error {
	cache.GetLogger().WithField("module", "mdb").Info(s)
	return nil
}

// ConnectDB connects to mongodb and stores the session
func ConnectMDB(url string, database string) {
	var err error

	log := cache.GetLogger()
	log.WithField("module", "mdb").Info("Connecting to " + url)

	mgoL := new(mgoLogger)
	mgo.SetLogger(mgoL)

	cache.GetLogger()

	newUrl := strings.TrimSuffix(url, "?ssl=true")
	newUrl = strings.Replace(newUrl, "ssl=true&", "", -1)

	dialInfo, err := mgo.ParseURL(newUrl)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	// setup TLS if we use SSL
	if newUrl != url {
		tlsConfig := &tls.Config{}
		tlsConfig.InsecureSkipVerify = true

		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			conn, err := tls.Dial("tcp", addr.String(), tlsConfig)
			return conn, err
		}
	}

	mDbSession, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	mDbSession.SetMode(mgo.Monotonic, true)
	mDbSession.SetSafe(&mgo.Safe{WMode: "majority"})

	mDbDatabase = database

	log.WithField("module", "mdb").Info("Connected!")
}

// GetDB is a simple getter for the mongodb database.
func GetMDb() *mgo.Database {
	return mDbSession.DB(mDbDatabase)
}

// GetDB is a simple getter for the mongodb session.
func GetMDbSession() *mgo.Session {
	return mDbSession
}

func MDbInsert(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).Insert(temp.Interface())
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "insert",
				Method:     "MDbInsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbInsertWithoutLogging(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	err = GetMDb().C(collection.String()).Insert(temp.Interface())

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbUpdate(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).UpdateId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "update",
				Method:     "MDbUpdate()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpdateWithoutLogging(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	err = GetMDb().C(collection.String()).UpdateId(id, data)

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsertID(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	_, err = GetMDb().C(collection.String()).UpsertId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsertID()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MDbUpsert(collection models.MongoDbCollection, selector interface{}, data interface{}) (err error) {
	start := time.Now()
	_, err = GetMDb().C(collection.String()).Upsert(selector, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Query:      fmt.Sprintf("%+v", selector),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	return err
}

func MDbDelete(collection models.MongoDbCollection, id bson.ObjectId) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).RemoveId(id)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "remove",
				Method:     "MDbDelete()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         id.String(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MdbCollection(collection models.MongoDbCollection) (query *mgo.Collection) {
	return GetMDb().C(collection.String())
}

func MDbIter(query *mgo.Query) (iter *mgo.Iter) {
	start := time.Now()
	iter = query.Iter()
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbIter()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOne(query *mgo.Query, object interface{}) (err error) {
	start := time.Now()
	err = query.One(object)
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbOne()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOneWithoutLogging(query *mgo.Query, object interface{}) (err error) {
	err = query.One(object)
	return
}

// Returns a human readable ID version of a ObjectID
// id	: the ObjectID to convert
func MdbIdToHuman(id bson.ObjectId) (text string) {
	return fmt.Sprintf(`%x`, string(id))
}

// Returns an ObjectID from a human readable ID
// text	: the human readable ID
func HumanToMdbId(text string) (id bson.ObjectId) {
	return bson.ObjectIdHex(text)
}

func stripRobyulDatabaseFromCollection(input string) (output string) {
	return strings.TrimPrefix(input, mDbDatabase+".")
}

type KeenMongoDbEvent struct {
	Seconds    float64
	Collection string
	Type       string
	Method     string
	Query      string `json:",omitempty"`
	Skip       int64  `json:",omitempty"`
	Limit      int64  `json:",omitempty"`
	Id         string `json:",omitempty"`
	Data       string `json:",omitempty"`
}
<commit_msg>[mdb] improves keen logging 👀<commit_after>package helpers

import (
	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"

	"reflect"

	"strings"

	"crypto/tls"

	"net"

	"fmt"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/pkg/errors"
)

var (
	mDbSession  *mgo.Session
	mDbDatabase string
)

type mgoLogger struct {
}

func (mgol mgoLogger) Output(calldepth int, s string) error {
	cache.GetLogger().WithField("module", "mdb").Info(s)
	return nil
}

// ConnectDB connects to mongodb and stores the session
func ConnectMDB(url string, database string) {
	var err error

	log := cache.GetLogger()
	log.WithField("module", "mdb").Info("Connecting to " + url)

	mgoL := new(mgoLogger)
	mgo.SetLogger(mgoL)

	cache.GetLogger()

	newUrl := strings.TrimSuffix(url, "?ssl=true")
	newUrl = strings.Replace(newUrl, "ssl=true&", "", -1)

	dialInfo, err := mgo.ParseURL(newUrl)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	// setup TLS if we use SSL
	if newUrl != url {
		tlsConfig := &tls.Config{}
		tlsConfig.InsecureSkipVerify = true

		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			conn, err := tls.Dial("tcp", addr.String(), tlsConfig)
			return conn, err
		}
	}

	mDbSession, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithField("module", "mdb").Error(err.Error())
		panic(err)
	}

	mDbSession.SetMode(mgo.Monotonic, true)
	mDbSession.SetSafe(&mgo.Safe{WMode: "majority"})

	mDbDatabase = database

	log.WithField("module", "mdb").Info("Connected!")
}

// GetDB is a simple getter for the mongodb database.
func GetMDb() *mgo.Database {
	return mDbSession.DB(mDbDatabase)
}

// GetDB is a simple getter for the mongodb session.
func GetMDbSession() *mgo.Session {
	return mDbSession
}

func MDbInsert(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).Insert(temp.Interface())
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "insert",
				Method:     "MDbInsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbInsertWithoutLogging(collection models.MongoDbCollection, data interface{}) (rid bson.ObjectId, err error) {
	ptr := reflect.New(reflect.TypeOf(data))
	temp := ptr.Elem()
	temp.Set(reflect.ValueOf(data))

	v := temp.FieldByName("ID")

	if !v.IsValid() {
		return bson.ObjectId(""), errors.New("invalid data")
	}

	newID := v.String()
	if newID == "" {
		newID = string(bson.NewObjectId())
		v.SetString(newID)
	}

	err = GetMDb().C(collection.String()).Insert(temp.Interface())

	if err != nil {
		return bson.ObjectId(""), err
	}

	return bson.ObjectId(newID), nil
}

func MDbUpdate(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).UpdateId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "update",
				Method:     "MDbUpdate()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         MdbIdToHuman(id),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpdateWithoutLogging(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (rid bson.ObjectId, err error) {
	if !id.Valid() {
		return bson.ObjectId(""), errors.New("invalid id")
	}

	err = GetMDb().C(collection.String()).UpdateId(id, data)

	if err != nil {
		return bson.ObjectId(""), err
	}

	return id, nil
}

func MDbUpsertID(collection models.MongoDbCollection, id bson.ObjectId, data interface{}) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	_, err = GetMDb().C(collection.String()).UpsertId(id, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsertID()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         MdbIdToHuman(id),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MDbUpsert(collection models.MongoDbCollection, selector interface{}, data interface{}) (err error) {
	start := time.Now()
	_, err = GetMDb().C(collection.String()).Upsert(selector, data)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "upsert",
				Method:     "MDbUpsert()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Query:      fmt.Sprintf("%+v", selector),
				Data:       fmt.Sprintf("%+v", data),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	return err
}

func MDbDelete(collection models.MongoDbCollection, id bson.ObjectId) (err error) {
	if !id.Valid() {
		return errors.New("invalid id")
	}

	start := time.Now()
	err = GetMDb().C(collection.String()).RemoveId(id)
	took := time.Since(start)

	if cache.HasKeen() {
		go func() {
			defer Recover()

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "remove",
				Method:     "MDbDelete()",
				Collection: stripRobyulDatabaseFromCollection(collection.String()),
				Id:         MdbIdToHuman(id),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}

	if err != nil {
		return err
	}

	return nil
}

func MdbCollection(collection models.MongoDbCollection) (query *mgo.Collection) {
	return GetMDb().C(collection.String())
}

func MDbIter(query *mgo.Query) (iter *mgo.Iter) {
	start := time.Now()
	iter = query.Iter()
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbIter()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOne(query *mgo.Query, object interface{}) (err error) {
	start := time.Now()
	err = query.One(object)
	took := time.Since(start)
	if cache.HasKeen() {
		go func() {
			defer Recover()

			queryValue := reflect.ValueOf(*query)
			queryOp := queryValue.FieldByName("query").FieldByName("op")

			err := cache.GetKeen().AddEvent("Robyul_MongoDB", &KeenMongoDbEvent{
				Seconds:    took.Seconds(),
				Type:       "query",
				Method:     "MdbOne()",
				Collection: stripRobyulDatabaseFromCollection(queryOp.FieldByName("collection").String()),
				Query:      fmt.Sprintf("%+v", reflect.ValueOf(queryOp.FieldByName("query")).Interface()),
				Skip:       queryOp.FieldByName("skip").Int(),
				Limit:      queryOp.FieldByName("limit").Int(),
			})
			if err != nil {
				cache.GetLogger().WithField("module", "mdb").Error("Error logging MongoDB request to keen: ", err.Error())
			}
		}()
	}
	return
}

func MdbOneWithoutLogging(query *mgo.Query, object interface{}) (err error) {
	err = query.One(object)
	return
}

// Returns a human readable ID version of a ObjectID
// id	: the ObjectID to convert
func MdbIdToHuman(id bson.ObjectId) (text string) {
	return fmt.Sprintf(`%x`, string(id))
}

// Returns an ObjectID from a human readable ID
// text	: the human readable ID
func HumanToMdbId(text string) (id bson.ObjectId) {
	return bson.ObjectIdHex(text)
}

func stripRobyulDatabaseFromCollection(input string) (output string) {
	return strings.TrimPrefix(input, mDbDatabase+".")
}

type KeenMongoDbEvent struct {
	Seconds    float64
	Collection string
	Type       string
	Method     string
	Query      string `json:",omitempty"`
	Skip       int64  `json:",omitempty"`
	Limit      int64  `json:",omitempty"`
	Id         string `json:",omitempty"`
	Data       string `json:",omitempty"`
}
<|endoftext|>


<commit_before>package gocql

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

type nodeState int32

func (n nodeState) String() string {
	if n == NodeUp {
		return "UP"
	} else if n == NodeDown {
		return "DOWN"
	}
	return fmt.Sprintf("UNKNOWN_%d", n)
}

const (
	NodeUp nodeState = iota
	NodeDown
)

type cassVersion struct {
	Major, Minor, Patch int
}

func (c *cassVersion) UnmarshalCQL(info TypeInfo, data []byte) error {
	v := strings.Split(string(data), ".")
	if len(v) != 3 {
		return fmt.Errorf("invalid schema_version: %v", string(data))
	}

	var err error
	c.Major, err = strconv.Atoi(v[0])
	if err != nil {
		return fmt.Errorf("invalid major version %v: %v", v[0], err)
	}

	c.Minor, err = strconv.Atoi(v[1])
	if err != nil {
		return fmt.Errorf("invalid minor version %v: %v", v[1], err)
	}

	c.Patch, err = strconv.Atoi(v[2])
	if err != nil {
		return fmt.Errorf("invalid patch version %v: %v", v[2], err)
	}

	return nil
}

func (c cassVersion) String() string {
	return fmt.Sprintf("v%d.%d.%d", c.Major, c.Minor, c.Patch)
}

func (c cassVersion) nodeUpDelay() time.Duration {
	if c.Major >= 2 && c.Minor >= 2 {
		// CASSANDRA-8236
		return 0
	}

	return 10 * time.Second
}

type HostInfo struct {
	// TODO(zariel): reduce locking maybe, not all values will change, but to ensure
	// that we are thread safe use a mutex to access all fields.
	mu         sync.RWMutex
	peer       string
	port       int
	dataCenter string
	rack       string
	hostId     string
	version    cassVersion
	state      nodeState
	tokens     []string
}

func (h *HostInfo) Equal(host *HostInfo) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	host.mu.RLock()
	defer host.mu.RUnlock()

	return h.peer == host.peer && h.hostId == host.hostId
}

func (h *HostInfo) Peer() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.peer
}

func (h *HostInfo) setPeer(peer string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.peer = peer
	return h
}

func (h *HostInfo) DataCenter() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.dataCenter
}

func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.dataCenter = dataCenter
	return h
}

func (h *HostInfo) Rack() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.rack
}

func (h *HostInfo) setRack(rack string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.rack = rack
	return h
}

func (h *HostInfo) HostID() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.hostId
}

func (h *HostInfo) setHostID(hostID string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.hostId = hostID
	return h
}

func (h *HostInfo) Version() cassVersion {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.version
}

func (h *HostInfo) setVersion(major, minor, patch int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.version = cassVersion{major, minor, patch}
	return h
}

func (h *HostInfo) State() nodeState {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.state
}

func (h *HostInfo) setState(state nodeState) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.state = state
	return h
}

func (h *HostInfo) Tokens() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.tokens
}

func (h *HostInfo) setTokens(tokens []string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.tokens = tokens
	return h
}

func (h *HostInfo) Port() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.port
}

func (h *HostInfo) setPort(port int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.port = port
	return h
}

func (h *HostInfo) update(from *HostInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.tokens = from.tokens
	h.version = from.version
	h.hostId = from.hostId
	h.dataCenter = from.dataCenter
}

func (h *HostInfo) IsUp() bool {
	return h.State() == NodeUp
}

func (h *HostInfo) String() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return fmt.Sprintf("[hostinfo peer=%q port=%d data_centre=%q rack=%q host_id=%q version=%q state=%s num_tokens=%d]", h.peer, h.port, h.dataCenter, h.rack, h.hostId, h.version, h.state, len(h.tokens))
}

// Polls system.peers at a specific interval to find new hosts
type ringDescriber struct {
	dcFilter   string
	rackFilter string
	session    *Session
	closeChan  chan bool
	// indicates that we can use system.local to get the connections remote address
	localHasRpcAddr bool

	mu              sync.Mutex
	prevHosts       []*HostInfo
	prevPartitioner string
}

func checkSystemLocal(control *controlConn) (bool, error) {
	iter := control.query("SELECT broadcast_address FROM system.local")
	if err := iter.err; err != nil {
		if errf, ok := err.(*errorFrame); ok {
			if errf.code == errSyntax {
				return false, nil
			}
		}

		return false, err
	}

	return true, nil
}

func (r *ringDescriber) GetHosts() (hosts []*HostInfo, partitioner string, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	// we need conn to be the same because we need to query system.peers and system.local
	// on the same node to get the whole cluster

	const (
		legacyLocalQuery = "SELECT data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
		// only supported in 2.2.0, 2.1.6, 2.0.16
		localQuery = "SELECT broadcast_address, data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
	)

	localHost := &HostInfo{}
	if r.localHasRpcAddr {
		iter := r.session.control.query(localQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.peer, &localHost.dataCenter, &localHost.rack,
			&localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}
	} else {
		iter := r.session.control.query(legacyLocalQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.dataCenter, &localHost.rack, &localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}

		addr, _, err := net.SplitHostPort(r.session.control.addr())
		if err != nil {
			// this should not happen, ever, as this is the address that was dialed by conn, here
			// a panic makes sense, please report a bug if it occurs.
			panic(err)
		}

		localHost.peer = addr
	}

	localHost.port = r.session.cfg.Port

	hosts = []*HostInfo{localHost}

	iter := r.session.control.query("SELECT rpc_address, data_center, rack, host_id, tokens, release_version FROM system.peers")
	if iter == nil {
		return r.prevHosts, r.prevPartitioner, nil
	}

	host := &HostInfo{port: r.session.cfg.Port}
	for iter.Scan(&host.peer, &host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version) {
		if r.matchFilter(host) {
			hosts = append(hosts, host)
		}
		host = &HostInfo{
			port: r.session.cfg.Port,
		}
	}

	if err = iter.Close(); err != nil {
		return nil, "", err
	}

	r.prevHosts = hosts
	r.prevPartitioner = partitioner

	return hosts, partitioner, nil
}

func (r *ringDescriber) matchFilter(host *HostInfo) bool {
	if r.dcFilter != "" && r.dcFilter != host.DataCenter() {
		return false
	}

	if r.rackFilter != "" && r.rackFilter != host.Rack() {
		return false
	}

	return true
}

func (r *ringDescriber) refreshRing() error {
	// if we have 0 hosts this will return the previous list of hosts to
	// attempt to reconnect to the cluster otherwise we would never find
	// downed hosts again, could possibly have an optimisation to only
	// try to add new hosts if GetHosts didnt error and the hosts didnt change.
	hosts, partitioner, err := r.GetHosts()
	if err != nil {
		return err
	}

	// TODO: move this to session
	// TODO: handle removing hosts here
	for _, h := range hosts {
		if host, ok := r.session.ring.addHostIfMissing(h); !ok {
			r.session.pool.addHost(h)
		} else {
			host.update(h)
		}
	}

	r.session.pool.SetPartitioner(partitioner)
	return nil
}
<commit_msg>handle snapshot versions<commit_after>package gocql

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

type nodeState int32

func (n nodeState) String() string {
	if n == NodeUp {
		return "UP"
	} else if n == NodeDown {
		return "DOWN"
	}
	return fmt.Sprintf("UNKNOWN_%d", n)
}

const (
	NodeUp nodeState = iota
	NodeDown
)

type cassVersion struct {
	Major, Minor, Patch int
}

func (c *cassVersion) UnmarshalCQL(info TypeInfo, data []byte) error {
	version := strings.TrimSuffix(string(data), "-SNAPSHOT")
	v := strings.Split(version, ".")
	if len(v) != 3 {
		return fmt.Errorf("invalid schema_version: %v", string(data))
	}

	var err error
	c.Major, err = strconv.Atoi(v[0])
	if err != nil {
		return fmt.Errorf("invalid major version %v: %v", v[0], err)
	}

	c.Minor, err = strconv.Atoi(v[1])
	if err != nil {
		return fmt.Errorf("invalid minor version %v: %v", v[1], err)
	}

	c.Patch, err = strconv.Atoi(v[2])
	if err != nil {
		return fmt.Errorf("invalid patch version %v: %v", v[2], err)
	}

	return nil
}

func (c cassVersion) String() string {
	return fmt.Sprintf("v%d.%d.%d", c.Major, c.Minor, c.Patch)
}

func (c cassVersion) nodeUpDelay() time.Duration {
	if c.Major >= 2 && c.Minor >= 2 {
		// CASSANDRA-8236
		return 0
	}

	return 10 * time.Second
}

type HostInfo struct {
	// TODO(zariel): reduce locking maybe, not all values will change, but to ensure
	// that we are thread safe use a mutex to access all fields.
	mu         sync.RWMutex
	peer       string
	port       int
	dataCenter string
	rack       string
	hostId     string
	version    cassVersion
	state      nodeState
	tokens     []string
}

func (h *HostInfo) Equal(host *HostInfo) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	host.mu.RLock()
	defer host.mu.RUnlock()

	return h.peer == host.peer && h.hostId == host.hostId
}

func (h *HostInfo) Peer() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.peer
}

func (h *HostInfo) setPeer(peer string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.peer = peer
	return h
}

func (h *HostInfo) DataCenter() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.dataCenter
}

func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.dataCenter = dataCenter
	return h
}

func (h *HostInfo) Rack() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.rack
}

func (h *HostInfo) setRack(rack string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.rack = rack
	return h
}

func (h *HostInfo) HostID() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.hostId
}

func (h *HostInfo) setHostID(hostID string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.hostId = hostID
	return h
}

func (h *HostInfo) Version() cassVersion {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.version
}

func (h *HostInfo) setVersion(major, minor, patch int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.version = cassVersion{major, minor, patch}
	return h
}

func (h *HostInfo) State() nodeState {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.state
}

func (h *HostInfo) setState(state nodeState) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.state = state
	return h
}

func (h *HostInfo) Tokens() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.tokens
}

func (h *HostInfo) setTokens(tokens []string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.tokens = tokens
	return h
}

func (h *HostInfo) Port() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.port
}

func (h *HostInfo) setPort(port int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.port = port
	return h
}

func (h *HostInfo) update(from *HostInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.tokens = from.tokens
	h.version = from.version
	h.hostId = from.hostId
	h.dataCenter = from.dataCenter
}

func (h *HostInfo) IsUp() bool {
	return h.State() == NodeUp
}

func (h *HostInfo) String() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return fmt.Sprintf("[hostinfo peer=%q port=%d data_centre=%q rack=%q host_id=%q version=%q state=%s num_tokens=%d]", h.peer, h.port, h.dataCenter, h.rack, h.hostId, h.version, h.state, len(h.tokens))
}

// Polls system.peers at a specific interval to find new hosts
type ringDescriber struct {
	dcFilter   string
	rackFilter string
	session    *Session
	closeChan  chan bool
	// indicates that we can use system.local to get the connections remote address
	localHasRpcAddr bool

	mu              sync.Mutex
	prevHosts       []*HostInfo
	prevPartitioner string
}

func checkSystemLocal(control *controlConn) (bool, error) {
	iter := control.query("SELECT broadcast_address FROM system.local")
	if err := iter.err; err != nil {
		if errf, ok := err.(*errorFrame); ok {
			if errf.code == errSyntax {
				return false, nil
			}
		}

		return false, err
	}

	return true, nil
}

func (r *ringDescriber) GetHosts() (hosts []*HostInfo, partitioner string, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	// we need conn to be the same because we need to query system.peers and system.local
	// on the same node to get the whole cluster

	const (
		legacyLocalQuery = "SELECT data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
		// only supported in 2.2.0, 2.1.6, 2.0.16
		localQuery = "SELECT broadcast_address, data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
	)

	localHost := &HostInfo{}
	if r.localHasRpcAddr {
		iter := r.session.control.query(localQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.peer, &localHost.dataCenter, &localHost.rack,
			&localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}
	} else {
		iter := r.session.control.query(legacyLocalQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.dataCenter, &localHost.rack, &localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}

		addr, _, err := net.SplitHostPort(r.session.control.addr())
		if err != nil {
			// this should not happen, ever, as this is the address that was dialed by conn, here
			// a panic makes sense, please report a bug if it occurs.
			panic(err)
		}

		localHost.peer = addr
	}

	localHost.port = r.session.cfg.Port

	hosts = []*HostInfo{localHost}

	iter := r.session.control.query("SELECT rpc_address, data_center, rack, host_id, tokens, release_version FROM system.peers")
	if iter == nil {
		return r.prevHosts, r.prevPartitioner, nil
	}

	host := &HostInfo{port: r.session.cfg.Port}
	for iter.Scan(&host.peer, &host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version) {
		if r.matchFilter(host) {
			hosts = append(hosts, host)
		}
		host = &HostInfo{
			port: r.session.cfg.Port,
		}
	}

	if err = iter.Close(); err != nil {
		return nil, "", err
	}

	r.prevHosts = hosts
	r.prevPartitioner = partitioner

	return hosts, partitioner, nil
}

func (r *ringDescriber) matchFilter(host *HostInfo) bool {
	if r.dcFilter != "" && r.dcFilter != host.DataCenter() {
		return false
	}

	if r.rackFilter != "" && r.rackFilter != host.Rack() {
		return false
	}

	return true
}

func (r *ringDescriber) refreshRing() error {
	// if we have 0 hosts this will return the previous list of hosts to
	// attempt to reconnect to the cluster otherwise we would never find
	// downed hosts again, could possibly have an optimisation to only
	// try to add new hosts if GetHosts didnt error and the hosts didnt change.
	hosts, partitioner, err := r.GetHosts()
	if err != nil {
		return err
	}

	// TODO: move this to session
	// TODO: handle removing hosts here
	for _, h := range hosts {
		if host, ok := r.session.ring.addHostIfMissing(h); !ok {
			r.session.pool.addHost(h)
		} else {
			host.update(h)
		}
	}

	r.session.pool.SetPartitioner(partitioner)
	return nil
}
<|endoftext|>