<commit_before>package dht

import (
	"testing"

	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	ci "github.com/jbenet/go-ipfs/crypto"
	identify "github.com/jbenet/go-ipfs/identify"
	peer "github.com/jbenet/go-ipfs/peer"
	swarm "github.com/jbenet/go-ipfs/swarm"
	u "github.com/jbenet/go-ipfs/util"

	"fmt"
	"time"
)

func setupDHTS(n int, t *testing.T) ([]*ma.Multiaddr, []*peer.Peer, []*IpfsDHT) {
	var addrs []*ma.Multiaddr
	for i := 0; i < 4; i++ {
		a, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", 5000+i))
		if err != nil {
			t.Fatal(err)
		}
		addrs = append(addrs, a)
	}

	var peers []*peer.Peer
	for i := 0; i < 4; i++ {
		p := new(peer.Peer)
		p.AddAddress(addrs[i])
		sk, pk, err := ci.GenerateKeyPair(ci.RSA, 512)
		if err != nil {
			panic(err)
		}
		p.PubKey = pk
		p.PrivKey = sk
		id, err := identify.IDFromPubKey(pk)
		if err != nil {
			panic(err)
		}
		p.ID = id
		peers = append(peers, p)
	}

	var dhts []*IpfsDHT
	for i := 0; i < 4; i++ {
		net := swarm.NewSwarm(peers[i])
		err := net.Listen()
		if err != nil {
			t.Fatal(err)
		}
		d := NewDHT(peers[i], net, ds.NewMapDatastore())
		dhts = append(dhts, d)
		d.Start()
	}

	return addrs, peers, dhts
}

func makePeer(addr *ma.Multiaddr) *peer.Peer {
	p := new(peer.Peer)
	p.AddAddress(addr)
	sk, pk, err := ci.GenerateKeyPair(ci.RSA, 512)
	if err != nil {
		panic(err)
	}
	p.PrivKey = sk
	p.PubKey = pk
	id, err := identify.IDFromPubKey(pk)
	if err != nil {
		panic(err)
	}

	p.ID = id
	return p
}

func TestPing(t *testing.T) {
	u.Debug = true
	addrA, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/2222")
	if err != nil {
		t.Fatal(err)
	}
	addrB, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5678")
	if err != nil {
		t.Fatal(err)
	}

	peerA := makePeer(addrA)
	peerB := makePeer(addrB)

	neta := swarm.NewSwarm(peerA)
	err = neta.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtA := NewDHT(peerA, neta, ds.NewMapDatastore())

	netb := swarm.NewSwarm(peerB)
	err = netb.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtB := NewDHT(peerB, netb, ds.NewMapDatastore())

	dhtA.Start()
	dhtB.Start()

	_, err = dhtA.Connect(addrB)
	if err != nil {
		t.Fatal(err)
	}

	//Test that we can ping the node
	err = dhtA.Ping(peerB, time.Second*2)
	if err != nil {
		t.Fatal(err)
	}

	dhtA.Halt()
	dhtB.Halt()
}

func TestValueGetSet(t *testing.T) {
	u.Debug = false
	addrA, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/1235")
	if err != nil {
		t.Fatal(err)
	}
	addrB, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5679")
	if err != nil {
		t.Fatal(err)
	}

	peerA := makePeer(addrA)
	peerB := makePeer(addrB)

	neta := swarm.NewSwarm(peerA)
	err = neta.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtA := NewDHT(peerA, neta, ds.NewMapDatastore())

	netb := swarm.NewSwarm(peerB)
	err = netb.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtB := NewDHT(peerB, netb, ds.NewMapDatastore())

	dhtA.Start()
	dhtB.Start()

	errsa := dhtA.network.GetErrChan()
	errsb := dhtB.network.GetErrChan()
	go func() {
		select {
		case err := <-errsa:
			t.Fatal(err)
		case err := <-errsb:
			t.Fatal(err)
		}
	}()

	_, err = dhtA.Connect(addrB)
	if err != nil {
		t.Fatal(err)
	}

	dhtA.PutValue("hello", []byte("world"))

	val, err := dhtA.GetValue("hello", time.Second*2)
	if err != nil {
		t.Fatal(err)
	}

	if string(val) != "world" {
		t.Fatalf("Expected 'world' got '%s'", string(val))
	}

}

func TestProvides(t *testing.T) {
	u.Debug = false

	addrs, _, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].putLocal(u.Key("hello"), []byte("world"))
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[3].getLocal(u.Key("hello"))
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].Provide(u.Key("hello"))
	if err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 60)

	provs, err := dhts[0].FindProviders(u.Key("hello"), time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if len(provs) != 1 {
		t.Fatal("Didnt get back providers")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}

func TestLayeredGet(t *testing.T) {
	u.Debug = false
	addrs, _, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatalf("Failed to connect: %s", err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].putLocal(u.Key("hello"), []byte("world"))
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].Provide(u.Key("hello"))
	if err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 60)

	val, err := dhts[0].GetValue(u.Key("hello"), time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if string(val) != "world" {
		t.Fatal("Got incorrect value.")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}

func TestFindPeer(t *testing.T) {
	u.Debug = false

	addrs, peers, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	p, err := dhts[0].FindPeer(peers[2].ID, time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if p == nil {
		t.Fatal("Failed to find peer.")
	}

	if !p.ID.Equal(peers[2].ID) {
		t.Fatal("Didnt find expected peer.")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}
<commit_msg>checking returned value<commit_after>package dht

import (
	"testing"

	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	ci "github.com/jbenet/go-ipfs/crypto"
	identify "github.com/jbenet/go-ipfs/identify"
	peer "github.com/jbenet/go-ipfs/peer"
	swarm "github.com/jbenet/go-ipfs/swarm"
	u "github.com/jbenet/go-ipfs/util"

	"bytes"
	"fmt"
	"time"
)

func setupDHTS(n int, t *testing.T) ([]*ma.Multiaddr, []*peer.Peer, []*IpfsDHT) {
	var addrs []*ma.Multiaddr
	for i := 0; i < 4; i++ {
		a, err := ma.NewMultiaddr(fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", 5000+i))
		if err != nil {
			t.Fatal(err)
		}
		addrs = append(addrs, a)
	}

	var peers []*peer.Peer
	for i := 0; i < 4; i++ {
		p := new(peer.Peer)
		p.AddAddress(addrs[i])
		sk, pk, err := ci.GenerateKeyPair(ci.RSA, 512)
		if err != nil {
			panic(err)
		}
		p.PubKey = pk
		p.PrivKey = sk
		id, err := identify.IDFromPubKey(pk)
		if err != nil {
			panic(err)
		}
		p.ID = id
		peers = append(peers, p)
	}

	var dhts []*IpfsDHT
	for i := 0; i < 4; i++ {
		net := swarm.NewSwarm(peers[i])
		err := net.Listen()
		if err != nil {
			t.Fatal(err)
		}
		d := NewDHT(peers[i], net, ds.NewMapDatastore())
		dhts = append(dhts, d)
		d.Start()
	}

	return addrs, peers, dhts
}

func makePeer(addr *ma.Multiaddr) *peer.Peer {
	p := new(peer.Peer)
	p.AddAddress(addr)
	sk, pk, err := ci.GenerateKeyPair(ci.RSA, 512)
	if err != nil {
		panic(err)
	}
	p.PrivKey = sk
	p.PubKey = pk
	id, err := identify.IDFromPubKey(pk)
	if err != nil {
		panic(err)
	}

	p.ID = id
	return p
}

func TestPing(t *testing.T) {
	u.Debug = true
	addrA, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/2222")
	if err != nil {
		t.Fatal(err)
	}
	addrB, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5678")
	if err != nil {
		t.Fatal(err)
	}

	peerA := makePeer(addrA)
	peerB := makePeer(addrB)

	neta := swarm.NewSwarm(peerA)
	err = neta.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtA := NewDHT(peerA, neta, ds.NewMapDatastore())

	netb := swarm.NewSwarm(peerB)
	err = netb.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtB := NewDHT(peerB, netb, ds.NewMapDatastore())

	dhtA.Start()
	dhtB.Start()

	_, err = dhtA.Connect(addrB)
	if err != nil {
		t.Fatal(err)
	}

	//Test that we can ping the node
	err = dhtA.Ping(peerB, time.Second*2)
	if err != nil {
		t.Fatal(err)
	}

	dhtA.Halt()
	dhtB.Halt()
}

func TestValueGetSet(t *testing.T) {
	u.Debug = false
	addrA, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/1235")
	if err != nil {
		t.Fatal(err)
	}
	addrB, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/5679")
	if err != nil {
		t.Fatal(err)
	}

	peerA := makePeer(addrA)
	peerB := makePeer(addrB)

	neta := swarm.NewSwarm(peerA)
	err = neta.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtA := NewDHT(peerA, neta, ds.NewMapDatastore())

	netb := swarm.NewSwarm(peerB)
	err = netb.Listen()
	if err != nil {
		t.Fatal(err)
	}
	dhtB := NewDHT(peerB, netb, ds.NewMapDatastore())

	dhtA.Start()
	dhtB.Start()

	errsa := dhtA.network.GetErrChan()
	errsb := dhtB.network.GetErrChan()
	go func() {
		select {
		case err := <-errsa:
			t.Fatal(err)
		case err := <-errsb:
			t.Fatal(err)
		}
	}()

	_, err = dhtA.Connect(addrB)
	if err != nil {
		t.Fatal(err)
	}

	dhtA.PutValue("hello", []byte("world"))

	val, err := dhtA.GetValue("hello", time.Second*2)
	if err != nil {
		t.Fatal(err)
	}

	if string(val) != "world" {
		t.Fatalf("Expected 'world' got '%s'", string(val))
	}

}

func TestProvides(t *testing.T) {
	u.Debug = false

	addrs, _, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].putLocal(u.Key("hello"), []byte("world"))
	if err != nil {
		t.Fatal(err)
	}

	bits, err := dhts[3].getLocal(u.Key("hello"))
	if err != nil && bytes.Equal(bits, []byte("world")) {
		t.Fatal(err)
	}

	err = dhts[3].Provide(u.Key("hello"))
	if err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 60)

	provs, err := dhts[0].FindProviders(u.Key("hello"), time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if len(provs) != 1 {
		t.Fatal("Didnt get back providers")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}

func TestLayeredGet(t *testing.T) {
	u.Debug = false
	addrs, _, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatalf("Failed to connect: %s", err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].putLocal(u.Key("hello"), []byte("world"))
	if err != nil {
		t.Fatal(err)
	}

	err = dhts[3].Provide(u.Key("hello"))
	if err != nil {
		t.Fatal(err)
	}

	time.Sleep(time.Millisecond * 60)

	val, err := dhts[0].GetValue(u.Key("hello"), time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if string(val) != "world" {
		t.Fatal("Got incorrect value.")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}

func TestFindPeer(t *testing.T) {
	u.Debug = false

	addrs, peers, dhts := setupDHTS(4, t)

	_, err := dhts[0].Connect(addrs[1])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[2])
	if err != nil {
		t.Fatal(err)
	}

	_, err = dhts[1].Connect(addrs[3])
	if err != nil {
		t.Fatal(err)
	}

	p, err := dhts[0].FindPeer(peers[2].ID, time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if p == nil {
		t.Fatal("Failed to find peer.")
	}

	if !p.ID.Equal(peers[2].ID) {
		t.Fatal("Didnt find expected peer.")
	}

	for i := 0; i < 4; i++ {
		dhts[i].Halt()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/vektra/mockery/mockery"
)

var fName = flag.String("name", "", "name of interface to generate mock for")
var fPrint = flag.Bool("print", false, "print the generated mock to stdout")
var fOutput = flag.String("output", "./mocks", "directory to write mocks to")
var fDir = flag.String("dir", ".", "directory to search for interfaces")
var fAll = flag.Bool("all", false, "generates mocks for all found interfaces")

func checkDir(p *mockery.Parser, dir, name string) bool {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return false
	}

	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		path := filepath.Join(dir, file.Name())

		if file.IsDir() {
			ret := checkDir(p, path, name)
			if ret {
				return true
			}
		}

		if !strings.HasSuffix(path, ".go") {
			continue
		}

		err = p.Parse(path)
		if err != nil {
			continue
		}

		node, err := p.Find(name)
		if err != nil {
			continue
		}

		if node != nil {
			return true
		}
	}

	return false
}

func main() {
	flag.Parse()

	if *fAll {
		mockAll()
	} else {
		mockFor(*fName)
	}
}

func walkDir(dir string) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return
	}

	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		path := filepath.Join(dir, file.Name())

		if file.IsDir() {
			walkDir(path)
		}

		if !strings.HasSuffix(path, ".go") {
			continue
		}

		p := mockery.NewParser()

		err = p.Parse(path)
		if err != nil {
			continue
		}

		for _, iface := range p.Interfaces() {
			genMock(iface)
		}
	}

	return
}

func mockAll() {
	walkDir(*fDir)
}

func mockFor(name string) {
	if name == "" {
		fmt.Fprintf(os.Stderr, "Use -name to specify the name of the interface")
		os.Exit(1)
	}

	parser := mockery.NewParser()

	ret := checkDir(parser, *fDir, name)
	if !ret {
		fmt.Printf("Unable to find %s in any go files under this path\n", name)
		os.Exit(1)
	}

	iface, err := parser.Find(name)
	if err != nil {
		fmt.Printf("Error finding %s: %s\n", name, err)
		os.Exit(1)
	}

	genMock(iface)
}

func genMock(iface *mockery.Interface) {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("Unable to generated mock for '%s': %s\n", iface.Name, r)
			return
		}
	}()

	var out io.Writer

	name := iface.Name

	if *fPrint {
		out = os.Stdout
	} else {
		path := filepath.Join(*fOutput, name+".go")
		os.MkdirAll(filepath.Dir(path), 0755)

		f, err := os.Create(path)
		if err != nil {
			fmt.Printf("Unable to create output file for generated mock: %s\n", err)
			os.Exit(1)
		}

		defer f.Close()

		out = f

		fmt.Printf("Generating mock for: %s\n", name)
	}

	gen := mockery.NewGenerator(iface)

	gen.GeneratePrologue()
	err := gen.Generate()
	if err != nil {
		fmt.Printf("Error with %s: %s\n", name, err)
		os.Exit(1)
	}

	gen.Write(out)
}
<commit_msg>Short curcuit walkDir on a dir properly<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/vektra/mockery/mockery"
)

