<commit_before>package yasha

import (
	"compress/bzip2"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/dotabuff/yasha/dota"
	"github.com/stretchr/testify/assert"
)

type testCase struct {
	matchId int64
	url     string

	expectLastChatMessage string
	expectHeroKillCount   map[string]int
	expectHeroDeathCount  map[string]int
}

// Esports match, played on patch 6.83c
func TestEsportsPatch683b(t *testing.T) {
	c := &testCase{
		matchId: 1405240741,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1405240741.dem",
		expectLastChatMessage: "Gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_ember_spirit": 0,
			"npc_dota_hero_broodmother":  5,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_chen":        2,
			"npc_dota_hero_broodmother": 0,
			"npc_dota_hero_sniper":      2,
			"npc_dota_hero_phoenix":     2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p0
func TestEsportsPatch684p0(t *testing.T) {
	c := &testCase{
		matchId: 1450235906,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1450235906.dem",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_broodmother": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_broodmother": 7,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p1
func TestEsportsPatch684p1(t *testing.T) {
	c := &testCase{
		matchId: 1458895412,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1458895412.dem",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_faceless_void": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_faceless_void": 2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84c
func TestEsportsPatch684c(t *testing.T) {
	c := &testCase{
		matchId: 1483980562,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1483980562.dem",
		expectLastChatMessage: "gg wp",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_dragon_knight": 5,
			"npc_dota_hero_bristleback":   1,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_earthshaker": 6,
			"npc_dota_hero_bristleback": 3,
		},
	}

	testReplayCase(t, c)
}

// Manually scrutinised match, played on patch 6.84p1
func TestPublicMatchPatch684p1(t *testing.T) {
	assert := assert.New(t)

	data, err := getReplayData(1456774107, "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1456774107.dem")
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
	}

	earthshakerDeaths := 0
	spiritBreakerDeaths := 0
	parser.OnCombatLog = func(entry CombatLogEntry) {
		// t.Logf("OnCombatLog: %s: %+v", reflect.TypeOf(entry), entry)
		switch log := entry.(type) {
		case *CombatLogDeath:
			if log.Target == "npc_dota_hero_earthshaker" {
				earthshakerDeaths++
			}
			if log.Target == "npc_dota_hero_spirit_breaker" {
				spiritBreakerDeaths++
			}
		}
	}

	var now time.Duration
	var gameTime, preGameStarttime float64
	parser.OnEntityPreserved = func(pe *PacketEntity) {
		if pe.Name == "DT_DOTAGamerulesProxy" {
			gameTime = pe.Values["DT_DOTAGamerules.m_fGameTime"].(float64)
			preGameStarttime = pe.Values["DT_DOTAGamerules.m_flPreGameStartTime"].(float64)
			now = time.Duration(gameTime-preGameStarttime) * time.Second
		}
	}

	// entindex:3 order_type:1 units:349 position:<x:6953.3125 y:6920.8438 z:384 > queue:false
	unitOrderCount := 0
	unitOrderQueuedCount := 0
	specificUnitOrder := false
	parser.OnSpectatorPlayerUnitOrders = func(n int, o *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) {
		unitOrderCount++
		if *o.Queue == true {
			unitOrderQueuedCount++
		}
		if *o.Entindex == 3 && *o.OrderType == 1 && o.Units[0] == 349 && *o.Queue == false &&
			*o.Position.X == 6953.3125 && *o.Position.Y == 6920.8438 && *o.Position.Y == 384.0 {
			specificUnitOrder = true
		}
	}

	chatWheelMessagesCount := 0
	parser.OnChatWheel = func(n int, o *dota.CDOTAUserMsg_ChatWheel) {
		chatWheelMessagesCount++
	}

	parser.Parse()

	assert.Equal(8, earthshakerDeaths)
	assert.Equal(11, spiritBreakerDeaths)          // not actually right but verified in replay
	assert.Equal(55316, unitOrderCount)            // regression test
	assert.Equal(102, unitOrderQueuedCount)        // regression test
	assert.Equal(int64(2585000000000), int64(now)) // regression test
	assert.Equal(0, chatWheelMessagesCount)        // regression test
}

func testReplayCase(t *testing.T, c *testCase) {
	assert := assert.New(t)

	data, err := getReplayData(c.matchId, c.url)
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	worldMins := &Vector3{}
	worldMaxes := &Vector3{}
	lastChatMessage := ""
	heroKillCount := make(map[string]int)
	heroDeathCount := make(map[string]int)

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
		lastChatMessage = o.GetText()
	}

	parser.OnChatEvent = func(n int, o *dota.CDOTAUserMsg_ChatEvent) {
	}

	parser.OnCombatLog = func(entry CombatLogEntry) {
		switch log := entry.(type) {
		case *CombatLogDeath:
			if strings.HasPrefix(log.Target, "npc_dota_hero_") {
				if _, ok := heroKillCount[log.Source]; !ok {
					heroKillCount[log.Source] = 0
				}
				heroKillCount[log.Source] += 1
			}

			if _, ok := heroDeathCount[log.Target]; !ok {
				heroDeathCount[log.Target] = 0
			}
			heroDeathCount[log.Target] += 1
		}
	}

	parser.OnEntityCreated = func(ent *PacketEntity) {
		if ent.Tick == 0 && ent.Name == "DT_WORLD" {
			worldMins = ent.Values["DT_WORLD.m_WorldMins"].(*Vector3)
			worldMaxes = ent.Values["DT_WORLD.m_WorldMaxs"].(*Vector3)
		}
	}

	parser.Parse()

	// Make sure we have found the death counts for specified heroes
	if c.expectHeroDeathCount != nil {
		for hero, count := range c.expectHeroDeathCount {
			assert.Equal(count, heroDeathCount[hero], "expected hero %s to have death count %d", hero, count)
		}
	}

	// Make sure we have found the kill counts for specified heroes.
	if c.expectHeroKillCount != nil {
		for hero, count := range c.expectHeroKillCount {
			assert.Equal(count, heroKillCount[hero], "expected hero %s to have kill count %d", hero, count)
		}
	}

	// Make sure we find the DT_WORLD entity and it has the correct min and max dimensions.
	// This serves to help ensure our Float and Vector3 parsing is correct.
	assert.Equal(&Vector3{X: -8576.0, Y: -7680.0, Z: -1536.0}, worldMins)
	assert.Equal(&Vector3{X: 9216.0, Y: 8192.0, Z: 256.0}, worldMaxes)

	// Make sure we found the chat messages and have properly found the last one
	assert.Equal(c.expectLastChatMessage, lastChatMessage)
}

func getReplayData(matchId int64, url string) ([]byte, error) {
	path := fmt.Sprintf("replays/%d.dem", matchId)
	if data, err := ioutil.ReadFile(path); err == nil {
		fmt.Printf("read replay %d from %s\n", matchId, path)
		return data, nil
	}

	fmt.Printf("downloading replay %d from %s...\n", matchId, url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Return an error if we don't get a 200
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("invalid status %d", resp.StatusCode)
	}

	var data []byte
	if url[len(url)-3:] == "bz2" {
		data, err = ioutil.ReadAll(bzip2.NewReader(resp.Body))
	} else {
		data, err = ioutil.ReadAll(resp.Body)
	}

	if err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(path, data, 0644); err != nil {
		return nil, err
	}

	fmt.Printf("downloaded replay %d from %s to %s\n", matchId, url, path)

	return data, nil
}
<commit_msg>Adapted tests to new signature of OnCombatLog hook.<commit_after>package yasha

import (
	"compress/bzip2"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/dotabuff/yasha/dota"
	"github.com/stretchr/testify/assert"
)

type testCase struct {
	matchId int64
	url     string

	expectLastChatMessage string
	expectHeroKillCount   map[string]int
	expectHeroDeathCount  map[string]int
}

// Esports match, played on patch 6.83c
func TestEsportsPatch683b(t *testing.T) {
	c := &testCase{
		matchId: 1405240741,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1405240741.dem",
		expectLastChatMessage: "Gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_ember_spirit": 0,
			"npc_dota_hero_broodmother":  5,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_chen":        2,
			"npc_dota_hero_broodmother": 0,
			"npc_dota_hero_sniper":      2,
			"npc_dota_hero_phoenix":     2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p0
func TestEsportsPatch684p0(t *testing.T) {
	c := &testCase{
		matchId: 1450235906,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1450235906.dem",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_broodmother": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_broodmother": 7,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p1
func TestEsportsPatch684p1(t *testing.T) {
	c := &testCase{
		matchId: 1458895412,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1458895412.dem",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_faceless_void": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_faceless_void": 2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84c
func TestEsportsPatch684c(t *testing.T) {
	c := &testCase{
		matchId: 1483980562,
		url:     "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1483980562.dem",
		expectLastChatMessage: "gg wp",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_dragon_knight": 5,
			"npc_dota_hero_bristleback":   1,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_earthshaker": 6,
			"npc_dota_hero_bristleback": 3,
		},
	}

	testReplayCase(t, c)
}

// Manually scrutinised match, played on patch 6.84p1
func TestPublicMatchPatch684p1(t *testing.T) {
	assert := assert.New(t)

	data, err := getReplayData(1456774107, "https://s3-us-west-2.amazonaws.com/yasha.dotabuff/1456774107.dem")
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
	}

	earthshakerDeaths := 0
	spiritBreakerDeaths := 0
	parser.OnCombatLog = func(tick int, entry CombatLogEntry) {
		// t.Logf("OnCombatLog: %s: %+v", reflect.TypeOf(entry), entry)
		switch log := entry.(type) {
		case *CombatLogDeath:
			if log.Target == "npc_dota_hero_earthshaker" {
				earthshakerDeaths++
			}
			if log.Target == "npc_dota_hero_spirit_breaker" {
				spiritBreakerDeaths++
			}
		}
	}

	var now time.Duration
	var gameTime, preGameStarttime float64
	parser.OnEntityPreserved = func(pe *PacketEntity) {
		if pe.Name == "DT_DOTAGamerulesProxy" {
			gameTime = pe.Values["DT_DOTAGamerules.m_fGameTime"].(float64)
			preGameStarttime = pe.Values["DT_DOTAGamerules.m_flPreGameStartTime"].(float64)
			now = time.Duration(gameTime-preGameStarttime) * time.Second
		}
	}

	// entindex:3 order_type:1 units:349 position:<x:6953.3125 y:6920.8438 z:384 > queue:false
	unitOrderCount := 0
	unitOrderQueuedCount := 0
	specificUnitOrder := false
	parser.OnSpectatorPlayerUnitOrders = func(n int, o *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) {
		unitOrderCount++
		if *o.Queue == true {
			unitOrderQueuedCount++
		}
		if *o.Entindex == 3 && *o.OrderType == 1 && o.Units[0] == 349 && *o.Queue == false &&
			*o.Position.X == 6953.3125 && *o.Position.Y == 6920.8438 && *o.Position.Y == 384.0 {
			specificUnitOrder = true
		}
	}

	chatWheelMessagesCount := 0
	parser.OnChatWheel = func(n int, o *dota.CDOTAUserMsg_ChatWheel) {
		chatWheelMessagesCount++
	}

	parser.Parse()

	assert.Equal(8, earthshakerDeaths)
	assert.Equal(11, spiritBreakerDeaths)          // not actually right but verified in replay
	assert.Equal(55316, unitOrderCount)            // regression test
	assert.Equal(102, unitOrderQueuedCount)        // regression test
	assert.Equal(int64(2585000000000), int64(now)) // regression test
	assert.Equal(0, chatWheelMessagesCount)        // regression test
}

func testReplayCase(t *testing.T, c *testCase) {
	assert := assert.New(t)

	data, err := getReplayData(c.matchId, c.url)
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	worldMins := &Vector3{}
	worldMaxes := &Vector3{}
	lastChatMessage := ""
	heroKillCount := make(map[string]int)
	heroDeathCount := make(map[string]int)

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
		lastChatMessage = o.GetText()
	}

	parser.OnChatEvent = func(n int, o *dota.CDOTAUserMsg_ChatEvent) {
	}

	parser.OnCombatLog = func(tick int, entry CombatLogEntry) {
		switch log := entry.(type) {
		case *CombatLogDeath:
			if strings.HasPrefix(log.Target, "npc_dota_hero_") {
				if _, ok := heroKillCount[log.Source]; !ok {
					heroKillCount[log.Source] = 0
				}
				heroKillCount[log.Source] += 1
			}

			if _, ok := heroDeathCount[log.Target]; !ok {
				heroDeathCount[log.Target] = 0
			}
			heroDeathCount[log.Target] += 1
		}
	}

	parser.OnEntityCreated = func(ent *PacketEntity) {
		if ent.Tick == 0 && ent.Name == "DT_WORLD" {
			worldMins = ent.Values["DT_WORLD.m_WorldMins"].(*Vector3)
			worldMaxes = ent.Values["DT_WORLD.m_WorldMaxs"].(*Vector3)
		}
	}

	parser.Parse()

	// Make sure we have found the death counts for specified heroes
	if c.expectHeroDeathCount != nil {
		for hero, count := range c.expectHeroDeathCount {
			assert.Equal(count, heroDeathCount[hero], "expected hero %s to have death count %d", hero, count)
		}
	}

	// Make sure we have found the kill counts for specified heroes.
	if c.expectHeroKillCount != nil {
		for hero, count := range c.expectHeroKillCount {
			assert.Equal(count, heroKillCount[hero], "expected hero %s to have kill count %d", hero, count)
		}
	}

	// Make sure we find the DT_WORLD entity and it has the correct min and max dimensions.
	// This serves to help ensure our Float and Vector3 parsing is correct.
	assert.Equal(&Vector3{X: -8576.0, Y: -7680.0, Z: -1536.0}, worldMins)
	assert.Equal(&Vector3{X: 9216.0, Y: 8192.0, Z: 256.0}, worldMaxes)

	// Make sure we found the chat messages and have properly found the last one
	assert.Equal(c.expectLastChatMessage, lastChatMessage)
}

func getReplayData(matchId int64, url string) ([]byte, error) {
	path := fmt.Sprintf("replays/%d.dem", matchId)
	if data, err := ioutil.ReadFile(path); err == nil {
		fmt.Printf("read replay %d from %s\n", matchId, path)
		return data, nil
	}

	fmt.Printf("downloading replay %d from %s...\n", matchId, url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Return an error if we don't get a 200
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("invalid status %d", resp.StatusCode)
	}

	var data []byte
	if url[len(url)-3:] == "bz2" {
		data, err = ioutil.ReadAll(bzip2.NewReader(resp.Body))
	} else {
		data, err = ioutil.ReadAll(resp.Body)
	}

	if err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(path, data, 0644); err != nil {
		return nil, err
	}

	fmt.Printf("downloaded replay %d from %s to %s\n", matchId, url, path)

	return data, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"

	"github.com/peterbourgon/g2s"
	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/farm"
	"github.com/soundcloud/roshi/instrumentation"
	"github.com/soundcloud/roshi/instrumentation/statsd"
	"github.com/soundcloud/roshi/shard"
	"github.com/tsenart/tb"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	var (
		redisInstances      = flag.String("redis.instances", "", "Semicolon-separated list of comma-separated lists of Redis instances")
		redisConnectTimeout = flag.Duration("redis.connect.timeout", 3*time.Second, "Redis connect timeout")
		redisReadTimeout    = flag.Duration("redis.read.timeout", 3*time.Second, "Redis read timeout")
		redisWriteTimeout   = flag.Duration("redis.write.timeout", 3*time.Second, "Redis write timeout")
		redisMCPI           = flag.Int("redis.mcpi", 10, "Max connections per Redis instance")
		redisHash           = flag.String("redis.hash", "murmur3", "Redis hash function: murmur3, fnv, fnva")
		maxSize             = flag.Int("max.size", 10000, "Maximum number of events per key")
		batchSize           = flag.Int("batch.size", 100, "keys to select per request")
		maxKeysPerSecond    = flag.Int64("max.keys.per.second", 1000, "max keys per second to walk")
		scanLogInterval     = flag.Duration("scan.log.interval", 5*time.Second, "how often to report scan rates in log")
		once                = flag.Bool("once", false, "walk entire keyspace once and exit (default false, walk forever)")
		statsdAddress       = flag.String("statsd.address", "", "Statsd address (blank to disable)")
		statsdSampleRate    = flag.Float64("statsd.sample.rate", 0.1, "Statsd sample rate for normal metrics")
		statsdBucketPrefix  = flag.String("statsd.bucket.prefix", "myservice.", "Statsd bucket key prefix, including trailing period")
	)
	flag.Parse()

	// Validate integer arguments.
	if *maxKeysPerSecond < int64(*batchSize) {
		log.Fatal("max keys per second should be bigger than batch size")
	}

	// Set up statsd instrumentation, if it's specified.
	stats := g2s.Noop()
	if *statsdAddress != "" {
		var err error
		stats, err = g2s.Dial("udp", *statsdAddress)
		if err != nil {
			log.Fatal(err)
		}
	}
	instr := statsd.New(stats, float32(*statsdSampleRate), *statsdBucketPrefix)

	// Parse hash function.
	var hashFunc func(string) uint32
	switch strings.ToLower(*redisHash) {
	case "murmur3":
		hashFunc = shard.Murmur3
	case "fnv":
		hashFunc = shard.FNV
	case "fnva":
		hashFunc = shard.FNVa
	default:
		log.Fatalf("unknown hash '%s'", *redisHash)
	}

	// Set up the clusters.
	clusters, err := makeClusters(
		*redisInstances,
		*redisConnectTimeout, *redisReadTimeout, *redisWriteTimeout,
		*redisMCPI,
		hashFunc,
		*maxSize,
		instr,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Set up our rate limiter. Remember: it's per-key, not per-request.
	throttle := newThrottle(*maxKeysPerSecond)

	// Perform the walk.
	dst := farm.New(clusters, len(clusters), farm.SendAllReadAll, farm.AllRepairs, instr)
	for {
		src := scan(clusters, *batchSize, *scanLogInterval) // new key set
		walkOnce(dst, throttle, src, *maxSize, instr)
		if *once {
			return
		}
	}
}

func makeClusters(
	redisInstances string,
	connectTimeout, readTimeout, writeTimeout time.Duration,
	redisMCPI int,
	hashFunc func(string) uint32,
	maxSize int,
	instr instrumentation.Instrumentation,
) ([]cluster.Cluster, error) {
	clusters := []cluster.Cluster{}
	for i, clusterInstances := range strings.Split(redisInstances, ";") {
		addresses := stripBlank(strings.Split(clusterInstances, ","))
		if len(addresses) <= 0 {
			continue
		}
		clusters = append(clusters, cluster.New(
			shard.New(
				addresses,
				connectTimeout, readTimeout, writeTimeout,
				redisMCPI,
				hashFunc,
			),
			maxSize,
			instr,
		))
		log.Printf("Redis cluster %d: %d instance(s)", i+1, len(addresses))
	}
	if len(clusters) <= 0 {
		return []cluster.Cluster{}, fmt.Errorf("no cluster(s)")
	}
	return clusters, nil
}

func scan(clusters []cluster.Cluster, batchSize int, logInterval time.Duration) <-chan []string {
	c := make(chan []string)
	go func() {
		defer close(c)
		logTick := time.Tick(logInterval)
		batches, keys, prev, mark := 0, 0, 0, time.Now()

		for i, index := range rand.Perm(len(clusters)) {
			log.Printf("scan: %d/%d, cluster index %d: begin", i+1, len(clusters), index)
			for batch := range clusters[index].Keys(batchSize) {
				select {
				case c <- batch:
					//log.Printf(
					//	"scan: %d/%d, cluster index %d: forwarded batch of %d",
					//	i+1, len(clusters), index,
					//	len(batch),
					//)
					batches += 1
					keys += len(batch)

				case <-logTick:
					log.Printf(
						"scan: %d/%d, cluster index %d: %d batches, %d keys, %.2f keys/sec",
						i+1, len(clusters), index,
						batches,
						keys,
						float64(keys-prev)/(time.Since(mark).Seconds()),
					)
					prev = keys
					mark = time.Now()
				}
			}
		}
	}()
	return c
}

func walkOnce(
	dst farm.Selecter,
	throttle *throttle,
	src <-chan []string,
	maxSize int,
	instr instrumentation.WalkInstrumentation,
) {
	for batch := range src {
		throttle.wait(int64(len(batch)))
		dst.Select(batch, 0, maxSize)
		instr.WalkKeys(len(batch))
	}
}

type throttle struct {
	bucket       *tb.Bucket
	waitInterval time.Duration
}

func newThrottle(maxPerSecond int64) *throttle {
	return &throttle{
		bucket:       tb.NewBucket(maxPerSecond, -1),
		waitInterval: (1 * time.Second) / time.Duration(maxPerSecond),
	}
}

func (t *throttle) wait(n int64) {
	got := t.bucket.Take(n)
	for got < n {
		time.Sleep(t.waitInterval)
		got += t.bucket.Take(n - got)
	}
}

func stripBlank(src []string) []string {
	dst := []string{}
	for _, s := range src {
		if s == "" {
			continue
		}
		dst = append(dst, s)
	}
	return dst
}
<commit_msg>roshi-walker: fix unused import<commit_after>package main

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

	"github.com/peterbourgon/g2s"
	"github.com/soundcloud/roshi/cluster"
	"github.com/soundcloud/roshi/farm"
	"github.com/soundcloud/roshi/instrumentation"
	"github.com/soundcloud/roshi/instrumentation/statsd"
	"github.com/soundcloud/roshi/shard"
	"github.com/tsenart/tb"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	var (
		redisInstances      = flag.String("redis.instances", "", "Semicolon-separated list of comma-separated lists of Redis instances")
		redisConnectTimeout = flag.Duration("redis.connect.timeout", 3*time.Second, "Redis connect timeout")
		redisReadTimeout    = flag.Duration("redis.read.timeout", 3*time.Second, "Redis read timeout")
		redisWriteTimeout   = flag.Duration("redis.write.timeout", 3*time.Second, "Redis write timeout")
		redisMCPI           = flag.Int("redis.mcpi", 10, "Max connections per Redis instance")
		redisHash           = flag.String("redis.hash", "murmur3", "Redis hash function: murmur3, fnv, fnva")
		maxSize             = flag.Int("max.size", 10000, "Maximum number of events per key")
		batchSize           = flag.Int("batch.size", 100, "keys to select per request")
		maxKeysPerSecond    = flag.Int64("max.keys.per.second", 1000, "max keys per second to walk")
		scanLogInterval     = flag.Duration("scan.log.interval", 5*time.Second, "how often to report scan rates in log")
		once                = flag.Bool("once", false, "walk entire keyspace once and exit (default false, walk forever)")
		statsdAddress       = flag.String("statsd.address", "", "Statsd address (blank to disable)")
		statsdSampleRate    = flag.Float64("statsd.sample.rate", 0.1, "Statsd sample rate for normal metrics")
		statsdBucketPrefix  = flag.String("statsd.bucket.prefix", "myservice.", "Statsd bucket key prefix, including trailing period")
	)
	flag.Parse()

	// Validate integer arguments.
	if *maxKeysPerSecond < int64(*batchSize) {
		log.Fatal("max keys per second should be bigger than batch size")
	}

	// Set up statsd instrumentation, if it's specified.
	stats := g2s.Noop()
	if *statsdAddress != "" {
		var err error
		stats, err = g2s.Dial("udp", *statsdAddress)
		if err != nil {
			log.Fatal(err)
		}
	}
	instr := statsd.New(stats, float32(*statsdSampleRate), *statsdBucketPrefix)

	// Parse hash function.
	var hashFunc func(string) uint32
	switch strings.ToLower(*redisHash) {
	case "murmur3":
		hashFunc = shard.Murmur3
	case "fnv":
		hashFunc = shard.FNV
	case "fnva":
		hashFunc = shard.FNVa
	default:
		log.Fatalf("unknown hash '%s'", *redisHash)
	}

	// Set up the clusters.
	clusters, err := makeClusters(
		*redisInstances,
		*redisConnectTimeout, *redisReadTimeout, *redisWriteTimeout,
		*redisMCPI,
		hashFunc,
		*maxSize,
		instr,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Set up our rate limiter. Remember: it's per-key, not per-request.
	throttle := newThrottle(*maxKeysPerSecond)

	// Perform the walk.
	dst := farm.New(clusters, len(clusters), farm.SendAllReadAll, farm.AllRepairs, instr)
	for {
		src := scan(clusters, *batchSize, *scanLogInterval) // new key set
		walkOnce(dst, throttle, src, *maxSize, instr)
		if *once {
			return
		}
	}
}

func makeClusters(
	redisInstances string,
	connectTimeout, readTimeout, writeTimeout time.Duration,
	redisMCPI int,
	hashFunc func(string) uint32,
	maxSize int,
	instr instrumentation.Instrumentation,
) ([]cluster.Cluster, error) {
	clusters := []cluster.Cluster{}
	for i, clusterInstances := range strings.Split(redisInstances, ";") {
		addresses := stripBlank(strings.Split(clusterInstances, ","))
		if len(addresses) <= 0 {
			continue
		}
		clusters = append(clusters, cluster.New(
			shard.New(
				addresses,
				connectTimeout, readTimeout, writeTimeout,
				redisMCPI,
				hashFunc,
			),
			maxSize,
			instr,
		))
		log.Printf("Redis cluster %d: %d instance(s)", i+1, len(addresses))
	}
	if len(clusters) <= 0 {
		return []cluster.Cluster{}, fmt.Errorf("no cluster(s)")
	}
	return clusters, nil
}

func scan(clusters []cluster.Cluster, batchSize int, logInterval time.Duration) <-chan []string {
	c := make(chan []string)
	go func() {
		defer close(c)
		logTick := time.Tick(logInterval)
		batches, keys, prev, mark := 0, 0, 0, time.Now()

		for i, index := range rand.Perm(len(clusters)) {
			log.Printf("scan: %d/%d, cluster index %d: begin", i+1, len(clusters), index)
			for batch := range clusters[index].Keys(batchSize) {
				select {
				case c <- batch:
					//log.Printf(
					//	"scan: %d/%d, cluster index %d: forwarded batch of %d",
					//	i+1, len(clusters), index,
					//	len(batch),
					//)
					batches += 1
					keys += len(batch)

				case <-logTick:
					log.Printf(
						"scan: %d/%d, cluster index %d: %d batches, %d keys, %.2f keys/sec",
						i+1, len(clusters), index,
						batches,
						keys,
						float64(keys-prev)/(time.Since(mark).Seconds()),
					)
					prev = keys
					mark = time.Now()
				}
			}
		}
	}()
	return c
}

func walkOnce(
	dst farm.Selecter,
	throttle *throttle,
	src <-chan []string,
	maxSize int,
	instr instrumentation.WalkInstrumentation,
) {
	for batch := range src {
		throttle.wait(int64(len(batch)))
		dst.Select(batch, 0, maxSize)
		instr.WalkKeys(len(batch))
	}
}

type throttle struct {
	bucket       *tb.Bucket
	waitInterval time.Duration
}

func newThrottle(maxPerSecond int64) *throttle {
	return &throttle{
		bucket:       tb.NewBucket(maxPerSecond, -1),
		waitInterval: (1 * time.Second) / time.Duration(maxPerSecond),
	}
}

func (t *throttle) wait(n int64) {
	got := t.bucket.Take(n)
	for got < n {
		time.Sleep(t.waitInterval)
		got += t.bucket.Take(n - got)
	}
}

func stripBlank(src []string) []string {
	dst := []string{}
	for _, s := range src {
		if s == "" {
			continue
		}
		dst = append(dst, s)
	}
	return dst
}
<|endoftext|>


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

package repo

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/Unknwon/com"
	log "gopkg.in/clog.v1"

	"github.com/gogits/git-module"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/auth"
	"github.com/gogits/gogs/modules/base"
	"github.com/gogits/gogs/modules/context"
	"github.com/gogits/gogs/modules/setting"
)

const (
	CREATE  base.TplName = "repo/create"
	MIGRATE base.TplName = "repo/migrate"
)

func MustBeNotBare(ctx *context.Context) {
	if ctx.Repo.Repository.IsBare {
		ctx.Handle(404, "MustBeNotBare", nil)
	}
}

func checkContextUser(ctx *context.Context, uid int64) *models.User {
	orgs, err := models.GetOwnedOrgsByUserIDDesc(ctx.User.ID, "updated_unix")
	if err != nil {
		ctx.Handle(500, "GetOwnedOrgsByUserIDDesc", err)
		return nil
	}
	ctx.Data["Orgs"] = orgs

	// Not equal means current user is an organization.
	if uid == ctx.User.ID || uid == 0 {
		return ctx.User
	}

	org, err := models.GetUserByID(uid)
	if models.IsErrUserNotExist(err) {
		return ctx.User
	}

	if err != nil {
		ctx.Handle(500, "GetUserByID", fmt.Errorf("[%d]: %v", uid, err))
		return nil
	}

	// Check ownership of organization.
	if !org.IsOrganization() || !(ctx.User.IsAdmin || org.IsOwnedBy(ctx.User.ID)) {
		ctx.Error(403)
		return nil
	}
	return org
}

func Create(ctx *context.Context) {
	ctx.Data["Title"] = ctx.Tr("new_repo")

	// Give default value for template to render.
	ctx.Data["Gitignores"] = models.Gitignores
	ctx.Data["Licenses"] = models.Licenses
	ctx.Data["Readmes"] = models.Readmes
	ctx.Data["readme"] = "Default"
	ctx.Data["private"] = ctx.User.LastRepoVisibility
	ctx.Data["IsForcedPrivate"] = setting.Repository.ForcePrivate

	ctxUser := checkContextUser(ctx, ctx.QueryInt64("org"))
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	ctx.HTML(200, CREATE)
}

func handleCreateError(ctx *context.Context, owner *models.User, err error, name string, tpl base.TplName, form interface{}) {
	switch {
	case models.IsErrReachLimitOfRepo(err):
		ctx.RenderWithErr(ctx.Tr("repo.form.reach_limit_of_creation", owner.RepoCreationNum()), tpl, form)
	case models.IsErrRepoAlreadyExist(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("form.repo_name_been_taken"), tpl, form)
	case models.IsErrNameReserved(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("repo.form.name_reserved", err.(models.ErrNameReserved).Name), tpl, form)
	case models.IsErrNamePatternNotAllowed(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("repo.form.name_pattern_not_allowed", err.(models.ErrNamePatternNotAllowed).Pattern), tpl, form)
	default:
		ctx.Handle(500, name, err)
	}
}

func CreatePost(ctx *context.Context, form auth.CreateRepoForm) {
	ctx.Data["Title"] = ctx.Tr("new_repo")

	ctx.Data["Gitignores"] = models.Gitignores
	ctx.Data["Licenses"] = models.Licenses
	ctx.Data["Readmes"] = models.Readmes

	ctxUser := checkContextUser(ctx, form.Uid)
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	if ctx.HasError() {
		ctx.HTML(200, CREATE)
		return
	}

	repo, err := models.CreateRepository(ctxUser, models.CreateRepoOptions{
		Name:        form.RepoName,
		Description: form.Description,
		Gitignores:  form.Gitignores,
		License:     form.License,
		Readme:      form.Readme,
		IsPrivate:   form.Private || setting.Repository.ForcePrivate,
		AutoInit:    form.AutoInit,
	})
	if err == nil {
		log.Trace("Repository created [%d]: %s/%s", repo.ID, ctxUser.Name, repo.Name)
		ctx.Redirect(setting.AppSubUrl + "/" + ctxUser.Name + "/" + repo.Name)
		return
	}

	if repo != nil {
		if errDelete := models.DeleteRepository(ctxUser.ID, repo.ID); errDelete != nil {
			log.Error(4, "DeleteRepository: %v", errDelete)
		}
	}

	handleCreateError(ctx, ctxUser, err, "CreatePost", CREATE, &form)
}

func Migrate(ctx *context.Context) {
	ctx.Data["Title"] = ctx.Tr("new_migrate")
	ctx.Data["private"] = ctx.User.LastRepoVisibility
	ctx.Data["IsForcedPrivate"] = setting.Repository.ForcePrivate
	ctx.Data["mirror"] = ctx.Query("mirror") == "1"

	ctxUser := checkContextUser(ctx, ctx.QueryInt64("org"))
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	ctx.HTML(200, MIGRATE)
}

func MigratePost(ctx *context.Context, form auth.MigrateRepoForm) {
	ctx.Data["Title"] = ctx.Tr("new_migrate")

	ctxUser := checkContextUser(ctx, form.Uid)
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	if ctx.HasError() {
		ctx.HTML(200, MIGRATE)
		return
	}

	remoteAddr, err := form.ParseRemoteAddr(ctx.User)
	if err != nil {
		if models.IsErrInvalidCloneAddr(err) {
			ctx.Data["Err_CloneAddr"] = true
			addrErr := err.(models.ErrInvalidCloneAddr)
			switch {
			case addrErr.IsURLError:
				ctx.RenderWithErr(ctx.Tr("form.url_error"), MIGRATE, &form)
			case addrErr.IsPermissionDenied:
				ctx.RenderWithErr(ctx.Tr("repo.migrate.permission_denied"), MIGRATE, &form)
			case addrErr.IsInvalidPath:
				ctx.RenderWithErr(ctx.Tr("repo.migrate.invalid_local_path"), MIGRATE, &form)
			default:
				ctx.Handle(500, "Unknown error", err)
			}
		} else {
			ctx.Handle(500, "ParseRemoteAddr", err)
		}
		return
	}

	repo, err := models.MigrateRepository(ctxUser, models.MigrateRepoOptions{
		Name:        form.RepoName,
		Description: form.Description,
		IsPrivate:   form.Private || setting.Repository.ForcePrivate,
		IsMirror:    form.Mirror,
		RemoteAddr:  remoteAddr,
	})
	if err == nil {
		log.Trace("Repository migrated [%d]: %s/%s", repo.ID, ctxUser.Name, form.RepoName)
		ctx.Redirect(setting.AppSubUrl + "/" + ctxUser.Name + "/" + form.RepoName)
		return
	}

	if repo != nil {
		if errDelete := models.DeleteRepository(ctxUser.ID, repo.ID); errDelete != nil {
			log.Error(4, "DeleteRepository: %v", errDelete)
		}
	}

	if strings.Contains(err.Error(), "Authentication failed") ||
		strings.Contains(err.Error(), "could not read Username") {
		ctx.Data["Err_Auth"] = true
		ctx.RenderWithErr(ctx.Tr("form.auth_failed", models.HandleCloneUserCredentials(err.Error(), true)), MIGRATE, &form)
		return
	} else if strings.Contains(err.Error(), "fatal:") {
		ctx.Data["Err_CloneAddr"] = true
		ctx.RenderWithErr(ctx.Tr("repo.migrate.failed", models.HandleCloneUserCredentials(err.Error(), true)), MIGRATE, &form)
		return
	}

	handleCreateError(ctx, ctxUser, err, "MigratePost", MIGRATE, &form)
}

func Action(ctx *context.Context) {
	var err error
	switch ctx.Params(":action") {
	case "watch":
		err = models.WatchRepo(ctx.User.ID, ctx.Repo.Repository.ID, true)
	case "unwatch":
		err = models.WatchRepo(ctx.User.ID, ctx.Repo.Repository.ID, false)
	case "star":
		err = models.StarRepo(ctx.User.ID, ctx.Repo.Repository.ID, true)
	case "unstar":
		err = models.StarRepo(ctx.User.ID, ctx.Repo.Repository.ID, false)
	case "desc": // FIXME: this is not used
		if !ctx.Repo.IsOwner() {
			ctx.Error(404)
			return
		}

		ctx.Repo.Repository.Description = ctx.Query("desc")
		ctx.Repo.Repository.Website = ctx.Query("site")
		err = models.UpdateRepository(ctx.Repo.Repository, false)
	}

	if err != nil {
		ctx.Handle(500, fmt.Sprintf("Action (%s)", ctx.Params(":action")), err)
		return
	}

	redirectTo := ctx.Query("redirect_to")
	if len(redirectTo) == 0 {
		redirectTo = ctx.Repo.RepoLink
	}
	ctx.Redirect(redirectTo)
}

func Download(ctx *context.Context) {
	var (
		uri         = ctx.Params("*")
		refName     string
		ext         string
		archivePath string
		archiveType git.ArchiveType
	)

	switch {
	case strings.HasSuffix(uri, ".zip"):
		ext = ".zip"
		archivePath = path.Join(ctx.Repo.GitRepo.Path, "archives/zip")
		archiveType = git.ZIP
	case strings.HasSuffix(uri, ".tar.gz"):
		ext = ".tar.gz"
		archivePath = path.Join(ctx.Repo.GitRepo.Path, "archives/targz")
		archiveType = git.TARGZ
	default:
		log.Trace("Unknown format: %s", uri)
		ctx.Error(404)
		return
	}
	refName = strings.TrimSuffix(uri, ext)

	if !com.IsDir(archivePath) {
		if err := os.MkdirAll(archivePath, os.ModePerm); err != nil {
			ctx.Handle(500, "Download -> os.MkdirAll(archivePath)", err)
			return
		}
	}

	// Get corresponding commit.
	var (
		commit *git.Commit
		err    error
	)
	gitRepo := ctx.Repo.GitRepo
	if gitRepo.IsBranchExist(refName) {
		commit, err = gitRepo.GetBranchCommit(refName)
		if err != nil {
			ctx.Handle(500, "GetBranchCommit", err)
			return
		}
	} else if gitRepo.IsTagExist(refName) {
		commit, err = gitRepo.GetTagCommit(refName)
		if err != nil {
			ctx.Handle(500, "GetTagCommit", err)
			return
		}
	} else if len(refName) == 40 {
		commit, err = gitRepo.GetCommit(refName)
		if err != nil {
			ctx.Handle(404, "GetCommit", nil)
			return
		}
	} else {
		ctx.Handle(404, "Download", nil)
		return
	}

	archivePath = path.Join(archivePath, base.ShortSha(commit.ID.String())+ext)
	if !com.IsFile(archivePath) {
		if err := commit.CreateArchive(archivePath, archiveType); err != nil {
			ctx.Handle(500, "Download -> CreateArchive "+archivePath, err)
			return
		}
	}

	ctx.ServeFile(archivePath, ctx.Repo.Repository.Name+"-"+refName+ext)
}
<commit_msg>routers/repo: allow shorter SHA to download archive (#3834)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package repo

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/Unknwon/com"
	log "gopkg.in/clog.v1"

	"github.com/gogits/git-module"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/auth"
	"github.com/gogits/gogs/modules/base"
	"github.com/gogits/gogs/modules/context"
	"github.com/gogits/gogs/modules/setting"
)

const (
	CREATE  base.TplName = "repo/create"
	MIGRATE base.TplName = "repo/migrate"
)

func MustBeNotBare(ctx *context.Context) {
	if ctx.Repo.Repository.IsBare {
		ctx.Handle(404, "MustBeNotBare", nil)
	}
}

func checkContextUser(ctx *context.Context, uid int64) *models.User {
	orgs, err := models.GetOwnedOrgsByUserIDDesc(ctx.User.ID, "updated_unix")
	if err != nil {
		ctx.Handle(500, "GetOwnedOrgsByUserIDDesc", err)
		return nil
	}
	ctx.Data["Orgs"] = orgs

	// Not equal means current user is an organization.
	if uid == ctx.User.ID || uid == 0 {
		return ctx.User
	}

	org, err := models.GetUserByID(uid)
	if models.IsErrUserNotExist(err) {
		return ctx.User
	}

	if err != nil {
		ctx.Handle(500, "GetUserByID", fmt.Errorf("[%d]: %v", uid, err))
		return nil
	}

	// Check ownership of organization.
	if !org.IsOrganization() || !(ctx.User.IsAdmin || org.IsOwnedBy(ctx.User.ID)) {
		ctx.Error(403)
		return nil
	}
	return org
}

func Create(ctx *context.Context) {
	ctx.Data["Title"] = ctx.Tr("new_repo")

	// Give default value for template to render.
	ctx.Data["Gitignores"] = models.Gitignores
	ctx.Data["Licenses"] = models.Licenses
	ctx.Data["Readmes"] = models.Readmes
	ctx.Data["readme"] = "Default"
	ctx.Data["private"] = ctx.User.LastRepoVisibility
	ctx.Data["IsForcedPrivate"] = setting.Repository.ForcePrivate

	ctxUser := checkContextUser(ctx, ctx.QueryInt64("org"))
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	ctx.HTML(200, CREATE)
}

func handleCreateError(ctx *context.Context, owner *models.User, err error, name string, tpl base.TplName, form interface{}) {
	switch {
	case models.IsErrReachLimitOfRepo(err):
		ctx.RenderWithErr(ctx.Tr("repo.form.reach_limit_of_creation", owner.RepoCreationNum()), tpl, form)
	case models.IsErrRepoAlreadyExist(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("form.repo_name_been_taken"), tpl, form)
	case models.IsErrNameReserved(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("repo.form.name_reserved", err.(models.ErrNameReserved).Name), tpl, form)
	case models.IsErrNamePatternNotAllowed(err):
		ctx.Data["Err_RepoName"] = true
		ctx.RenderWithErr(ctx.Tr("repo.form.name_pattern_not_allowed", err.(models.ErrNamePatternNotAllowed).Pattern), tpl, form)
	default:
		ctx.Handle(500, name, err)
	}
}

func CreatePost(ctx *context.Context, form auth.CreateRepoForm) {
	ctx.Data["Title"] = ctx.Tr("new_repo")

	ctx.Data["Gitignores"] = models.Gitignores
	ctx.Data["Licenses"] = models.Licenses
	ctx.Data["Readmes"] = models.Readmes

	ctxUser := checkContextUser(ctx, form.Uid)
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	if ctx.HasError() {
		ctx.HTML(200, CREATE)
		return
	}

	repo, err := models.CreateRepository(ctxUser, models.CreateRepoOptions{
		Name:        form.RepoName,
		Description: form.Description,
		Gitignores:  form.Gitignores,
		License:     form.License,
		Readme:      form.Readme,
		IsPrivate:   form.Private || setting.Repository.ForcePrivate,
		AutoInit:    form.AutoInit,
	})
	if err == nil {
		log.Trace("Repository created [%d]: %s/%s", repo.ID, ctxUser.Name, repo.Name)
		ctx.Redirect(setting.AppSubUrl + "/" + ctxUser.Name + "/" + repo.Name)
		return
	}

	if repo != nil {
		if errDelete := models.DeleteRepository(ctxUser.ID, repo.ID); errDelete != nil {
			log.Error(4, "DeleteRepository: %v", errDelete)
		}
	}

	handleCreateError(ctx, ctxUser, err, "CreatePost", CREATE, &form)
}

func Migrate(ctx *context.Context) {
	ctx.Data["Title"] = ctx.Tr("new_migrate")
	ctx.Data["private"] = ctx.User.LastRepoVisibility
	ctx.Data["IsForcedPrivate"] = setting.Repository.ForcePrivate
	ctx.Data["mirror"] = ctx.Query("mirror") == "1"

	ctxUser := checkContextUser(ctx, ctx.QueryInt64("org"))
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	ctx.HTML(200, MIGRATE)
}

func MigratePost(ctx *context.Context, form auth.MigrateRepoForm) {
	ctx.Data["Title"] = ctx.Tr("new_migrate")

	ctxUser := checkContextUser(ctx, form.Uid)
	if ctx.Written() {
		return
	}
	ctx.Data["ContextUser"] = ctxUser

	if ctx.HasError() {
		ctx.HTML(200, MIGRATE)
		return
	}

	remoteAddr, err := form.ParseRemoteAddr(ctx.User)
	if err != nil {
		if models.IsErrInvalidCloneAddr(err) {
			ctx.Data["Err_CloneAddr"] = true
			addrErr := err.(models.ErrInvalidCloneAddr)
			switch {
			case addrErr.IsURLError:
				ctx.RenderWithErr(ctx.Tr("form.url_error"), MIGRATE, &form)
			case addrErr.IsPermissionDenied:
				ctx.RenderWithErr(ctx.Tr("repo.migrate.permission_denied"), MIGRATE, &form)
			case addrErr.IsInvalidPath:
				ctx.RenderWithErr(ctx.Tr("repo.migrate.invalid_local_path"), MIGRATE, &form)
			default:
				ctx.Handle(500, "Unknown error", err)
			}
		} else {
			ctx.Handle(500, "ParseRemoteAddr", err)
		}
		return
	}

	repo, err := models.MigrateRepository(ctxUser, models.MigrateRepoOptions{
		Name:        form.RepoName,
		Description: form.Description,
		IsPrivate:   form.Private || setting.Repository.ForcePrivate,
		IsMirror:    form.Mirror,
		RemoteAddr:  remoteAddr,
	})
	if err == nil {
		log.Trace("Repository migrated [%d]: %s/%s", repo.ID, ctxUser.Name, form.RepoName)
		ctx.Redirect(setting.AppSubUrl + "/" + ctxUser.Name + "/" + form.RepoName)
		return
	}

	if repo != nil {
		if errDelete := models.DeleteRepository(ctxUser.ID, repo.ID); errDelete != nil {
			log.Error(4, "DeleteRepository: %v", errDelete)
		}
	}

	if strings.Contains(err.Error(), "Authentication failed") ||
		strings.Contains(err.Error(), "could not read Username") {
		ctx.Data["Err_Auth"] = true
		ctx.RenderWithErr(ctx.Tr("form.auth_failed", models.HandleCloneUserCredentials(err.Error(), true)), MIGRATE, &form)
		return
	} else if strings.Contains(err.Error(), "fatal:") {
		ctx.Data["Err_CloneAddr"] = true
		ctx.RenderWithErr(ctx.Tr("repo.migrate.failed", models.HandleCloneUserCredentials(err.Error(), true)), MIGRATE, &form)
		return
	}

	handleCreateError(ctx, ctxUser, err, "MigratePost", MIGRATE, &form)
}

func Action(ctx *context.Context) {
	var err error
	switch ctx.Params(":action") {
	case "watch":
		err = models.WatchRepo(ctx.User.ID, ctx.Repo.Repository.ID, true)
	case "unwatch":
		err = models.WatchRepo(ctx.User.ID, ctx.Repo.Repository.ID, false)
	case "star":
		err = models.StarRepo(ctx.User.ID, ctx.Repo.Repository.ID, true)
	case "unstar":
		err = models.StarRepo(ctx.User.ID, ctx.Repo.Repository.ID, false)
	case "desc": // FIXME: this is not used
		if !ctx.Repo.IsOwner() {
			ctx.Error(404)
			return
		}

		ctx.Repo.Repository.Description = ctx.Query("desc")
		ctx.Repo.Repository.Website = ctx.Query("site")
		err = models.UpdateRepository(ctx.Repo.Repository, false)
	}

	if err != nil {
		ctx.Handle(500, fmt.Sprintf("Action (%s)", ctx.Params(":action")), err)
		return
	}

	redirectTo := ctx.Query("redirect_to")
	if len(redirectTo) == 0 {
		redirectTo = ctx.Repo.RepoLink
	}
	ctx.Redirect(redirectTo)
}

func Download(ctx *context.Context) {
	var (
		uri         = ctx.Params("*")
		refName     string
		ext         string
		archivePath string
		archiveType git.ArchiveType
	)

	switch {
	case strings.HasSuffix(uri, ".zip"):
		ext = ".zip"
		archivePath = path.Join(ctx.Repo.GitRepo.Path, "archives/zip")
		archiveType = git.ZIP
	case strings.HasSuffix(uri, ".tar.gz"):
		ext = ".tar.gz"
		archivePath = path.Join(ctx.Repo.GitRepo.Path, "archives/targz")
		archiveType = git.TARGZ
	default:
		log.Trace("Unknown format: %s", uri)
		ctx.Error(404)
		return
	}
	refName = strings.TrimSuffix(uri, ext)

	if !com.IsDir(archivePath) {
		if err := os.MkdirAll(archivePath, os.ModePerm); err != nil {
			ctx.Handle(500, "Download -> os.MkdirAll(archivePath)", err)
			return
		}
	}

	// Get corresponding commit.
	var (
		commit *git.Commit
		err    error
	)
	gitRepo := ctx.Repo.GitRepo
	if gitRepo.IsBranchExist(refName) {
		commit, err = gitRepo.GetBranchCommit(refName)
		if err != nil {
			ctx.Handle(500, "GetBranchCommit", err)
			return
		}
	} else if gitRepo.IsTagExist(refName) {
		commit, err = gitRepo.GetTagCommit(refName)
		if err != nil {
			ctx.Handle(500, "GetTagCommit", err)
			return
		}
	} else if len(refName) >= 7 && len(refName) <= 40 {
		commit, err = gitRepo.GetCommit(refName)
		if err != nil {
			ctx.NotFound()
			return
		}
	} else {
		ctx.NotFound()
		return
	}

	archivePath = path.Join(archivePath, base.ShortSha(commit.ID.String())+ext)
	if !com.IsFile(archivePath) {
		if err := commit.CreateArchive(archivePath, archiveType); err != nil {
			ctx.Handle(500, "Download -> CreateArchive "+archivePath, err)
			return
		}
	}

	ctx.ServeFile(archivePath, ctx.Repo.Repository.Name+"-"+refName+ext)
}
<|endoftext|>


<commit_before>package zipfs

import (
	"archive/zip"
	"io"
	"os"
	"path/filepath"
	"syscall"

	"github.com/spf13/afero"
)

type File struct {
	fs            *Fs
	zipfile       *zip.File
	reader        io.ReadCloser
	offset        int64
	isdir, closed bool
	buf           []byte
}

func (f *File) fillBuffer(offset int64) (err error) {
	if f.reader == nil {
		if f.reader, err = f.zipfile.Open(); err != nil {
			return
		}
	}
	if offset > int64(f.zipfile.UncompressedSize64) {
		offset = int64(f.zipfile.UncompressedSize64)
		err = io.EOF
	}
	if len(f.buf) >= int(offset) {
		return
	}
	buf := make([]byte, int(offset)-len(f.buf))
	if n, readErr := io.ReadFull(f.reader, buf); n > 0 {
		f.buf = append(f.buf, buf[:n]...)
	} else if readErr != nil {
		err = readErr
	}
	return
}

func (f *File) Close() (err error) {
	f.zipfile = nil
	f.closed = true
	f.buf = nil
	if f.reader != nil {
		err = f.reader.Close()
		f.reader = nil
	}
	return
}

func (f *File) Read(p []byte) (n int, err error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	err = f.fillBuffer(f.offset + int64(len(p)))
	n = copy(p, f.buf[f.offset:])
	f.offset += int64(len(p))
	return
}

func (f *File) ReadAt(p []byte, off int64) (n int, err error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	err = f.fillBuffer(off + int64(len(p)))
	n = copy(p, f.buf[int(off):])
	return
}

func (f *File) Seek(offset int64, whence int) (int64, error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	switch whence {
	case os.SEEK_SET:
	case os.SEEK_CUR:
		offset += f.offset
	case os.SEEK_END:
		offset += int64(f.zipfile.UncompressedSize64)
	default:
		return 0, syscall.EINVAL
	}
	if offset < 0 || offset > int64(f.zipfile.UncompressedSize64) {
		return 0, afero.ErrOutOfRange
	}
	f.offset = offset
	return offset, nil
}

func (f *File) Write(p []byte) (n int, err error) { return 0, syscall.EPERM }

func (f *File) WriteAt(p []byte, off int64) (n int, err error) { return 0, syscall.EPERM }

func (f *File) Name() string {
	if f.zipfile == nil {
		return string(filepath.Separator)
	}
	return filepath.Join(splitpath(f.zipfile.Name))
}

func (f *File) getDirEntries() (map[string]*zip.File, error) {
	if !f.isdir {
		return nil, syscall.ENOTDIR
	}
	name := f.Name()
	entries, ok := f.fs.files[name]
	if !ok {
		return nil, &os.PathError{Op: "readdir", Path: name, Err: syscall.ENOENT}
	}
	return entries, nil
}

func (f *File) Readdir(count int) (fi []os.FileInfo, err error) {
	zipfiles, err := f.getDirEntries()
	if err != nil {
		return nil, err
	}
	for _, zipfile := range zipfiles {
		fi = append(fi, zipfile.FileInfo())
		if count > 0 && len(fi) >= count {
			break
		}
	}
	return
}

func (f *File) Readdirnames(count int) (names []string, err error) {
	zipfiles, err := f.getDirEntries()
	if err != nil {
		return nil, err
	}
	for filename := range zipfiles {
		names = append(names, filename)
		if count > 0 && len(names) >= count {
			break
		}
	}
	return
}

func (f *File) Stat() (os.FileInfo, error) {
	if f.zipfile == nil {
		return &pseudoRoot{}, nil
	}
	return f.zipfile.FileInfo(), nil
}

func (f *File) Sync() error { return nil }

func (f *File) Truncate(size int64) error { return syscall.EPERM }

func (f *File) WriteString(s string) (ret int, err error) { return 0, syscall.EPERM }
<commit_msg>Fix panic when not filling up zipfs's read buffer<commit_after>package zipfs

import (
	"archive/zip"
	"io"
	"os"
	"path/filepath"
	"syscall"

	"github.com/spf13/afero"
)

type File struct {
	fs            *Fs
	zipfile       *zip.File
	reader        io.ReadCloser
	offset        int64
	isdir, closed bool
	buf           []byte
}

func (f *File) fillBuffer(offset int64) (err error) {
	if f.reader == nil {
		if f.reader, err = f.zipfile.Open(); err != nil {
			return
		}
	}
	if offset > int64(f.zipfile.UncompressedSize64) {
		offset = int64(f.zipfile.UncompressedSize64)
		err = io.EOF
	}
	if len(f.buf) >= int(offset) {
		return
	}
	buf := make([]byte, int(offset)-len(f.buf))
	if n, readErr := io.ReadFull(f.reader, buf); n > 0 {
		f.buf = append(f.buf, buf[:n]...)
	} else if readErr != nil {
		err = readErr
	}
	return
}

func (f *File) Close() (err error) {
	f.zipfile = nil
	f.closed = true
	f.buf = nil
	if f.reader != nil {
		err = f.reader.Close()
		f.reader = nil
	}
	return
}

func (f *File) Read(p []byte) (n int, err error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	err = f.fillBuffer(f.offset + int64(len(p)))
	n = copy(p, f.buf[f.offset:])
	f.offset += int64(n)
	return
}

func (f *File) ReadAt(p []byte, off int64) (n int, err error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	err = f.fillBuffer(off + int64(len(p)))
	n = copy(p, f.buf[int(off):])
	return
}

func (f *File) Seek(offset int64, whence int) (int64, error) {
	if f.isdir {
		return 0, syscall.EISDIR
	}
	if f.closed {
		return 0, afero.ErrFileClosed
	}
	switch whence {
	case os.SEEK_SET:
	case os.SEEK_CUR:
		offset += f.offset
	case os.SEEK_END:
		offset += int64(f.zipfile.UncompressedSize64)
	default:
		return 0, syscall.EINVAL
	}
	if offset < 0 || offset > int64(f.zipfile.UncompressedSize64) {
		return 0, afero.ErrOutOfRange
	}
	f.offset = offset
	return offset, nil
}

func (f *File) Write(p []byte) (n int, err error) { return 0, syscall.EPERM }

func (f *File) WriteAt(p []byte, off int64) (n int, err error) { return 0, syscall.EPERM }

func (f *File) Name() string {
	if f.zipfile == nil {
		return string(filepath.Separator)
	}
	return filepath.Join(splitpath(f.zipfile.Name))
}

func (f *File) getDirEntries() (map[string]*zip.File, error) {
	if !f.isdir {
		return nil, syscall.ENOTDIR
	}
	name := f.Name()
	entries, ok := f.fs.files[name]
	if !ok {
		return nil, &os.PathError{Op: "readdir", Path: name, Err: syscall.ENOENT}
	}
	return entries, nil
}

func (f *File) Readdir(count int) (fi []os.FileInfo, err error) {
	zipfiles, err := f.getDirEntries()
	if err != nil {
		return nil, err
	}
	for _, zipfile := range zipfiles {
		fi = append(fi, zipfile.FileInfo())
		if count > 0 && len(fi) >= count {
			break
		}
	}
	return
}

func (f *File) Readdirnames(count int) (names []string, err error) {
	zipfiles, err := f.getDirEntries()
	if err != nil {
		return nil, err
	}
	for filename := range zipfiles {
		names = append(names, filename)
		if count > 0 && len(names) >= count {
			break
		}
	}
	return
}

func (f *File) Stat() (os.FileInfo, error) {
	if f.zipfile == nil {
		return &pseudoRoot{}, nil
	}
	return f.zipfile.FileInfo(), nil
}

func (f *File) Sync() error { return nil }

func (f *File) Truncate(size int64) error { return syscall.EPERM }

func (f *File) WriteString(s string) (ret int, err error) { return 0, syscall.EPERM }
<|endoftext|>


<commit_before>package managesystemagent

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/rancher/fleet/pkg/apis/fleet.cattle.io/v1alpha1"
	rancherv1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	namespaces "github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/wrangler"
	upgradev1 "github.com/rancher/system-upgrade-controller/pkg/apis/upgrade.cattle.io/v1"
	"github.com/rancher/wrangler/pkg/name"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

var (
	workerSelector = metav1.LabelSelector{
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "node-role.kubernetes.io/etcd",
				Operator: metav1.LabelSelectorOpNotIn,
				Values:   []string{"true"},
			},
			{
				Key:      "node-role.kubernetes.io/control-plane",
				Operator: metav1.LabelSelectorOpNotIn,
				Values:   []string{"true"},
			},
			{
				Key:      "beta.kubernetes.io/os",
				Operator: metav1.LabelSelectorOpNotIn,
				Values:   []string{"windows"},
			},
		},
	}
	controlPlaneSelector = metav1.LabelSelector{
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "node-role.kubernetes.io/etcd",
				Operator: metav1.LabelSelectorOpNotIn,
				Values:   []string{"true"},
			},
			{
				Key:      "node-role.kubernetes.io/control-plane",
				Operator: metav1.LabelSelectorOpIn,
				Values:   []string{"true"},
			},
		},
	}
	etcdSelector = metav1.LabelSelector{
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "node-role.kubernetes.io/etcd",
				Operator: metav1.LabelSelectorOpIn,
				Values:   []string{"true"},
			},
			{
				Key:      "node-role.kubernetes.io/control-plane",
				Operator: metav1.LabelSelectorOpNotIn,
				Values:   []string{"true"},
			},
		},
	}
	controlPlaneAndEtcdSelector = metav1.LabelSelector{
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "node-role.kubernetes.io/etcd",
				Operator: metav1.LabelSelectorOpIn,
				Values:   []string{"true"},
			},
			{
				Key:      "node-role.kubernetes.io/control-plane",
				Operator: metav1.LabelSelectorOpIn,
				Values:   []string{"true"},
			},
		},
	}
)

