<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 Romain LÉTENDART
//
// See LICENSE file.

// Main package for the goxxx project
//
// For the details see the file goxxx.go, as godoc won't show the documentation for the main package.
package main

import (
	"flag"
	"fmt"
	"github.com/vaz-ar/cfg_flags"
	"github.com/vaz-ar/goxxx/core"
	"github.com/vaz-ar/goxxx/database"
	"github.com/vaz-ar/goxxx/help"
	"github.com/vaz-ar/goxxx/invoke"
	"github.com/vaz-ar/goxxx/memo"
	"github.com/vaz-ar/goxxx/search"
	"github.com/vaz-ar/goxxx/webinfo"
	"github.com/vaz-ar/goxxx/xkcd"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

const (
	// Application version
	GLOBAL_VERSION string = "0.0.2"

	// Equivalent to enums (cf. https://golang.org/ref/spec#Iota)
	FLAGS_EXIT     = iota //  == 0
	FLAGS_SUCCESS         //  == 1
	FLAGS_FAILURE         //  == 2
	FLAGS_ADD_USER        //  == 3
)

// Config struct
type configData struct {
	channel       string
	channelKey    string
	nick          string
	server        string
	modules       []string
	debug         bool
	emailServer   string
	emailPort     int
	emailSender   string
	emailAccount  string
	emailPassword string
}

// Process the command line arguments
func getOptions() (config configData, returnCode int) {
	// IRC
	flag.StringVar(&config.channel, "channel", "", "IRC channel name")
	flag.StringVar(&config.channelKey, "key", "", "IRC channel key (optional)")
	flag.StringVar(&config.nick, "nick", "goxxx", "the bot's nickname (optional)")
	flag.StringVar(&config.server, "server", "chat.freenode.net:6697", "IRC_SERVER[:PORT] (optional)")
	modules := flag.String("modules", "memo,webinfo,invoke,search,xkcd", "Modules to enable (separated by commas)")
	// Email
	flag.StringVar(&config.emailServer, "email_server", "", "SMTP server address")
	flag.IntVar(&config.emailPort, "email_port", 0, "SMTP server port")
	flag.StringVar(&config.emailSender, "email_sender", "", "Email address to use in the \"From\" part of the header")
	flag.StringVar(&config.emailAccount, "email_account", "", "Email address from which to send emails")
	flag.StringVar(&config.emailPassword, "email_pwd", "", "password for the SMTP server")
	// Application
	flag.BoolVar(&config.debug, "debug", false, "Debug mode")
	version := flag.Bool("version", false, "Display goxxx version")

	flag.Usage = func() {
		fmt.Println("Usage:", os.Args[0], "-channel CHANNEL [ARGUMENTS]")
		fmt.Println()
		fmt.Println("Arguments description:")
		flag.PrintDefaults()
		fmt.Println("\nCommands description:")
		fmt.Println("  add_user <nick> <email>: Add an user to the database\n")
	}

	// Hybrid config: use flags and INI file
	// Command line flags take precedence on INI values
	if err := cfg_flags.Parse("goxxx.ini"); err != nil {
		flag.Usage()
		log.Fatal(err)
	}

	config.modules = strings.Split(*modules, ",")

	if *version {
		fmt.Printf("\nGoxxx version: %s\n\n", GLOBAL_VERSION)
		returnCode = FLAGS_EXIT
		return
	}

	lenArgs := len(flag.Args())
	// add_user command
	if lenArgs > 0 && flag.Args()[0] == "add_user" {
		if lenArgs != 3 {
			flag.Usage()
			returnCode = FLAGS_FAILURE
			return
		}
		returnCode = FLAGS_ADD_USER
	} else if config.channel == "" {
		flag.Usage()
		returnCode = FLAGS_FAILURE
	} else {
		returnCode = FLAGS_SUCCESS
	}
	return
}

func main() {
	// Set log output to a file
	logFile, err := os.OpenFile("./goxxx_logs.txt", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("Error opening file: %v", err)
	}
	defer logFile.Close()
	log.SetOutput(logFile)

	config, returnCode := getOptions()
	if returnCode == FLAGS_EXIT {
		return
	} else if returnCode == FLAGS_FAILURE {
		log.Fatal("Initialisation failed (getOptions())")
	}
	if config.debug {
		// In debug mode we show the file name and the line from where the log come from
		log.SetFlags(log.LstdFlags | log.Lshortfile)
	}

	// Create the database
	db := database.NewDatabase("", false)
	defer db.Close()

	// Process commands if necessary
	if returnCode == FLAGS_ADD_USER {
		if err := database.AddUser(flag.Args()[1], flag.Args()[2]); err == nil {
			fmt.Println("User added to the database")
		} else {
			fmt.Printf("\nadd_user error: %s\n", err)
		}
		return
	}

	// Create the bot
	bot := core.NewBot(config.nick, config.server, config.channel, config.channelKey)

	// Initialise packages
	for _, module := range config.modules {
		switch strings.TrimSpace(module) {
		case "invoke":
			if !invoke.Init(db, config.emailSender, config.emailAccount, config.emailPassword, config.emailServer, config.emailPort) {
				log.Println("Error while initialising invoke package")
				return
			}
			bot.AddCmdHandler(invoke.HandleInvokeCmd, bot.ReplyToNick)
			help.AddMessages(invoke.HELP_INVOKE)
			log.Println("invoke module loaded")

		case "memo":
			memo.Init(db)
			bot.AddMsgHandler(memo.SendMemo, bot.ReplyToNick)
			bot.AddCmdHandler(memo.HandleMemoCmd, bot.ReplyToAll)
			bot.AddCmdHandler(memo.HandleMemoStatusCmd, bot.ReplyToNick)
			help.AddMessages(memo.HELP_MEMO, memo.HELP_MEMOSTAT)
			log.Println("memo module loaded")

		case "search":
			bot.AddCmdHandler(search.HandleSearchCmd, bot.Reply)
			help.AddMessages(
				search.HELP_DUCKDUCKGO,
				search.HELP_WIKIPEDIA,
				search.HELP_WIKIPEDIA_FR,
				search.HELP_URBANDICTIONNARY)
			log.Println("search module loaded")

		case "webinfo":
			webinfo.Init(db)
			bot.AddMsgHandler(webinfo.HandleUrls, bot.ReplyToAll)
			log.Println("webinfo module loaded")

		case "xkcd":
			bot.AddCmdHandler(xkcd.HandleXKCDCmd, bot.ReplyToAll)
			help.AddMessages(xkcd.HELP_XKCD, xkcd.HELP_XKCD_NUM)
			log.Println("xkcd module loaded")

		default:
		}
	}
	bot.AddCmdHandler(help.HandleHelpCmd, bot.ReplyToNick)

	log.Println("Goxxx started")

	// Go signal notification works by sending os.Signal values on a channel.
	// We'll create a channel to receive these notifications
	// (we'll also make one to notify us when the program can exit).
	interruptSignals := make(chan os.Signal, 1)
	done := make(chan bool, 1)

	// signal.Notify registers the given channel to receive notifications of the specified signals.
	signal.Notify(interruptSignals, syscall.SIGINT, syscall.SIGTERM)

	// This goroutine executes a blocking receive for signals.
	// When it gets one it'll print it out and then notify the program that it can finish.
	go func() {
		sig := <-interruptSignals
		log.Printf("System signal received: %s\n", sig)
		done <- true
	}()

	// Start the bot
	go bot.Run()

	// The current routine will be blocked here until done is true
	<-done

	// Close the bot connection and the database
	bot.Stop()
	db.Close()

	log.Println("Goxxx exiting")
}
<commit_msg>Remove the configuration file path as a parameter of the cfg_flags.Parse() function<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 Romain LÉTENDART
//
// See LICENSE file.

// Main package for the goxxx project
//
// For the details see the file goxxx.go, as godoc won't show the documentation for the main package.
package main

import (
	"flag"
	"fmt"
	"github.com/vaz-ar/cfg_flags"
	"github.com/vaz-ar/goxxx/core"
	"github.com/vaz-ar/goxxx/database"
	"github.com/vaz-ar/goxxx/help"
	"github.com/vaz-ar/goxxx/invoke"
	"github.com/vaz-ar/goxxx/memo"
	"github.com/vaz-ar/goxxx/search"
	"github.com/vaz-ar/goxxx/webinfo"
	"github.com/vaz-ar/goxxx/xkcd"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

const (
	// Application version
	GLOBAL_VERSION string = "0.0.2"

	// Equivalent to enums (cf. https://golang.org/ref/spec#Iota)
	FLAGS_EXIT     = iota //  == 0
	FLAGS_SUCCESS         //  == 1
	FLAGS_FAILURE         //  == 2
	FLAGS_ADD_USER        //  == 3
)

// Config struct
type configData struct {
	channel       string
	channelKey    string
	nick          string
	server        string
	modules       []string
	debug         bool
	emailServer   string
	emailPort     int
	emailSender   string
	emailAccount  string
	emailPassword string
}

// Process the command line arguments
func getOptions() (config configData, returnCode int) {
	// IRC
	flag.StringVar(&config.channel, "channel", "", "IRC channel name")
	flag.StringVar(&config.channelKey, "key", "", "IRC channel key (optional)")
	flag.StringVar(&config.nick, "nick", "goxxx", "the bot's nickname (optional)")
	flag.StringVar(&config.server, "server", "chat.freenode.net:6697", "IRC_SERVER[:PORT] (optional)")
	modules := flag.String("modules", "memo,webinfo,invoke,search,xkcd", "Modules to enable (separated by commas)")
	// Email
	flag.StringVar(&config.emailServer, "email_server", "", "SMTP server address")
	flag.IntVar(&config.emailPort, "email_port", 0, "SMTP server port")
	flag.StringVar(&config.emailSender, "email_sender", "", "Email address to use in the \"From\" part of the header")
	flag.StringVar(&config.emailAccount, "email_account", "", "Email address from which to send emails")
	flag.StringVar(&config.emailPassword, "email_pwd", "", "password for the SMTP server")
	// Application
	flag.BoolVar(&config.debug, "debug", false, "Debug mode")
	version := flag.Bool("version", false, "Display goxxx version")

	flag.Usage = func() {
		fmt.Println("Usage:", os.Args[0], "-channel CHANNEL [ARGUMENTS]")
		fmt.Println()
		fmt.Println("Arguments description:")
		flag.PrintDefaults()
		fmt.Println("\nCommands description:")
		fmt.Println("  add_user <nick> <email>: Add an user to the database\n")
	}

	// Hybrid config: use flags and INI file
	// Command line flags take precedence on INI values
	if err := cfg_flags.Parse(); err != nil {
		flag.Usage()
		log.Fatal(err)
	}

	config.modules = strings.Split(*modules, ",")

	if *version {
		fmt.Printf("\nGoxxx version: %s\n\n", GLOBAL_VERSION)
		returnCode = FLAGS_EXIT
		return
	}

	lenArgs := len(flag.Args())
	// add_user command
	if lenArgs > 0 && flag.Args()[0] == "add_user" {
		if lenArgs != 3 {
			flag.Usage()
			returnCode = FLAGS_FAILURE
			return
		}
		returnCode = FLAGS_ADD_USER
	} else if config.channel == "" {
		flag.Usage()
		returnCode = FLAGS_FAILURE
	} else {
		returnCode = FLAGS_SUCCESS
	}
	return
}

func main() {
	// Set log output to a file
	logFile, err := os.OpenFile("./goxxx_logs.txt", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("Error opening file: %v", err)
	}
	defer logFile.Close()
	log.SetOutput(logFile)

	config, returnCode := getOptions()
	if returnCode == FLAGS_EXIT {
		return
	} else if returnCode == FLAGS_FAILURE {
		log.Fatal("Initialisation failed (getOptions())")
	}
	if config.debug {
		// In debug mode we show the file name and the line from where the log come from
		log.SetFlags(log.LstdFlags | log.Lshortfile)
	}

	// Create the database
	db := database.NewDatabase("", false)
	defer db.Close()

	// Process commands if necessary
	if returnCode == FLAGS_ADD_USER {
		if err := database.AddUser(flag.Args()[1], flag.Args()[2]); err == nil {
			fmt.Println("User added to the database")
		} else {
			fmt.Printf("\nadd_user error: %s\n", err)
		}
		return
	}

	// Create the bot
	bot := core.NewBot(config.nick, config.server, config.channel, config.channelKey)

	// Initialise packages
	for _, module := range config.modules {
		switch strings.TrimSpace(module) {
		case "invoke":
			if !invoke.Init(db, config.emailSender, config.emailAccount, config.emailPassword, config.emailServer, config.emailPort) {
				log.Println("Error while initialising invoke package")
				continue
			}
			bot.AddCmdHandler(invoke.HandleInvokeCmd, bot.ReplyToNick)
			help.AddMessages(invoke.HELP_INVOKE)
			log.Println("invoke module loaded")

		case "memo":
			memo.Init(db)
			bot.AddMsgHandler(memo.SendMemo, bot.ReplyToNick)
			bot.AddCmdHandler(memo.HandleMemoCmd, bot.ReplyToAll)
			bot.AddCmdHandler(memo.HandleMemoStatusCmd, bot.ReplyToNick)
			help.AddMessages(memo.HELP_MEMO, memo.HELP_MEMOSTAT)
			log.Println("memo module loaded")

		case "search":
			bot.AddCmdHandler(search.HandleSearchCmd, bot.Reply)
			help.AddMessages(
				search.HELP_DUCKDUCKGO,
				search.HELP_WIKIPEDIA,
				search.HELP_WIKIPEDIA_FR,
				search.HELP_URBANDICTIONNARY)
			log.Println("search module loaded")

		case "webinfo":
			webinfo.Init(db)
			bot.AddMsgHandler(webinfo.HandleUrls, bot.ReplyToAll)
			log.Println("webinfo module loaded")

		case "xkcd":
			bot.AddCmdHandler(xkcd.HandleXKCDCmd, bot.ReplyToAll)
			help.AddMessages(xkcd.HELP_XKCD, xkcd.HELP_XKCD_NUM)
			log.Println("xkcd module loaded")

		default:
		}
	}
	bot.AddCmdHandler(help.HandleHelpCmd, bot.ReplyToNick)

	log.Println("Goxxx started")

	// Go signal notification works by sending os.Signal values on a channel.
	// We'll create a channel to receive these notifications
	// (we'll also make one to notify us when the program can exit).
	interruptSignals := make(chan os.Signal, 1)
	done := make(chan bool, 1)

	// signal.Notify registers the given channel to receive notifications of the specified signals.
	signal.Notify(interruptSignals, syscall.SIGINT, syscall.SIGTERM)

	// This goroutine executes a blocking receive for signals.
	// When it gets one it'll print it out and then notify the program that it can finish.
	go func() {
		sig := <-interruptSignals
		log.Printf("System signal received: %s\n", sig)
		done <- true
	}()

	// Start the bot
	go bot.Run()

	// The current routine will be blocked here until done is true
	<-done

	// Close the bot connection and the database
	bot.Stop()
	db.Close()

	log.Println("Goxxx exiting")
}
<|endoftext|>


<commit_before>package cobe

import (
	"fmt"
	"log"
	"math/rand"
	"strings"
	"time"
)

type Brain struct {
	graph  *graph
	tok    tokenizer
	scorer scorer
}

const spaceTokenID tokenID = -1

func OpenBrain(path string) (*Brain, error) {
	graph, err := openGraph(path)
	if err != nil {
		return nil, err
	}

	version, err := graph.GetInfoString("version")
	if err != nil {
		return nil, err
	}

	if version != "2" {
		return nil, fmt.Errorf("cannot read version %s brain", version)
	}

	tokenizer, err := graph.GetInfoString("tokenizer")
	if err != nil {
		return nil, err
	}

	return &Brain{graph, getTokenizer(tokenizer), &cobeScorer{}}, nil
}

func (b *Brain) Close() {
	if b.graph != nil {
		b.graph.Close()
		b.graph = nil
	}
}

func getTokenizer(name string) tokenizer {
	switch strings.ToLower(name) {
	case "cobe":
		return newCobeTokenizer()
	case "megahal":
		return newMegaHALTokenizer()
	}

	return nil
}

func (b *Brain) Learn(text string) {
	tokens := b.tok.Split(text)

	// skip learning if too few tokens (but don't count spaces)
	if countGoodTokens(tokens) <= b.graph.getOrder() {
		return
	}

	var tokenIds []tokenID
	for _, text := range tokens {
		var tokenID tokenID
		if text == " " {
			tokenID = spaceTokenID
		} else {
			tokenID = b.graph.GetOrCreateToken(text)
		}

		tokenIds = append(tokenIds, tokenID)
	}

	var prevNode nodeID
	b.forEdges(tokenIds, func(prev, next []tokenID, hasSpace bool) {
		if prevNode == 0 {
			prevNode = b.graph.GetOrCreateNode(prev)
		}
		nextNode := b.graph.GetOrCreateNode(next)

		b.graph.addEdge(prevNode, nextNode, hasSpace)
		prevNode = nextNode
	})
}

func countGoodTokens(tokens []string) int {
	var count int
	for _, token := range tokens {
		if token != " " {
			count++
		}
	}

	return count
}

func (b *Brain) forEdges(tokenIds []tokenID, f func([]tokenID, []tokenID, bool)) {
	// Call f() on every N-gram (N = brain order) in tokenIds.
	order := b.graph.getOrder()

	chain := b.toChain(order, tokenIds)
	edges := toEdges(order, chain)

	for _, e := range edges {
		f(e.prev, e.next, e.hasSpace)
	}
}

func (b *Brain) toChain(order int, tokenIds []tokenID) []tokenID {
	var chain []tokenID
	for i := 0; i < order; i++ {
		chain = append(chain, b.graph.endTokenID)
	}

	chain = append(chain, tokenIds...)

	for i := 0; i < order; i++ {
		chain = append(chain, b.graph.endTokenID)
	}

	return chain
}

type edge struct {
	prev     []tokenID
	next     []tokenID
	hasSpace bool
}

func toEdges(order int, tokenIds []tokenID) []edge {
	var tokens []tokenID
	var spaces []int

	// Turn tokenIds (containing some SPACE_TOKEN_ID) into a list
	// of tokens and a list of positions in the tokens slice after
	// which spaces were found.

	for i := 0; i < len(tokenIds); i++ {
		tokens = append(tokens, tokenIds[i])

		if i < len(tokenIds)-1 && tokenIds[i+1] == spaceTokenID {
			spaces = append(spaces, len(tokens))
			i++
		}
	}

	var ret []edge

	prev := tokens[0:order]
	for i := 1; i < len(tokens)-order+1; i++ {
		next := tokens[i : i+order]

		var hasSpace bool
		if len(spaces) > 0 && spaces[0] == i+order-1 {
			hasSpace = true
			spaces = spaces[1:]
		}

		ret = append(ret, edge{prev, next, hasSpace})
		prev = next
	}

	return ret
}

func (b *Brain) Reply(text string) string {
	tokens := b.tok.Split(text)
	tokenIds := b.graph.filterPivots(unique(tokens))

	stemTokenIds := b.conflateStems(tokens)
	tokenIds = uniqueIds(append(tokenIds, stemTokenIds...))

	if len(tokenIds) == 0 {
		tokenIds = b.babble()
	}

	if len(tokenIds) == 0 {
		return "I don't know enough to answer you yet!"
	}

	var count int

	var bestReply *reply
	var bestScore float64 = -1

	stop := make(chan bool)
	replies := b.replySearch(tokenIds, stop)

	timeout := time.After(500 * time.Millisecond)
loop:
	for {
		select {
		case edges := <-replies:
			if edges == nil {
				// Channel was closed: run another search
				replies = b.replySearch(tokenIds, stop)
				continue
			}

			reply := newReply(b.graph, edges)
			score := b.scorer.Score(reply)

			if score > bestScore {
				bestReply = reply
				bestScore = score
			}

			count++
		case <-timeout:
			break loop
		}
	}

	fmt.Printf("Got %d total replies\n", count)
	if bestReply == nil {
		return "I don't know enough to answer you yet!"
	}

	return bestReply.ToString()
}

func (b *Brain) conflateStems(tokens []string) []tokenID {
	var ret []tokenID

	for _, token := range tokens {
		tokenIds := b.graph.getTokensByStem(token)
		ret = append(ret, tokenIds...)
	}

	return ret
}

func (b *Brain) babble() []tokenID {
	var tokenIds []tokenID

	for i := 0; i < 5; i++ {
		t := b.graph.getRandomToken()
		if t > 0 {
			tokenIds = append(tokenIds, tokenID(t))
		}
	}

	return tokenIds
}

// replySearch combines a forward and a reverse search over the graph
// into a series of replies.
func (b *Brain) replySearch(tokenIds []tokenID, stop chan bool) <-chan []edgeID {
	pivotID := b.pickPivot(tokenIds)
	pivotNode := b.graph.getRandomNodeWithToken(pivotID)

	endNode := b.graph.endContextID

	revIter := &history{b.graph.search(pivotNode, endNode, Reverse), nil}
	fwdIter := &history{b.graph.search(pivotNode, endNode, Forward), nil}

	replies := make(chan []edgeID)

	go func() {
	loop:
		for {
			rev := revIter.Next()
			if rev {
				// combine new rev with all fwds
				for _, f := range fwdIter.h {
					select {
					case replies <- join(revIter.Result(), f):
						// nothing
					case <-stop:
						break loop
					}
				}
			}

			fwd := fwdIter.Next()
			if fwd {
				// combine new rev with all fwds
				for _, r := range revIter.h {
					select {
					case replies <- join(r, fwdIter.Result()):
						// nothing
					case <-stop:
						break loop
					}
				}
			}

			if !rev && !fwd {
				break
			}
		}

		close(replies)
	}()

	return replies
}

type history struct {
	s *search
	h [][]edgeID
}

func (h *history) Next() bool {
	ret := h.s.Next()
	if ret {
		h.h = append(h.h, h.s.Result())
	}

	return ret
}

func (h *history) Result() []edgeID {
	return h.s.Result()
}

func join(rev []edgeID, fwd []edgeID) []edgeID {
	edges := make([]edgeID, 0, len(rev)+len(fwd))

	// rev is a path from the pivot node to the beginning of a
	// reply: join its edges in reverse order.
	for i := len(rev) - 1; i >= 0; i-- {
		edges = append(edges, rev[i])
	}

	return append(edges, fwd...)
}

func (b *Brain) pickPivot(tokenIds []tokenID) tokenID {
	return tokenIds[rand.Intn(len(tokenIds))]
}

func unique(tokens []string) []string {
	// Reduce tokens to a unique set by sending them through a map.
	m := make(map[string]int)
	for _, token := range tokens {
		m[token]++
	}

	ret := make([]string, 0, len(m))
	for token := range m {
		ret = append(ret, token)
	}

	return ret
}

func uniqueIds(ids []tokenID) []tokenID {
	// Reduce token ids to a unique set by sending them through a map.
	m := make(map[tokenID]int)
	for _, id := range ids {
		m[id]++
	}

	ret := make([]tokenID, 0, len(m))
	for id := range m {
		ret = append(ret, id)
	}

	return ret
}

type reply struct {
	graph   *graph
	edges   []edgeID
	hasText bool
	text    string
}

func newReply(graph *graph, edges []edgeID) *reply {
	return &reply{graph, edges, false, ""}
}

func (r *reply) ToString() string {
	if !r.hasText {
		var parts []string

		for _, edge := range r.edges {
			word, hasSpace, err := r.graph.getTextByEdge(edge)
			if err != nil {
				log.Printf("ERROR: %s\n", err)
			}

			parts = append(parts, word)
			if hasSpace {
				parts = append(parts, " ")
			}
		}

		r.hasText = true
		r.text = strings.Join(parts, "")
	}

	return r.text
}
<commit_msg>Don't break the main reply loop until we've gotten one reply<commit_after>package cobe

import (
	"fmt"
	"log"
	"math/rand"
	"strings"
	"time"
)

type Brain struct {
	graph  *graph
	tok    tokenizer
	scorer scorer
}

const spaceTokenID tokenID = -1

func OpenBrain(path string) (*Brain, error) {
	graph, err := openGraph(path)
	if err != nil {
		return nil, err
	}

	version, err := graph.GetInfoString("version")
	if err != nil {
		return nil, err
	}

	if version != "2" {
		return nil, fmt.Errorf("cannot read version %s brain", version)
	}

	tokenizer, err := graph.GetInfoString("tokenizer")
	if err != nil {
		return nil, err
	}

	return &Brain{graph, getTokenizer(tokenizer), &cobeScorer{}}, nil
}

func (b *Brain) Close() {
	if b.graph != nil {
		b.graph.Close()
		b.graph = nil
	}
}

func getTokenizer(name string) tokenizer {
	switch strings.ToLower(name) {
	case "cobe":
		return newCobeTokenizer()
	case "megahal":
		return newMegaHALTokenizer()
	}

	return nil
}

func (b *Brain) Learn(text string) {
	tokens := b.tok.Split(text)

	// skip learning if too few tokens (but don't count spaces)
	if countGoodTokens(tokens) <= b.graph.getOrder() {
		return
	}

	var tokenIds []tokenID
	for _, text := range tokens {
		var tokenID tokenID
		if text == " " {
			tokenID = spaceTokenID
		} else {
			tokenID = b.graph.GetOrCreateToken(text)
		}

		tokenIds = append(tokenIds, tokenID)
	}

	var prevNode nodeID
	b.forEdges(tokenIds, func(prev, next []tokenID, hasSpace bool) {
		if prevNode == 0 {
			prevNode = b.graph.GetOrCreateNode(prev)
		}
		nextNode := b.graph.GetOrCreateNode(next)

		b.graph.addEdge(prevNode, nextNode, hasSpace)
		prevNode = nextNode
	})
}

func countGoodTokens(tokens []string) int {
	var count int
	for _, token := range tokens {
		if token != " " {
			count++
		}
	}

	return count
}

func (b *Brain) forEdges(tokenIds []tokenID, f func([]tokenID, []tokenID, bool)) {
	// Call f() on every N-gram (N = brain order) in tokenIds.
	order := b.graph.getOrder()

	chain := b.toChain(order, tokenIds)
	edges := toEdges(order, chain)

	for _, e := range edges {
		f(e.prev, e.next, e.hasSpace)
	}
}

func (b *Brain) toChain(order int, tokenIds []tokenID) []tokenID {
	var chain []tokenID
	for i := 0; i < order; i++ {
		chain = append(chain, b.graph.endTokenID)
	}

	chain = append(chain, tokenIds...)

	for i := 0; i < order; i++ {
		chain = append(chain, b.graph.endTokenID)
	}

	return chain
}

type edge struct {
	prev     []tokenID
	next     []tokenID
	hasSpace bool
}

func toEdges(order int, tokenIds []tokenID) []edge {
	var tokens []tokenID
	var spaces []int

	// Turn tokenIds (containing some SPACE_TOKEN_ID) into a list
	// of tokens and a list of positions in the tokens slice after
	// which spaces were found.

	for i := 0; i < len(tokenIds); i++ {
		tokens = append(tokens, tokenIds[i])

		if i < len(tokenIds)-1 && tokenIds[i+1] == spaceTokenID {
			spaces = append(spaces, len(tokens))
			i++
		}
	}

	var ret []edge

	prev := tokens[0:order]
	for i := 1; i < len(tokens)-order+1; i++ {
		next := tokens[i : i+order]

		var hasSpace bool
		if len(spaces) > 0 && spaces[0] == i+order-1 {
			hasSpace = true
			spaces = spaces[1:]
		}

		ret = append(ret, edge{prev, next, hasSpace})
		prev = next
	}

	return ret
}

func (b *Brain) Reply(text string) string {
	tokens := b.tok.Split(text)
	tokenIds := b.graph.filterPivots(unique(tokens))

	stemTokenIds := b.conflateStems(tokens)
	tokenIds = uniqueIds(append(tokenIds, stemTokenIds...))

	if len(tokenIds) == 0 {
		tokenIds = b.babble()
	}

	if len(tokenIds) == 0 {
		return "I don't know enough to answer you yet!"
	}

	var count int

	var bestReply *reply
	var bestScore float64 = -1

	stop := make(chan bool)
	replies := b.replySearch(tokenIds, stop)

	timeout := time.After(500 * time.Millisecond)
loop:
	for {
		select {
		case edges := <-replies:
			if edges == nil {
				// Channel was closed: run another search
				replies = b.replySearch(tokenIds, stop)
				continue
			}

			reply := newReply(b.graph, edges)
			score := b.scorer.Score(reply)

			if score > bestScore {
				bestReply = reply
				bestScore = score
			}

			count++
		case <-timeout:
			if bestReply != nil {
				break loop
			}
		}
	}

	fmt.Printf("Got %d total replies\n", count)
	if bestReply == nil {
		return "I don't know enough to answer you yet!"
	}

	return bestReply.ToString()
}

func (b *Brain) conflateStems(tokens []string) []tokenID {
	var ret []tokenID

	for _, token := range tokens {
		tokenIds := b.graph.getTokensByStem(token)
		ret = append(ret, tokenIds...)
	}

	return ret
}

func (b *Brain) babble() []tokenID {
	var tokenIds []tokenID

	for i := 0; i < 5; i++ {
		t := b.graph.getRandomToken()
		if t > 0 {
			tokenIds = append(tokenIds, tokenID(t))
		}
	}

	return tokenIds
}

// replySearch combines a forward and a reverse search over the graph
// into a series of replies.
func (b *Brain) replySearch(tokenIds []tokenID, stop chan bool) <-chan []edgeID {
	pivotID := b.pickPivot(tokenIds)
	pivotNode := b.graph.getRandomNodeWithToken(pivotID)

	endNode := b.graph.endContextID

	revIter := &history{b.graph.search(pivotNode, endNode, Reverse), nil}
	fwdIter := &history{b.graph.search(pivotNode, endNode, Forward), nil}

	replies := make(chan []edgeID)

	go func() {
	loop:
		for {
			rev := revIter.Next()
			if rev {
				// combine new rev with all fwds
				for _, f := range fwdIter.h {
					select {
					case replies <- join(revIter.Result(), f):
						// nothing
					case <-stop:
						break loop
					}
				}
			}

			fwd := fwdIter.Next()
			if fwd {
				// combine new rev with all fwds
				for _, r := range revIter.h {
					select {
					case replies <- join(r, fwdIter.Result()):
						// nothing
					case <-stop:
						break loop
					}
				}
			}

			if !rev && !fwd {
				break
			}
		}

		close(replies)
	}()

	return replies
}

type history struct {
	s *search
	h [][]edgeID
}

func (h *history) Next() bool {
	ret := h.s.Next()
	if ret {
		h.h = append(h.h, h.s.Result())
	}

	return ret
}

func (h *history) Result() []edgeID {
	return h.s.Result()
}

func join(rev []edgeID, fwd []edgeID) []edgeID {
	edges := make([]edgeID, 0, len(rev)+len(fwd))

	// rev is a path from the pivot node to the beginning of a
	// reply: join its edges in reverse order.
	for i := len(rev) - 1; i >= 0; i-- {
		edges = append(edges, rev[i])
	}

	return append(edges, fwd...)
}

func (b *Brain) pickPivot(tokenIds []tokenID) tokenID {
	return tokenIds[rand.Intn(len(tokenIds))]
}

func unique(tokens []string) []string {
	// Reduce tokens to a unique set by sending them through a map.
	m := make(map[string]int)
	for _, token := range tokens {
		m[token]++
	}

	ret := make([]string, 0, len(m))
	for token := range m {
		ret = append(ret, token)
	}

	return ret
}

func uniqueIds(ids []tokenID) []tokenID {
	// Reduce token ids to a unique set by sending them through a map.
	m := make(map[tokenID]int)
	for _, id := range ids {
		m[id]++
	}

	ret := make([]tokenID, 0, len(m))
	for id := range m {
		ret = append(ret, id)
	}

	return ret
}

type reply struct {
	graph   *graph
	edges   []edgeID
	hasText bool
	text    string
}

func newReply(graph *graph, edges []edgeID) *reply {
	return &reply{graph, edges, false, ""}
}

func (r *reply) ToString() string {
	if !r.hasText {
		var parts []string

		for _, edge := range r.edges {
			word, hasSpace, err := r.graph.getTextByEdge(edge)
			if err != nil {
				log.Printf("ERROR: %s\n", err)
			}

			parts = append(parts, word)
			if hasSpace {
				parts = append(parts, " ")
			}
		}

		r.hasText = true
		r.text = strings.Join(parts, "")
	}

	return r.text
}
<|endoftext|>


<commit_before>package pool

import (
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/ubclaunchpad/cumulus/blockchain"
	"github.com/ubclaunchpad/cumulus/common/util"
	"github.com/ubclaunchpad/cumulus/consensus"
	"github.com/ubclaunchpad/cumulus/miner"
)

// PooledTransaction is a Transaction with a timestamp.
type PooledTransaction struct {
	Transaction *blockchain.Transaction
	Time        time.Time
}

// Pool is a set of valid Transactions.
type Pool struct {
	Order             []*PooledTransaction
	ValidTransactions map[blockchain.Hash]*PooledTransaction
}

// New initializes a new pool.
func New() *Pool {
	return &Pool{
		Order:             []*PooledTransaction{},
		ValidTransactions: map[blockchain.Hash]*PooledTransaction{},
	}
}

// Len returns the number of transactions in the Pool.
func (p *Pool) Len() int {
	return len(p.ValidTransactions)
}

// Empty returns true if the pool has exactly zero transactions in it.
func (p *Pool) Empty() bool {
	return p.Len() == 0
}

// Get returns the tranasction with input transaction Hash h.
func (p *Pool) Get(h blockchain.Hash) *blockchain.Transaction {
	return p.ValidTransactions[h].Transaction
}

// GetN returns the Nth transaction in the pool.
func (p *Pool) GetN(N int) *blockchain.Transaction {
	return p.Order[N].Transaction
}

// GetIndex returns the index of the transaction in the ordering.
func (p *Pool) GetIndex(t *blockchain.Transaction) int {
	hash := blockchain.HashSum(t)
	target := p.ValidTransactions[hash].Time
	return getIndex(p.Order, target, 0, p.Len()-1)
}

// getIndex does a binary search for a PooledTransaction by timestamp.
func getIndex(a []*PooledTransaction, target time.Time, low, high int) int {
	mid := (low + high) / 2
	if a[mid].Time == target {
		return mid
	} else if target.Before(a[mid].Time) {
		return getIndex(a, target, low, mid-1)
	} else {
		return getIndex(a, target, mid+1, high)
	}
}

// Push inserts a transaction into the pool, returning
// true if the Transaction was inserted (was valid).
// TODO: This should return an error if could not add.
func (p *Pool) Push(t *blockchain.Transaction, bc *blockchain.BlockChain) bool {
	if ok, err := consensus.VerifyTransaction(bc, t); ok {
		p.set(t)
		return true
	} else {
		log.Debug(err)
		return false
	}
}

// PushUnsafe adds a transaction to the pool without validation.
func (p *Pool) PushUnsafe(t *blockchain.Transaction) {
	p.set(t)
}

// Silently adds a transaction to the pool.
// Deletes a transaction if it exists from the input hash.
func (p *Pool) set(t *blockchain.Transaction) {
	hash := blockchain.HashSum(t)
	if txn, ok := p.ValidTransactions[hash]; ok {
		p.Delete(txn.Transaction)
	}
	vt := &PooledTransaction{
		Transaction: t,
		Time:        time.Now(),
	}
	p.Order = append(p.Order, vt)
	p.ValidTransactions[hash] = vt
}

// Delete removes a transaction from the Pool.
func (p *Pool) Delete(t *blockchain.Transaction) {
	hash := blockchain.HashSum(t)
	vt, ok := p.ValidTransactions[hash]
	if ok {
		i := p.GetIndex(vt.Transaction)
		p.Order = append(p.Order[0:i], p.Order[i+1:]...)
		delete(p.ValidTransactions, hash)
	}
}

// Update updates the Pool by removing the Transactions found in the
// Block. If the Block is found invalid wrt bc, then false is returned and no
// Transactions are removed from the Pool.
func (p *Pool) Update(b *blockchain.Block, bc *blockchain.BlockChain) bool {
	if ok, _ := consensus.VerifyBlock(bc, b); !ok {
		return false
	}
	for _, t := range b.Transactions {
		p.Delete(t)
	}
	return true
}

// Pop returns the next transaction and removes it from the pool.
func (p *Pool) Pop() *blockchain.Transaction {
	if p.Len() > 0 {
		next := p.GetN(0)
		p.Delete(next)
		return next
	}
	return nil
}

// Peek returns the next transaction and does not remove it from the pool.
func (p *Pool) Peek() *blockchain.Transaction {
	if p.Len() > 0 {
		return p.GetN(0)
	}
	return nil
}

// NextBlock produces a new block from the pool for mining.
func (p *Pool) NextBlock(chain *blockchain.BlockChain,
	address blockchain.Address, size uint32) *blockchain.Block {
	var txns []*blockchain.Transaction

	// Hash the last block in the chain.
	ix := len(chain.Blocks) - 1
	lastHash := blockchain.HashSum(chain.Blocks[ix])

	// Build a new block for mining.
	b := &blockchain.Block{
		BlockHeader: blockchain.BlockHeader{
			BlockNumber: uint32(len(chain.Blocks)),
			LastBlock:   lastHash,
			Time:        util.UnixNow(),
			Nonce:       0,
		}, Transactions: txns,
	}

	// Prepend the cloudbase transaction for this miner.
	miner.CloudBase(b, chain, address)

	// Try to grab as many transactions as the block will allow.
	// Test each transaction to see if we break size before adding.
	for p.Len() > 0 {
		nextSize := p.Peek().Len()
		if b.Len()+nextSize < int(size) {
			b.Transactions = append(b.Transactions, p.Pop())
		} else {
			break
		}
	}

	return b
}
<commit_msg>Replace Len method with Size<commit_after>package pool

import (
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/ubclaunchpad/cumulus/blockchain"
	"github.com/ubclaunchpad/cumulus/common/util"
	"github.com/ubclaunchpad/cumulus/consensus"
	"github.com/ubclaunchpad/cumulus/miner"
)

// PooledTransaction is a Transaction with a timestamp.
type PooledTransaction struct {
	Transaction *blockchain.Transaction
	Time        time.Time
}

// Pool is a set of valid Transactions.
type Pool struct {
	Order             []*PooledTransaction
	ValidTransactions map[blockchain.Hash]*PooledTransaction
}

// New initializes a new pool.
func New() *Pool {
	return &Pool{
		Order:             []*PooledTransaction{},
		ValidTransactions: map[blockchain.Hash]*PooledTransaction{},
	}
}

// Size returns the number of transactions in the Pool.
func (p *Pool) Size() int {
	return len(p.ValidTransactions)
}

// Empty returns true if the pool has exactly zero transactions in it.
func (p *Pool) Empty() bool {
	return p.Size() == 0
}

// Get returns the tranasction with input transaction Hash h.
func (p *Pool) Get(h blockchain.Hash) *blockchain.Transaction {
	return p.ValidTransactions[h].Transaction
}

// GetN returns the Nth transaction in the pool.
func (p *Pool) GetN(N int) *blockchain.Transaction {
	return p.Order[N].Transaction
}

// GetIndex returns the index of the transaction in the ordering.
func (p *Pool) GetIndex(t *blockchain.Transaction) int {
	hash := blockchain.HashSum(t)
	target := p.ValidTransactions[hash].Time
	return getIndex(p.Order, target, 0, p.Size()-1)
}

// getIndex does a binary search for a PooledTransaction by timestamp.
func getIndex(a []*PooledTransaction, target time.Time, low, high int) int {
	mid := (low + high) / 2
	if a[mid].Time == target {
		return mid
	} else if target.Before(a[mid].Time) {
		return getIndex(a, target, low, mid-1)
	} else {
		return getIndex(a, target, mid+1, high)
	}
}

// Push inserts a transaction into the pool, returning
// true if the Transaction was inserted (was valid).
// TODO: This should return an error if could not add.
func (p *Pool) Push(t *blockchain.Transaction, bc *blockchain.BlockChain) bool {
	if ok, err := consensus.VerifyTransaction(bc, t); ok {
		p.set(t)
		return true
	} else {
		log.Debug(err)
		return false
	}
}

// PushUnsafe adds a transaction to the pool without validation.
func (p *Pool) PushUnsafe(t *blockchain.Transaction) {
	p.set(t)
}

// Silently adds a transaction to the pool.
// Deletes a transaction if it exists from the input hash.
func (p *Pool) set(t *blockchain.Transaction) {
	hash := blockchain.HashSum(t)
	if txn, ok := p.ValidTransactions[hash]; ok {
		p.Delete(txn.Transaction)
	}
	vt := &PooledTransaction{
		Transaction: t,
		Time:        time.Now(),
	}
	p.Order = append(p.Order, vt)
	p.ValidTransactions[hash] = vt
}

// Delete removes a transaction from the Pool.
func (p *Pool) Delete(t *blockchain.Transaction) {
	hash := blockchain.HashSum(t)
	vt, ok := p.ValidTransactions[hash]
	if ok {
		i := p.GetIndex(vt.Transaction)
		p.Order = append(p.Order[0:i], p.Order[i+1:]...)
		delete(p.ValidTransactions, hash)
	}
}

// Update updates the Pool by removing the Transactions found in the
// Block. If the Block is found invalid wrt bc, then false is returned and no
// Transactions are removed from the Pool.
func (p *Pool) Update(b *blockchain.Block, bc *blockchain.BlockChain) bool {
	if ok, _ := consensus.VerifyBlock(bc, b); !ok {
		return false
	}
	for _, t := range b.Transactions {
		p.Delete(t)
	}
	return true
}

// Pop returns the next transaction and removes it from the pool.
func (p *Pool) Pop() *blockchain.Transaction {
	if p.Size() > 0 {
		next := p.GetN(0)
		p.Delete(next)
		return next
	}
	return nil
}

// Peek returns the next transaction and does not remove it from the pool.
func (p *Pool) Peek() *blockchain.Transaction {
	if p.Size() > 0 {
		return p.GetN(0)
	}
	return nil
}

// NextBlock produces a new block from the pool for mining.
func (p *Pool) NextBlock(chain *blockchain.BlockChain,
	address blockchain.Address, size uint32) *blockchain.Block {
	var txns []*blockchain.Transaction

	// Hash the last block in the chain.
	ix := len(chain.Blocks) - 1
	lastHash := blockchain.HashSum(chain.Blocks[ix])

	// Build a new block for mining.
	b := &blockchain.Block{
		BlockHeader: blockchain.BlockHeader{
			BlockNumber: uint32(len(chain.Blocks)),
			LastBlock:   lastHash,
			Time:        util.UnixNow(),
			Nonce:       0,
		}, Transactions: txns,
	}

	// Prepend the cloudbase transaction for this miner.
	miner.CloudBase(b, chain, address)

	// Try to grab as many transactions as the block will allow.
	// Test each transaction to see if we break size before adding.
	for p.Size() > 0 {
		nextSize := p.Peek().Len()
		if b.Len()+nextSize < int(size) {
			b.Transactions = append(b.Transactions, p.Pop())
		} else {
			break
		}
	}

	return b
}
<|endoftext|>


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

package gofs

// TODO
// create() can either create a folder or a file.
// When a folder/file is created, make all subdirectories in the map as well

/* TEST5
 * Supports:
 *  [+] UTF=8 file names <- not yet
 *  [+] 2^128 files
 *  [+] o(1) seek/write time for metadata
 *  [+] There can be two files with the same name, but only if one is a directory
 */

import (
    "fmt"
    "crypto/md5"
    "encoding/hex"
    "compress/gzip"
    "bytes"
	"sync"
    "strings"
)

const STATUS_OK                 int = 0
const STATUS_ERROR              int = -1
const STATUS_EXISTS		        int = -2
const STATUS_NOT_FOUND          int = -3
const STATUS_NOT_READABLE       int = -4

const IRP_PURGE                 int = 2 /* Flush the entire database and all files */
const IRP_DELETE                int = 3 /* Delete a file/folder */
const IRP_WRITE                 int = 4 /* Write data to a file */
const IRP_CREATE                int = 5 /* Create a new file or folder */

const FLAG_FILE                 int = 1
const FLAG_DIRECTORY            int = 2

type gofs_header struct {
    filename    string
    key         [16]byte
    meta        map[string]*gofs_file
    t_size      uint /* Total size of all files */
    io_in       chan *gofs_io_block
	create_sync	sync.Mutex
}

type gofs_file struct {
    filename    string
    filetype    int /* FLAG_FILE, FLAG_DIRECTORY */
    datasum     string
    data        []byte
	lock		sync.Mutex
}

type gofs_io_block struct {
    file        *gofs_file
    name        string
    data        []byte
    status      int /* 0 == fail, 1 == ok, 2 == purge, 3 == delete, 4 == write */
    flags       int
    io_out      chan *gofs_io_block
}

func create_db(filename string) *gofs_header {
    var header                      = new(gofs_header)
    header.filename                 = filename
    header.meta                     = make(map[string]*gofs_file)
    header.meta[s("/")]             = new(gofs_file)
    header.meta[s("/")].filename    = "/"

    /* i/o channel processor. Performs i/o to the filesystem */
    header.io_in = make(chan *gofs_io_block)
    go func (f *gofs_header) {
        for {
            var io = <- header.io_in

            switch io.status {
            case IRP_PURGE:
                /* PURGE */
                out("ERROR: PURGING")
                close(header.io_in)
                return
            case IRP_DELETE:
                /* DELETE */
                // FIXME/ADDME
                io.status = STATUS_ERROR
                if io.file.filename == "/" { /* Cannot delete the root file */
                    io.status = STATUS_ERROR
                    io.io_out <- io
                } else {
                    if i := f.check(io.name); i != nil {
                        delete(f.meta, s(io.name))
                        f.meta[s(io.name)] = nil
                        io.status = STATUS_OK
                    }
                    io.io_out <- io
                }
            case IRP_WRITE:
                /* WRITE */
                if i := f.check(io.name); i != nil {
					io.file.lock.Lock()
                    if f.write_internal(i, io.data) == len(io.data) {
                        io.status = STATUS_OK
						io.file.lock.Unlock()
                        io.io_out <- io
                    } else {
                        io.status = STATUS_ERROR
						io.file.lock.Unlock()
                        io.io_out <- io
                    }
                }
            case IRP_CREATE:
                f.meta[s(io.name)] = new(gofs_file)
                io.file = f.meta[s(io.name)]
                io.file.filename = io.name

                if string(io.name[len(io.name) - 1:]) == "/" {
                    io.file.filetype = FLAG_DIRECTORY
                } else {
                    io.file.filetype = FLAG_FILE
                }

                /* Recursively create all subdirectory files */
                sub_strings := strings.Split(io.name, "/")
                sub_array := make([]string, len(sub_strings) - 2)
                copy(sub_array, sub_strings[1:len(sub_strings) - 1]) /* We do not need the first/last file */
                var tmp string = ""
                for e := range sub_array {
                    tmp += "/" + sub_array[e]

                    /* Create a subdirectory header */
                    func (sub_directory string, f *gofs_header) {
                        if f := f.check(sub_directory); f != nil {
                            return /* There can exist two files with the same name,
                                       as long as one is a directory and the other is a file */
                        }

                        f.meta[s(tmp)] = new(gofs_file)
                        f.meta[s(tmp)].filename = sub_directory + "/" /* Explicit directory name */
                        f.meta[s(tmp)].filetype = FLAG_DIRECTORY
                    } (tmp, f)
                }

                io.status = STATUS_OK
                io.io_out <- io
            }
        }
    } (header)

    return header
}

func (f *gofs_header) unmount_db(filename *string) int {
    type comp_data struct {
        file *gofs_file
        data_compressed bytes.Buffer
    }

    commit_ch := make(chan *comp_data)
    for k := range f.meta {
        header := new(comp_data)
        header.file = f.meta[k]

        go func (d *comp_data) {
            if d.file.filename == "/" {
                return
            }

            /*
             * Perform compression of the file, and store it in 'd'
             */
            if d.file.filetype == FLAG_FILE /* File */ && len(d.file.data) > 0 {
                /* Compression required since this is a file, and it's length is > 0 */
                w := gzip.NewWriter(&d.data_compressed)
                w.Write(d.file.data)
                w.Close()
            }
            commit_ch <- d
        }(header)
    }

    /* Do not count "/" as a file, since it is not sent in channel */
    total_files := f.get_file_count() - 1
    for total_files != 0 {
        var header = <- commit_ch
		out("inbound: " + header.file.filename)
        total_files -= 1
    }

    close(commit_ch)
    return STATUS_OK
}

func (f *gofs_header) get_file_count() uint {
    var total uint = 0
    for range f.meta {
        total += 1
    }

    return total
}

func (f *gofs_header) check(name string) *gofs_file {
    if sum := s(name); f.meta[sum] != nil {
        return f.meta[sum]
    }

    return nil
}

func (f *gofs_header) generate_irp(name string, data []byte, irp_type int) *gofs_io_block {
    switch irp_type {
    case IRP_DELETE:
        /* DELETE */
        var file_header = f.check(name)
        if file_header == nil {
            return nil /* ERROR -- deleting non-existant file */
        }

        irp := new(gofs_io_block)
        irp.file = file_header
        irp.name = name
		irp.io_out = make(chan *gofs_io_block)

        irp.status = IRP_DELETE

        return irp
    case IRP_WRITE:
        /* WRITE */
        var file_header = f.check(name)
        if file_header == nil {
            return nil
        }

        irp := new(gofs_io_block)
        irp.file = file_header
        irp.name = name
        irp.data = make([]byte, len(data))
		irp.io_out = make(chan *gofs_io_block)
        copy(irp.data, data)

        irp.status = IRP_WRITE /* write IRP request */

        return irp

    case IRP_CREATE:
        /* CREATE IRP */
        irp := new(gofs_io_block)
        irp.name = name
        irp.status = IRP_CREATE
        irp.io_out = make(chan *gofs_io_block)

        return irp
    }

    return nil
}

func (f *gofs_header) create(name string) (*gofs_file, int) {
    if file := f.check(name); file != nil {
        return nil, STATUS_EXISTS
    }
	f.create_sync.Lock()
    var irp *gofs_io_block = f.generate_irp(name, nil, IRP_CREATE)

    f.io_in <- irp
    output_irp := <- irp.io_out
	f.create_sync.Unlock()
    if output_irp.file == nil {
        return nil, STATUS_ERROR
    }
    close(output_irp.io_out)

    return output_irp.file, STATUS_OK
}

func (f *gofs_header) read(name string) ([]byte, int) {
    var file_header = f.check(name)
    if file_header == nil {
        return nil, STATUS_NOT_FOUND
    }

    if file_header.filetype == FLAG_DIRECTORY {
        return nil, STATUS_NOT_READABLE
    }

    output := make([]byte, len(file_header.data))
    copy(output, file_header.data)
    return output, STATUS_OK
}

func (f *gofs_header) delete(name string) int {
    irp := f.generate_irp(name, nil, IRP_DELETE)
    if irp == nil {
        return STATUS_ERROR /* ERROR -- File does not exist */
    }

    f.io_in <- irp
    var output_irp = <- irp.io_out

    close(irp.io_out)
    if output_irp.status != STATUS_OK {
        return STATUS_ERROR /* failed */
    }

    return STATUS_OK
}

func (f *gofs_header) write(name string, d []byte) int {
	if i := f.check(name); i == nil {
		return STATUS_ERROR
	}

    irp := f.generate_irp(name, d, IRP_WRITE)
    if irp == nil {
        return STATUS_ERROR /* FAILURE */
    }

    /*
     * Send the write request IRP and receive the response
     *  IRP indicating the write status of the request
     */
    f.io_in <- irp
    var output_irp = <- irp.io_out

    close(irp.io_out)
    if output_irp.status != STATUS_OK {
        return STATUS_ERROR /* failed */
    }

    return STATUS_OK
}

func (f *gofs_header) write_internal(d *gofs_file, data []byte) int {
    if len(data) == 0 {
        return len(data)
    }

    if uint(len(data)) >= uint(len(d.data)) {
        f.t_size += uint(len(data)) - uint(len(d.data))
    } else {
        f.t_size -= uint(len(d.data)) - uint(len(data))
    }

    d.data = make([]byte, len(data))
    copy(d.data, data)
    d.datasum = s(string(data))

    datalen := len(d.data)

    return datalen
}

func (f *gofs_header) get_total_filesizes() uint {
    return f.t_size
}

func (f *gofs_header) get_file_list() []string {
    var output []string

    for k := range f.meta {
        file := f.meta[k]
        if file.filetype == FLAG_DIRECTORY {
            output = append(output, "(DIR)  " + file.filename)
            continue
        }
        output = append(output, "(FILE) " + file.filename)
    }

    return output
}

/* Returns an md5sum of a string */
func s(name string) string {
    name_seeded := name + "gofs_magic"
    d := make([]byte, len(name_seeded))
    copy(d, name_seeded)
    sum := md5.Sum(d)
    return hex.EncodeToString(sum[:])
}

func out(debug string) {
    fmt.Println(debug)
}

func out_hex(debug []byte) {
    fmt.Printf("%v\r\n", debug)
}
<commit_msg>Added MAX_FILENAME_LENGTH. Fixed some issues with unmount().<commit_after>/*
 * Copyright (c) 2017 AlexRuzin (stan.ruzin@gmail.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package gofs

// TODO
// create() can either create a folder or a file.
// When a folder/file is created, make all subdirectories in the map as well

/* TEST5
 * Supports:
 *  [+] UTF=8 file names <- not yet
 *  [+] 2^128 files
 *  [+] o(1) seek/write time for metadata
 *  [+] There can be two files with the same name, but only if one is a directory
 */

import (
    "fmt"
    "crypto/md5"
    "encoding/hex"
    "compress/gzip"
    "bytes"
	"sync"
    "strings"
)

/*
 * Configurable constants
 */
const MAX_FILENAME_LENGTH		int = 256

const STATUS_OK                 int = 0
const STATUS_ERROR              int = -1
const STATUS_EXISTS		        int = -2
const STATUS_NOT_FOUND          int = -3
const STATUS_NOT_READABLE       int = -4
const STATUS_NAME_EXCEEDED		int = -5 /* Input name is too long for create() */

const IRP_PURGE                 int = 2 /* Flush the entire database and all files */
const IRP_DELETE                int = 3 /* Delete a file/folder */
const IRP_WRITE                 int = 4 /* Write data to a file */
const IRP_CREATE                int = 5 /* Create a new file or folder */

const FLAG_FILE                 int = 1
const FLAG_DIRECTORY            int = 2

type gofs_header struct {
    filename    string
    key         [16]byte
    meta        map[string]*gofs_file
    t_size      uint /* Total size of all files */
    io_in       chan *gofs_io_block
	create_sync	sync.Mutex
}

type gofs_file struct {
    filename    string
    filetype    int /* FLAG_FILE, FLAG_DIRECTORY */
    datasum     string
    data        []byte
	lock		sync.Mutex
}

type gofs_io_block struct {
    file        *gofs_file
    name        string
    data        []byte
    status      int /* 0 == fail, 1 == ok, 2 == purge, 3 == delete, 4 == write */
    flags       int
    io_out      chan *gofs_io_block
}

func create_db(filename string) *gofs_header {
    var header                      = new(gofs_header)
    header.filename                 = filename
    header.meta                     = make(map[string]*gofs_file)
    header.meta[s("/")]             = new(gofs_file)
    header.meta[s("/")].filename    = "/"

    /* i/o channel processor. Performs i/o to the filesystem */
    header.io_in = make(chan *gofs_io_block)
    go func (f *gofs_header) {
        for {
            var io = <- header.io_in

            switch io.status {
            case IRP_PURGE:
                /* PURGE */
                out("ERROR: PURGING")
                close(header.io_in)
                return
            case IRP_DELETE:
                /* DELETE */
                // FIXME/ADDME
                io.status = STATUS_ERROR
                if io.file.filename == "/" { /* Cannot delete the root file */
                    io.status = STATUS_ERROR
                    io.io_out <- io
                } else {
                    if i := f.check(io.name); i != nil {
                        delete(f.meta, s(io.name))
                        f.meta[s(io.name)] = nil
                        io.status = STATUS_OK
                    }
                    io.io_out <- io
                }
            case IRP_WRITE:
                /* WRITE */
                if i := f.check(io.name); i != nil {
					io.file.lock.Lock()
                    if f.write_internal(i, io.data) == len(io.data) {
                        io.status = STATUS_OK
						io.file.lock.Unlock()
                        io.io_out <- io
                    } else {
                        io.status = STATUS_ERROR
						io.file.lock.Unlock()
                        io.io_out <- io
                    }
                }
            case IRP_CREATE:
                f.meta[s(io.name)] = new(gofs_file)
                io.file = f.meta[s(io.name)]
                io.file.filename = io.name

                if string(io.name[len(io.name) - 1:]) == "/" {
                    io.file.filetype = FLAG_DIRECTORY
                } else {
                    io.file.filetype = FLAG_FILE
                }

                /* Recursively create all subdirectory files */
                sub_strings := strings.Split(io.name, "/")
                sub_array := make([]string, len(sub_strings) - 2)
                copy(sub_array, sub_strings[1:len(sub_strings) - 1]) /* We do not need the first/last file */
                var tmp string = ""
                for e := range sub_array {
                    tmp += "/" + sub_array[e]

                    /* Create a subdirectory header */
                    func (sub_directory string, f *gofs_header) {
                        if f := f.check(sub_directory); f != nil {
                            return /* There can exist two files with the same name,
                                       as long as one is a directory and the other is a file */
                        }

                        f.meta[s(tmp)] = new(gofs_file)
                        f.meta[s(tmp)].filename = sub_directory + "/" /* Explicit directory name */
                        f.meta[s(tmp)].filetype = FLAG_DIRECTORY
                    } (tmp, f)
                }

                io.status = STATUS_OK
                io.io_out <- io
            }
        }
    } (header)

    return header
}

func (f *gofs_header) get_file_count() uint {
    var total uint = 0
    for range f.meta {
        total += 1
    }

    return total
}

func (f *gofs_header) check(name string) *gofs_file {
    if sum := s(name); f.meta[sum] != nil {
        return f.meta[sum]
    }

    return nil
}

func (f *gofs_header) generate_irp(name string, data []byte, irp_type int) *gofs_io_block {
    switch irp_type {
    case IRP_DELETE:
        /* DELETE */
        var file_header = f.check(name)
        if file_header == nil {
            return nil /* ERROR -- deleting non-existant file */
        }

        irp := new(gofs_io_block)
        irp.file = file_header
        irp.name = name
		irp.io_out = make(chan *gofs_io_block)

        irp.status = IRP_DELETE

        return irp
    case IRP_WRITE:
        /* WRITE */
        var file_header = f.check(name)
        if file_header == nil {
            return nil
        }

        irp := new(gofs_io_block)
        irp.file = file_header
        irp.name = name
        irp.data = make([]byte, len(data))
		irp.io_out = make(chan *gofs_io_block)
        copy(irp.data, data)

        irp.status = IRP_WRITE /* write IRP request */

        return irp

    case IRP_CREATE:
        /* CREATE IRP */
        irp := new(gofs_io_block)
        irp.name = name
        irp.status = IRP_CREATE
        irp.io_out = make(chan *gofs_io_block)

        return irp
    }

    return nil
}

func (f *gofs_header) create(name string) (*gofs_file, int) {
    if file := f.check(name); file != nil {
        return nil, STATUS_EXISTS
    }

    if len(name) > MAX_FILENAME_LENGTH {
    	return nil, STATUS_NAME_EXCEEDED
	}

	f.create_sync.Lock()
    var irp *gofs_io_block = f.generate_irp(name, nil, IRP_CREATE)

    f.io_in <- irp
    output_irp := <- irp.io_out
	f.create_sync.Unlock()
    if output_irp.file == nil {
        return nil, STATUS_ERROR
    }
    close(output_irp.io_out)

    return output_irp.file, STATUS_OK
}

func (f *gofs_header) read(name string) ([]byte, int) {
    var file_header = f.check(name)
    if file_header == nil {
        return nil, STATUS_NOT_FOUND
    }

    if file_header.filetype == FLAG_DIRECTORY {
        return nil, STATUS_NOT_READABLE
    }

    output := make([]byte, len(file_header.data))
    copy(output, file_header.data)
    return output, STATUS_OK
}

func (f *gofs_header) delete(name string) int {
    irp := f.generate_irp(name, nil, IRP_DELETE)
    if irp == nil {
        return STATUS_ERROR /* ERROR -- File does not exist */
    }

    f.io_in <- irp
    var output_irp = <- irp.io_out

    close(irp.io_out)
    if output_irp.status != STATUS_OK {
        return STATUS_ERROR /* failed */
    }

    return STATUS_OK
}

func (f *gofs_header) write(name string, d []byte) int {
	if i := f.check(name); i == nil {
		return STATUS_ERROR
	}

    irp := f.generate_irp(name, d, IRP_WRITE)
    if irp == nil {
        return STATUS_ERROR /* FAILURE */
    }

    /*
     * Send the write request IRP and receive the response
     *  IRP indicating the write status of the request
     */
    f.io_in <- irp
    var output_irp = <- irp.io_out

    close(irp.io_out)
    if output_irp.status != STATUS_OK {
        return STATUS_ERROR /* failed */
    }

    return STATUS_OK
}

func (f *gofs_header) write_internal(d *gofs_file, data []byte) int {
    if len(data) == 0 {
        return len(data)
    }

    if uint(len(data)) >= uint(len(d.data)) {
        f.t_size += uint(len(data)) - uint(len(d.data))
    } else {
        f.t_size -= uint(len(d.data)) - uint(len(data))
    }

    d.data = make([]byte, len(data))
    copy(d.data, data)
    d.datasum = s(string(data))

    datalen := len(d.data)

    return datalen
}

func (f *gofs_header) unmount_db(filename *string) int {
	type raw_file struct {
		raw_sum [16]byte
		gzip_size uint
		filetype int
		name [MAX_FILENAME_LENGTH]byte
	}

	type comp_data struct {
		file *gofs_file
		data_compressed bytes.Buffer
		raw raw_file
	}

	commit_ch := make(chan *comp_data)
	for k := range f.meta {
		header := new(comp_data)
		header.file = f.meta[k]

		go func (d *comp_data) {
			if d.file.filename == "/" {
				return
			}

			/*
			 * Perform compression of the file, and store it in 'd'
			 */
			if d.file.filetype == FLAG_FILE /* File */ && len(d.file.data) > 0 {
				/* Compression required since this is a file, and it's length is > 0 */
				w := gzip.NewWriter(&d.data_compressed)
				w.Write(d.file.data)
				w.Close()
			}
			commit_ch <- d
		}(header)
	}

	/* Do not count "/" as a file, since it is not sent in channel */
	total_files := f.get_file_count() - 1
	for total_files != 0 {
		var _ = <- commit_ch
		//out("inbound: " + header.file.filename)
		total_files -= 1
	}

	close(commit_ch)
	return STATUS_OK
}

func (f *gofs_header) get_total_filesizes() uint {
    return f.t_size
}

func (f *gofs_header) get_file_list() []string {
    var output []string

    for k := range f.meta {
        file := f.meta[k]
        if file.filetype == FLAG_DIRECTORY {
            output = append(output, "(DIR)  " + file.filename)
            continue
        }
        output = append(output, "(FILE) " + file.filename)
    }

    return output
}

/* Returns an md5sum of a string */
func s(name string) string {
    name_seeded := name + "gofs_magic"
    d := make([]byte, len(name_seeded))
    copy(d, name_seeded)
    sum := md5.Sum(d)
    return hex.EncodeToString(sum[:])
}

func out(debug string) {
    fmt.Println(debug)
}

func out_hex(debug []byte) {
    fmt.Printf("%v\r\n", debug)
}
<|endoftext|>


<commit_before>// Tideland Go REST Server Library - REST - Path
//
// Copyright (C) 2009-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package rest

//--------------------
// IMPORTS
//--------------------

import (
	"github.com/tideland/golib/stringex"
)

//--------------------
// CONSTANTS
//--------------------

// Path indexes for the different parts.
const (
	PathDomain = 0
	PathResource = 1
	PathResourceID = 2
)

//--------------------
// PATH
//--------------------

// Path provides access to the parts of a
// request path interesting for handling a
// job.
type Path interface {
	// Length returns the number of parts of the path.
	Length() int

	// Part returns the parts of the URL path based on the
	// index or an empty string.
	Part(index int) string

	// Domain returns the requests domain.
	Domain() string

	// Resource returns the requests resource.
	Resource() string

	// ResourceID return the requests resource ID.
	ResourceID() string
}

// path implements Path.
type path struct {
	path []string
}

func newPath(url ) *path {
}

// EOF<commit_msg>Continued path implementation<commit_after>// Tideland Go REST Server Library - REST - Path
//
// Copyright (C) 2009-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package rest

//--------------------
// IMPORTS
//--------------------

import (
	"net/http"

	"github.com/tideland/golib/stringex"
)

//--------------------
// CONSTANTS
//--------------------

// Path indexes for the different parts.
const (
	PathDomain     = 0
	PathResource   = 1
	PathResourceID = 2
)

//--------------------
// PATH
//--------------------

// Path provides access to the parts of a
// request path interesting for handling a
// job.
type Path interface {
	// Length returns the number of parts of the path.
	Length() int

	// Part returns the parts of the URL path based on the
	// index or an empty string.
	Part(index int) string

	// Domain returns the requests domain.
	Domain() string

	// Resource returns the requests resource.
	Resource() string

	// ResourceID return the requests resource ID.
	ResourceID() string
}

// path implements Path.
type path struct {
	parts []string
}

// newPath returns the analyzed path.
func newPath(env *environment, r *http.Request) *path {
	parts := stringex.SplitMap(r.URL.Path, "/", func(p string) (string, bool) {
		if part == "" {
			return "", false
		}
		return part, true
	})[env.basepartsLen:]
	switch len(parts) {
	case 1:
		parts = append(parts, env.defaultResource)
	case 0:
		parts = append(parts, env.defaultDomain, nev.defaultResource)
	}
	return &path{
		parts: parts,
	}
}

// EOF
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	bzl "github.com/bazelbuild/buildifier/core"
	"github.com/golang/glog"
)

var kubeRoot = os.Getenv("KUBE_ROOT")
var dryRun = flag.Bool("dry-run", false, "run in dry mode")

func main() {
	flag.Parse()
	flag.Set("alsologtostderr", "true")
	v := Venderor{
		ctx: &build.Default,
	}
	if len(flag.Args()) == 2 {
		v.updateSinglePkg(flag.Args()[1])
	} else {
		v.walkVendor()
		if err := v.walkRepo(); err != nil {
			glog.Fatalf("err: %v", err)
		}
	}
}

type Venderor struct {
	ctx *build.Context
}

func writeHeaders(file *bzl.File) {
	pkgRule := bzl.Rule{
		&bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: "package"},
		},
	}
	pkgRule.SetAttr("default_visibility", asExpr([]string{"//visibility:public"}))

	file.Stmt = append(file.Stmt,
		[]bzl.Expr{
			pkgRule.Call,
			&bzl.CallExpr{
				X:    &bzl.LiteralExpr{Token: "licenses"},
				List: []bzl.Expr{asExpr([]string{"notice"})},
			},
			&bzl.CallExpr{
				X: &bzl.LiteralExpr{Token: "load"},
				List: asExpr([]string{
					"@io_bazel_rules_go//go:def.bzl",
					"go_binary",
					"go_library",
					"go_test",
					"cgo_library",
				}).(*bzl.ListExpr).List,
			},
		}...,
	)
}