var fName = flag.String("name", "", "name of interface to generate mock for")
var fPrint = flag.Bool("print", false, "print the generated mock to stdout")
var fOutput = flag.String("output", "./mocks", "directory to write mocks to")
var fDir = flag.String("dir", ".", "directory to search for interfaces")
var fAll = flag.Bool("all", false, "generates mocks for all found interfaces")

func checkDir(p *mockery.Parser, dir, name string) bool {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return false
	}

	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		path := filepath.Join(dir, file.Name())

		if file.IsDir() {
			ret := checkDir(p, path, name)
			if ret {
				return true
			}
		}

		if !strings.HasSuffix(path, ".go") {
			continue
		}

		err = p.Parse(path)
		if err != nil {
			continue
		}

		node, err := p.Find(name)
		if err != nil {
			continue
		}

		if node != nil {
			return true
		}
	}

	return false
}

func main() {
	flag.Parse()

	if *fAll {
		mockAll()
	} else {
		mockFor(*fName)
	}
}

func walkDir(dir string) {
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return
	}

	for _, file := range files {
		if strings.HasPrefix(file.Name(), ".") {
			continue
		}

		path := filepath.Join(dir, file.Name())

		if file.IsDir() {
			walkDir(path)
			continue
		}

		if !strings.HasSuffix(path, ".go") {
			continue
		}

		p := mockery.NewParser()

		err = p.Parse(path)
		if err != nil {
			continue
		}

		for _, iface := range p.Interfaces() {
			genMock(iface)
		}
	}

	return
}

func mockAll() {
	walkDir(*fDir)
}

func mockFor(name string) {
	if name == "" {
		fmt.Fprintf(os.Stderr, "Use -name to specify the name of the interface")
		os.Exit(1)
	}

	parser := mockery.NewParser()

	ret := checkDir(parser, *fDir, name)
	if !ret {
		fmt.Printf("Unable to find %s in any go files under this path\n", name)
		os.Exit(1)
	}

	iface, err := parser.Find(name)
	if err != nil {
		fmt.Printf("Error finding %s: %s\n", name, err)
		os.Exit(1)
	}

	genMock(iface)
}

func genMock(iface *mockery.Interface) {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("Unable to generated mock for '%s': %s\n", iface.Name, r)
			return
		}
	}()

	var out io.Writer

	name := iface.Name

	if *fPrint {
		out = os.Stdout
	} else {
		path := filepath.Join(*fOutput, name+".go")
		os.MkdirAll(filepath.Dir(path), 0755)

		f, err := os.Create(path)
		if err != nil {
			fmt.Printf("Unable to create output file for generated mock: %s\n", err)
			os.Exit(1)
		}

		defer f.Close()

		out = f

		fmt.Printf("Generating mock for: %s\n", name)
	}

	gen := mockery.NewGenerator(iface)

	gen.GeneratePrologue()
	err := gen.Generate()
	if err != nil {
		fmt.Printf("Error with %s: %s\n", name, err)
		os.Exit(1)
	}

	gen.Write(out)
}
<|endoftext|>


<commit_before>package main

import "log"
import "strconv"
import "time"
import "golang.org/x/net/context"
import "google.golang.org/grpc"
import pb "github.com/brotherlogic/cardserver/card"

func Build() pb.CardList {
     // Generate 10 cards from the current time onwards
     now := time.Now().Truncate(time.Minute)
     cards := pb.CardList{}
     for i:=0; i < 10 ; i++ {
          card := pb.Card{}
	  card.Text = strconv.Itoa(now.Hour()) + ":" + strconv.Itoa(now.Minute())
	  card.ApplicationDate = now.Unix()

	  //Add a minute
	  now = now.Add(time.Minute)
	  card.ExpirationDate = now.Unix()

	  cards.Cards = append(cards.Cards, &card)
     }

     return cards
}

func main() {
     cards := Build()
     conn, err := grpc.Dial("localhost:50051", grpc.WithInsecure())

     defer conn.Close()
     client := pb.NewCardServiceClient(conn)
     _, err = client.AddCards(context.Background(), &cards)
     if err != nil {
     	log.Printf("Problem adding cards %v", err)
     }
}<commit_msg>Prepends the time to show as proper 24 hour clock. This closes #19.<commit_after>package main

import "log"
import "strconv"
import "time"
import "golang.org/x/net/context"
import "google.golang.org/grpc"
import pb "github.com/brotherlogic/cardserver/card"

func Prepend(str string) string {
	if len(str) == 1 {
		return "0" + str
	} else {
		return str
	}
}

func Build() pb.CardList {
	// Generate 10 cards from the current time onwards
	now := time.Now().Truncate(time.Minute)
	cards := pb.CardList{}
	for i := 0; i < 10; i++ {
		card := pb.Card{}
		card.Text = Prepend(strconv.Itoa(now.Hour())) + ":" + Prepend(strconv.Itoa(now.Minute()))
		card.ApplicationDate = now.Unix()

		//Add a minute
		now = now.Add(time.Minute)
		card.ExpirationDate = now.Unix()

		cards.Cards = append(cards.Cards, &card)
	}

	return cards
}

func main() {
	cards := Build()
	conn, err := grpc.Dial("localhost:50051", grpc.WithInsecure())

	defer conn.Close()
	client := pb.NewCardServiceClient(conn)
	_, err = client.AddCards(context.Background(), &cards)
	if err != nil {
		log.Printf("Problem adding cards %v", err)
	}
}
<|endoftext|>


<commit_before>package git_pipeline_test

import (
	"fmt"
	"net/http"
	"os"
	"os/exec"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/testflight/bosh"
	"github.com/concourse/testflight/gitserver"
	"github.com/concourse/testflight/guidserver"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"testing"
	"time"
)

// has ruby, curl
const guidServerRootfs = "/var/vcap/packages/bosh_deployment_resource"

// has git, curl
const gitServerRootfs = "/var/vcap/packages/git_resource"

var flyBin string

var (
	gardenClient garden.Client

	gitServer *gitserver.Server

	successGitServer       *gitserver.Server
	failureGitServer       *gitserver.Server
	noUpdateGitServer      *gitserver.Server
	ensureSuccessGitServer *gitserver.Server
	ensureFailureGitServer *gitserver.Server

	atcURL string
)

type DeploymentTemplateData struct {
	DirectorUUID       string
	GardenLinuxVersion string
}

var _ = BeforeSuite(func() {
	gardenLinuxVersion := os.Getenv("GARDEN_LINUX_VERSION")
	Ω(gardenLinuxVersion).ShouldNot(BeEmpty(), "must set $GARDEN_LINUX_VERSION")

	var err error

	flyBin, err = gexec.Build("github.com/concourse/fly", "-race")
	Ω(err).ShouldNot(HaveOccurred())

	directorUUID := bosh.DirectorUUID()

	bosh.DeleteDeployment("concourse-testflight")

	deploymentData := DeploymentTemplateData{
		DirectorUUID:       directorUUID,
		GardenLinuxVersion: gardenLinuxVersion,
	}

	bosh.Deploy("deployment.yml.tmpl", deploymentData)

	gardenClient = client.New(connection.New("tcp", "10.244.15.2:7777"))
	Eventually(gardenClient.Ping, 10*time.Second).ShouldNot(HaveOccurred())

	guidserver.Start(guidServerRootfs, gardenClient)

	gitServer = gitserver.Start(gitServerRootfs, gardenClient)
	successGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	failureGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	noUpdateGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	ensureSuccessGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	ensureFailureGitServer = gitserver.Start(gitServerRootfs, gardenClient)

	atcURL = "http://10.244.15.2:8080"

	Eventually(errorPolling(atcURL), 1*time.Minute).ShouldNot(HaveOccurred())

	configureCmd := exec.Command(
		flyBin,
		"-t", atcURL,
		"configure",
		"pipeline-name",
		"-c", "pipeline.yml",
		"-v", "failure-git-server="+failureGitServer.URI(),
		"-v", "guid-server-curl-command="+guidserver.CurlCommand(),
		"-v", "no-update-git-server="+noUpdateGitServer.URI(),
		"-v", "origin-git-server="+gitServer.URI(),
		"-v", "success-git-server="+successGitServer.URI(),
		"-v", "ensure-success-git-server="+ensureSuccessGitServer.URI(),
		"-v", "ensure-failure-git-server="+ensureFailureGitServer.URI(),
		"-v", "testflight-helper-image="+guidServerRootfs,
		"--paused=false",
	)

	stdin, err := configureCmd.StdinPipe()
	Ω(err).ShouldNot(HaveOccurred())

	defer stdin.Close()

	configure, err := gexec.Start(configureCmd, GinkgoWriter, GinkgoWriter)
	Ω(err).ShouldNot(HaveOccurred())

	Eventually(configure, 10).Should(gbytes.Say("apply configuration?"))

	fmt.Fprintln(stdin, "y")

	Eventually(configure, 10).Should(gexec.Exit(0))
})