type handler struct{}

func Register(ctx context.Context, clients *wrangler.Context) {
	h := &handler{}
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Fleet.Bundle(),
				clients.Provisioning.Cluster()),
		"", "manage-system-agent", h.OnChange, nil)
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Mgmt.ManagedChart(),
				clients.Provisioning.Cluster()),
		"", "manage-system-upgrade-controller", h.OnChangeInstallSUC, nil)
}

func (h *handler) OnChange(cluster *rancherv1.Cluster, status rancherv1.ClusterStatus) ([]runtime.Object, rancherv1.ClusterStatus, error) {
	if cluster.Spec.RKEConfig == nil || settings.SystemAgentUpgradeImage.Get() == "" {
		return nil, status, nil
	}

	bundle := &v1alpha1.Bundle{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: cluster.Namespace,
			Name:      name.SafeConcatName(cluster.Name, "managed", "system", "agent"),
		},
		Spec: v1alpha1.BundleSpec{
			BundleDeploymentOptions: v1alpha1.BundleDeploymentOptions{
				DefaultNamespace: namespaces.System,
			},
			Resources: []v1alpha1.BundleResource{
				{
					Name:    "cp.yaml",
					Content: installer("cp", "false", "false", "true", cluster.Spec.AgentEnvVars, &controlPlaneSelector),
				},
				{
					Name:    "etcd.yaml",
					Content: installer("etcd", "false", "true", "false", cluster.Spec.AgentEnvVars, &etcdSelector),
				},
				{
					Name:    "cp-and-etcd.yaml",
					Content: installer("cp-and-etcd", "false", "true", "true", cluster.Spec.AgentEnvVars, &controlPlaneAndEtcdSelector),
				},
				{
					Name:    "worker.yaml",
					Content: installer("worker", "true", "false", "false", cluster.Spec.AgentEnvVars, &workerSelector),
				},
			},
			Targets: []v1alpha1.BundleTarget{
				{
					ClusterName: cluster.Name,
					ClusterSelector: &metav1.LabelSelector{
						MatchExpressions: []metav1.LabelSelectorRequirement{
							{
								Key:      "provisioning.cattle.io/unmanaged-system-agent",
								Operator: metav1.LabelSelectorOpDoesNotExist,
							},
						},
					},
				},
			},
		},
	}

	return []runtime.Object{
		bundle,
	}, status, nil
}