func writeRules(file *bzl.File, rules []*bzl.Rule) {
	for _, rule := range rules {
		file.Stmt = append(file.Stmt, rule.Call)
	}
}

func (v *Venderor) resolve(ipath string) Label {
	if strings.HasPrefix(ipath, "k8s.io/kubernetes") {
		return Label{
			pkg: strings.TrimPrefix(ipath, "k8s.io/kubernetes/"),
			tag: "go_default_library",
		}
	}
	return Label{
		pkg: "vendor",
		tag: ipath,
	}
}

func (v *Venderor) walk(root string, f func(path, ipath string, pkg *build.Package) error) error {
	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		ipath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if !info.IsDir() {
			return nil
		}
		pkg, err := v.ctx.ImportDir(filepath.Join(kubeRoot, path), build.ImportComment)
		if err != nil {
			if _, ok := err.(*build.NoGoError); err != nil && ok {
				return nil
			} else {
				return err
			}
		}

		return f(path, ipath, pkg)
	})
}

func (v *Venderor) walkRepo() error {
	for _, root := range []string{
		"./pkg",
		"./cmd",
		"./third_party",
		"./plugin",
		"./test",
		"./federation",
	} {
		if err := v.walk(root, v.updatePkg); err != nil {
			return err
		}
	}
	return nil
}