var _ = AfterSuite(func() {
	gitServer.Stop()
	successGitServer.Stop()
	failureGitServer.Stop()
	noUpdateGitServer.Stop()
	ensureSuccessGitServer.Stop()
	ensureFailureGitServer.Stop()

	guidserver.Stop(gardenClient)
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Git Pipeline Suite")
}

func errorPolling(url string) func() error {
	return func() error {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
		}

		return err
	}
}
<commit_msg>don't clean up git/guid servers<commit_after>package git_pipeline_test

import (
	"fmt"
	"net/http"
	"os"
	"os/exec"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/testflight/bosh"
	"github.com/concourse/testflight/gitserver"
	"github.com/concourse/testflight/guidserver"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"testing"
	"time"
)

// has ruby, curl
const guidServerRootfs = "/var/vcap/packages/bosh_deployment_resource"

// has git, curl
const gitServerRootfs = "/var/vcap/packages/git_resource"

var flyBin string

var (
	gardenClient garden.Client

	gitServer *gitserver.Server

	successGitServer       *gitserver.Server
	failureGitServer       *gitserver.Server
	noUpdateGitServer      *gitserver.Server
	ensureSuccessGitServer *gitserver.Server
	ensureFailureGitServer *gitserver.Server

	atcURL string
)

type DeploymentTemplateData struct {
	DirectorUUID       string
	GardenLinuxVersion string
}

var _ = BeforeSuite(func() {
	gardenLinuxVersion := os.Getenv("GARDEN_LINUX_VERSION")
	Ω(gardenLinuxVersion).ShouldNot(BeEmpty(), "must set $GARDEN_LINUX_VERSION")

	var err error

	flyBin, err = gexec.Build("github.com/concourse/fly", "-race")
	Ω(err).ShouldNot(HaveOccurred())

	directorUUID := bosh.DirectorUUID()

	bosh.DeleteDeployment("concourse-testflight")

	deploymentData := DeploymentTemplateData{
		DirectorUUID:       directorUUID,
		GardenLinuxVersion: gardenLinuxVersion,
	}

	bosh.Deploy("deployment.yml.tmpl", deploymentData)

	gardenClient = client.New(connection.New("tcp", "10.244.15.2:7777"))
	Eventually(gardenClient.Ping, 10*time.Second).ShouldNot(HaveOccurred())

	guidserver.Start(guidServerRootfs, gardenClient)

	gitServer = gitserver.Start(gitServerRootfs, gardenClient)
	successGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	failureGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	noUpdateGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	ensureSuccessGitServer = gitserver.Start(gitServerRootfs, gardenClient)
	ensureFailureGitServer = gitserver.Start(gitServerRootfs, gardenClient)

	atcURL = "http://10.244.15.2:8080"

	Eventually(errorPolling(atcURL), 1*time.Minute).ShouldNot(HaveOccurred())

	configureCmd := exec.Command(
		flyBin,
		"-t", atcURL,
		"configure",
		"pipeline-name",
		"-c", "pipeline.yml",
		"-v", "failure-git-server="+failureGitServer.URI(),
		"-v", "guid-server-curl-command="+guidserver.CurlCommand(),
		"-v", "no-update-git-server="+noUpdateGitServer.URI(),
		"-v", "origin-git-server="+gitServer.URI(),
		"-v", "success-git-server="+successGitServer.URI(),
		"-v", "ensure-success-git-server="+ensureSuccessGitServer.URI(),
		"-v", "ensure-failure-git-server="+ensureFailureGitServer.URI(),
		"-v", "testflight-helper-image="+guidServerRootfs,
		"--paused=false",
	)

	stdin, err := configureCmd.StdinPipe()
	Ω(err).ShouldNot(HaveOccurred())

	defer stdin.Close()

	configure, err := gexec.Start(configureCmd, GinkgoWriter, GinkgoWriter)
	Ω(err).ShouldNot(HaveOccurred())

	Eventually(configure, 10).Should(gbytes.Say("apply configuration?"))

	fmt.Fprintln(stdin, "y")

	Eventually(configure, 10).Should(gexec.Exit(0))
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Git Pipeline Suite")
}

func errorPolling(url string) func() error {
	return func() error {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
		}

		return err
	}
}
<|endoftext|>


<commit_before>package neo4j

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"koding/tools/config"
	"koding/tools/logger"
	"koding/tools/statsd"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo/bson"
)

var (
	BASE_URL         string
	CYPHER_URL       string
	INDEX_NODE_PATH  = "/db/data/index/node/koding"
	UNIQUE_NODE_PATH = "/db/data/index/node/koding?unique"
	INDEX_PATH       = "/db/data/index/node"
	NODE_URL         = "/db/data/node"
	MAX_RETRIES      = 5
	TIMEOUT          = 20
	DEADLINE         = 40
	CYPHER_PATH      = "db/data/cypher"
)

func init() {
	statsd.SetAppName("neo4j")
}

var log = logger.New("neo4jfeeder")

type Relationship struct {
	Id         bson.ObjectId `bson:"_id,omitempty"`
	TargetId   bson.ObjectId `bson:"targetId,omitempty"`
	TargetName string        `bson:"targetName"`
	SourceId   bson.ObjectId `bson:"sourceId,omitempty"`
	SourceName string        `bson:"sourceName"`
	As         string        `bson:"as"`
	Timestamp  time.Time     `bson:"timestamp"`
	Data       bson.Binary
}

func SetupNeo4j(c *config.Config) {
	BASE_URL = c.Neo4j.Write + ":" + strconv.Itoa(c.Neo4j.Port)
	CYPHER_URL = fmt.Sprintf("%v/%v", BASE_URL, CYPHER_PATH)
}

func GetBaseURL() string {
	if BASE_URL == "" {
		log.Fatal("Base url is not set. Please call SetupNeo4j() before you use this pkg.")
	}

	return BASE_URL
}

// Setup the dial timeout
func dialTimeout(timeout time.Duration, deadline time.Duration) func(network, addr string) (c net.Conn, err error) {
	return func(netw, addr string) (net.Conn, error) {
		conn, err := net.DialTimeout(netw, addr, timeout)
		if err != nil {
			return nil, err
		}
		conn.SetDeadline(time.Now().Add(deadline))
		return conn, nil
	}
}

// Gets URL and string data to be sent and makes POST request
// reads response body and returns as string
func sendRequest(requestType, url, data string, attempt int) string {
	sTimer := statsd.StartTimer("sendRequest")

	// Set the timeout & deadline
	timeOut := time.Duration(TIMEOUT) * time.Second
	deadLine := time.Duration(DEADLINE) * time.Second

	transport := http.Transport{
		Dial: dialTimeout(timeOut, deadLine),
	}

	client := http.Client{
		Transport: &transport,
	}

	//convert string into bytestream
	dataByte := strings.NewReader(data)
	req, err := http.NewRequest(requestType, url, dataByte)

	// read response body
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	res, err := client.Do(req)
	if err != nil && attempt <= MAX_RETRIES {
		sTimer.Failed()
		log.Error("Request timed out: %v", err)
		attempt++
		sendRequest(requestType, url, data, attempt)
	}
	if err != nil && attempt > MAX_RETRIES {
		log.Error("req to %v timed out after %v retries", url, attempt)
	}

	body, _ := ioutil.ReadAll(res.Body)

	defer res.Body.Close()

	sTimer.Success()

	return string(body)
}

// connect source and target with relation property
// response will be object
func CreateRelationship(relation, source, target string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateRelationship")

	relationshipData := fmt.Sprintf(`{"to" : "%s", "type" : "%s" }`, target, relation)
	relRes := sendRequest("POST", fmt.Sprintf("%s", source), relationshipData, 1)

	relNode, err := jsonDecode(relRes)
	if err != nil {
		log.Error("Problem with relation response %v", relRes)
		sTimer.Failed()

		return relNode
	}

	sTimer.Success()

	return relNode
}

// connect source and target with relation property
// response will be object
func CreateRelationshipWithData(relation, source, target, data string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateRelationshipWithData")

	relationshipData := fmt.Sprintf(`{"to" : "%s", "type" : "%s", "data" : %s }`, target, relation, data)
	relRes := sendRequest("POST", fmt.Sprintf("%s", source), relationshipData, 1)

	relNode, err := jsonDecode(relRes)
	if err != nil {
		sTimer.Failed()
		log.Error("Problem with relation response %v", relRes)

		return relNode
	}

	sTimer.Success()

	return relNode
}

// creates a unique node with given id and node name
// response will be Object
func CreateUniqueNode(id string, name string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateUniqueNode")

	url := GetBaseURL() + UNIQUE_NODE_PATH

	postData := generatePostJsonData(id, name)

	response := sendRequest("POST", url, postData, 1)

	node, err := jsonDecode(response)
	if err != nil {
		log.Error("Problem with unique node creation response %v", response)
		sTimer.Failed()
	} else {
		sTimer.Success()
	}

	return node
}

// deletes a relation between two node using relationship info
func DeleteRelationship(sourceId, targetId, relationship string) bool {
	sTimer := statsd.StartTimer("DeleteRelationship")

	//get source node information
	sourceInfo := GetNode(sourceId)

	//get target node information
	targetInfo := GetNode(targetId)

	if len(sourceInfo) < 1 || len(targetInfo) < 1 {
		return false
	}

	if _, ok := sourceInfo[0]["self"]; !ok {
		return false
	}

	if _, ok := targetInfo[0]["self"]; !ok {
		return false
	}

	// create  url to get relationship information of source node
	relationshipsURL := fmt.Sprintf("%s", sourceInfo[0]["self"]) + "/relationships/all/" + relationship

	//this request returns objects in an array
	response := sendRequest("GET", relationshipsURL, "", 1)
	//so use json array decoder
	relationships, err := jsonArrayDecode(response)
	if err != nil {
		log.Error("Problem with unique node creation response %v", response)
		return false
	}

	if len(relationships) < 1 {
		return false
	}

	if _, ok := relationships[0]["self"]; !ok {
		return false
	}

	foundNode := false

	for _, relation := range relationships {
		if relation["end"] == targetInfo[0]["self"] {
			toBeDeletedRelationURL := fmt.Sprintf("%s", relation["self"])
			sendRequest("DELETE", toBeDeletedRelationURL, "", 1)
			foundNode = true

			break
		}
	}

	if !foundNode {
		sTimer.Failed()
		log.Error("not found! %v", relationships[0]["self"])
	} else {
		sTimer.Success()
	}

	return true
}

// gets node from neo4j with given unique node id
//response will be object
func GetNode(id string) []map[string]interface{} {

	url := GetBaseURL() + INDEX_NODE_PATH + "/id/" + id

	response := sendRequest("GET", url, "", 1)

	nodeData, err := jsonArrayDecode(response)
	if err != nil {
		log.Error("Problem with response %v", response)
	}

	return nodeData
}

// updates node with given data
// response will be object
func UpdateNode(id, propertiesJSON string) map[string]interface{} {

	node := GetNode(id)

	if len(node) < 1 {
		return nil
	}

	//if self is not there!
	if _, ok := node[0]["self"]; !ok {
		return nil
	}

	// create  url to get relationship information of source node
	propertiesURL := fmt.Sprintf("%s", node[0]["self"]) + "/properties"

	response := sendRequest("PUT", propertiesURL, propertiesJSON, 1)
	if response != "" {
		res, err := jsonDecode(response)
		if err != nil {
			log.Error("Problem with response %v, %v", err, res)
		}
	}

	return make(map[string]interface{})
}