func installer(name, worker, etcd, controlPlane string, envs []corev1.EnvVar, selector *metav1.LabelSelector) string {
	image := strings.SplitN(settings.SystemAgentUpgradeImage.Get(), ":", 2)
	version := "latest"
	if len(image) == 2 {
		version = image[1]
	}

	plan := &upgradev1.Plan{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Plan",
			APIVersion: "upgrade.cattle.io/v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "system-agent-upgrader-" + name,
			Namespace: namespaces.System,
		},
		Spec: upgradev1.PlanSpec{
			Concurrency: 10,
			Version:     version,
			Tolerations: []corev1.Toleration{{
				Operator: corev1.TolerationOpExists,
			}},
			NodeSelector: selector,
			Upgrade: &upgradev1.ContainerSpec{
				Image:   settings.PrefixPrivateRegistry(image[0]),
				Command: nil,
				Args:    nil,
				Env: append(envs, []corev1.EnvVar{
					{
						Name:  "CATTLE_ROLE_WORKER",
						Value: worker,
					},
					{
						Name:  "CATTLE_ROLE_ETCD",
						Value: etcd,
					},
					{
						Name:  "CATTLE_ROLE_CONTROL_PLANE",
						Value: controlPlane,
					},
				}...),
				EnvFrom: []corev1.EnvFromSource{{
					SecretRef: &corev1.SecretEnvSource{
						LocalObjectReference: corev1.LocalObjectReference{
							Name: "steve-aggregation",
						},
					},
				}},
			},
		},
	}

	file, err := json.Marshal(plan)
	if err != nil {
		panic(err)
	}
	return string(file)
}
<commit_msg>Dynamically determine node roles for system-agent suc plan<commit_after>package managesystemagent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/rancher/fleet/pkg/apis/fleet.cattle.io/v1alpha1"
	rancherv1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	v3 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	namespaces "github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/systemtemplate"
	"github.com/rancher/rancher/pkg/wrangler"
	upgradev1 "github.com/rancher/system-upgrade-controller/pkg/apis/upgrade.cattle.io/v1"
	"github.com/rancher/wrangler/pkg/generic"
	"github.com/rancher/wrangler/pkg/gvk"
	"github.com/rancher/wrangler/pkg/name"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