func (v *Venderor) updateSinglePkg(path string) error {
	pkg, err := v.ctx.ImportDir("./"+path, build.ImportComment)
	if err != nil {
		if _, ok := err.(*build.NoGoError); err != nil && ok {
			return nil
		} else {
			return err
		}
	}
	return v.updatePkg(path, "", pkg)
}

func (v *Venderor) updatePkg(path, _ string, pkg *build.Package) error {
	var rules []*bzl.Rule

	var attrs Attrs = make(Attrs)
	srcs := asExpr(merge(pkg.GoFiles, pkg.SFiles)).(*bzl.ListExpr)

	deps := v.extractDeps(pkg.Imports)

	if len(srcs.List) == 0 {
		return nil
	}
	attrs.Set("srcs", srcs)

	if len(deps.List) > 0 {
		attrs.Set("deps", deps)
	}

	if pkg.IsCommand() {
		rules = append(rules, newRule("go_binary", filepath.Base(pkg.Dir), attrs))
	} else {
		rules = append(rules, newRule("go_library", "go_default_library", attrs))
		if len(pkg.TestGoFiles) != 0 {
			rules = append(rules, newRule("go_test", "go_default_test", map[string]bzl.Expr{
				"srcs":    asExpr(pkg.TestGoFiles),
				"deps":    v.extractDeps(pkg.TestImports),
				"library": asExpr("go_default_library"),
			}))
		}
	}

	if len(pkg.XTestGoFiles) != 0 {
		rules = append(rules, newRule("go_test", "go_default_xtest", map[string]bzl.Expr{
			"srcs": asExpr(pkg.XTestGoFiles),
			"deps": v.extractDeps(pkg.XTestImports),
		}))
	}

	wrote, err := ReconcileRules(filepath.Join(path, "BUILD"), rules)
	if err != nil {
		return err
	}
	if wrote {
		fmt.Fprintf(os.Stderr, "wrote BUILD for %q\n", pkg.ImportPath)
	}
	return nil
}