func DeleteNode(id string) bool {
	sTimer := statsd.StartTimer("DeleteNode")

	node := GetNode(id)

	if len(node) < 1 {
		sTimer.Failed()
		return false
	}

	//if self is not there!
	selfUrl, ok := node[0]["self"]
	if !ok {
		sTimer.Failed()
		return false
	}

	splitStrings := strings.Split(selfUrl.(string), "/")
	nodeId := splitStrings[len(splitStrings)-1]

	query := fmt.Sprintf(`
    {"query" : "START n=node(%v) MATCH n-[r?]-items DELETE r, n"}
  `, nodeId)

	response := sendRequest("POST", CYPHER_URL, query, 1)

	var result map[string][]interface{}
	err := json.Unmarshal([]byte(response), &result)
	if err != nil {
		sTimer.Failed()
		log.Error("Deleting node Marshalling error: %v", err)
		return false
	}

	sTimer.Success()

	return true
}

// creates a unique tree head node to hold all nodes
// it is called once during runtime while initializing
func CreateUniqueIndex(name string) {
	//create unique index
	url := GetBaseURL() + INDEX_PATH

	bd := sendRequest("POST", url, `{"name":"`+name+`"}`, 1)

	log.Info("Created unique index for data: %v", bd)
}

// This is a custom json string generator as http request body to neo4j
func generatePostJsonData(id, name string) string {
	return fmt.Sprintf(`{ "key" : "id", "value" : "%s", "properties" : { "id" : "%s", "name" : "%s" } }`, id, id, name)
}

//here, mapping of decoded json
func jsonArrayDecode(data string) ([]map[string]interface{}, error) {
	sTimer := statsd.StartTimer("jsonArrayDecode")

	var source []map[string]interface{}

	err := json.Unmarshal([]byte(data), &source)
	if err != nil {
		sTimer.Failed()
		log.Error("Marshalling error: %v", err)
		return nil, err
	}

	sTimer.Success()

	return source, nil
}

//here, mapping of decoded json
func jsonDecode(data string) (map[string]interface{}, error) {
	var source map[string]interface{}

	err := json.Unmarshal([]byte(data), &source)
	if err != nil {
		log.Error("Marshalling error: %v", err)
		return nil, err
	}

	return source, nil
}

var NotAllowedNames = []string{
	"CStatusActivity",
	"CFolloweeBucketActivity",
	"CFollowerBucketActivity",
	"CCodeSnipActivity",
	"CDiscussionActivity",
	"CReplieeBucketActivity",
	"CReplierBucketActivity",
	"CBlogPostActivity",
	"CNewMemberBucketActivity",
	"CTutorialActivity",
	"CLikeeBucketActivity",
	"CLikerBucketActivity",
	"CInstalleeBucketActivity",
	"CInstallerBucketActivity",
	"CActivity",
	"CRunnableActivity",
	"JAppStorage",
	"JFeed",
}
<commit_msg>Neo4j: add more blacklisted items<commit_after>package neo4j

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"koding/tools/config"
	"koding/tools/logger"
	"koding/tools/statsd"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo/bson"
)

var (
	BASE_URL         string
	CYPHER_URL       string
	INDEX_NODE_PATH  = "/db/data/index/node/koding"
	UNIQUE_NODE_PATH = "/db/data/index/node/koding?unique"
	INDEX_PATH       = "/db/data/index/node"
	NODE_URL         = "/db/data/node"
	MAX_RETRIES      = 5
	TIMEOUT          = 20
	DEADLINE         = 40
	CYPHER_PATH      = "db/data/cypher"
)

func init() {
	statsd.SetAppName("neo4j")
}

var log = logger.New("neo4jfeeder")

type Relationship struct {
	Id         bson.ObjectId `bson:"_id,omitempty"`
	TargetId   bson.ObjectId `bson:"targetId,omitempty"`
	TargetName string        `bson:"targetName"`
	SourceId   bson.ObjectId `bson:"sourceId,omitempty"`
	SourceName string        `bson:"sourceName"`
	As         string        `bson:"as"`
	Timestamp  time.Time     `bson:"timestamp"`
	Data       bson.Binary
}

func SetupNeo4j(c *config.Config) {
	BASE_URL = c.Neo4j.Write + ":" + strconv.Itoa(c.Neo4j.Port)
	CYPHER_URL = fmt.Sprintf("%v/%v", BASE_URL, CYPHER_PATH)
}

func GetBaseURL() string {
	if BASE_URL == "" {
		log.Fatal("Base url is not set. Please call SetupNeo4j() before you use this pkg.")
	}

	return BASE_URL
}

// Setup the dial timeout
func dialTimeout(timeout time.Duration, deadline time.Duration) func(network, addr string) (c net.Conn, err error) {
	return func(netw, addr string) (net.Conn, error) {
		conn, err := net.DialTimeout(netw, addr, timeout)
		if err != nil {
			return nil, err
		}
		conn.SetDeadline(time.Now().Add(deadline))
		return conn, nil
	}
}

// Gets URL and string data to be sent and makes POST request
// reads response body and returns as string
func sendRequest(requestType, url, data string, attempt int) string {
	sTimer := statsd.StartTimer("sendRequest")

	// Set the timeout & deadline
	timeOut := time.Duration(TIMEOUT) * time.Second
	deadLine := time.Duration(DEADLINE) * time.Second

	transport := http.Transport{
		Dial: dialTimeout(timeOut, deadLine),
	}

	client := http.Client{
		Transport: &transport,
	}

	//convert string into bytestream
	dataByte := strings.NewReader(data)
	req, err := http.NewRequest(requestType, url, dataByte)

	// read response body
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	res, err := client.Do(req)
	if err != nil && attempt <= MAX_RETRIES {
		sTimer.Failed()
		log.Error("Request timed out: %v", err)
		attempt++
		sendRequest(requestType, url, data, attempt)
	}
	if err != nil && attempt > MAX_RETRIES {
		log.Error("req to %v timed out after %v retries", url, attempt)
	}

	body, _ := ioutil.ReadAll(res.Body)

	defer res.Body.Close()

	sTimer.Success()

	return string(body)
}

// connect source and target with relation property
// response will be object
func CreateRelationship(relation, source, target string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateRelationship")

	relationshipData := fmt.Sprintf(`{"to" : "%s", "type" : "%s" }`, target, relation)
	relRes := sendRequest("POST", fmt.Sprintf("%s", source), relationshipData, 1)

	relNode, err := jsonDecode(relRes)
	if err != nil {
		log.Error("Problem with relation response %v", relRes)
		sTimer.Failed()

		return relNode
	}

	sTimer.Success()

	return relNode
}

// connect source and target with relation property
// response will be object
func CreateRelationshipWithData(relation, source, target, data string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateRelationshipWithData")

	relationshipData := fmt.Sprintf(`{"to" : "%s", "type" : "%s", "data" : %s }`, target, relation, data)
	relRes := sendRequest("POST", fmt.Sprintf("%s", source), relationshipData, 1)

	relNode, err := jsonDecode(relRes)
	if err != nil {
		sTimer.Failed()
		log.Error("Problem with relation response %v", relRes)

		return relNode
	}

	sTimer.Success()

	return relNode
}

// creates a unique node with given id and node name
// response will be Object
func CreateUniqueNode(id string, name string) map[string]interface{} {
	sTimer := statsd.StartTimer("CreateUniqueNode")

	url := GetBaseURL() + UNIQUE_NODE_PATH

	postData := generatePostJsonData(id, name)

	response := sendRequest("POST", url, postData, 1)

	node, err := jsonDecode(response)
	if err != nil {
		log.Error("Problem with unique node creation response %v", response)
		sTimer.Failed()
	} else {
		sTimer.Success()
	}

	return node
}

// deletes a relation between two node using relationship info
func DeleteRelationship(sourceId, targetId, relationship string) bool {
	sTimer := statsd.StartTimer("DeleteRelationship")

	//get source node information
	sourceInfo := GetNode(sourceId)

	//get target node information
	targetInfo := GetNode(targetId)

	if len(sourceInfo) < 1 || len(targetInfo) < 1 {
		return false
	}

	if _, ok := sourceInfo[0]["self"]; !ok {
		return false
	}

	if _, ok := targetInfo[0]["self"]; !ok {
		return false
	}

	// create  url to get relationship information of source node
	relationshipsURL := fmt.Sprintf("%s", sourceInfo[0]["self"]) + "/relationships/all/" + relationship

	//this request returns objects in an array
	response := sendRequest("GET", relationshipsURL, "", 1)
	//so use json array decoder
	relationships, err := jsonArrayDecode(response)
	if err != nil {
		log.Error("Problem with unique node creation response %v", response)
		return false
	}

	if len(relationships) < 1 {
		return false
	}

	if _, ok := relationships[0]["self"]; !ok {
		return false
	}

	foundNode := false

	for _, relation := range relationships {
		if relation["end"] == targetInfo[0]["self"] {
			toBeDeletedRelationURL := fmt.Sprintf("%s", relation["self"])
			sendRequest("DELETE", toBeDeletedRelationURL, "", 1)
			foundNode = true

			break
		}
	}

	if !foundNode {
		sTimer.Failed()
		log.Error("not found! %v", relationships[0]["self"])
	} else {
		sTimer.Success()
	}

	return true
}

// gets node from neo4j with given unique node id
//response will be object
func GetNode(id string) []map[string]interface{} {

	url := GetBaseURL() + INDEX_NODE_PATH + "/id/" + id

	response := sendRequest("GET", url, "", 1)

	nodeData, err := jsonArrayDecode(response)
	if err != nil {
		log.Error("Problem with response %v", response)
	}

	return nodeData
}

// updates node with given data
// response will be object
func UpdateNode(id, propertiesJSON string) map[string]interface{} {

	node := GetNode(id)

	if len(node) < 1 {
		return nil
	}

	//if self is not there!
	if _, ok := node[0]["self"]; !ok {
		return nil
	}

	// create  url to get relationship information of source node
	propertiesURL := fmt.Sprintf("%s", node[0]["self"]) + "/properties"

	response := sendRequest("PUT", propertiesURL, propertiesJSON, 1)
	if response != "" {
		res, err := jsonDecode(response)
		if err != nil {
			log.Error("Problem with response %v, %v", err, res)
		}
	}

	return make(map[string]interface{})
}

func DeleteNode(id string) bool {
	sTimer := statsd.StartTimer("DeleteNode")

	node := GetNode(id)

	if len(node) < 1 {
		sTimer.Failed()
		return false
	}

	//if self is not there!
	selfUrl, ok := node[0]["self"]
	if !ok {
		sTimer.Failed()
		return false
	}

	splitStrings := strings.Split(selfUrl.(string), "/")
	nodeId := splitStrings[len(splitStrings)-1]

	query := fmt.Sprintf(`
    {"query" : "START n=node(%v) MATCH n-[r?]-items DELETE r, n"}
  `, nodeId)

	response := sendRequest("POST", CYPHER_URL, query, 1)

	var result map[string][]interface{}
	err := json.Unmarshal([]byte(response), &result)
	if err != nil {
		sTimer.Failed()
		log.Error("Deleting node Marshalling error: %v", err)
		return false
	}

	sTimer.Success()

	return true
}

// creates a unique tree head node to hold all nodes
// it is called once during runtime while initializing
func CreateUniqueIndex(name string) {
	//create unique index
	url := GetBaseURL() + INDEX_PATH

	bd := sendRequest("POST", url, `{"name":"`+name+`"}`, 1)

	log.Info("Created unique index for data: %v", bd)
}

// This is a custom json string generator as http request body to neo4j
func generatePostJsonData(id, name string) string {
	return fmt.Sprintf(`{ "key" : "id", "value" : "%s", "properties" : { "id" : "%s", "name" : "%s" } }`, id, id, name)
}