type handler struct {
	clusterRegistrationTokens v3.ClusterRegistrationTokenCache
}

func Register(ctx context.Context, clients *wrangler.Context) {
	h := &handler{
		clusterRegistrationTokens: clients.Mgmt.ClusterRegistrationToken().Cache(),
	}
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Fleet.Bundle(),
				clients.Provisioning.Cluster(),
				clients.Core.Secret(),
				clients.RBAC.RoleBinding(),
				clients.RBAC.Role()),
		"", "manage-system-agent", h.OnChange, &generic.GeneratingHandlerOptions{
			AllowCrossNamespace: true,
		})
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Mgmt.ManagedChart(),
				clients.Provisioning.Cluster()),
		"", "manage-system-upgrade-controller", h.OnChangeInstallSUC, nil)
}

func (h *handler) OnChange(cluster *rancherv1.Cluster, status rancherv1.ClusterStatus) ([]runtime.Object, rancherv1.ClusterStatus, error) {
	if cluster.Spec.RKEConfig == nil || settings.SystemAgentUpgradeImage.Get() == "" {
		return nil, status, nil
	}

	var (
		secretName = "steve-aggregation"
		result     []runtime.Object
	)

	if cluster.Status.ClusterName == "local" && cluster.Namespace == "fleet-local" {
		secretName += "-local-"

		token, err := h.clusterRegistrationTokens.Get(cluster.Status.ClusterName, "default-token")
		if err != nil {
			return nil, status, err
		}
		if token.Status.Token == "" {
			return nil, status, fmt.Errorf("token not yet generated for %s/%s", token.Namespace, token.Name)
		}

		digest := sha256.New()
		digest.Write([]byte(settings.InternalServerURL.Get()))
		digest.Write([]byte(token.Status.Token))
		digest.Write([]byte(systemtemplate.InternalCAChecksum()))
		d := digest.Sum(nil)
		secretName += hex.EncodeToString(d[:])[:12]

		result = append(result, &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      secretName,
				Namespace: namespaces.System,
			},
			Data: map[string][]byte{
				"CATTLE_SERVER":      []byte(settings.InternalServerURL.Get()),
				"CATTLE_TOKEN":       []byte(token.Status.Token),
				"CATTLE_CA_CHECKSUM": []byte(systemtemplate.InternalCAChecksum()),
			},
		})
	}

	resources, err := ToResources(installer(len(cluster.Spec.RKEConfig.NodeConfig) == 0, secretName))
	if err != nil {
		return nil, status, err
	}

	result = append(result, &v1alpha1.Bundle{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: cluster.Namespace,
			Name:      name.SafeConcatName(cluster.Name, "managed", "system", "agent"),
		},
		Spec: v1alpha1.BundleSpec{
			BundleDeploymentOptions: v1alpha1.BundleDeploymentOptions{
				DefaultNamespace: namespaces.System,
			},
			Resources: resources,
			Targets: []v1alpha1.BundleTarget{
				{
					ClusterName: cluster.Name,
					ClusterSelector: &metav1.LabelSelector{
						MatchExpressions: []metav1.LabelSelectorRequirement{
							{
								Key:      "provisioning.cattle.io/unmanaged-system-agent",
								Operator: metav1.LabelSelectorOpDoesNotExist,
							},
						},
					},
				},
			},
		},
	})

	return result, status, nil
}