func (v *Venderor) walkVendor() {
	var rules []*bzl.Rule
	if err := v.walk("./vendor", func(path, ipath string, pkg *build.Package) error {
		var attrs Attrs = make(Attrs)

		srcs := asExpr(
			apply(
				merge(pkg.GoFiles, pkg.SFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		cgoSrcs := asExpr(
			apply(
				merge(pkg.CgoFiles, pkg.CFiles, pkg.CXXFiles, pkg.HFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		deps := v.extractDeps(pkg.Imports)
		attrs.Set("srcs", srcs)

		if len(deps.List) > 0 {
			attrs.Set("deps", deps)
		}

		if pkg.IsCommand() {
			rules = append(rules, newRule("go_binary", v.resolve(ipath).tag, attrs))
		} else {
			if len(cgoSrcs.List) != 0 {
				cgoPname := v.resolve(ipath).tag + "_cgo"
				cgoDeps := v.extractDeps(pkg.TestImports)
				cgoRule := newRule("cgo_library", cgoPname, map[string]bzl.Expr{
					"srcs":      cgoSrcs,
					"clinkopts": asExpr([]string{"-ldl", "-lz", "-lm", "-lpthread", "-ldl"}),
				})
				rules = append(rules, cgoRule)
				if len(cgoDeps.List) != 0 {
					cgoRule.SetAttr("deps", cgoDeps)
				}
				attrs["library"] = asExpr(cgoPname)
			}
			rules = append(rules, newRule("go_library", v.resolve(ipath).tag, attrs))
		}
		return nil
	}); err != nil {
		glog.Fatalf("err: %v", err)
	}
	if _, err := ReconcileRules("./vendor/BUILD", rules); err != nil {
		glog.Fatalf("err: %v", err)
	}
}

func (v *Venderor) extractDeps(deps []string) *bzl.ListExpr {
	return asExpr(
		apply(
			merge(deps),
			filterer(func(s string) bool {
				pkg, err := v.ctx.Import(s, kubeRoot, build.ImportComment)
				if err != nil {
					return false
				}
				if pkg.Goroot {
					return false
				}
				return true
			}),
			mapper(func(s string) string {
				return v.resolve(s).String()
			}),
		),
	).(*bzl.ListExpr)
}

type Attrs map[string]bzl.Expr

func (a Attrs) Set(name string, expr bzl.Expr) {
	a[name] = expr
}

type Label struct {
	pkg, tag string
}

func (l Label) String() string {
	return fmt.Sprintf("//%v:%v", l.pkg, l.tag)
}

func asExpr(e interface{}) bzl.Expr {
	rv := reflect.ValueOf(e)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%d", e)}
	case reflect.Float32, reflect.Float64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%f", e)}
	case reflect.String:
		return &bzl.StringExpr{Value: e.(string)}
	case reflect.Slice, reflect.Array:
		var list []bzl.Expr
		for i := 0; i < rv.Len(); i++ {
			list = append(list, asExpr(rv.Index(i).Interface()))
		}
		return &bzl.ListExpr{List: list}
	default:
		glog.Fatalf("Uh oh")
		return nil
	}
}

type Sed func(s []string) []string

func mapString(in []string, f func(string) string) []string {
	var out []string
	for _, s := range in {
		out = append(out, f(s))
	}
	return out
}

func mapper(f func(string) string) Sed {
	return func(in []string) []string {
		return mapString(in, f)
	}
}

func filterString(in []string, f func(string) bool) []string {
	var out []string
	for _, s := range in {
		if f(s) {
			out = append(out, s)
		}
	}
	return out
}

func filterer(f func(string) bool) Sed {
	return func(in []string) []string {
		return filterString(in, f)
	}
}

func apply(stream []string, seds ...Sed) []string {
	for _, sed := range seds {
		stream = sed(stream)
	}
	return stream
}

func merge(streams ...[]string) []string {
	var out []string
	for _, stream := range streams {
		out = append(out, stream...)
	}
	return out
}

func newRule(kind, name string, attrs map[string]bzl.Expr) *bzl.Rule {
	rule := &bzl.Rule{
		Call: &bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: kind},
		},
	}
	rule.SetAttr("name", asExpr(name))
	for k, v := range attrs {
		rule.SetAttr(k, v)
	}
	rule.SetAttr("tags", asExpr([]string{"automanaged"}))
	return rule
}

func ReconcileRules(path string, rules []*bzl.Rule) (bool, error) {
	info, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		f := &bzl.File{}
		writeHeaders(f)
		writeRules(f, rules)
		return writeFile(path, f, false)
	} else if err != nil {
		return false, err
	}
	if info.IsDir() {
		return false, fmt.Errorf("%q cannot be a directory", path)
	}
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return false, err
	}
	f, err := bzl.Parse(path, b)
	if err != nil {
		return false, err
	}
	oldRules := make(map[string]*bzl.Rule)
	for _, r := range f.Rules("") {
		oldRules[r.Name()] = r
	}
	for _, r := range rules {
		o, ok := oldRules[r.Name()]
		if !ok {
			f.Stmt = append(f.Stmt, r.Call)
			continue
		}
		if !RuleIsManaged(o) {
			continue
		}
		reconcileAttr := func(o, n *bzl.Rule, name string) {
			if e := n.Attr(name); e != nil {
				o.SetAttr(name, e)
			} else {
				o.DelAttr(name)
			}
		}
		reconcileAttr(o, r, "srcs")
		reconcileAttr(o, r, "deps")
		reconcileAttr(o, r, "library")
		delete(oldRules, r.Name())
	}
	for _, r := range oldRules {
		if !RuleIsManaged(r) {
			continue
		}
		f.DelRules(r.Kind(), r.Name())
	}
	return writeFile(path, f, true)
}

func RuleIsManaged(r *bzl.Rule) bool {
	var automanaged bool
	for _, tag := range r.AttrStrings("tags") {
		if tag == "automanaged" {
			automanaged = true
			break
		}
	}
	return automanaged
}

func writeFile(path string, f *bzl.File, exists bool) (bool, error) {
	var info bzl.RewriteInfo
	bzl.Rewrite(f, &info)
	out := bzl.Format(f)
	if exists {
		orig, err := ioutil.ReadFile(path)
		if err != nil {
			return false, err
		}
		if bytes.Compare(out, orig) == 0 {
			return false, nil
		}
	}
	if *dryRun {
		return true, nil
	}
	return true, ioutil.WriteFile(path, out, 0644)

}
<commit_msg>better error output<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	bzl "github.com/bazelbuild/buildifier/core"
	"github.com/golang/glog"
)