//here, mapping of decoded json
func jsonArrayDecode(data string) ([]map[string]interface{}, error) {
	sTimer := statsd.StartTimer("jsonArrayDecode")

	var source []map[string]interface{}

	err := json.Unmarshal([]byte(data), &source)
	if err != nil {
		sTimer.Failed()
		log.Error("Marshalling error: %v", err)
		return nil, err
	}

	sTimer.Success()

	return source, nil
}

//here, mapping of decoded json
func jsonDecode(data string) (map[string]interface{}, error) {
	var source map[string]interface{}

	err := json.Unmarshal([]byte(data), &source)
	if err != nil {
		log.Error("Marshalling error: %v", err)
		return nil, err
	}

	return source, nil
}

var NotAllowedNames = []string{
	"CStatusActivity",
	"CStatus",

	"CFolloweeBucketActivity",
	"CFolloweeBucket",

	"CFollowerBucketActivity",
	"CFollowerBucket",

	"GroupJoineeBucketActivity",
	"GroupJoineeBucket",

	"GroupJoinerBucketActivity",
	"GroupJoinerBucket",

	"CInstalleeBucketActivity",
	"CInstalleeBucket",

	"CInstallerBucketActivity",
	"CInstallerBucket",

	"CLikeeBucketActivity",
	"CLikeeBucket",

	"CLikerBucketActivity",
	"CLikerBucket",

	"CReplieeBucketActivity",
	"CReplieeBucket",

	"CReplierBucketActivity",
	"CReplierBucket",

	"CCodeSnipActivity",
	"CCodeSnip",

	"CDiscussionActivity",
	"CDiscussion",

	"CBlogPostActivity",
	"CBlogPost",

	"CNewMemberBucketActivity",
	"CNewMemberBucket",

	"CRunnableActivity",
	"CRunnable",

	"CTutorialActivity",
	"CTutorial",

	"CActivity",
	"JAppStorage",
	"JFeed",

	"JBlogPost",
	"JChatConversation",
	"JCodeShare",
	"JCodeSnip",
	"JConversationSlice",
	"JDiscussion",
	"JDomainStat",
	"JDomain",
	"JEmailConfirmation",
	"JEmailNotification",
	"JEnvironment",
	"JGroupBundle",
	"JGuest",
	"JInvitationRequest",
	"JInvitation",
	"JKodingKey",
	"JLimit",
	"JLocationStates",
	"JLocation",
	"JMailNotification",
	"JMails",
	"JMarkdownDoc",
	"JMembershipPolicy",
	"JMessage",
	"JName",
	"JOpinion",
	"JPasswordRecovery",
	"JPrivateMessage",
	"JReferrableEmail",
	"JReferral",
	"JStatusUpdate",
	"JStorage",
	"JVM",
	"JApp",
}
<|endoftext|>


<commit_before>package fsutils

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"time"
)

type FileEntry struct {
	Name     string      `json:"name"`
	FullPath string      `json:"fullPath"`
	IsDir    bool        `json:"isDir"`
	Size     int64       `json:"size"`
	Mode     os.FileMode `json:"mode"`
	Time     time.Time   `json:"time"`
	IsBroken bool        `json:"isBroken"`
	Readable bool        `json:"readable"`
	Writable bool        `json:"writable"`
}

func NewFileEntry(name string, fullPath string) *FileEntry {
	return &FileEntry{Name: name, FullPath: fullPath}
}

func ReadDirectory(p string) ([]FileEntry, error) {
	files, err := ioutil.ReadDir(p)
	if err != nil {
		return nil, err
	}

	ls := make([]FileEntry, len(files))
	for i, info := range files {
		ls[i] = makeFileEntry(path.Join(p, info.Name()), info)
	}

	return ls, nil
}

func Glob(glob string) ([]string, error) {
	files, err := filepath.Glob(glob)
	if err != nil {
		return nil, err
	}

	return files, nil
}

func ReadFile(path string) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return nil, err
	}

	if fi.Size() > 10*1024*1024 {
		return nil, fmt.Errorf("File larger than 10MiB.")
	}

	buf := make([]byte, fi.Size())
	if _, err := io.ReadFull(file, buf); err != nil {
		return nil, err
	}

	return buf, nil
}

func WriteFile(filename string, data []byte, DoNotOverwrite, Append bool) error {
	flags := os.O_RDWR | os.O_CREATE
	if DoNotOverwrite {
		flags |= os.O_EXCL
	}

	if !Append {
		flags |= os.O_TRUNC
	} else {
		flags |= os.O_APPEND
	}

	file, err := os.OpenFile(filename, flags, 0666)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write(data)
	if err != nil {
		return err
	}

	return nil
}

var suffixRegexp = regexp.MustCompile(`.((_\d+)?)(\.\w*)?$`)

func EnsureNonexistentPath(name string) (string, error) {
	index := 1
	for {
		_, err := os.Stat(name)
		if err != nil {
			if os.IsNotExist(err) {
				break
			}
			return "", err
		}

		loc := suffixRegexp.FindStringSubmatchIndex(name)
		name = name[:loc[2]] + "_" + strconv.Itoa(index) + name[loc[3]:]
		index++
	}

	return name, nil
}

func GetInfo(path string) (*FileEntry, error) {
	fi, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, errors.New("file does not exist")
		}
		return nil, err
	}

	fileEntry := makeFileEntry(path, fi)

	return &fileEntry, nil
}

func makeFileEntry(fullPath string, fi os.FileInfo) FileEntry {
	entry := FileEntry{
		Name:     fi.Name(),
		FullPath: fullPath,
		IsDir:    fi.IsDir(),
		Size:     fi.Size(),
		Mode:     fi.Mode(),
		Time:     fi.ModTime(),
	}

	if fi.Mode()&os.ModeSymlink != 0 {
		symlinkInfo, err := os.Stat(path.Dir(fullPath) + "/" + fi.Name())
		if err != nil {
			entry.IsBroken = true
			return entry
		}
		entry.IsDir = symlinkInfo.IsDir()
		entry.Size = symlinkInfo.Size()
		entry.Mode = symlinkInfo.Mode()
		entry.Time = symlinkInfo.ModTime()
	}

	return entry
}

func SetPermissions(name string, mode os.FileMode, recursive bool) error {
	var doChange func(name string) error

	doChange = func(name string) error {
		if err := os.Chmod(name, mode); err != nil {
			return err
		}

		if !recursive {
			return nil
		}

		fi, err := os.Stat(name)
		if err != nil {
			return err
		}

		if !fi.IsDir() {
			return nil
		}

		dir, err := os.Open(name)
		if err != nil {
			return err
		}
		defer dir.Close()

		entries, err := dir.Readdirnames(0)
		if err != nil {
			return err
		}
		var firstErr error
		for _, entry := range entries {
			err := doChange(name + "/" + entry)
			if err != nil && firstErr == nil {
				firstErr = err
			}
		}
		return firstErr
	}

	return doChange(name)
}

func Remove(path string) error {
	return os.Remove(path)
}

func Rename(oldname, newname string) error {
	return os.Rename(oldname, newname)
}

func CreateDirectory(name string, recursive bool) error {
	if recursive {
		return os.MkdirAll(name, 0755)
	}

	return os.Mkdir(name, 0755)
}
<commit_msg>kite: add isReadable and isWritable methods<commit_after>package fsutils

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"time"
)

type FileEntry struct {
	Name     string      `json:"name"`
	FullPath string      `json:"fullPath"`
	IsDir    bool        `json:"isDir"`
	Size     int64       `json:"size"`
	Mode     os.FileMode `json:"mode"`
	Time     time.Time   `json:"time"`
	IsBroken bool        `json:"isBroken"`
	Readable bool        `json:"readable"`
	Writable bool        `json:"writable"`
}

func NewFileEntry(name string, fullPath string) *FileEntry {
	return &FileEntry{Name: name, FullPath: fullPath}
}

func ReadDirectory(p string) ([]FileEntry, error) {
	files, err := ioutil.ReadDir(p)
	if err != nil {
		return nil, err
	}

	ls := make([]FileEntry, len(files))
	for i, info := range files {
		ls[i] = makeFileEntry(path.Join(p, info.Name()), info)
	}

	return ls, nil
}

func Glob(glob string) ([]string, error) {
	files, err := filepath.Glob(glob)
	if err != nil {
		return nil, err
	}

	return files, nil
}

func ReadFile(path string) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return nil, err
	}

	if fi.Size() > 10*1024*1024 {
		return nil, fmt.Errorf("File larger than 10MiB.")
	}

	buf := make([]byte, fi.Size())
	if _, err := io.ReadFull(file, buf); err != nil {
		return nil, err
	}

	return buf, nil
}

func WriteFile(filename string, data []byte, DoNotOverwrite, Append bool) error {
	flags := os.O_RDWR | os.O_CREATE
	if DoNotOverwrite {
		flags |= os.O_EXCL
	}

	if !Append {
		flags |= os.O_TRUNC
	} else {
		flags |= os.O_APPEND
	}

	file, err := os.OpenFile(filename, flags, 0666)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write(data)
	if err != nil {
		return err
	}

	return nil
}

var suffixRegexp = regexp.MustCompile(`.((_\d+)?)(\.\w*)?$`)

func EnsureNonexistentPath(name string) (string, error) {
	index := 1
	for {
		_, err := os.Stat(name)
		if err != nil {
			if os.IsNotExist(err) {
				break
			}
			return "", err
		}

		loc := suffixRegexp.FindStringSubmatchIndex(name)
		name = name[:loc[2]] + "_" + strconv.Itoa(index) + name[loc[3]:]
		index++
	}

	return name, nil
}

func GetInfo(path string) (*FileEntry, error) {
	fi, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, errors.New("file does not exist")
		}
		return nil, err
	}

	fileEntry := makeFileEntry(path, fi)

	return &fileEntry, nil
}

func makeFileEntry(fullPath string, fi os.FileInfo) FileEntry {
	entry := FileEntry{
		Name:     fi.Name(),
		FullPath: fullPath,
		IsDir:    fi.IsDir(),
		Size:     fi.Size(),
		Mode:     fi.Mode(),
		Time:     fi.ModTime(),
		Readable: isReadable(fi.Mode()),
		Writable: isReadable(fi.Mode()),
	}

	if fi.Mode()&os.ModeSymlink != 0 {
		symlinkInfo, err := os.Stat(path.Dir(fullPath) + "/" + fi.Name())
		if err != nil {
			entry.IsBroken = true
			return entry
		}
		entry.IsDir = symlinkInfo.IsDir()
		entry.Size = symlinkInfo.Size()
		entry.Mode = symlinkInfo.Mode()
		entry.Time = symlinkInfo.ModTime()
	}

	return entry
}

// check for owner permission
func isReadable(mode os.FileMode) bool { return mode&0400 != 0 }

// check for owner permission
func isWritable(mode os.FileMode) bool { return mode&0200 != 0 }

func SetPermissions(name string, mode os.FileMode, recursive bool) error {
	var doChange func(name string) error

	doChange = func(name string) error {
		if err := os.Chmod(name, mode); err != nil {
			return err
		}

		if !recursive {
			return nil
		}

		fi, err := os.Stat(name)
		if err != nil {
			return err
		}

		if !fi.IsDir() {
			return nil
		}

		dir, err := os.Open(name)
		if err != nil {
			return err
		}
		defer dir.Close()

		entries, err := dir.Readdirnames(0)
		if err != nil {
			return err
		}
		var firstErr error
		for _, entry := range entries {
			err := doChange(name + "/" + entry)
			if err != nil && firstErr == nil {
				firstErr = err
			}
		}
		return firstErr
	}

	return doChange(name)
}