func installer(allWorkers bool, secretName string) []runtime.Object {
	image := strings.SplitN(settings.SystemAgentUpgradeImage.Get(), ":", 2)
	version := "latest"
	if len(image) == 2 {
		version = image[1]
	}

	env := []corev1.EnvVar{{
		Name:  "CATTLE_SERVER_QUERY",
		Value: "?internal=true",
	}}

	if allWorkers {
		env = append(env, corev1.EnvVar{
			Name:  "CATTLE_ROLE_WORKER",
			Value: "true",
		})
	}

	return []runtime.Object{
		&upgradev1.Plan{
			TypeMeta: metav1.TypeMeta{
				Kind:       "Plan",
				APIVersion: "upgrade.cattle.io/v1",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			},
			Spec: upgradev1.PlanSpec{
				Concurrency: 10,
				Version:     version,
				Tolerations: []corev1.Toleration{{
					Operator: corev1.TolerationOpExists,
				}},
				NodeSelector:       &metav1.LabelSelector{},
				ServiceAccountName: "system-agent-upgrader",
				Upgrade: &upgradev1.ContainerSpec{
					Image: settings.PrefixPrivateRegistry(image[0]),
					Env:   env,
					EnvFrom: []corev1.EnvFromSource{{
						SecretRef: &corev1.SecretEnvSource{
							LocalObjectReference: corev1.LocalObjectReference{
								Name: secretName,
							},
						},
					}},
				},
			},
		},
		&corev1.ServiceAccount{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			},
		},
		&rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Rules: []rbacv1.PolicyRule{{
				Verbs:     []string{"get"},
				APIGroups: []string{""},
				Resources: []string{"nodes"},
			}},
		},
		&rbacv1.ClusterRoleBinding{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Subjects: []rbacv1.Subject{{
				Kind:      "ServiceAccount",
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			}},
			RoleRef: rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "ClusterRole",
				Name:     "system-agent-upgrader",
			},
		},
	}
}