var kubeRoot = os.Getenv("KUBE_ROOT")
var dryRun = flag.Bool("dry-run", false, "run in dry mode")

func main() {
	flag.Parse()
	flag.Set("alsologtostderr", "true")
	v := Venderor{
		ctx: &build.Default,
	}
	if len(flag.Args()) == 1 {
		v.updateSinglePkg(flag.Args()[0])
	} else {
		v.walkVendor()
		if err := v.walkRepo(); err != nil {
			glog.Fatalf("err: %v", err)
		}
	}
}

type Venderor struct {
	ctx *build.Context
}

func writeHeaders(file *bzl.File) {
	pkgRule := bzl.Rule{
		&bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: "package"},
		},
	}
	pkgRule.SetAttr("default_visibility", asExpr([]string{"//visibility:public"}))

	file.Stmt = append(file.Stmt,
		[]bzl.Expr{
			pkgRule.Call,
			&bzl.CallExpr{
				X:    &bzl.LiteralExpr{Token: "licenses"},
				List: []bzl.Expr{asExpr([]string{"notice"})},
			},
			&bzl.CallExpr{
				X: &bzl.LiteralExpr{Token: "load"},
				List: asExpr([]string{
					"@io_bazel_rules_go//go:def.bzl",
					"go_binary",
					"go_library",
					"go_test",
					"cgo_library",
				}).(*bzl.ListExpr).List,
			},
		}...,
	)
}