func Remove(path string) error {
	return os.Remove(path)
}

func Rename(oldname, newname string) error {
	return os.Rename(oldname, newname)
}

func CreateDirectory(name string, recursive bool) error {
	if recursive {
		return os.MkdirAll(name, 0755)
	}

	return os.Mkdir(name, 0755)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package topotests

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"context"

	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/memorytopo"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// This file tests the CellInfo part of the topo.Server API.

func TestCellInfo(t *testing.T) {
	cell := "cell1"
	ctx := context.Background()
	ts := memorytopo.NewServer(cell)

	// Check GetCellInfo returns what memorytopo created.
	ci, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.Root != "" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	var cells []string
	cells, err = ts.ExpandCells(ctx, cell)
	require.NoError(t, err)
	require.EqualValues(t, []string{"cell1"}, cells)

	// Update the Server Address.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "new address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test update with no change.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "bad address"
		return topo.NewError(topo.NoUpdateNeeded, cell)
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test failing update.
	updateErr := fmt.Errorf("inside error")
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		return updateErr
	}); err != updateErr {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}

	// Test update on non-existing object.
	newCell := "new_cell"
	if err := ts.UpdateCellInfoFields(ctx, newCell, func(ci *topodatapb.CellInfo) error {
		ci.Root = "/"
		ci.ServerAddress = "good address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, newCell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "good address" || ci.Root != "/" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Add a record that should block CellInfo deletion for safety reasons.
	if err := ts.UpdateSrvKeyspace(ctx, cell, "keyspace", &topodatapb.SrvKeyspace{}); err != nil {
		t.Fatalf("UpdateSrvKeyspace failed: %v", err)
	}
	srvKeyspaces, err := ts.GetSrvKeyspaceNames(ctx, cell)
	if err != nil {
		t.Fatalf("GetSrvKeyspaceNames failed: %v", err)
	}
	if len(srvKeyspaces) == 0 {
		t.Fatalf("UpdateSrvKeyspace did not add SrvKeyspace.")
	}

	// Try to delete without force; it should fail.
	if err := ts.DeleteCellInfo(ctx, cell, false); err == nil {
		t.Fatalf("DeleteCellInfo should have failed without -force")
	}

	// Use the force.
	if err := ts.DeleteCellInfo(ctx, cell, true); err != nil {
		t.Fatalf("DeleteCellInfo failed even with -force: %v", err)
	}
	if _, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/); !topo.IsErrType(err, topo.NoNode) {
		t.Fatalf("GetCellInfo(non-existing cell) failed: %v", err)
	}
}

func TestExpandCells(t *testing.T) {
	ctx := context.Background()
	var cells []string
	var err error
	var allCells = "cell1,cell2,cell3"
	type testCase struct {
		name      string
		cellsIn   string
		cellsOut  []string
		errString string
	}

	testCases := []testCase{
		{"single", "cell1", []string{"cell1"}, ""},
		{"multiple", "cell1,cell2,cell3", []string{"cell1", "cell2", "cell3"}, ""},
		{"empty", "", []string{"cell1", "cell2", "cell3"}, ""},
		{"bad", "unknown", nil, "node doesn't exist"},
	}

	for _, tCase := range testCases {
		t.Run(tCase.name, func(t *testing.T) {
			cellsIn := tCase.cellsIn
			if cellsIn == "" {
				cellsIn = allCells
			}
			topoCells := strings.Split(cellsIn, ",")
			var ts *topo.Server
			if tCase.name == "bad" {
				ts = memorytopo.NewServer()
			} else {
				ts = memorytopo.NewServer(topoCells...)
			}
			cells, err = ts.ExpandCells(ctx, cellsIn)
			if tCase.errString != "" {
				require.Error(t, err)
				require.Contains(t, err.Error(), tCase.errString)
			} else {
				require.NoError(t, err)
			}
			require.EqualValues(t, tCase.cellsOut, cells)
		})
	}

	t.Run("aliases", func(t *testing.T) {
		cells := []string{"cell1", "cell2", "cell3"}
		ts := memorytopo.NewServer(cells...)
		err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: cells})
		require.NoError(t, err)

		tests := []struct {
			name      string
			in        string
			out       []string
			shouldErr bool
		}{
			{
				name: "alias only",
				in:   "alias",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "alias and cell in alias", // test deduping logic
				in:   "alias,cell1",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "just cells",
				in:   "cell1",
				out:  []string{"cell1"},
			},
			{
				name:      "missing alias",
				in:        "not_an_alias",
				shouldErr: true,
			},
		}

		for _, tt := range tests {
			tt := tt
			t.Run(tt.name, func(t *testing.T) {
				expanded, err := ts.ExpandCells(ctx, tt.in)
				if tt.shouldErr {
					assert.Error(t, err)
					return
				}

				require.NoError(t, err)
				assert.ElementsMatch(t, expanded, tt.out)
			})
		}
	})
}
<commit_msg>Add unit tests for DeleteCellInfo(force={true,false}) with unreachable local cell<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package topotests

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"context"

	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/memorytopo"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// This file tests the CellInfo part of the topo.Server API.

func TestCellInfo(t *testing.T) {
	cell := "cell1"
	ctx := context.Background()
	ts := memorytopo.NewServer(cell)

	// Check GetCellInfo returns what memorytopo created.
	ci, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.Root != "" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	var cells []string
	cells, err = ts.ExpandCells(ctx, cell)
	require.NoError(t, err)
	require.EqualValues(t, []string{"cell1"}, cells)

	// Update the Server Address.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "new address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test update with no change.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "bad address"
		return topo.NewError(topo.NoUpdateNeeded, cell)
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test failing update.
	updateErr := fmt.Errorf("inside error")
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		return updateErr
	}); err != updateErr {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}

	// Test update on non-existing object.
	newCell := "new_cell"
	if err := ts.UpdateCellInfoFields(ctx, newCell, func(ci *topodatapb.CellInfo) error {
		ci.Root = "/"
		ci.ServerAddress = "good address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, newCell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "good address" || ci.Root != "/" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Add a record that should block CellInfo deletion for safety reasons.
	if err := ts.UpdateSrvKeyspace(ctx, cell, "keyspace", &topodatapb.SrvKeyspace{}); err != nil {
		t.Fatalf("UpdateSrvKeyspace failed: %v", err)
	}
	srvKeyspaces, err := ts.GetSrvKeyspaceNames(ctx, cell)
	if err != nil {
		t.Fatalf("GetSrvKeyspaceNames failed: %v", err)
	}
	if len(srvKeyspaces) == 0 {
		t.Fatalf("UpdateSrvKeyspace did not add SrvKeyspace.")
	}

	// Try to delete without force; it should fail.
	if err := ts.DeleteCellInfo(ctx, cell, false); err == nil {
		t.Fatalf("DeleteCellInfo should have failed without -force")
	}

	// Use the force.
	if err := ts.DeleteCellInfo(ctx, cell, true); err != nil {
		t.Fatalf("DeleteCellInfo failed even with -force: %v", err)
	}
	if _, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/); !topo.IsErrType(err, topo.NoNode) {
		t.Fatalf("GetCellInfo(non-existing cell) failed: %v", err)
	}
}

func TestExpandCells(t *testing.T) {
	ctx := context.Background()
	var cells []string
	var err error
	var allCells = "cell1,cell2,cell3"
	type testCase struct {
		name      string
		cellsIn   string
		cellsOut  []string
		errString string
	}

	testCases := []testCase{
		{"single", "cell1", []string{"cell1"}, ""},
		{"multiple", "cell1,cell2,cell3", []string{"cell1", "cell2", "cell3"}, ""},
		{"empty", "", []string{"cell1", "cell2", "cell3"}, ""},
		{"bad", "unknown", nil, "node doesn't exist"},
	}

	for _, tCase := range testCases {
		t.Run(tCase.name, func(t *testing.T) {
			cellsIn := tCase.cellsIn
			if cellsIn == "" {
				cellsIn = allCells
			}
			topoCells := strings.Split(cellsIn, ",")
			var ts *topo.Server
			if tCase.name == "bad" {
				ts = memorytopo.NewServer()
			} else {
				ts = memorytopo.NewServer(topoCells...)
			}
			cells, err = ts.ExpandCells(ctx, cellsIn)
			if tCase.errString != "" {
				require.Error(t, err)
				require.Contains(t, err.Error(), tCase.errString)
			} else {
				require.NoError(t, err)
			}
			require.EqualValues(t, tCase.cellsOut, cells)
		})
	}

	t.Run("aliases", func(t *testing.T) {
		cells := []string{"cell1", "cell2", "cell3"}
		ts := memorytopo.NewServer(cells...)
		err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: cells})
		require.NoError(t, err)

		tests := []struct {
			name      string
			in        string
			out       []string
			shouldErr bool
		}{
			{
				name: "alias only",
				in:   "alias",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "alias and cell in alias", // test deduping logic
				in:   "alias,cell1",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "just cells",
				in:   "cell1",
				out:  []string{"cell1"},
			},
			{
				name:      "missing alias",
				in:        "not_an_alias",
				shouldErr: true,
			},
		}

		for _, tt := range tests {
			tt := tt
			t.Run(tt.name, func(t *testing.T) {
				expanded, err := ts.ExpandCells(ctx, tt.in)
				if tt.shouldErr {
					assert.Error(t, err)
					return
				}

				require.NoError(t, err)
				assert.ElementsMatch(t, expanded, tt.out)
			})
		}
	})
}

func TestDeleteCellInfo(t *testing.T) {
	ctx := context.Background()
	ts := memorytopo.NewServer("zone1", "unreachable")

	err := ts.UpdateCellInfoFields(ctx, "unreachable", func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = memorytopo.UnreachableServerAddr
		return nil
	})
	require.NoError(t, err, "failed to update cell to point at unreachable addr")

	tests := []struct {
		force       bool
		shouldErr   bool
		shouldExist bool
	}{
		{
			force:       false,
			shouldErr:   true,
			shouldExist: true,
		},
		{
			force:       true,
			shouldErr:   false,
			shouldExist: false,
		},
	}
	for _, tt := range tests {
		func() {
			ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
			defer cancel()

			err := ts.DeleteCellInfo(ctx, "unreachable", tt.force)
			if tt.shouldErr {
				assert.Error(t, err, "force=%t", tt.force)
			} else {
				assert.NoError(t, err, "force=%t", tt.force)
			}

			ci, err := ts.GetCellInfo(ctx, "unreachable", true /* strongRead */)
			if tt.shouldExist {
				assert.NoError(t, err)
				assert.NotNil(t, ci)
			} else {
				assert.True(t, topo.IsErrType(err, topo.NoNode), "expected cell %q to not exist", "unreachable")
			}
		}()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/getlantern/golog"
	"github.com/getlantern/tdb"
)

var (
	log = golog.LoggerFor("tdbdemo")
)