func ToResources(objs []runtime.Object) (result []v1alpha1.BundleResource, err error) {
	for _, obj := range objs {
		obj = obj.DeepCopyObject()
		if err := gvk.Set(obj); err != nil {
			return nil, fmt.Errorf("failed to set gvk: %w", err)
		}

		typeMeta, err := meta.TypeAccessor(obj)
		if err != nil {
			return nil, err
		}

		meta, err := meta.Accessor(obj)
		if err != nil {
			return nil, err
		}

		data, err := json.Marshal(obj)
		if err != nil {
			return nil, err
		}

		digest := sha256.Sum256(data)
		filename := name.SafeConcatName(typeMeta.GetKind(), meta.GetNamespace(), meta.GetName(), hex.EncodeToString(digest[:])[:12]) + ".yaml"
		result = append(result, v1alpha1.BundleResource{
			Name:    filename,
			Content: string(data),
		})
	}
	return
}
<|endoftext|>


<commit_before>package persist

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"

	"github.com/dancannon/gorethink"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
	"google.golang.org/grpc"
)

// A Table is a rethinkdb table name.
type Table string

// A PrimaryKey is a rethinkdb primary key identifier.
type PrimaryKey string

// An Index is a rethinkdb index.
type Index string

const (
	repoTable   Table = "Repos"
	branchTable Table = "Branches"

	commitTable       Table = "Commits"
	commitClocksIndex Index = "CommitClocksIndex"
	commitRepoIndex   Index = "CommitRepoIndex"

	diffTable         Table = "Diffs"
	diffCommitIDIndex Index = "DiffCommitIDIndex"
	diffPathIndex     Index = "DiffPathIndex"

	connectTimeoutSeconds = 5
)

var (
	tables = []Table{
		repoTable,
		branchTable,
		commitTable,
		diffTable,
	}

	tableToTableCreateOpts = map[Table][]gorethink.TableCreateOpts{
		repoTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "Name",
			},
		},
		branchTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		commitTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		diffTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
	}
)

type driver struct {
	blockAddress string
	blockClient  pfs.BlockAPIClient

	dbAddress string
	dbName    string
	dbClient  *gorethink.Session
}