func writeRules(file *bzl.File, rules []*bzl.Rule) {
	for _, rule := range rules {
		file.Stmt = append(file.Stmt, rule.Call)
	}
}

func (v *Venderor) resolve(ipath string) Label {
	if strings.HasPrefix(ipath, "k8s.io/kubernetes") {
		return Label{
			pkg: strings.TrimPrefix(ipath, "k8s.io/kubernetes/"),
			tag: "go_default_library",
		}
	}
	return Label{
		pkg: "vendor",
		tag: ipath,
	}
}

func (v *Venderor) walk(root string, f func(path, ipath string, pkg *build.Package) error) error {
	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		ipath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if !info.IsDir() {
			return nil
		}
		pkg, err := v.ctx.ImportDir(filepath.Join(kubeRoot, path), build.ImportComment)
		if err != nil {
			if _, ok := err.(*build.NoGoError); err != nil && ok {
				return nil
			} else {
				return err
			}
		}

		return f(path, ipath, pkg)
	})
}

func (v *Venderor) walkRepo() error {
	for _, root := range []string{
		"./pkg",
		"./cmd",
		"./third_party",
		"./plugin",
		"./test",
		"./federation",
	} {
		if err := v.walk(root, v.updatePkg); err != nil {
			return err
		}
	}
	return nil
}