func main() {
	go func() {
		log.Error(http.ListenAndServe("localhost:4000", nil))
	}()

	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "tdbtest")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	log.Debugf("Writing data to %v", tmpDir)

	numReporters := 5000
	uniquesPerReporter := 100
	uniquesPerPeriod := 20
	reportingPeriods := 1000
	reportingInterval := time.Millisecond
	resolution := reportingInterval * 5
	retainPeriods := 20
	retentionPeriod := time.Duration(retainPeriods) * reportingInterval * 100
	targetPointsPerSecond := 20000
	numWriters := 4
	targetPointsPerSecondPerWriter := targetPointsPerSecond / numWriters
	targetDeltaFor1000Points := 1000 * time.Second / time.Duration(targetPointsPerSecondPerWriter)
	log.Debugf("Target delta for 1000 points: %v", targetDeltaFor1000Points)

	db, err := tdb.NewDB(&tdb.DBOpts{
		Dir:                  tmpDir,
		BatchSize:            100000,
		RocksDBStatsInterval: 60 * time.Second,
	})
	if err != nil {
		log.Fatal(err)
	}
	err = db.CreateTable("test", retentionPeriod, fmt.Sprintf(`
SELECT
	SUM(i) AS i,
	SUM(ii) AS ii,
	AVG(ii / i) AS iii
FROM inbound
GROUP BY period(%v)`, resolution))
	if err != nil {
		log.Fatal(err)
	}

	inserts := int64(0)
	start := time.Now()

	report := func() {
		delta := time.Now().Sub(start)
		start = time.Now()
		i := atomic.SwapInt64(&inserts, 0)
		var ms runtime.MemStats
		runtime.ReadMemStats(&ms)
		preGC := float64(ms.HeapAlloc) / 1024.0 / 1024.0
		runtime.GC()
		runtime.ReadMemStats(&ms)
		postGC := float64(ms.HeapAlloc) / 1024.0 / 1024.0
		fmt.Printf(`
%s inserts at %s inserts per second
%v
HeapAlloc pre/post GC %f/%f MiB
`,
			humanize.Comma(i), humanize.Comma(i/int64(delta.Seconds())),
			db.PrintTableStats("test"),
			preGC, postGC)
	}

	go func() {
		for {
			tk := time.NewTicker(30 * time.Second)
			for range tk.C {
				report()
			}
		}
	}()

	go func() {
		for {
			tk := time.NewTicker(30 * time.Second)
			for range tk.C {
				now := db.Now("test")
				q, err := db.SQLQuery(`
SELECT COUNT(i) AS the_count
FROM test
GROUP BY r, u, period(168h)
`)
				if err != nil {
					log.Errorf("Unable to build query: %v", err)
					continue
				}
				start := time.Now()
				result, err := q.Run()
				delta := time.Now().Sub(start)
				if err != nil {
					log.Errorf("Unable to run query: %v", err)
					continue
				}
				count := float64(0)
				if len(result.Entries) > 0 {
					count = result.Entries[0].Fields["the_count"][0].Get()
				}
				fmt.Printf("\nQuery at %v returned %v in %v\n", now, humanize.Comma(int64(count)), delta)
			}
		}
	}()

	var wg sync.WaitGroup
	wg.Add(numWriters)
	for _w := 0; _w < numWriters; _w++ {
		go func() {
			defer wg.Done()
			c := 0
			start := time.Now()
			for i := 0; i < reportingPeriods; i++ {
				ts := epoch.Add(time.Duration(i) * reportingInterval)
				for r := 0; r < numReporters/numWriters; r++ {
					for u := 0; u < uniquesPerPeriod; u++ {
						p := &tdb.Point{
							Ts: ts,
							Dims: map[string]interface{}{
								"r": rand.Intn(numReporters),
								"u": rand.Intn(uniquesPerReporter),
								"b": rand.Float64() > 0.99,
							},
							Vals: map[string]float64{
								"i":  float64(rand.Intn(100000)),
								"ii": float64(rand.Intn(100)),
							},
						}
						ierr := db.Insert("inbound", p)
						if ierr != nil {
							log.Errorf("Unable to insert: %v", err)
							return
						}
						atomic.AddInt64(&inserts, 1)
						c++

						// Control rate
						if c > 0 && c%1000 == 0 {
							delta := time.Now().Sub(start)
							if delta < targetDeltaFor1000Points {
								time.Sleep(targetDeltaFor1000Points - delta)
							}
							start = time.Now()
						}
					}
				}
				fmt.Print(".")
			}
		}()
	}

	wg.Wait()
	report()
}
<commit_msg>tdbdemo updates<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/getlantern/golog"
	"github.com/getlantern/tdb"
)

var (
	log = golog.LoggerFor("tdbdemo")
)

func main() {
	go func() {
		log.Error(http.ListenAndServe("localhost:4000", nil))
	}()

	epoch := time.Date(2015, time.January, 1, 0, 0, 0, 0, time.UTC)

	tmpDir, err := ioutil.TempDir("", "tdbtest")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)

	log.Debugf("Writing data to %v", tmpDir)

	numReporters := 5000
	uniquesPerReporter := 100
	uniquesPerPeriod := 20
	reportingPeriods := 1000
	reportingInterval := time.Millisecond
	resolution := reportingInterval * 5
	retainPeriods := 20
	retentionPeriod := time.Duration(retainPeriods) * reportingInterval * 100
	targetPointsPerSecond := 20000
	numWriters := 4
	targetPointsPerSecondPerWriter := targetPointsPerSecond / numWriters
	targetDeltaFor1000Points := 1000 * time.Second / time.Duration(targetPointsPerSecondPerWriter)
	log.Debugf("Target delta for 1000 points: %v", targetDeltaFor1000Points)

	db, err := tdb.NewDB(&tdb.DBOpts{
		Dir:                  tmpDir,
		BatchSize:            100000,
		RocksDBStatsInterval: 60 * time.Second,
	})
	if err != nil {
		log.Fatal(err)
	}
	err = db.CreateTable("test", retentionPeriod, fmt.Sprintf(`
SELECT
	SUM(i) AS i,
	SUM(ii) AS ii,
	AVG(ii / i) AS iii
FROM inbound
GROUP BY period(%v)`, resolution))
	if err != nil {
		log.Fatal(err)
	}

	inserts := int64(0)
	start := time.Now()

	report := func() {
		delta := time.Now().Sub(start)
		start = time.Now()
		i := atomic.SwapInt64(&inserts, 0)
		var ms runtime.MemStats
		runtime.ReadMemStats(&ms)
		preGC := float64(ms.HeapAlloc) / 1024.0 / 1024.0
		runtime.GC()
		runtime.ReadMemStats(&ms)
		postGC := float64(ms.HeapAlloc) / 1024.0 / 1024.0
		fmt.Printf(`
%s inserts at %s inserts per second
%v
HeapAlloc pre/post GC %f/%f MiB
`,
			humanize.Comma(i), humanize.Comma(i/int64(delta.Seconds())),
			db.PrintTableStats("test"),
			preGC, postGC)
	}

	go func() {
		for {
			tk := time.NewTicker(30 * time.Second)
			for range tk.C {
				report()
			}
		}
	}()

	go func() {
		for {
			tk := time.NewTicker(30 * time.Second)
			for range tk.C {
				now := db.Now("test")
				q, err := db.SQLQuery(`
SELECT COUNT(i) AS the_count
FROM test
GROUP BY x, period(168h)
`)
				if err != nil {
					log.Errorf("Unable to build query: %v", err)
					continue
				}
				start := time.Now()
				result, err := q.Run()
				delta := time.Now().Sub(start)
				if err != nil {
					log.Errorf("Unable to run query: %v", err)
					continue
				}
				count := float64(0)
				if len(result.Entries) > 0 {
					count = result.Entries[0].Fields["the_count"][0].Get()
				}
				fmt.Printf("\nQuery at %v returned %v in %v\n", now, humanize.Comma(int64(count)), delta)
			}
		}
	}()

	var wg sync.WaitGroup
	wg.Add(numWriters)
	for _w := 0; _w < numWriters; _w++ {
		go func() {
			defer wg.Done()
			c := 0
			start := time.Now()
			for i := 0; i < reportingPeriods; i++ {
				ts := epoch.Add(time.Duration(i) * reportingInterval)
				for r := 0; r < numReporters/numWriters; r++ {
					for u := 0; u < uniquesPerPeriod; u++ {
						p := &tdb.Point{
							Ts: ts,
							Dims: map[string]interface{}{
								"r": rand.Intn(numReporters),
								"u": rand.Intn(uniquesPerReporter),
								"b": rand.Float64() > 0.99,
								"x": 1,
							},
							Vals: map[string]float64{
								"i":  float64(rand.Intn(100000)),
								"ii": float64(rand.Intn(100)),
							},
						}
						ierr := db.Insert("inbound", p)
						if ierr != nil {
							log.Errorf("Unable to insert: %v", err)
							return
						}
						atomic.AddInt64(&inserts, 1)
						c++

						// Control rate
						if c > 0 && c%1000 == 0 {
							delta := time.Now().Sub(start)
							if delta < targetDeltaFor1000Points {
								time.Sleep(targetDeltaFor1000Points - delta)
							}
							start = time.Now()
						}
					}
				}
				fmt.Print(".")
			}
		}()
	}

	wg.Wait()
	report()
}
<|endoftext|>


<commit_before>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
/*
 *  Filename:    godirs.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description:
 *  Usage:       godirs [options] ARGUMENT ...
 */

//  Package dispatch provides goroutine dispatch and concurrency limiting.
//  It provides an object Dispatch which is a queueing system for concurrent
//  functions. It implements a dynamic limit on the number of routines that
//  it runs simultaneously. It also uses a Queue interface, allowing for
//  alternate queue implementations.
//
//  See github.com/bmatsuo/dispatch/queues for more about the Queue interface.
//
//  See github.com/bmatsuo/dispatch/examples for usage examples.
package dispatch
import (
    "sync"
    "log"
    "github.com/bmatsuo/dispatch/queues"
)

//  A Dispatch is an automated function dispatch queue with a limited
//  number of concurrent gorountines.
type Dispatch struct {
    // The maximum number of goroutines can be changed while the queue is
    // processing.
    MaxGo      int

    // Handle waiting when the limit of concurrent goroutines has been reached.
    waitingToRun bool
    //nextWake     chan bool
    nextWait     *sync.WaitGroup

    // Handle waiting when function queue is empty.
    waitingOnQ   bool
    restart      *sync.WaitGroup

    // Manage the Start()'ing of a Dispatch, avoiding race conditions.
    startLock    *sync.Mutex
    started      bool

    // Handle goroutine-safe queue operations.
    qLock        *sync.Mutex
    queue        queues.Queue

    // Handle goroutine-safe limiting and identifier operations.
    pLock        *sync.Mutex
    processing   int         // Number of QueueTasks running
    idcount      int64       // pid counter

    // Handle stopping of the Start() method.
    kill         chan bool
}

//  Create a new queue object with a specified limit on concurrency.
func New(maxroutines int) *Dispatch {
    return NewCustom(maxroutines, queues.NewFIFO())
}
func NewCustom(maxroutines int, queue queues.Queue) *Dispatch {
    var rl = new(Dispatch)
    rl.startLock = new(sync.Mutex)
    rl.qLock     = new(sync.Mutex)
    rl.pLock     = new(sync.Mutex)
    rl.restart   = new(sync.WaitGroup)
    rl.kill      = make(chan bool)
    //rl.nextWake  = make(chan bool)
    rl.nextWait  = new(sync.WaitGroup)
    rl.queue     = queue
    rl.MaxGo     = maxroutines
    rl.idcount   = 0
    return rl
}

//  Goroutines called from a Dispatch are given an int identifier unique
//  to that routine.
type StdTask struct {
    F func(id int64)
}
func (dt StdTask) Type() string {
    return "StdTask"
}
func (dt StdTask) SetFunc(f func(id int64)) {
    dt.F = f
}
func (dt StdTask) Func() func(id int64) {
    return dt.F
}
type dispatchTaskWrapper struct {
    id int64
    t  queues.Task
}
func (dtw dispatchTaskWrapper) Func() func(id int64) {
    return dtw.t.Func()
}
func (dtw dispatchTaskWrapper) Id() int64 {
    return dtw.id
}
func (dtw dispatchTaskWrapper) Task() queues.Task {
    return dtw.t
}