func NewDriver(blockAddress string, dbAddress string, dbName string) (drive.Driver, error) {
	clientConn, err := grpc.Dial(blockAddress, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	dbClient, err := dbConnect(dbAddress)
	if err != nil {
		return nil, err
	}

	return &driver{
		blockAddress: blockAddress,
		blockClient:  pfs.NewBlockAPIClient(clientConn),
		dbAddress:    dbAddress,
		dbName:       dbName,
		dbClient:     dbClient,
	}, nil
}

func InitDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBCreate(databaseName).RunWrite(session); err != nil {
		return err
	}

	// Create tables
	for _, table := range tables {
		tableCreateOpts := tableToTableCreateOpts[table]
		if _, err := gorethink.DB(databaseName).TableCreate(table, tableCreateOpts...).RunWrite(session); err != nil {
			return err
		}
	}

	// Create indexes
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexCreate(commitClocksIndex).RunWrite(session); err != nil {
		return err
	}
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexCreate(commitRepoIndex).RunWrite(session); err != nil {
		return err
	}
	if _, err := gorethink.DB(databaseName).Table(diffTable).IndexCreate(diffCommitIDIndex).RunWrite(session); err != nil {
		return err
	}
	if _, err := gorethink.DB(databaseName).Table(diffTable).IndexCreate(diffPathIndex).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func RemoveDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBDrop(databaseName).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func dbConnect(address string) (*gorethink.Session, error) {
	return gorethink.Connect(gorethink.ConnectOpts{
		Address: address,
		Timeout: connectTimeoutSeconds * time.Second,
	})
}

func validateRepoName(name string) error {
	match, _ := regexp.MatchString("^[a-zA-Z0-9_]+$", name)

	if !match {
		return fmt.Errorf("repo name (%v) invalid: only alphanumeric and underscore characters allowed", name)
	}

	return nil
}

func (d *driver) getTerm(table Table) gorethink.Term {
	return gorethink.DB(d.dbName).Table(table)
}

func (d *driver) CreateRepo(repo *pfs.Repo, created *google_protobuf.Timestamp,
	provenance []*pfs.Repo, shards map[uint64]bool) error {

	err := validateRepoName(repo.Name)
	if err != nil {
		return err
	}

	_, err = d.getTerm(repoTable).Insert(&Repo{
		Name:    repo.Name,
		Created: created,
	}).RunWrite(d.dbClient)
	return err
}

func (d *driver) InspectRepo(repo *pfs.Repo, shards map[uint64]bool) (repoInfo *pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Get(repo.Name).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	rawRepo := &Repo{}
	cursor.Next(rawRepo)
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	repoInfo = &pfs.RepoInfo{
		Repo:    &pfs.Repo{rawRepo.Name},
		Created: rawRepo.Created,
	}
	return repoInfo, nil
}

func (d *driver) ListRepo(provenance []*pfs.Repo, shards map[uint64]bool) (repoInfos []*pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		repo := &Repo{}
		if !cursor.Next(repo) {
			break
		}
		repoInfos = append(repoInfos, &pfs.RepoInfo{
			Repo:    &pfs.Repo{repo.Name},
			Created: repo.Created,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return repoInfos, nil
}

func (d *driver) DeleteRepo(repo *pfs.Repo, shards map[uint64]bool, force bool) error {
	_, err := d.getTerm(repoTable).Get(repo.Name).Delete().RunWrite(d.dbClient)
	return err
}

func (d *driver) StartCommit(repo *pfs.Repo, commitID string, parentID string, branch string,
	started *google_protobuf.Timestamp, provenance []*pfs.Commit, shards map[uint64]bool) error {
	if commitID == "" {
		commitID = uuid.NewWithoutDashes()
	}
	_, err := d.getTerm(commitTable).Insert(&Commit{
		ID:         commitID,
		Repo:       repo.Name,
		Started:    prototime.TimeToTimestamp(time.Now()),
		Provenance: []string{parentID}, // Incorrect. Need all ancestors
	}).RunWrite(d.dbClient)

	return err
}

// FinishCommit blocks until its parent has been finished/cancelled
func (d *driver) FinishCommit(commit *pfs.Commit, finished *google_protobuf.Timestamp, cancel bool, shards map[uint64]bool) error {
	_, err := d.getTerm(commitTable).Get(commit.ID).Update(
		map[string]interface{}{
			"Finished": finished,
		},
	).RunWrite(d.dbClient)

	return err
}

func (d *driver) InspectCommit(commit *pfs.Commit, shards map[uint64]bool) (commitInfo *pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Get(commit.ID).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()

	rawCommit := &Commit{}
	cursor.Next(rawCommit)
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	return rawCommitToCommitInfo(rawCommit), nil
}

func rawCommitToCommitInfo(rawCommit *Commit) *pfs.CommitInfo {
	commitType := pfs.CommitType_COMMIT_TYPE_READ
	if rawCommit.Finished == nil {
		commitType = pfs.CommitType_COMMIT_TYPE_WRITE
	}
	return &pfs.CommitInfo{
		Commit: &pfs.Commit{
			Repo: &pfs.Repo{rawCommit.Repo},
			ID:   rawCommit.ID,
		},
		Started:    rawCommit.Started,
		Finished:   rawCommit.Finished,
		CommitType: commitType,
	}
}

func (d *driver) ListCommit(repos []*pfs.Repo, commitType pfs.CommitType, fromCommit []*pfs.Commit,
	provenance []*pfs.Commit, all bool, shards map[uint64]bool) (commitInfos []*pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Filter(func(commit gorethink.Term) gorethink.Term {
		var predicates []interface{}
		for _, repo := range repos {
			predicates = append(predicates, commit.Field("Repo").Eq(repo.Name))
		}
		return gorethink.Or(predicates...)
	}).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		rawCommit := &Commit{}
		if !cursor.Next(rawCommit) {
			break
		}
		commitInfos = append(commitInfos, rawCommitToCommitInfo(rawCommit))
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return commitInfos, nil
}

func (d *driver) ListBranch(repo *pfs.Repo, shards map[uint64]bool) ([]*pfs.CommitInfo, error) {
	return nil, nil
}

func (d *driver) DeleteCommit(commit *pfs.Commit, shards map[uint64]bool) error {
	return nil
}

func (d *driver) PutFile(file *pfs.File, handle string,
	delimiter pfs.Delimiter, shard uint64, reader io.Reader) (retErr error) {
	return nil
}

func (d *driver) MakeDirectory(file *pfs.File, shard uint64) (retErr error) {
	return nil
}

func (d *driver) GetFile(file *pfs.File, filterShard *pfs.Shard, offset int64,
	size int64, from *pfs.Commit, shard uint64, unsafe bool, handle string) (io.ReadCloser, error) {
	return nil, nil
}

func (d *driver) InspectFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, unsafe bool, handle string) (*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) ListFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, recurse bool, unsafe bool, handle string) ([]*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) DeleteFile(file *pfs.File, shard uint64, unsafe bool, handle string) error {
	return nil
}

func (d *driver) DeleteAll(shards map[uint64]bool) error {
	return nil
}

func (d *driver) AddShard(shard uint64) error {
	return nil
}

func (d *driver) DeleteShard(shard uint64) error {
	return nil
}

func (d *driver) Dump() {
}
<commit_msg>Create commitBranchIndex<commit_after>package persist

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"

	"github.com/dancannon/gorethink"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
	"google.golang.org/grpc"
)

// A Table is a rethinkdb table name.
type Table string

// A PrimaryKey is a rethinkdb primary key identifier.
type PrimaryKey string

// An Index is a rethinkdb index.
type Index string

const (
	repoTable   Table = "Repos"
	branchTable Table = "Branches"

	commitTable Table = "Commits"
	// commitBranchIndex maps commits to branches
	commitBranchIndex Index = "CommitBranchIndex"

	diffTable Table = "Diffs"

	connectTimeoutSeconds = 5
)

var (
	tables = []Table{
		repoTable,
		branchTable,
		commitTable,
		diffTable,
	}

	tableToTableCreateOpts = map[Table][]gorethink.TableCreateOpts{
		repoTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "Name",
			},
		},
		branchTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		commitTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		diffTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
	}
)

type driver struct {
	blockAddress string
	blockClient  pfs.BlockAPIClient

	dbAddress string
	dbName    string
	dbClient  *gorethink.Session
}