func (v *Venderor) updateSinglePkg(path string) error {
	pkg, err := v.ctx.ImportDir("./"+path, build.ImportComment)
	if err != nil {
		if _, ok := err.(*build.NoGoError); err != nil && ok {
			return nil
		} else {
			return err
		}
	}
	return v.updatePkg(path, "", pkg)
}

func (v *Venderor) updatePkg(path, _ string, pkg *build.Package) error {
	var rules []*bzl.Rule

	var attrs Attrs = make(Attrs)
	srcs := asExpr(merge(pkg.GoFiles, pkg.SFiles)).(*bzl.ListExpr)

	deps := v.extractDeps(pkg.Imports)

	if len(srcs.List) == 0 {
		return nil
	}
	attrs.Set("srcs", srcs)

	if len(deps.List) > 0 {
		attrs.Set("deps", deps)
	}

	if pkg.IsCommand() {
		rules = append(rules, newRule("go_binary", filepath.Base(pkg.Dir), attrs))
	} else {
		rules = append(rules, newRule("go_library", "go_default_library", attrs))
		if len(pkg.TestGoFiles) != 0 {
			rules = append(rules, newRule("go_test", "go_default_test", map[string]bzl.Expr{
				"srcs":    asExpr(pkg.TestGoFiles),
				"deps":    v.extractDeps(pkg.TestImports),
				"library": asExpr("go_default_library"),
			}))
		}
	}

	if len(pkg.XTestGoFiles) != 0 {
		rules = append(rules, newRule("go_test", "go_default_xtest", map[string]bzl.Expr{
			"srcs": asExpr(pkg.XTestGoFiles),
			"deps": v.extractDeps(pkg.XTestImports),
		}))
	}

	wrote, err := ReconcileRules(filepath.Join(path, "BUILD"), rules)
	if err != nil {
		return err
	}
	if wrote {
		fmt.Fprintf(os.Stderr, "wrote BUILD for %q\n", pkg.Dir)
	}
	return nil
}