//  Enqueue a task for execution as a goroutine.
func (gq *Dispatch) Enqueue(t queues.Task) int64 {
    // Wrap the function so it works with the goroutine limiting code.
    var f = t.Func()
    var dtFunc = func (id int64) {
        // Run the given function.
        f(id)

        // Decrement the process counter.
        gq.pLock.Lock()
        log.Printf("processing: %d, waiting: %v", gq.processing, gq.waitingToRun)
        gq.processing--
        if gq.waitingToRun {
            gq.waitingToRun = false
            gq.nextWait.Done()
        }
        gq.pLock.Unlock()
    }
    t.SetFunc(dtFunc)

    // Lock the queue and enqueue a new task.
    gq.qLock.Lock()
    gq.idcount++
    var id = gq.idcount
    gq.queue.Enqueue(dispatchTaskWrapper{id, t})
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    gq.qLock.Unlock()

    return id
}

//  Stop the queue after gq.Start() has been called. Any goroutines which
//  have not already been dequeued will not be executed until gq.Start()
//  is called again.
func (gq *Dispatch) Stop() {
    // Lock out Start() and queue ops for the entire call.
    gq.startLock.Lock()
    defer gq.startLock.Unlock()
    gq.qLock.Lock()
    defer gq.qLock.Unlock()

    if !gq.started {
        return
    }

    // Clear channel flags and close channels, stoping further processing.
    close(gq.kill)
    gq.started = false
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.waitingToRun {
        gq.waitingToRun = false
        gq.nextWait.Done()
    }
    //close(gq.nextWake)
}

//  Start the next task in the queue. It's assumed that the queue is non-
//  empty. Furthermore, there should only be one goroutine in this method
//  (for this object) at a time. Both conditions are enforced in
//  gq.Start(), which calls gq.next() exclusively.
func (gq *Dispatch) next() {
    for true {
        // Attempt to start processing the file.
        gq.pLock.Lock()
        if gq.processing >= gq.MaxGo {
            gq.waitingToRun = true
            gq.nextWait.Add(1)
            gq.pLock.Unlock()
            gq.nextWait.Wait()
            /*
            var cont, ok =<-gq.nextWake
            if !ok {
                gq.nextWake = make(chan bool)
                return
            }
            if !cont {
                return
            }
            */
            continue
        }
        // Keep the books and reset wait time before unlocking.
        gq.processing++
        gq.pLock.Unlock()

        // Get an element from the queue.
        gq.qLock.Lock()
        var wrapper = gq.queue.Dequeue().(queues.RegisteredTask)
        gq.qLock.Unlock()

        // Begin processing and asyncronously return.
        //var task = taskelm.Value.(dispatchTaskWrapper)
        var task = wrapper.Func()
        go task(wrapper.Id())
        return
    }
}

//  Start executing goroutines. Don't stop until gq.Stop() is called.
func (gq *Dispatch) Start() {
    // Avoid multiple gq.Start() methods and avoid race conditions.
    gq.startLock.Lock()
    if gq.started {
        panic("already started")
    }
    gq.started = true
    gq.startLock.Unlock()

    // Recreate any channels that were closed by a previous Stop().
    var inited = false
    for !inited {
        select {
        case _, okKill :=<-gq.kill:
            if !okKill {
                gq.kill = make(chan bool)
            }
        /*
        case _, okWake :=<-gq.nextWake:
            if !okWake {
                gq.nextWake = make(chan bool)
            }
        */
        default:
            inited = true
        }
    }

    // Process the queue
    for true {
        select {
        case die, ok :=<-gq.kill:
            // If something came out of this channel, we must stop.
            if !ok {
                // Recreate the channel on a closure.
                gq.kill = make(chan bool)
                return
            }
            if die {
                return
            }
        default:
            // Check the queue size and determine if we need to wait.
            gq.qLock.Lock()
            var wait = gq.queue.Len() == 0
            if gq.waitingOnQ = wait ; wait {
                gq.restart.Add(1)
            }
            gq.qLock.Unlock()

            if wait {
                // Wait for a restart signal from gq.Enqueue
                gq.restart.Wait()
            } else {
                // Process the head of the queue and start the loop again.
                gq.next()
                continue
            }
        }
    }
}
<commit_msg>Figured it out. Fix bug setting StdTask functions.<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
/*
 *  Filename:    godirs.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Tue Jul  5 22:13:49 PDT 2011
 *  Description:
 *  Usage:       godirs [options] ARGUMENT ...
 */

//  Package dispatch provides goroutine dispatch and concurrency limiting.
//  It provides an object Dispatch which is a queueing system for concurrent
//  functions. It implements a dynamic limit on the number of routines that
//  it runs simultaneously. It also uses a Queue interface, allowing for
//  alternate queue implementations.
//
//  See github.com/bmatsuo/dispatch/queues for more about the Queue interface.
//
//  See github.com/bmatsuo/dispatch/examples for usage examples.
package dispatch
import (
    "sync"
    "log"
    "github.com/bmatsuo/dispatch/queues"
)

//  A Dispatch is an automated function dispatch queue with a limited
//  number of concurrent gorountines.
type Dispatch struct {
    // The maximum number of goroutines can be changed while the queue is
    // processing.
    MaxGo      int

    // Handle waiting when the limit of concurrent goroutines has been reached.
    waitingToRun bool
    //nextWake     chan bool
    nextWait     *sync.WaitGroup

    // Handle waiting when function queue is empty.
    waitingOnQ   bool
    restart      *sync.WaitGroup

    // Manage the Start()'ing of a Dispatch, avoiding race conditions.
    startLock    *sync.Mutex
    started      bool

    // Handle goroutine-safe queue operations.
    qLock        *sync.Mutex
    queue        queues.Queue

    // Handle goroutine-safe limiting and identifier operations.
    pLock        *sync.Mutex
    processing   int         // Number of QueueTasks running
    idcount      int64       // pid counter

    // Handle stopping of the Start() method.
    kill         chan bool
}

//  Create a new queue object with a specified limit on concurrency.
func New(maxroutines int) *Dispatch {
    return NewCustom(maxroutines, queues.NewFIFO())
}
func NewCustom(maxroutines int, queue queues.Queue) *Dispatch {
    var rl = new(Dispatch)
    rl.startLock = new(sync.Mutex)
    rl.qLock     = new(sync.Mutex)
    rl.pLock     = new(sync.Mutex)
    rl.restart   = new(sync.WaitGroup)
    rl.kill      = make(chan bool)
    //rl.nextWake  = make(chan bool)
    rl.nextWait  = new(sync.WaitGroup)
    rl.queue     = queue
    rl.MaxGo     = maxroutines
    rl.idcount   = 0
    return rl
}

//  Goroutines called from a Dispatch are given an int identifier unique
//  to that routine.
type StdTask struct {
    F func(id int64)
}
func (dt *StdTask) Type() string {
    return "StdTask"
}
func (dt *StdTask) SetFunc(f func(id int64)) {
    dt.F = f
}
func (dt *StdTask) Func() func(id int64) {
    return dt.F
}
type dispatchTaskWrapper struct {
    id int64
    t  queues.Task
}
func (dtw dispatchTaskWrapper) Func() func(id int64) {
    return dtw.t.Func()
}
func (dtw dispatchTaskWrapper) Id() int64 {
    return dtw.id
}
func (dtw dispatchTaskWrapper) Task() queues.Task {
    return dtw.t
}

//  Enqueue a task for execution as a goroutine.
func (gq *Dispatch) Enqueue(t queues.Task) int64 {
    // Wrap the function so it works with the goroutine limiting code.
    var f = t.Func()
    var dtFunc = func (id int64) {
        // Run the given function.
        f(id)

        // Decrement the process counter.
        gq.pLock.Lock()
        log.Printf("processing: %d, waiting: %v", gq.processing, gq.waitingToRun)
        gq.processing--
        if gq.waitingToRun {
            gq.waitingToRun = false
            gq.nextWait.Done()
        }
        gq.pLock.Unlock()
    }
    t.SetFunc(dtFunc)

    // Lock the queue and enqueue a new task.
    gq.qLock.Lock()
    gq.idcount++
    var id = gq.idcount
    gq.queue.Enqueue(dispatchTaskWrapper{id, t})
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    gq.qLock.Unlock()

    return id
}

//  Stop the queue after gq.Start() has been called. Any goroutines which
//  have not already been dequeued will not be executed until gq.Start()
//  is called again.
func (gq *Dispatch) Stop() {
    // Lock out Start() and queue ops for the entire call.
    gq.startLock.Lock()
    defer gq.startLock.Unlock()
    gq.qLock.Lock()
    defer gq.qLock.Unlock()

    if !gq.started {
        return
    }

    // Clear channel flags and close channels, stoping further processing.
    close(gq.kill)
    gq.started = false
    if gq.waitingOnQ {
        gq.waitingOnQ = false
        gq.restart.Done()
    }
    if gq.waitingToRun {
        gq.waitingToRun = false
        gq.nextWait.Done()
    }
    //close(gq.nextWake)
}

//  Start the next task in the queue. It's assumed that the queue is non-
//  empty. Furthermore, there should only be one goroutine in this method
//  (for this object) at a time. Both conditions are enforced in
//  gq.Start(), which calls gq.next() exclusively.
func (gq *Dispatch) next() {
    for true {
        // Attempt to start processing the file.
        gq.pLock.Lock()
        if gq.processing >= gq.MaxGo {
            gq.waitingToRun = true
            gq.nextWait.Add(1)
            gq.pLock.Unlock()
            gq.nextWait.Wait()
            /*
            var cont, ok =<-gq.nextWake
            if !ok {
                gq.nextWake = make(chan bool)
                return
            }
            if !cont {
                return
            }
            */
            continue
        }
        // Keep the books and reset wait time before unlocking.
        gq.processing++
        gq.pLock.Unlock()

        // Get an element from the queue.
        gq.qLock.Lock()
        var wrapper = gq.queue.Dequeue().(queues.RegisteredTask)
        gq.qLock.Unlock()

        // Begin processing and asyncronously return.
        //var task = taskelm.Value.(dispatchTaskWrapper)
        var task = wrapper.Func()
        go task(wrapper.Id())
        return
    }
}

//  Start executing goroutines. Don't stop until gq.Stop() is called.
func (gq *Dispatch) Start() {
    // Avoid multiple gq.Start() methods and avoid race conditions.
    gq.startLock.Lock()
    if gq.started {
        panic("already started")
    }
    gq.started = true
    gq.startLock.Unlock()

    // Recreate any channels that were closed by a previous Stop().
    var inited = false
    for !inited {
        select {
        case _, okKill :=<-gq.kill:
            if !okKill {
                gq.kill = make(chan bool)
            }
        /*
        case _, okWake :=<-gq.nextWake:
            if !okWake {
                gq.nextWake = make(chan bool)
            }
        */
        default:
            inited = true
        }
    }

    // Process the queue
    for true {
        select {
        case die, ok :=<-gq.kill:
            // If something came out of this channel, we must stop.
            if !ok {
                // Recreate the channel on a closure.
                gq.kill = make(chan bool)
                return
            }
            if die {
                return
            }
        default:
            // Check the queue size and determine if we need to wait.
            gq.qLock.Lock()
            var wait = gq.queue.Len() == 0
            if gq.waitingOnQ = wait ; wait {
                gq.restart.Add(1)
            }
            gq.qLock.Unlock()

            if wait {
                // Wait for a restart signal from gq.Enqueue
                gq.restart.Wait()
            } else {
                // Process the head of the queue and start the loop again.
                gq.next()
                continue
            }
        }
    }
}
<|endoftext|>