func NewDriver(blockAddress string, dbAddress string, dbName string) (drive.Driver, error) {
	clientConn, err := grpc.Dial(blockAddress, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	dbClient, err := dbConnect(dbAddress)
	if err != nil {
		return nil, err
	}

	return &driver{
		blockAddress: blockAddress,
		blockClient:  pfs.NewBlockAPIClient(clientConn),
		dbAddress:    dbAddress,
		dbName:       dbName,
		dbClient:     dbClient,
	}, nil
}

func InitDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBCreate(databaseName).RunWrite(session); err != nil {
		return err
	}

	// Create tables
	for _, table := range tables {
		tableCreateOpts := tableToTableCreateOpts[table]
		if _, err := gorethink.DB(databaseName).TableCreate(table, tableCreateOpts...).RunWrite(session); err != nil {
			return err
		}
	}

	// Create indexes
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexCreateFunc(commitBranchIndex, func(row gorethink.Term) interface{} {
		return row.Field("BranchClocks").Map(func(branchClock gorethink.Term) interface{} {
			lastClock := branchClock.Field("Clocks").Nth(-1)
			return []interface{}{
				lastClock.Field("Branch"),
				lastClock.Field("Clock"),
			}
		})
	}, gorethink.IndexCreateOpts{
		Multi: true,
	}).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func RemoveDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBDrop(databaseName).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func dbConnect(address string) (*gorethink.Session, error) {
	return gorethink.Connect(gorethink.ConnectOpts{
		Address: address,
		Timeout: connectTimeoutSeconds * time.Second,
	})
}

func validateRepoName(name string) error {
	match, _ := regexp.MatchString("^[a-zA-Z0-9_]+$", name)

	if !match {
		return fmt.Errorf("repo name (%v) invalid: only alphanumeric and underscore characters allowed", name)
	}

	return nil
}

func (d *driver) getTerm(table Table) gorethink.Term {
	return gorethink.DB(d.dbName).Table(table)
}

func (d *driver) CreateRepo(repo *pfs.Repo, created *google_protobuf.Timestamp,
	provenance []*pfs.Repo, shards map[uint64]bool) error {

	err := validateRepoName(repo.Name)
	if err != nil {
		return err
	}

	_, err = d.getTerm(repoTable).Insert(&Repo{
		Name:    repo.Name,
		Created: created,
	}).RunWrite(d.dbClient)
	return err
}

func (d *driver) InspectRepo(repo *pfs.Repo, shards map[uint64]bool) (repoInfo *pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Get(repo.Name).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	rawRepo := &Repo{}
	cursor.Next(rawRepo)
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	repoInfo = &pfs.RepoInfo{
		Repo:    &pfs.Repo{rawRepo.Name},
		Created: rawRepo.Created,
	}
	return repoInfo, nil
}

func (d *driver) ListRepo(provenance []*pfs.Repo, shards map[uint64]bool) (repoInfos []*pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		repo := &Repo{}
		if !cursor.Next(repo) {
			break
		}
		repoInfos = append(repoInfos, &pfs.RepoInfo{
			Repo:    &pfs.Repo{repo.Name},
			Created: repo.Created,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return repoInfos, nil
}

func (d *driver) DeleteRepo(repo *pfs.Repo, shards map[uint64]bool, force bool) error {
	_, err := d.getTerm(repoTable).Get(repo.Name).Delete().RunWrite(d.dbClient)
	return err
}

func (d *driver) StartCommit(repo *pfs.Repo, commitID string, parentID string, branch string,
	started *google_protobuf.Timestamp, provenance []*pfs.Commit, shards map[uint64]bool) error {
	if commitID == "" {
		commitID = uuid.NewWithoutDashes()
	}
	_, err := d.getTerm(commitTable).Insert(&Commit{
		ID:         commitID,
		Repo:       repo.Name,
		Started:    prototime.TimeToTimestamp(time.Now()),
		Provenance: []string{parentID}, // Incorrect. Need all ancestors
	}).RunWrite(d.dbClient)

	return err
}

// FinishCommit blocks until its parent has been finished/cancelled
func (d *driver) FinishCommit(commit *pfs.Commit, finished *google_protobuf.Timestamp, cancel bool, shards map[uint64]bool) error {
	_, err := d.getTerm(commitTable).Get(commit.ID).Update(
		map[string]interface{}{
			"Finished": finished,
		},
	).RunWrite(d.dbClient)

	return err
}

func (d *driver) InspectCommit(commit *pfs.Commit, shards map[uint64]bool) (commitInfo *pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Get(commit.ID).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()

	rawCommit := &Commit{}
	cursor.Next(rawCommit)
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	return rawCommitToCommitInfo(rawCommit), nil
}

func rawCommitToCommitInfo(rawCommit *Commit) *pfs.CommitInfo {
	commitType := pfs.CommitType_COMMIT_TYPE_READ
	if rawCommit.Finished == nil {
		commitType = pfs.CommitType_COMMIT_TYPE_WRITE
	}
	return &pfs.CommitInfo{
		Commit: &pfs.Commit{
			Repo: &pfs.Repo{rawCommit.Repo},
			ID:   rawCommit.ID,
		},
		Started:    rawCommit.Started,
		Finished:   rawCommit.Finished,
		CommitType: commitType,
	}
}

func (d *driver) ListCommit(repos []*pfs.Repo, commitType pfs.CommitType, fromCommit []*pfs.Commit,
	provenance []*pfs.Commit, all bool, shards map[uint64]bool) (commitInfos []*pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Filter(func(commit gorethink.Term) gorethink.Term {
		var predicates []interface{}
		for _, repo := range repos {
			predicates = append(predicates, commit.Field("Repo").Eq(repo.Name))
		}
		return gorethink.Or(predicates...)
	}).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		rawCommit := &Commit{}
		if !cursor.Next(rawCommit) {
			break
		}
		commitInfos = append(commitInfos, rawCommitToCommitInfo(rawCommit))
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return commitInfos, nil
}

func (d *driver) ListBranch(repo *pfs.Repo, shards map[uint64]bool) ([]*pfs.CommitInfo, error) {
	return nil, nil
}

func (d *driver) DeleteCommit(commit *pfs.Commit, shards map[uint64]bool) error {
	return nil
}

func (d *driver) PutFile(file *pfs.File, handle string,
	delimiter pfs.Delimiter, shard uint64, reader io.Reader) (retErr error) {
	return nil
}

func (d *driver) MakeDirectory(file *pfs.File, shard uint64) (retErr error) {
	return nil
}

func (d *driver) GetFile(file *pfs.File, filterShard *pfs.Shard, offset int64,
	size int64, from *pfs.Commit, shard uint64, unsafe bool, handle string) (io.ReadCloser, error) {
	return nil, nil
}

func (d *driver) InspectFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, unsafe bool, handle string) (*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) ListFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, recurse bool, unsafe bool, handle string) ([]*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) DeleteFile(file *pfs.File, shard uint64, unsafe bool, handle string) error {
	return nil
}

func (d *driver) DeleteAll(shards map[uint64]bool) error {
	return nil
}

func (d *driver) AddShard(shard uint64) error {
	return nil
}

func (d *driver) DeleteShard(shard uint64) error {
	return nil
}

func (d *driver) Dump() {
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 structs

import (
	"reflect"
	"strings"
)

// FindField compares the pointers (pointerToAField with all fields in pointerToAStruct)
func FindField(pointerToAField interface{}, pointerToAStruct interface{}) (field *reflect.StructField, found bool) {
	fieldVal := reflect.ValueOf(pointerToAField)

	if fieldVal.Kind() != reflect.Ptr {
		panic("pointerToAField must be a pointer")
	}

	strct := reflect.Indirect(reflect.ValueOf(pointerToAStruct))
	numField := strct.NumField()
	for i := 0; i < numField; i++ {
		sf := strct.Field(i)

		if sf.CanAddr() {
			if fieldVal.Pointer() == sf.Addr().Pointer() {
				field := strct.Type().Field(i)
				return &field, true
			}
		}
	}

	return nil, false
}

// ListExportedFields returns all fields of a structure that starts wit uppercase letter
func ListExportedFields(val interface{}, predicates ...ExportedPredicate) []*reflect.StructField {
	valType := reflect.Indirect(reflect.ValueOf(val)).Type()
	len := valType.NumField()
	ret := []*reflect.StructField{}
	for i := 0; i < len; i++ {
		structField := valType.Field(i)

		if FieldExported(&structField, predicates...) {
			ret = append(ret, &structField)
		}
	}

	return ret
}

// ListExportedFieldsWithVals returns all fields of a structure that starts wit uppercase letter with values
func ListExportedFieldsWithVals(val interface{}, predicates ...ExportedPredicate) (fields []*reflect.StructField, values []interface{}) {
	valRefl := reflect.Indirect(reflect.ValueOf(val))
	valType := valRefl.Type()
	len := valType.NumField()
	fields = []*reflect.StructField{}
	values = []interface{}{}
	for i := 0; i < len; i++ {
		structField := valType.Field(i)

		if FieldExported(&structField, predicates...) {
			// if exported
			fields = append(fields, &structField)
			values = append(values, valRefl.Field(i).Interface())
		}
	}

	return fields, values
}

// ExportedPredicate defines a callback (used in func FieldExported)
type ExportedPredicate func(field *reflect.StructField) bool

// FieldExported returns true if field name starts with uppercase
func FieldExported(field *reflect.StructField, predicates ...ExportedPredicate) (exported bool) {
	if field.Name[0] == strings.ToUpper(string(field.Name[0]))[0] {
		expPredic := true
		for _, predicate := range predicates {
			if !predicate(field) {
				expPredic = false
				break
			}
		}

		return expPredic
	}

	return false
}

// ListExportedFieldsPtrs iterates struct fields and return slice of pointers to field values
func ListExportedFieldsPtrs(val interface{}, predicates ...ExportedPredicate) []interface{} {
	rVal := reflect.Indirect(reflect.ValueOf(val))
	ptrs := []interface{}{}
	for i := 0; i < rVal.NumField(); i++ {
		field := rVal.Field(i)
		structField := rVal.Type().Field(i)
		if !FieldExported(&structField, predicates...) {
			continue
		}

		switch field.Kind() {
		case reflect.Ptr, reflect.Interface:
			if field.IsNil() {
				p := reflect.New(field.Type().Elem())
				field.Set(p)
				ptrs = append(ptrs, p.Interface())
			} else {
				ptrs = append(ptrs, field.Interface())
			}
		case reflect.Slice, reflect.Chan, reflect.Map:
			if field.IsNil() {
				p := reflect.New(field.Type())
				field.Set(p.Elem())
				ptrs = append(ptrs, field.Addr().Interface())
			} else {
				ptrs = append(ptrs, field.Interface())
			}
		default:
			if field.CanAddr() {
				ptrs = append(ptrs, field.Addr().Interface())
			} else if field.IsValid() {
				ptrs = append(ptrs, field.Interface())
			} else {
				panic("invalid field")
			}
		}
	}

	return ptrs
}
<commit_msg>SPOPT-1175 fix for insert statement fix for session<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 structs

import (
	"github.com/satori/go.uuid"
	"reflect"
	"strings"
)

// FindField compares the pointers (pointerToAField with all fields in pointerToAStruct)
func FindField(pointerToAField interface{}, pointerToAStruct interface{}) (field *reflect.StructField, found bool) {
	fieldVal := reflect.ValueOf(pointerToAField)

	if fieldVal.Kind() != reflect.Ptr {
		panic("pointerToAField must be a pointer")
	}

	strct := reflect.Indirect(reflect.ValueOf(pointerToAStruct))
	numField := strct.NumField()
	for i := 0; i < numField; i++ {
		sf := strct.Field(i)

		if sf.CanAddr() {
			if fieldVal.Pointer() == sf.Addr().Pointer() {
				field := strct.Type().Field(i)
				return &field, true
			}
		}
	}

	return nil, false
}

// ListExportedFields returns all fields of a structure that starts wit uppercase letter
func ListExportedFields(val interface{}, predicates ...ExportedPredicate) []*reflect.StructField {
	valType := reflect.Indirect(reflect.ValueOf(val)).Type()
	len := valType.NumField()
	ret := []*reflect.StructField{}
	for i := 0; i < len; i++ {
		structField := valType.Field(i)

		if FieldExported(&structField, predicates...) {
			ret = append(ret, &structField)
		}
	}

	return ret
}

// ListExportedFieldsWithVals returns all fields of a structure that starts wit uppercase letter with values
func ListExportedFieldsWithVals(val interface{}, predicates ...ExportedPredicate) (fields []*reflect.StructField, values []interface{}) {
	valRefl := reflect.Indirect(reflect.ValueOf(val))
	valType := valRefl.Type()
	len := valType.NumField()
	fields = []*reflect.StructField{}
	values = []interface{}{}
	for i := 0; i < len; i++ {
		structField := valType.Field(i)

		if FieldExported(&structField, predicates...) {
			// if exported
			fields = append(fields, &structField)
			values = append(values, valRefl.Field(i).Interface())
		}
	}

	return fields, values
}

// ExportedPredicate defines a callback (used in func FieldExported)
type ExportedPredicate func(field *reflect.StructField) bool

// FieldExported returns true if field name starts with uppercase
func FieldExported(field *reflect.StructField, predicates ...ExportedPredicate) (exported bool) {
	if field.Name[0] == strings.ToUpper(string(field.Name[0]))[0] {
		expPredic := true
		for _, predicate := range predicates {
			if !predicate(field) {
				expPredic = false
				break
			}
		}

		return expPredic
	}

	return false
}

// ListExportedFieldsPtrs iterates struct fields and return slice of pointers to field values
func ListExportedFieldsPtrs(val interface{}, predicates ...ExportedPredicate) []interface{} {
	rVal := reflect.Indirect(reflect.ValueOf(val))
	ptrs := []interface{}{}
	for i := 0; i < rVal.NumField(); i++ {
		field := rVal.Field(i)
		structField := rVal.Type().Field(i)
		if !FieldExported(&structField, predicates...) {
			continue
		}

		switch field.Kind() {
		case reflect.Ptr, reflect.Interface:
			if field.IsNil() {
				p := reflect.New(field.Type().Elem())
				field.Set(p)
				ptrs = append(ptrs, p.Interface())
			} else {
				ptrs = append(ptrs, field.Interface())
			}
		case reflect.Slice, reflect.Chan, reflect.Map:
			if field.IsNil() {
				p := reflect.New(field.Type())
				field.Set(p.Elem())
				if field.Type() != reflect.TypeOf(uuid.UUID{}) {
					ptrs = append(ptrs, field.Addr().Interface())
				}
			} else {
				if field.Type() != reflect.TypeOf(uuid.UUID{}) {
					ptrs = append(ptrs, field.Interface())
				}
			}
		default:
			if field.CanAddr() {
				ptrs = append(ptrs, field.Addr().Interface())
			} else if field.IsValid() {
				ptrs = append(ptrs, field.Interface())
			} else {
				panic("invalid field")
			}
		}
	}

	return ptrs
}
<|endoftext|>


<commit_before>package hapi

import (
    "net/http"
    "testing"

    "github.com/gorilla/context"

//     "net/http"
)

func (h *HypermediaAPI) DoRegisterTest(name, requestedType, expectedType, expectedID string, t *testing.T) {
    negType,typeHandler := h.TypeAndHandler("GET","/",requestedType)
    if negType != expectedType {
        t.Fatalf("%s: TypeAndHandler returned '%s', expected '%s'\n", name, negType, expectedType)
    }
    w := new(mockResponseWriter)
    r,_ := http.NewRequest("GET","/",nil)

    if typeHandler != nil {
        typeHandler( w, r )
    }
    if id,ok := context.GetOk(r,"id"); ! ok {
        if ( len(expectedID) > 0 ) {
            t.Fatalf("%s: Error reading id after request\n", name)
        }
    } else if id != expectedID {
        t.Fatalf("%s: Handler identified itself as '%s', expected '%s'\n", name, id, expectedID)
    }
}

func TestRegister1(t *testing.T) {
    router := New()
    router.TestRegister("text/html","1")
    router.DoRegisterTest("text/html 1","text/html","text/html","1",t)
    router.DoRegisterTest("text/html 2","text/*","text/html","1",t)
    router.DoRegisterTest("text/html 3","*/*","text/html","1",t)
    router.DoRegisterTest("text/html 4","text/plain","","",t)
}
<commit_msg>Add tests for httprouter.Handle() handling<commit_after>package hapi

import (
    "net/http"
    "testing"

    "github.com/gorilla/context"
    "github.com/julienschmidt/httprouter"   /* HTTP router */
)

func (h *HypermediaAPI) DoRegisterTest(name, requestedType, expectedType, expectedID string, t *testing.T) {
    negType,typeHandler := h.TypeAndHandler("GET","/",requestedType)
    if negType != expectedType {
        t.Fatalf("%s: TypeAndHandler returned '%s', expected '%s'\n", name, negType, expectedType)
    }
    w := new(mockResponseWriter)
    r,_ := http.NewRequest("GET","/",nil)

    if typeHandler != nil {
        typeHandler( w, r )
    }
    if id,ok := context.GetOk(r,"id"); ! ok {
        if ( len(expectedID) > 0 ) {
            t.Fatalf("%s: Error reading id after request\n", name)
        }
    } else if id != expectedID {
        t.Fatalf("%s: Handler identified itself as '%s', expected '%s'\n", name, id, expectedID)
    }
}

func TestRegister1(t *testing.T) {
    router := New()
    router.TestRegister("text/html","1")
    router.DoRegisterTest("text/html 1","text/html","text/html","1",t)
    router.DoRegisterTest("text/html 2","text/*","text/html","1",t)
    router.DoRegisterTest("text/html 3","*/*","text/html","1",t)
    router.DoRegisterTest("text/html 4","text/plain","","",t)
}

func TestHandle(t *testing.T) {
    var negotiatedType string
    var foo string
    handler := func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
        negotiatedType = context.Get(r,"Content-Type").(string)
        foo = p.ByName("foo")
    }

    router := New()
    router.Handle("GET","/:foo","text/html",handler)
    w := new(mockResponseWriter)
    r,_ := http.NewRequest("GET","/bar",nil)
    router.ServeHTTP(w,r)
    if negotiatedType != "text/html" {
        t.Fatalf("httprouter.Handle set Content-Type to '%s', expected 'text/html'. httprouter context not working.\n", negotiatedType)
    }
    if foo != "bar" {
        t.Fatalf("httprouter.Handle set foo to '%s', expected 'bar'. httprouter.Params handling not working.\n", foo)
    }
}
<|endoftext|>


<commit_before>package hash

import (
    "math"
    "math/rand"
    "testing"
)

func TestCountMin(t *testing.T) {
    sketch := New(210, 1300)
    freq := make(map[uint32]uint32)

    rng := rand.New(rand.NewSource(42))
    for i := 0; i < 10000; i++ {
        h := rng.Uint32()
        sketch.Add(h, 1)
        freq[h] += 1
    }

    // XXX Should test if error is within margin with some probability.
    for k, v := range freq {
        if math.Abs(float64(sketch.Get(k) - v)) > 4 {
            t.Errorf("difference too big: got %d, want %d", sketch.Get(k), v)
        }
    }
}
<commit_msg>CountMin benchmark<commit_after>package hash

import (
    "math"
    "math/rand"
    "testing"
)

func TestCountMin(t *testing.T) {
    sketch := New(210, 1300)
    freq := make(map[uint32]uint32)

    rng := rand.New(rand.NewSource(42))
    for i := 0; i < 10000; i++ {
        h := rng.Uint32()
        sketch.Add(h, 1)
        freq[h] += 1
    }

    // XXX Should test if error is within margin with some probability.
    for k, v := range freq {
        if math.Abs(float64(sketch.Get(k) - v)) > 4 {
            t.Errorf("difference too big: got %d, want %d", sketch.Get(k), v)
        }
    }
}

func BenchmarkCountMinAdd(b *testing.B) {
    sketch := New(256, 256)

    rng := rand.New(rand.NewSource(42))
    for i := 0; i < b.N; i++ {
        for j := 0; j < 2000000; j++ {
            sketch.Add(rng.Uint32(), 1)
        }
    }
}
<|endoftext|>