func (v *Venderor) walkVendor() {
	var rules []*bzl.Rule
	if err := v.walk("./vendor", func(path, ipath string, pkg *build.Package) error {
		var attrs Attrs = make(Attrs)

		srcs := asExpr(
			apply(
				merge(pkg.GoFiles, pkg.SFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		cgoSrcs := asExpr(
			apply(
				merge(pkg.CgoFiles, pkg.CFiles, pkg.CXXFiles, pkg.HFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		deps := v.extractDeps(pkg.Imports)
		attrs.Set("srcs", srcs)

		if len(deps.List) > 0 {
			attrs.Set("deps", deps)
		}

		if pkg.IsCommand() {
			rules = append(rules, newRule("go_binary", v.resolve(ipath).tag, attrs))
		} else {
			if len(cgoSrcs.List) != 0 {
				cgoPname := v.resolve(ipath).tag + "_cgo"
				cgoDeps := v.extractDeps(pkg.TestImports)
				cgoRule := newRule("cgo_library", cgoPname, map[string]bzl.Expr{
					"srcs":      cgoSrcs,
					"clinkopts": asExpr([]string{"-ldl", "-lz", "-lm", "-lpthread", "-ldl"}),
				})
				rules = append(rules, cgoRule)
				if len(cgoDeps.List) != 0 {
					cgoRule.SetAttr("deps", cgoDeps)
				}
				attrs["library"] = asExpr(cgoPname)
			}
			rules = append(rules, newRule("go_library", v.resolve(ipath).tag, attrs))
		}
		return nil
	}); err != nil {
		glog.Fatalf("err: %v", err)
	}
	if _, err := ReconcileRules("./vendor/BUILD", rules); err != nil {
		glog.Fatalf("err: %v", err)
	}
}

func (v *Venderor) extractDeps(deps []string) *bzl.ListExpr {
	return asExpr(
		apply(
			merge(deps),
			filterer(func(s string) bool {
				pkg, err := v.ctx.Import(s, kubeRoot, build.ImportComment)
				if err != nil {
					return false
				}
				if pkg.Goroot {
					return false
				}
				return true
			}),
			mapper(func(s string) string {
				return v.resolve(s).String()
			}),
		),
	).(*bzl.ListExpr)
}

type Attrs map[string]bzl.Expr

func (a Attrs) Set(name string, expr bzl.Expr) {
	a[name] = expr
}

type Label struct {
	pkg, tag string
}

func (l Label) String() string {
	return fmt.Sprintf("//%v:%v", l.pkg, l.tag)
}

func asExpr(e interface{}) bzl.Expr {
	rv := reflect.ValueOf(e)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%d", e)}
	case reflect.Float32, reflect.Float64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%f", e)}
	case reflect.String:
		return &bzl.StringExpr{Value: e.(string)}
	case reflect.Slice, reflect.Array:
		var list []bzl.Expr
		for i := 0; i < rv.Len(); i++ {
			list = append(list, asExpr(rv.Index(i).Interface()))
		}
		return &bzl.ListExpr{List: list}
	default:
		glog.Fatalf("Uh oh")
		return nil
	}
}

type Sed func(s []string) []string

func mapString(in []string, f func(string) string) []string {
	var out []string
	for _, s := range in {
		out = append(out, f(s))
	}
	return out
}

func mapper(f func(string) string) Sed {
	return func(in []string) []string {
		return mapString(in, f)
	}
}

func filterString(in []string, f func(string) bool) []string {
	var out []string
	for _, s := range in {
		if f(s) {
			out = append(out, s)
		}
	}
	return out
}

func filterer(f func(string) bool) Sed {
	return func(in []string) []string {
		return filterString(in, f)
	}
}

func apply(stream []string, seds ...Sed) []string {
	for _, sed := range seds {
		stream = sed(stream)
	}
	return stream
}

func merge(streams ...[]string) []string {
	var out []string
	for _, stream := range streams {
		out = append(out, stream...)
	}
	return out
}

func newRule(kind, name string, attrs map[string]bzl.Expr) *bzl.Rule {
	rule := &bzl.Rule{
		Call: &bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: kind},
		},
	}
	rule.SetAttr("name", asExpr(name))
	for k, v := range attrs {
		rule.SetAttr(k, v)
	}
	rule.SetAttr("tags", asExpr([]string{"automanaged"}))
	return rule
}

func ReconcileRules(path string, rules []*bzl.Rule) (bool, error) {
	info, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		f := &bzl.File{}
		writeHeaders(f)
		writeRules(f, rules)
		return writeFile(path, f, false)
	} else if err != nil {
		return false, err
	}
	if info.IsDir() {
		return false, fmt.Errorf("%q cannot be a directory", path)
	}
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return false, err
	}
	f, err := bzl.Parse(path, b)
	if err != nil {
		return false, err
	}
	oldRules := make(map[string]*bzl.Rule)
	for _, r := range f.Rules("") {
		oldRules[r.Name()] = r
	}
	for _, r := range rules {
		o, ok := oldRules[r.Name()]
		if !ok {
			f.Stmt = append(f.Stmt, r.Call)
			continue
		}
		if !RuleIsManaged(o) {
			continue
		}
		reconcileAttr := func(o, n *bzl.Rule, name string) {
			if e := n.Attr(name); e != nil {
				o.SetAttr(name, e)
			} else {
				o.DelAttr(name)
			}
		}
		reconcileAttr(o, r, "srcs")
		reconcileAttr(o, r, "deps")
		reconcileAttr(o, r, "library")
		delete(oldRules, r.Name())
	}
	for _, r := range oldRules {
		if !RuleIsManaged(r) {
			continue
		}
		f.DelRules(r.Kind(), r.Name())
	}
	return writeFile(path, f, true)
}

func RuleIsManaged(r *bzl.Rule) bool {
	var automanaged bool
	for _, tag := range r.AttrStrings("tags") {
		if tag == "automanaged" {
			automanaged = true
			break
		}
	}
	return automanaged
}

func writeFile(path string, f *bzl.File, exists bool) (bool, error) {
	var info bzl.RewriteInfo
	bzl.Rewrite(f, &info)
	out := bzl.Format(f)
	if exists {
		orig, err := ioutil.ReadFile(path)
		if err != nil {
			return false, err
		}
		if bytes.Compare(out, orig) == 0 {
			return false, nil
		}
	}
	if *dryRun {
		return true, nil
	}
	return true, ioutil.WriteFile(path, out, 0644)

}
<|endoftext|>


<commit_before>package retriable

import (
	"github.com/ganmacs/retriable/backoff"
	"time"
)

const (
	defaultRetries = 3
)

type Operation func() error

type Options struct {
	operation Operation
	backoff   backoff.BackOff
	retries   int
}

func Retry(op Operation) error {
	opt := &Options{
		operation: op,
		retries:   defaultRetries,
		backoff:   backoff.NewExponentialBackOff(),
	}

	return doRetry(opt)
}

func doRetry(opt *Options) error {
	var bo = opt.backoff
	var retries = opt.retries

	var next time.Duration
	var err error

	for i := 0; i < retries; i++ {
		if err = opt.operation(); err == nil {
			return nil
		}

		// TODO fix -1
		if next = bo.Next(); next == -1 {
			return err
		}

		time.Sleep(next)
	}

	return err
}
<commit_msg>Extract backoff from options<commit_after>package retriable

import (
	"github.com/ganmacs/retriable/backoff"
	"time"
)

const (
	defaultRetries = 3
)

type Operation func() error

type Options struct {
	operation Operation
	retries   int
}

func Retry(op Operation) error {
	opt := &Options{
		operation: op,
		retries:   defaultRetries,
	}

	return doRetry(backoff.NewExponentialBackOff(), opt)
}

func doRetry(bo backoff.BackOff, opt *Options) error {
	var retries = opt.retries

	var next time.Duration
	var err error

	for i := 0; i < retries; i++ {
		if err = opt.operation(); err == nil {
			return nil
		}

		// TODO fix -1
		if next = bo.Next(); next == -1 {
			return err
		}

		time.Sleep(next)
	}

	return err
}
<|endoftext|>


<commit_before>package emath

import (
	"fmt"
)

func main() {
	fmt.Println("Hello github")
}
<commit_msg>Signed-off-by: Adam Wu <adam@ewings.cc><commit_after>package emath

import (
	"fmt"
)

func main() {
	fmt.Println("Hello github!!!")
}
<|endoftext|>


<commit_before>package balancer

import (
	"context"

	g "github.com/aukbit/pluto/client/grpc"
	"google.golang.org/grpc"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

type Connector interface {
	Client() interface{}
	Connector() *connector
	Health() bool
}

const (
	// DefaultName prefix connector name
	DefaultName    = "connector"
	defaultVersion = "v1.0.0"
)

// NewConnector returns a new connector with cfg passed in
func NewConnector(cfgs ...ConfigFn) *connector {
	return newConnector(cfgs...)
}

type ConnsCh chan *connector

// connector struct
type connector struct {
	cfg        *Config
	requestsCh chan Request          // requests channel to receive requests from balancer
	pending    int                   // count pending tasks
	index      int                   // index in the heap
	conn       *grpc.ClientConn      // grpc connection to communicate with the server
	client     interface{}           // grpc client stub to perform RPCs
	stopCh     chan bool             // receive a stop call
	doneCh     chan bool             // guarantees has beeen stopped correctly
	health     healthpb.HealthClient // Client API for Health service
	// logger     *zap.Logger
}

// newConnector ...
func newConnector(cfgs ...ConfigFn) *connector {
	c := newConfig(cfgs...)
	conn := &connector{
		cfg:        c,
		requestsCh: make(chan Request),
		stopCh:     make(chan bool),
		doneCh:     make(chan bool),
	}
	// conn.logger, _ = zap.NewProduction()

	conn.initLogger()
	return conn
}

// dial establish grpc client connection with the grpc server
func (c *connector) dial() error {
	// c.logger.Info("dial")
	// append logger
	c.cfg.UnaryClientInterceptors = append(c.cfg.UnaryClientInterceptors, loggerUnaryClientInterceptor(c))
	// dial
	// TODO use TLS
	conn, err := grpc.Dial(
		c.cfg.Target,
		grpc.WithInsecure(),
		grpc.WithUnaryInterceptor(g.WrapperUnaryClient(c.cfg.UnaryClientInterceptors...)))
	if err != nil {
		return err
	}
	// keep connection for later close
	c.conn = conn
	// register proto client to get a stub to perform RPCs
	c.client = c.cfg.GRPCRegister(conn)
	// register proto health client to get a stub to perform RPCs
	c.health = healthpb.NewHealthClient(conn)
	return nil
}

// watch waits for any call from balancer
func (c *connector) watch() {
	// c.logger.Info("watch")
	for {
		select {
		case req := <-c.requestsCh: // get request from balancer
			req.connsCh <- c
		case <-c.stopCh:
			close(c.doneCh)
			return
		}
	}
}

func (c *connector) Init() error {
	if err := c.dial(); err != nil {
		return err
	}
	go c.watch()
	return nil
}

func (c *connector) Client() interface{} {
	return c.client
}

func (c *connector) Connector() *connector {
	return c
}

// Close stops connector and close grpc connection
func (c *connector) Close() {
	// c.logger.Info("close")
	c.conn.Close()
	c.stopCh <- true
	<-c.doneCh
}

// Health check if a round trip with server is valid or not
func (c *connector) Health() bool {
	hcr, err := c.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{})
	if err != nil {
		return false
	}
	return hcr.Status.String() == "SERVING"
}

func (c *connector) initLogger() {
	// c.logger = c.logger.With(
	// zap.String("type", "connector"),
	// zap.String("id", c.cfg.ID),
	// zap.String("name", c.cfg.Name),
	// zap.String("target", c.cfg.Target),
	// zap.String("parent", c.cfg.ParentID))
}
<commit_msg>add logger to connector<commit_after>package balancer

import (
	"context"

	g "github.com/aukbit/pluto/client/grpc"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

type Connector interface {
	Client() interface{}
	Connector() *connector
	Health() bool
}

const (
	// DefaultName prefix connector name
	DefaultName    = "connector"
	defaultVersion = "v1.0.0"
)

// NewConnector returns a new connector with cfg passed in
func NewConnector(cfgs ...ConfigFn) *connector {
	return newConnector(cfgs...)
}

type ConnsCh chan *connector

// connector struct
type connector struct {
	cfg        *Config
	requestsCh chan Request          // requests channel to receive requests from balancer
	pending    int                   // count pending tasks
	index      int                   // index in the heap
	conn       *grpc.ClientConn      // grpc connection to communicate with the server
	client     interface{}           // grpc client stub to perform RPCs
	stopCh     chan bool             // receive a stop call
	doneCh     chan bool             // guarantees has beeen stopped correctly
	health     healthpb.HealthClient // Client API for Health service
	logger     *zap.Logger
}

// newConnector ...
func newConnector(cfgs ...ConfigFn) *connector {
	c := newConfig(cfgs...)
	conn := &connector{
		cfg:        c,
		requestsCh: make(chan Request),
		stopCh:     make(chan bool),
		doneCh:     make(chan bool),
	}
	conn.logger, _ = zap.NewProduction()
	conn.initLogger()
	return conn
}

// dial establish grpc client connection with the grpc server
func (c *connector) dial() error {
	c.logger.Info("dial")
	// append logger
	c.cfg.UnaryClientInterceptors = append(c.cfg.UnaryClientInterceptors, loggerUnaryClientInterceptor(c))
	// dial
	// TODO use TLS
	conn, err := grpc.Dial(
		c.cfg.Target,
		grpc.WithInsecure(),
		grpc.WithUnaryInterceptor(g.WrapperUnaryClient(c.cfg.UnaryClientInterceptors...)))
	if err != nil {
		return err
	}
	// keep connection for later close
	c.conn = conn
	// register proto client to get a stub to perform RPCs
	c.client = c.cfg.GRPCRegister(conn)
	// register proto health client to get a stub to perform RPCs
	c.health = healthpb.NewHealthClient(conn)
	return nil
}

// watch waits for any call from balancer
func (c *connector) watch() {
	c.logger.Info("watch")
	for {
		select {
		case req := <-c.requestsCh: // get request from balancer
			req.connsCh <- c
		case <-c.stopCh:
			close(c.doneCh)
			return
		}
	}
}

func (c *connector) Init() error {
	if err := c.dial(); err != nil {
		return err
	}
	go c.watch()
	return nil
}

func (c *connector) Client() interface{} {
	return c.client
}

func (c *connector) Connector() *connector {
	return c
}

// Close stops connector and close grpc connection
func (c *connector) Close() {
	c.logger.Info("close")
	c.conn.Close()
	c.stopCh <- true
	<-c.doneCh
}

// Health check if a round trip with server is valid or not
func (c *connector) Health() bool {
	hcr, err := c.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{})
	if err != nil {
		return false
	}
	return hcr.Status.String() == "SERVING"
}

func (c *connector) initLogger() {
	c.logger = c.logger.With(
		zap.String("type", "connector"),
		zap.String("id", c.cfg.ID),
		zap.String("name", c.cfg.Name),
		zap.String("target", c.cfg.Target),
		zap.String("parent", c.cfg.ParentID),
	)
}
<|endoftext|>


<commit_before>package aranGO

// Graph structure
type Graph struct {

}
<commit_msg>Manage graphs<commit_after>package aranGO

import(
  "errors"
)

// Graph structure
type Graph struct {
  Id        string              `json:"_id,omitempty"`
  Key       string              `json:"_key"`
  Name      string              `json:"name"`
  EdgesDef  []EdgeDefinition    `json:"edgeDefinitions"`
  Orphan    []string            `json:"orphanCollections"`
  db        *Database
}

// Remove vertex collections
func(g *Graph) RemoveVertex(col string) error{
  if g.db == nil{
    return errors.New("Invalid db")
  }
  if col == "" {
    return errors.New("Invalid collection")
  }

  res, err :=g.db.get("gharial",g.Key+"/vertex/"+col,"DELETE",nil,nil,nil)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 400:
      return errors.New("Cannot remove collection")
    default:
      return errors.New("Invalid graph")
  }
}

// Remove edge
func(g *Graph) RemoveEdge(col string) error{
  if g.db == nil{
    return errors.New("Invalid db")
  }
  if col == "" {
    return errors.New("Invalid edge")
  }

  res, err :=g.db.get("gharial",g.Key+"/edge/"+col,"DELETE",nil,nil,nil)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 400:
      return errors.New("Cannot remove edge")
    default:
      return errors.New("Invalid graph")
  }
}

// Add vertex collections
func (g *Graph) AddVertex(col string) error{
  if g.db == nil{
    return errors.New("Invalid db")
  }

  pay := map[string]string { "collection" : col}
  res, err := g.db.send("gharial",g.Name+"/edge","POST",pay,g,g)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 400:
      return errors.New("Unable to add vertex collection")
    default:
      return errors.New("Invalid graph")
  }
}

func (g *Graph) AddEdge(ed *EdgeDefinition) error{
  if g.db == nil{
    return errors.New("Invalid db")
  }

  res, err := g.db.send("gharial",g.Name+"/edge","POST",ed,g,g)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 400:
      return errors.New("Unable to add edge definition")
    default:
      return errors.New("Invalid graph")
  }
}

func (g *Graph) ReplaceEdge(name string,ed *EdgeDefinition) error{
  if g.db == nil{
    return errors.New("Invalid db")
  }

  res, err := g.db.send("gharial",g.Name+"/edge/"+name,"POST",ed,g,g)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 400:
      return errors.New("Unable to replace edge definition")
    default:
      return errors.New("Invalid graph")
  }
}

func (g *Graph) ListEdges() ([]string,error){
  if g.db == nil{
    return []string{},errors.New("Invalid db")
  }
  var gr graphResponse
  res, err := g.db.get("gharial",g.Name+"/edge","GET",nil,&gr,&gr)
  if err != nil {
    return []string{},err
  }

  switch res.Status(){
    case 200:
      return gr.Col,nil
    default:
      return []string{},errors.New("Invalid graph")
  }
}

func (g *Graph) ListVertex() ([]string,error){
  if g.db == nil{
    return []string{},errors.New("Invalid db")
  }
  var gr graphResponse
  res, err := g.db.get("gharial",g.Name+"/vertex","GET",nil,&gr,&gr)
  if err != nil {
    return []string{},err
  }

  switch res.Status(){
    case 200:
      return gr.Col,nil
    default:
      return []string{},errors.New("Invalid graph")
  }
}

func (g *Graph) AddEdgeDefinition(ed EdgeDefinition) error {
  if ed.Collection == "" {
    return errors.New("Invalid collection")
  }
  g.EdgesDef = append(g.EdgesDef,ed)
  return nil
}

type graphResponse struct {
  Error   bool    `json:"error"`
  Code    int     `json:"code"`
  Graph   Graph   `json:"graph"`
  Graphs  []Graph `json:"graphs"`
  Col     []string `json:"collections"`
}

type EdgeDefinition struct {
  Collection  string    `json:"collection"`
  From        []string  `json:"from"`
  To          []string  `json:"to"`
}

func NewEdgeDefinition(col string,from []string,to []string ) *EdgeDefinition{
  var e EdgeDefinition
  if col == "" {
    return nil
  }
  e.Collection = col
  e.From = from
  e.To = to
  return &e
}

// Creates graphs
func (db *Database) CreateGraph(name string,eds []EdgeDefinition) (*Graph,error) {
  var g Graph
  var gr graphResponse
  if name != "" {
    g.Name = name
    g.EdgesDef = eds
  }else{
    return nil,errors.New("Invalid graph name")
  }
  if eds == nil {
    return nil,errors.New("Invalid edges")
  }
  res,err := db.send("gharial","","POST",g,&gr,&gr)
  if err != nil {
    return nil,err
  }

  switch res.Status() {
    case 201:
      return &g,nil
    case 409:
      return nil,errors.New("Conflic creating graph")
    default:
      return nil,errors.New("Conflic creating graph")
  }
}

func (db *Database) DropGraph(name string) error{

  if name == ""{
    return errors.New("Invalid graph name")
  }
  res, err := db.get("gharial",name,"DELETE",nil,nil,nil)
  if err != nil {
    return err
  }

  switch res.Status() {
    case 200:
      return nil
    case 404:
      return errors.New("Graph not found")
    default:
      return nil
  }

}

func (db *Database) Graph(name string) (*Graph){
  var g graphResponse
  if name == ""{
    return nil
  }
  _, err := db.get("gharial",name,"GET",nil,&g,&g)
  if err != nil {
    return nil
  }
  // set DB
  g.Graph.db = db
  return &g.Graph
}

func (db *Database) ListGraphs() ([]Graph,error){
  var gr graphResponse

  res, err := db.get("gharial","","GET",nil,&gr,&gr)
  if err != nil {
    return nil,err
  }

  switch res.Status(){
    case 200:
      return gr.Graphs,nil
    default:
      return nil,errors.New("Unable to list graphs")
  }
}
<|endoftext|>