<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

type rodSlice []rod

func (rs rodSlice) Len() int           { return len(rs) }
func (rs rodSlice) Swap(i, j int)      { rs[i], rs[j] = rs[j], rs[i] }
func (rs rodSlice) Less(i, j int) bool { return int(rs[i]) < int(rs[j]) }

type consumableSlice []consumable

func (cs consumableSlice) Len() int           { return len(cs) }
func (cs consumableSlice) Swap(i, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs consumableSlice) Less(i, j int) bool { return cs[i].Int() < cs[j].Int() }

type statusSlice []status

func (sts statusSlice) Len() int           { return len(sts) }
func (sts statusSlice) Swap(i, j int)      { sts[i], sts[j] = sts[j], sts[i] }
func (sts statusSlice) Less(i, j int) bool { return sts[i] < sts[j] }

type monsSlice []monsterKind

func (ms monsSlice) Len() int      { return len(ms) }
func (ms monsSlice) Swap(i, j int) { ms[i], ms[j] = ms[j], ms[i] }
func (ms monsSlice) Less(i, j int) bool {
	return ms[i].Dangerousness() > ms[j].Dangerousness()
}

func (g *game) DumpAptitudes() string {
	apts := []string{}
	for apt, b := range g.Player.Aptitudes {
		if b {
			apts = append(apts, apt.String())
		}
	}
	sort.Strings(apts)
	if len(apts) == 0 {
		return "You do not have any special aptitudes."
	}
	return "Aptitudes:\n" + strings.Join(apts, "\n")
}

func (g *game) DumpStatuses() string {
	sts := sort.StringSlice{}
	for st, c := range g.Player.Statuses {
		if c > 0 {
			sts = append(sts, st.String())
		}
	}
	sort.Sort(sts)
	if len(sts) == 0 {
		return "You are free of any status effects."
	}
	return "Statuses:\n" + strings.Join(sts, "\n")
}

func (g *game) SortedRods() rodSlice {
	var rs rodSlice
	for k, p := range g.Player.Rods {
		if p == nil {
			continue
		}
		rs = append(rs, k)
	}
	sort.Sort(rs)
	return rs
}

func (g *game) SortedKilledMonsters() monsSlice {
	var ms monsSlice
	for mk, p := range g.KilledMons {
		if p == 0 {
			continue
		}
		ms = append(ms, mk)
	}
	sort.Sort(ms)
	return ms
}

func (g *game) SortedPotions() consumableSlice {
	var cs consumableSlice
	for k := range g.Player.Consumables {
		switch k := k.(type) {
		case potion:
			cs = append(cs, k)
		}
	}
	sort.Sort(cs)
	return cs
}

func (g *game) SortedProjectiles() consumableSlice {
	var cs consumableSlice
	for k := range g.Player.Consumables {
		switch k := k.(type) {
		case projectile:
			cs = append(cs, k)
		}
	}
	sort.Sort(cs)
	return cs
}

func (g *game) Dump() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, " -- Boohu version %s character file --\n\n", Version)
	if g.Wizard {
		fmt.Fprintf(buf, "**WIZARD MODE**\n")
	}
	if g.Player.HP > 0 && g.Depth > 12 {
		fmt.Fprintf(buf, "You escaped from Hareka's Underground alive!\n")
	} else if g.Player.HP <= 0 {
		fmt.Fprintf(buf, "You died while exploring depth %d of Hareka's Underground.\n", g.Depth)
	} else {
		fmt.Fprintf(buf, "You are exploring depth %d of Hareka's Underground.\n", g.Depth)
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "You have %d/%d HP, and %d/%d MP.\n", g.Player.HP, g.Player.HPMax(), g.Player.MP, g.Player.MPMax())
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, g.DumpAptitudes())
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, g.DumpStatuses())
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, "Equipment:\n")
	fmt.Fprintf(buf, "You are wearing %s.\n", Indefinite(g.Player.Armour.String(), false))
	fmt.Fprintf(buf, "You are wielding %s.\n", Indefinite(g.Player.Weapon.String(), false))
	if g.Player.Shield != NoShield {
		if g.Player.Weapon.TwoHanded() {
			fmt.Fprintf(buf, "You have %s (unused).\n", Indefinite(g.Player.Shield.String(), false))
		} else {
			fmt.Fprintf(buf, "You are wearing %s.\n", Indefinite(g.Player.Shield.String(), false))
		}
	}
	fmt.Fprintf(buf, "\n")
	rs := g.SortedRods()
	if len(rs) > 0 {
		fmt.Fprintf(buf, "Rods:\n")
		for _, c := range rs {
			fmt.Fprintf(buf, "- %s (%d/%d charges)\n",
				c, g.Player.Rods[c].Charge, c.MaxCharge())
		}
	} else {
		fmt.Fprintf(buf, "You do not have any rods.\n")
	}
	fmt.Fprintf(buf, "\n")
	ps := g.SortedPotions()
	if len(ps) > 0 {
		fmt.Fprintf(buf, "Potions:\n")
		for _, c := range ps {
			fmt.Fprintf(buf, "- %s (%d available)\n", c, g.Player.Consumables[c])
		}
	} else {
		fmt.Fprintf(buf, "You do not have any potions.\n")
	}
	fmt.Fprintf(buf, "\n")
	ps = g.SortedProjectiles()
	if len(ps) > 0 {
		fmt.Fprintf(buf, "Projectiles:\n")
		for _, c := range ps {
			fmt.Fprintf(buf, "- %s (%d available)\n", c, g.Player.Consumables[c])
		}
	} else {
		fmt.Fprintf(buf, "You do not have any projectiles.\n")
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Miscelaneous:\n")
	fmt.Fprintf(buf, "You collected %d gold coins.\n", g.Player.Gold)
	fmt.Fprintf(buf, "You killed %d monsters.\n", g.Killed)
	fmt.Fprintf(buf, "You spent %.1f turns in the Underground.\n", float64(g.Turn)/10)
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Last messages:\n")
	for i := len(g.Log) - 10; i < len(g.Log); i++ {
		if i >= 0 {
			fmt.Fprintf(buf, "%s\n", g.Log[i])
		}
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Dungeon:\n")
	fmt.Fprintf(buf, "┌%s┐\n", strings.Repeat("─", g.Dungeon.Width))
	buf.WriteString(g.DumpDungeon())
	fmt.Fprintf(buf, "└%s┘\n", strings.Repeat("─", g.Dungeon.Width))
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, g.DumpedKilledMonsters())
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Timeline:\n")
	fmt.Fprintf(buf, g.DumpStory())
	return buf.String()
}

func (g *game) DumpStory() string {
	return strings.Join(g.Story, "\n")
}

func (g *game) DumpDungeon() string {
	buf := bytes.Buffer{}
	for i, c := range g.Dungeon.Cells {
		if i%g.Dungeon.Width == 0 {
			if i == 0 {
				buf.WriteRune('│')
			} else {
				buf.WriteString("│\n│")
			}
		}
		pos := g.Dungeon.CellPosition(i)
		if !c.Explored {
			buf.WriteRune(' ')
			if i == len(g.Dungeon.Cells)-1 {
				buf.WriteString("│\n")
			}
			continue
		}
		var r rune
		switch c.T {
		case WallCell:
			r = '#'
		case FreeCell:
			switch {
			case pos == g.Player.Pos:
				r = '@'
			default:
				r = '.'
				if _, ok := g.Clouds[pos]; ok && g.Player.LOS[pos] {
					r = '§'
				}
				if c, ok := g.Collectables[pos]; ok {
					r = c.Consumable.Letter()
				} else if eq, ok := g.Equipables[pos]; ok {
					r = eq.Letter()
				} else if rod, ok := g.Rods[pos]; ok {
					r = rod.Letter()
				} else if _, ok := g.Stairs[pos]; ok {
					r = '>'
				} else if _, ok := g.Gold[pos]; ok {
					r = '$'
				}
				m, _ := g.MonsterAt(pos)
				if m.Exists() && (g.Player.LOS[m.Pos] || g.Wizard) {
					r = m.Kind.Letter()
				}
			}
		}
		buf.WriteRune(r)
		if i == len(g.Dungeon.Cells)-1 {
			buf.WriteString("│\n")
		}
	}
	return buf.String()
}

func (g *game) DumpedKilledMonsters() string {
	buf := &bytes.Buffer{}
	fmt.Fprint(buf, "Killed Monsters:\n")
	ms := g.SortedKilledMonsters()
	for _, mk := range ms {
		fmt.Fprintf(buf, "- %s: %d\n", mk, g.KilledMons[mk])
	}
	return buf.String()
}

func (g *game) SimplifedDump() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, " ♣ Boohu version %s play summary ♣\n\n", Version)
	if g.Wizard {
		fmt.Fprintf(buf, "**WIZARD MODE**\n")
	}
	if g.Player.HP > 0 && g.Depth > 12 {
		fmt.Fprintf(buf, "You escaped from Hareka's Underground alive!\n")
	} else if g.Player.HP <= 0 {
		fmt.Fprintf(buf, "You died while exploring depth %d of Hareka's Underground.\n", g.Depth)
	} else {
		fmt.Fprintf(buf, "You are exploring depth %d of Hareka's Underground.\n", g.Depth)
	}
	fmt.Fprintf(buf, "You collected %d gold coins.\n", g.Player.Gold)
	fmt.Fprintf(buf, "You killed %d monsters.\n", g.Killed)
	fmt.Fprintf(buf, "You spent %.1f turns in the Underground.\n", float64(g.Turn)/10)
	fmt.Fprintf(buf, "\n")
	dataDir, err := g.DataDir()
	if err == nil {
		fmt.Fprintf(buf, "Full dump written to %s.\n", filepath.Join(dataDir, "dump"))
	}
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, "───Press esc or space to quit───")
	return buf.String()
}

func (g *game) WriteDump() error {
	dataDir, err := g.DataDir()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing dump: %s", err)
		return err
	}
	err = ioutil.WriteFile(filepath.Join(dataDir, "dump"), []byte(g.Dump()), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing dump: %s", err)
		return err
	}
	return nil
}
<commit_msg>draw doors in dump too<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

type rodSlice []rod

func (rs rodSlice) Len() int           { return len(rs) }
func (rs rodSlice) Swap(i, j int)      { rs[i], rs[j] = rs[j], rs[i] }
func (rs rodSlice) Less(i, j int) bool { return int(rs[i]) < int(rs[j]) }

type consumableSlice []consumable

func (cs consumableSlice) Len() int           { return len(cs) }
func (cs consumableSlice) Swap(i, j int)      { cs[i], cs[j] = cs[j], cs[i] }
func (cs consumableSlice) Less(i, j int) bool { return cs[i].Int() < cs[j].Int() }

type statusSlice []status

func (sts statusSlice) Len() int           { return len(sts) }
func (sts statusSlice) Swap(i, j int)      { sts[i], sts[j] = sts[j], sts[i] }
func (sts statusSlice) Less(i, j int) bool { return sts[i] < sts[j] }

type monsSlice []monsterKind

func (ms monsSlice) Len() int      { return len(ms) }
func (ms monsSlice) Swap(i, j int) { ms[i], ms[j] = ms[j], ms[i] }
func (ms monsSlice) Less(i, j int) bool {
	return ms[i].Dangerousness() > ms[j].Dangerousness()
}

func (g *game) DumpAptitudes() string {
	apts := []string{}
	for apt, b := range g.Player.Aptitudes {
		if b {
			apts = append(apts, apt.String())
		}
	}
	sort.Strings(apts)
	if len(apts) == 0 {
		return "You do not have any special aptitudes."
	}
	return "Aptitudes:\n" + strings.Join(apts, "\n")
}

func (g *game) DumpStatuses() string {
	sts := sort.StringSlice{}
	for st, c := range g.Player.Statuses {
		if c > 0 {
			sts = append(sts, st.String())
		}
	}
	sort.Sort(sts)
	if len(sts) == 0 {
		return "You are free of any status effects."
	}
	return "Statuses:\n" + strings.Join(sts, "\n")
}

func (g *game) SortedRods() rodSlice {
	var rs rodSlice
	for k, p := range g.Player.Rods {
		if p == nil {
			continue
		}
		rs = append(rs, k)
	}
	sort.Sort(rs)
	return rs
}

func (g *game) SortedKilledMonsters() monsSlice {
	var ms monsSlice
	for mk, p := range g.KilledMons {
		if p == 0 {
			continue
		}
		ms = append(ms, mk)
	}
	sort.Sort(ms)
	return ms
}

func (g *game) SortedPotions() consumableSlice {
	var cs consumableSlice
	for k := range g.Player.Consumables {
		switch k := k.(type) {
		case potion:
			cs = append(cs, k)
		}
	}
	sort.Sort(cs)
	return cs
}

func (g *game) SortedProjectiles() consumableSlice {
	var cs consumableSlice
	for k := range g.Player.Consumables {
		switch k := k.(type) {
		case projectile:
			cs = append(cs, k)
		}
	}
	sort.Sort(cs)
	return cs
}

func (g *game) Dump() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, " -- Boohu version %s character file --\n\n", Version)
	if g.Wizard {
		fmt.Fprintf(buf, "**WIZARD MODE**\n")
	}
	if g.Player.HP > 0 && g.Depth > 12 {
		fmt.Fprintf(buf, "You escaped from Hareka's Underground alive!\n")
	} else if g.Player.HP <= 0 {
		fmt.Fprintf(buf, "You died while exploring depth %d of Hareka's Underground.\n", g.Depth)
	} else {
		fmt.Fprintf(buf, "You are exploring depth %d of Hareka's Underground.\n", g.Depth)
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "You have %d/%d HP, and %d/%d MP.\n", g.Player.HP, g.Player.HPMax(), g.Player.MP, g.Player.MPMax())
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, g.DumpAptitudes())
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, g.DumpStatuses())
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, "Equipment:\n")
	fmt.Fprintf(buf, "You are wearing %s.\n", Indefinite(g.Player.Armour.String(), false))
	fmt.Fprintf(buf, "You are wielding %s.\n", Indefinite(g.Player.Weapon.String(), false))
	if g.Player.Shield != NoShield {
		if g.Player.Weapon.TwoHanded() {
			fmt.Fprintf(buf, "You have %s (unused).\n", Indefinite(g.Player.Shield.String(), false))
		} else {
			fmt.Fprintf(buf, "You are wearing %s.\n", Indefinite(g.Player.Shield.String(), false))
		}
	}
	fmt.Fprintf(buf, "\n")
	rs := g.SortedRods()
	if len(rs) > 0 {
		fmt.Fprintf(buf, "Rods:\n")
		for _, c := range rs {
			fmt.Fprintf(buf, "- %s (%d/%d charges)\n",
				c, g.Player.Rods[c].Charge, c.MaxCharge())
		}
	} else {
		fmt.Fprintf(buf, "You do not have any rods.\n")
	}
	fmt.Fprintf(buf, "\n")
	ps := g.SortedPotions()
	if len(ps) > 0 {
		fmt.Fprintf(buf, "Potions:\n")
		for _, c := range ps {
			fmt.Fprintf(buf, "- %s (%d available)\n", c, g.Player.Consumables[c])
		}
	} else {
		fmt.Fprintf(buf, "You do not have any potions.\n")
	}
	fmt.Fprintf(buf, "\n")
	ps = g.SortedProjectiles()
	if len(ps) > 0 {
		fmt.Fprintf(buf, "Projectiles:\n")
		for _, c := range ps {
			fmt.Fprintf(buf, "- %s (%d available)\n", c, g.Player.Consumables[c])
		}
	} else {
		fmt.Fprintf(buf, "You do not have any projectiles.\n")
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Miscelaneous:\n")
	fmt.Fprintf(buf, "You collected %d gold coins.\n", g.Player.Gold)
	fmt.Fprintf(buf, "You killed %d monsters.\n", g.Killed)
	fmt.Fprintf(buf, "You spent %.1f turns in the Underground.\n", float64(g.Turn)/10)
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Last messages:\n")
	for i := len(g.Log) - 10; i < len(g.Log); i++ {
		if i >= 0 {
			fmt.Fprintf(buf, "%s\n", g.Log[i])
		}
	}
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Dungeon:\n")
	fmt.Fprintf(buf, "┌%s┐\n", strings.Repeat("─", g.Dungeon.Width))
	buf.WriteString(g.DumpDungeon())
	fmt.Fprintf(buf, "└%s┘\n", strings.Repeat("─", g.Dungeon.Width))
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, g.DumpedKilledMonsters())
	fmt.Fprintf(buf, "\n")
	fmt.Fprintf(buf, "Timeline:\n")
	fmt.Fprintf(buf, g.DumpStory())
	return buf.String()
}

func (g *game) DumpStory() string {
	return strings.Join(g.Story, "\n")
}

func (g *game) DumpDungeon() string {
	buf := bytes.Buffer{}
	for i, c := range g.Dungeon.Cells {
		if i%g.Dungeon.Width == 0 {
			if i == 0 {
				buf.WriteRune('│')
			} else {
				buf.WriteString("│\n│")
			}
		}
		pos := g.Dungeon.CellPosition(i)
		if !c.Explored {
			buf.WriteRune(' ')
			if i == len(g.Dungeon.Cells)-1 {
				buf.WriteString("│\n")
			}
			continue
		}
		var r rune
		switch c.T {
		case WallCell:
			r = '#'
		case FreeCell:
			switch {
			case pos == g.Player.Pos:
				r = '@'
			default:
				r = '.'
				if _, ok := g.Clouds[pos]; ok && g.Player.LOS[pos] {
					r = '§'
				}
				if c, ok := g.Collectables[pos]; ok {
					r = c.Consumable.Letter()
				} else if eq, ok := g.Equipables[pos]; ok {
					r = eq.Letter()
				} else if rod, ok := g.Rods[pos]; ok {
					r = rod.Letter()
				} else if _, ok := g.Stairs[pos]; ok {
					r = '>'
				} else if _, ok := g.Gold[pos]; ok {
					r = '$'
				} else if _, ok := g.Doors[pos]; ok {
					r = '+'
				}
				m, _ := g.MonsterAt(pos)
				if m.Exists() && (g.Player.LOS[m.Pos] || g.Wizard) {
					r = m.Kind.Letter()
				}
			}
		}
		buf.WriteRune(r)
		if i == len(g.Dungeon.Cells)-1 {
			buf.WriteString("│\n")
		}
	}
	return buf.String()
}

func (g *game) DumpedKilledMonsters() string {
	buf := &bytes.Buffer{}
	fmt.Fprint(buf, "Killed Monsters:\n")
	ms := g.SortedKilledMonsters()
	for _, mk := range ms {
		fmt.Fprintf(buf, "- %s: %d\n", mk, g.KilledMons[mk])
	}
	return buf.String()
}

func (g *game) SimplifedDump() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, " ♣ Boohu version %s play summary ♣\n\n", Version)
	if g.Wizard {
		fmt.Fprintf(buf, "**WIZARD MODE**\n")
	}
	if g.Player.HP > 0 && g.Depth > 12 {
		fmt.Fprintf(buf, "You escaped from Hareka's Underground alive!\n")
	} else if g.Player.HP <= 0 {
		fmt.Fprintf(buf, "You died while exploring depth %d of Hareka's Underground.\n", g.Depth)
	} else {
		fmt.Fprintf(buf, "You are exploring depth %d of Hareka's Underground.\n", g.Depth)
	}
	fmt.Fprintf(buf, "You collected %d gold coins.\n", g.Player.Gold)
	fmt.Fprintf(buf, "You killed %d monsters.\n", g.Killed)
	fmt.Fprintf(buf, "You spent %.1f turns in the Underground.\n", float64(g.Turn)/10)
	fmt.Fprintf(buf, "\n")
	dataDir, err := g.DataDir()
	if err == nil {
		fmt.Fprintf(buf, "Full dump written to %s.\n", filepath.Join(dataDir, "dump"))
	}
	fmt.Fprintf(buf, "\n\n")
	fmt.Fprintf(buf, "───Press esc or space to quit───")
	return buf.String()
}

func (g *game) WriteDump() error {
	dataDir, err := g.DataDir()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing dump: %s", err)
		return err
	}
	err = ioutil.WriteFile(filepath.Join(dataDir, "dump"), []byte(g.Dump()), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing dump: %s", err)
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// EDNS0
//
// EDNS0 is an extension mechanism for the DNS defined in RFC 2671 and updated
// by RFC 6891. It defines a
// standard RR type, the OPT RR, which is then completely abused.
// Basic use pattern for creating an (empty) OPT RR:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "." // MUST be the root zone, per definition.
//	o.Hdr.Rrtype = dns.TypeOPT
//
// The rdata of an OPT RR consists out of a slice of EDNS0 interfaces. Currently
// only a few have been standardized: EDNS0_NSID (RFC 5001) and EDNS0_SUBNET (draft). Note that
// these options may be combined in an OPT RR.
// Basic use pattern for a server to check if (and which) options are set:
//
//	// o is a dns.OPT
//	for _, s := range o.Option {
//		switch e := s.(type) {
//		case *dns.EDNS0_NSID:
//			// do stuff with e.Nsid
//		case *dns.EDNS0_SUBNET:
//			// access e.Family, e.Address, etc.
//		}
//	}
package dns

import (
	"encoding/hex"
	"errors"
	"net"
	"strconv"
)

// EDNS0 Option codes.
const (
	EDNS0LLQ    = 0x1    // long lived queries: http://tools.ietf.org/html/draft-sekar-dns-llq-01
	EDNS0UL     = 0x2    // update lease draft: http://files.dns-sd.org/draft-sekar-dns-ul.txt
	EDNS0NSID   = 0x3    // nsid (RFC5001)
	EDNS0SUBNET = 0x50fa // client-subnet draft: http://tools.ietf.org/html/draft-vandergaast-edns-client-subnet-01
	_DO         = 1 << 7 // dnssec ok
)

type OPT struct {
	Hdr    RR_Header
	Option []EDNS0 `dns:"opt"`
}

func (rr *OPT) Header() *RR_Header {
	return &rr.Hdr
}

func (rr *OPT) String() string {
	s := "\n;; OPT PSEUDOSECTION:\n; EDNS: version " + strconv.Itoa(int(rr.Version())) + "; "
	if rr.Do() {
		s += "flags: do; "
	} else {
		s += "flags: ; "
	}
	s += "udp: " + strconv.Itoa(int(rr.UDPSize()))

	for _, o := range rr.Option {
		switch o.(type) {
		case *EDNS0_NSID:
			s += "\n; NSID: " + o.String()
			h, e := o.pack()
			var r string
			if e == nil {
				for _, c := range h {
					r += "(" + string(c) + ")"
				}
				s += "  " + r
			}
		case *EDNS0_SUBNET:
			s += "\n; SUBNET: " + o.String()
		case *EDNS0_UL:
			s += "\n; LEASE: " + o.String()
		}
	}
	return s
}

func (rr *OPT) len() int {
	l := rr.Hdr.len()
	for i := 0; i < len(rr.Option); i++ {
		lo, _ := rr.Option[i].pack()
		l += 2 + len(lo)
	}
	return l
}

func (rr *OPT) copy() RR {
	return &OPT{*rr.Hdr.copyHeader(), rr.Option}
}

// Version returns the EDNS version used. Only zero is defined.
func (rr *OPT) Version() uint8 {
	return uint8(rr.Hdr.Ttl & 0x00FF00FFFF)
}

// SetVersion sets the version of EDNS. This is usually zero.
func (rr *OPT) SetVersion(v uint8) {
	rr.Hdr.Ttl = rr.Hdr.Ttl&0xFF00FFFF | uint32(v)
}

// UDPSize returns the UDP buffer size.
func (rr *OPT) UDPSize() uint16 {
	return rr.Hdr.Class
}

// SetUDPSize sets the UDP buffer size.
func (rr *OPT) SetUDPSize(size uint16) {
	rr.Hdr.Class = size
}

// Do returns the value of the DO (DNSSEC OK) bit.
func (rr *OPT) Do() bool {
	return byte(rr.Hdr.Ttl>>8)&_DO == _DO
}

// SetDo sets the DO (DNSSEC OK) bit.
func (rr *OPT) SetDo() {
	b1 := byte(rr.Hdr.Ttl >> 24)
	b2 := byte(rr.Hdr.Ttl >> 16)
	b3 := byte(rr.Hdr.Ttl >> 8)
	b4 := byte(rr.Hdr.Ttl)
	b3 |= _DO // Set it
	rr.Hdr.Ttl = uint32(b1)<<24 | uint32(b2)<<16 | uint32(b3)<<8 | uint32(b4)
}

// EDNS0 defines an EDNS0 Option. An OPT RR can have multiple options appended to
// it. Basic use pattern for adding an option to and OPT RR:
//
//	// o is the OPT RR, e is the EDNS0 option
//	o.Option = append(o.Option, e)
type EDNS0 interface {
	// Option returns the option code for the option.
	Option() uint16
	// pack returns the bytes of the option data.
	pack() ([]byte, error)
	// unpack sets the data as found in the buffer. Is also sets
	// the length of the slice as the length of the option data.
	unpack([]byte)
	// String returns the string representation of the option.
	String() string
}

// The nsid EDNS0 option is used to retrieve some sort of nameserver
// identifier. When seding a request Nsid must be set to the empty string
// The identifier is an opaque string encoded as hex.
// Basic use pattern for creating an nsid option:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_NSID)
//	e.Code = dns.EDNS0NSID
//	o.Option = append(o.Option, e)
type EDNS0_NSID struct {
	Code uint16 // Always EDNS0NSID
	Nsid string // This string needs to be hex encoded
}

func (e *EDNS0_NSID) Option() uint16 {
	return EDNS0NSID
}

func (e *EDNS0_NSID) pack() ([]byte, error) {
	h, err := hex.DecodeString(e.Nsid)
	if err != nil {
		return nil, err
	}
	return h, nil
}

func (e *EDNS0_NSID) unpack(b []byte) {
	e.Nsid = hex.EncodeToString(b)
}

func (e *EDNS0_NSID) String() string {
	return string(e.Nsid)
}

// The subnet EDNS0 option is used to give the remote nameserver
// an idea of where the client lives. It can then give back a different
// answer depending on the location or network topology.
// Basic use pattern for creating an subnet option:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_SUBNET)
//	e.Code = dns.EDNS0SUBNET
//	e.Family = 1	// 1 for IPv4 source address, 2 for IPv6
//	e.NetMask = 32	// 32 for IPV4, 128 for IPv6
//	e.SourceScope = 0
//	e.Address = net.ParseIP("127.0.0.1").To4()	// for IPv4
//	// e.Address = net.ParseIP("2001:7b8:32a::2")	// for IPV6
//	o.Option = append(o.Option, e)
type EDNS0_SUBNET struct {
	Code          uint16 // Always EDNS0SUBNET
	Family        uint16 // 1 for IP, 2 for IP6
	SourceNetmask uint8
	SourceScope   uint8
	Address       net.IP
}

func (e *EDNS0_SUBNET) Option() uint16 {
	return EDNS0SUBNET
}

func (e *EDNS0_SUBNET) pack() ([]byte, error) {
	b := make([]byte, 4)
	b[0], b[1] = packUint16(e.Family)
	b[2] = e.SourceNetmask
	b[3] = e.SourceScope
	switch e.Family {
	case 1:
		if e.SourceNetmask > net.IPv4len*8 {
			return nil, errors.New("dns: bad netmask")
		}
		ip := make([]byte, net.IPv4len)
		a := e.Address.To4().Mask(net.CIDRMask(int(e.SourceNetmask), net.IPv4len*8))
		for i := 0; i < net.IPv4len; i++ {
			if i+1 > len(e.Address) {
				break
			}
			ip[i] = a[i]
		}
		needLength := e.SourceNetmask / 8
		if e.SourceNetmask%8 > 0 {
			needLength++
		}
		ip = ip[:needLength]
		b = append(b, ip...)
	case 2:
		if e.SourceNetmask > net.IPv6len*8 {
			return nil, errors.New("dns: bad netmask")
		}
		ip := make([]byte, net.IPv6len)
		a := e.Address.Mask(net.CIDRMask(int(e.SourceNetmask), net.IPv6len*8))
		for i := 0; i < net.IPv6len; i++ {
			if i+1 > len(e.Address) {
				break
			}
			ip[i] = a[i]
		}
		needLength := e.SourceNetmask / 8
		if e.SourceNetmask%8 > 0 {
			needLength++
		}
		ip = ip[:needLength]
		b = append(b, ip...)
	default:
		return nil, errors.New("dns: bad address family")
	}
	return b, nil
}

func (e *EDNS0_SUBNET) unpack(b []byte) {
	lb := len(b)
	if lb < 4 {
		return
	}
	e.Family, _ = unpackUint16(b, 0)
	e.SourceNetmask = b[2]
	e.SourceScope = b[3]
	switch e.Family {
	case 1:
		addr := make([]byte, 4)
		for i := 0; i < int(e.SourceNetmask/8); i++ {
			if 4+i > len(b) {
				break
			}
			addr[i] = b[4+i]
		}
		e.Address = net.IPv4(addr[0], addr[1], addr[2], addr[3])
	case 2:
		addr := make([]byte, 16)
		for i := 0; i < int(e.SourceNetmask/8); i++ {
			if 4+i > len(b) {
				break
			}
			addr[i] = b[4+i]
		}
		e.Address = net.IP{addr[0], addr[1], addr[2], addr[3], addr[4],
			addr[5], addr[6], addr[7], addr[8], addr[9], addr[10],
			addr[11], addr[12], addr[13], addr[14], addr[15]}
	}
	return
}

func (e *EDNS0_SUBNET) String() (s string) {
	if e.Address == nil {
		s = "<nil>"
	} else if e.Address.To4() != nil {
		s = e.Address.String()
	} else {
		s = "[" + e.Address.String() + "]"
	}
	s += "/" + strconv.Itoa(int(e.SourceNetmask)) + "/" + strconv.Itoa(int(e.SourceScope))
	return
}

// The UL (Update Lease) EDNS0 (draft RFC) option is used to tell the server to set
// an expiration on an update RR. This is helpful for clients that cannot clean
// up after themselves. This is a draft RFC and more information can be found at
// http://files.dns-sd.org/draft-sekar-dns-ul.txt
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_UL)
//	e.Code = dns.EDNS0UL
//	e.Lease = 120 // in seconds
//	o.Option = append(o.Option, e)

type EDNS0_UL struct {
	Code  uint16 // Always EDNS0UL
	Lease uint32
}

func (e *EDNS0_UL) Option() uint16 {
	return EDNS0UL
}

// Copied: http://golang.org/src/pkg/net/dnsmsg.go
func (e *EDNS0_UL) pack() ([]byte, error) {
	b := make([]byte, 4)
	b[0] = byte(e.Lease >> 24)
	b[1] = byte(e.Lease >> 16)
	b[2] = byte(e.Lease >> 8)
	b[3] = byte(e.Lease)
	return b, nil
}

func (e *EDNS0_UL) unpack(b []byte) {
	e.Lease = uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
}

func (e *EDNS0_UL) String() string {
	return strconv.Itoa(int(e.Lease))
}
<commit_msg>Preliminary support for EDNS0_LLQ<commit_after>// EDNS0
//
// EDNS0 is an extension mechanism for the DNS defined in RFC 2671 and updated
// by RFC 6891. It defines a
// standard RR type, the OPT RR, which is then completely abused.
// Basic use pattern for creating an (empty) OPT RR:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "." // MUST be the root zone, per definition.
//	o.Hdr.Rrtype = dns.TypeOPT
//
// The rdata of an OPT RR consists out of a slice of EDNS0 interfaces. Currently
// only a few have been standardized: EDNS0_NSID (RFC 5001) and EDNS0_SUBNET (draft). Note that
// these options may be combined in an OPT RR.
// Basic use pattern for a server to check if (and which) options are set:
//
//	// o is a dns.OPT
//	for _, s := range o.Option {
//		switch e := s.(type) {
//		case *dns.EDNS0_NSID:
//			// do stuff with e.Nsid
//		case *dns.EDNS0_SUBNET:
//			// access e.Family, e.Address, etc.
//		}
//	}
package dns

import (
	"encoding/hex"
	"errors"
	"net"
	"strconv"
)

// EDNS0 Option codes.
const (
	EDNS0LLQ    = 0x1    // long lived queries: http://tools.ietf.org/html/draft-sekar-dns-llq-01
	EDNS0UL     = 0x2    // update lease draft: http://files.dns-sd.org/draft-sekar-dns-ul.txt
	EDNS0NSID   = 0x3    // nsid (RFC5001)
	EDNS0SUBNET = 0x50fa // client-subnet draft: http://tools.ietf.org/html/draft-vandergaast-edns-client-subnet-01
	_DO         = 1 << 7 // dnssec ok
)

type OPT struct {
	Hdr    RR_Header
	Option []EDNS0 `dns:"opt"`
}

func (rr *OPT) Header() *RR_Header {
	return &rr.Hdr
}

func (rr *OPT) String() string {
	s := "\n;; OPT PSEUDOSECTION:\n; EDNS: version " + strconv.Itoa(int(rr.Version())) + "; "
	if rr.Do() {
		s += "flags: do; "
	} else {
		s += "flags: ; "
	}
	s += "udp: " + strconv.Itoa(int(rr.UDPSize()))

	for _, o := range rr.Option {
		switch o.(type) {
		case *EDNS0_NSID:
			s += "\n; NSID: " + o.String()
			h, e := o.pack()
			var r string
			if e == nil {
				for _, c := range h {
					r += "(" + string(c) + ")"
				}
				s += "  " + r
			}
		case *EDNS0_SUBNET:
			s += "\n; SUBNET: " + o.String()
		case *EDNS0_UL:
			s += "\n; LEASE: " + o.String()
		}
	}
	return s
}

func (rr *OPT) len() int {
	l := rr.Hdr.len()
	for i := 0; i < len(rr.Option); i++ {
		lo, _ := rr.Option[i].pack()
		l += 2 + len(lo)
	}
	return l
}

func (rr *OPT) copy() RR {
	return &OPT{*rr.Hdr.copyHeader(), rr.Option}
}

// Version returns the EDNS version used. Only zero is defined.
func (rr *OPT) Version() uint8 {
	return uint8(rr.Hdr.Ttl & 0x00FF00FFFF)
}

// SetVersion sets the version of EDNS. This is usually zero.
func (rr *OPT) SetVersion(v uint8) {
	rr.Hdr.Ttl = rr.Hdr.Ttl&0xFF00FFFF | uint32(v)
}

// UDPSize returns the UDP buffer size.
func (rr *OPT) UDPSize() uint16 {
	return rr.Hdr.Class
}

// SetUDPSize sets the UDP buffer size.
func (rr *OPT) SetUDPSize(size uint16) {
	rr.Hdr.Class = size
}

// Do returns the value of the DO (DNSSEC OK) bit.
func (rr *OPT) Do() bool {
	return byte(rr.Hdr.Ttl>>8)&_DO == _DO
}

// SetDo sets the DO (DNSSEC OK) bit.
func (rr *OPT) SetDo() {
	b1 := byte(rr.Hdr.Ttl >> 24)
	b2 := byte(rr.Hdr.Ttl >> 16)
	b3 := byte(rr.Hdr.Ttl >> 8)
	b4 := byte(rr.Hdr.Ttl)
	b3 |= _DO // Set it
	rr.Hdr.Ttl = uint32(b1)<<24 | uint32(b2)<<16 | uint32(b3)<<8 | uint32(b4)
}

// EDNS0 defines an EDNS0 Option. An OPT RR can have multiple options appended to
// it. Basic use pattern for adding an option to and OPT RR:
//
//	// o is the OPT RR, e is the EDNS0 option
//	o.Option = append(o.Option, e)
type EDNS0 interface {
	// Option returns the option code for the option.
	Option() uint16
	// pack returns the bytes of the option data.
	pack() ([]byte, error)
	// unpack sets the data as found in the buffer. Is also sets
	// the length of the slice as the length of the option data.
	unpack([]byte)
	// String returns the string representation of the option.
	String() string
}

// The nsid EDNS0 option is used to retrieve some sort of nameserver
// identifier. When seding a request Nsid must be set to the empty string
// The identifier is an opaque string encoded as hex.
// Basic use pattern for creating an nsid option:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_NSID)
//	e.Code = dns.EDNS0NSID
//	o.Option = append(o.Option, e)
type EDNS0_NSID struct {
	Code uint16 // Always EDNS0NSID
	Nsid string // This string needs to be hex encoded
}

func (e *EDNS0_NSID) Option() uint16 {
	return EDNS0NSID
}

func (e *EDNS0_NSID) pack() ([]byte, error) {
	h, err := hex.DecodeString(e.Nsid)
	if err != nil {
		return nil, err
	}
	return h, nil
}

func (e *EDNS0_NSID) unpack(b []byte) {
	e.Nsid = hex.EncodeToString(b)
}

func (e *EDNS0_NSID) String() string {
	return string(e.Nsid)
}

// The subnet EDNS0 option is used to give the remote nameserver
// an idea of where the client lives. It can then give back a different
// answer depending on the location or network topology.
// Basic use pattern for creating an subnet option:
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_SUBNET)
//	e.Code = dns.EDNS0SUBNET
//	e.Family = 1	// 1 for IPv4 source address, 2 for IPv6
//	e.NetMask = 32	// 32 for IPV4, 128 for IPv6
//	e.SourceScope = 0
//	e.Address = net.ParseIP("127.0.0.1").To4()	// for IPv4
//	// e.Address = net.ParseIP("2001:7b8:32a::2")	// for IPV6
//	o.Option = append(o.Option, e)
type EDNS0_SUBNET struct {
	Code          uint16 // Always EDNS0SUBNET
	Family        uint16 // 1 for IP, 2 for IP6
	SourceNetmask uint8
	SourceScope   uint8
	Address       net.IP
}

func (e *EDNS0_SUBNET) Option() uint16 {
	return EDNS0SUBNET
}

func (e *EDNS0_SUBNET) pack() ([]byte, error) {
	b := make([]byte, 4)
	b[0], b[1] = packUint16(e.Family)
	b[2] = e.SourceNetmask
	b[3] = e.SourceScope
	switch e.Family {
	case 1:
		if e.SourceNetmask > net.IPv4len*8 {
			return nil, errors.New("dns: bad netmask")
		}
		ip := make([]byte, net.IPv4len)
		a := e.Address.To4().Mask(net.CIDRMask(int(e.SourceNetmask), net.IPv4len*8))
		for i := 0; i < net.IPv4len; i++ {
			if i+1 > len(e.Address) {
				break
			}
			ip[i] = a[i]
		}
		needLength := e.SourceNetmask / 8
		if e.SourceNetmask%8 > 0 {
			needLength++
		}
		ip = ip[:needLength]
		b = append(b, ip...)
	case 2:
		if e.SourceNetmask > net.IPv6len*8 {
			return nil, errors.New("dns: bad netmask")
		}
		ip := make([]byte, net.IPv6len)
		a := e.Address.Mask(net.CIDRMask(int(e.SourceNetmask), net.IPv6len*8))
		for i := 0; i < net.IPv6len; i++ {
			if i+1 > len(e.Address) {
				break
			}
			ip[i] = a[i]
		}
		needLength := e.SourceNetmask / 8
		if e.SourceNetmask%8 > 0 {
			needLength++
		}
		ip = ip[:needLength]
		b = append(b, ip...)
	default:
		return nil, errors.New("dns: bad address family")
	}
	return b, nil
}

func (e *EDNS0_SUBNET) unpack(b []byte) {
	lb := len(b)
	if lb < 4 {
		return
	}
	e.Family, _ = unpackUint16(b, 0)
	e.SourceNetmask = b[2]
	e.SourceScope = b[3]
	switch e.Family {
	case 1:
		addr := make([]byte, 4)
		for i := 0; i < int(e.SourceNetmask/8); i++ {
			if 4+i > len(b) {
				break
			}
			addr[i] = b[4+i]
		}
		e.Address = net.IPv4(addr[0], addr[1], addr[2], addr[3])
	case 2:
		addr := make([]byte, 16)
		for i := 0; i < int(e.SourceNetmask/8); i++ {
			if 4+i > len(b) {
				break
			}
			addr[i] = b[4+i]
		}
		e.Address = net.IP{addr[0], addr[1], addr[2], addr[3], addr[4],
			addr[5], addr[6], addr[7], addr[8], addr[9], addr[10],
			addr[11], addr[12], addr[13], addr[14], addr[15]}
	}
	return
}

func (e *EDNS0_SUBNET) String() (s string) {
	if e.Address == nil {
		s = "<nil>"
	} else if e.Address.To4() != nil {
		s = e.Address.String()
	} else {
		s = "[" + e.Address.String() + "]"
	}
	s += "/" + strconv.Itoa(int(e.SourceNetmask)) + "/" + strconv.Itoa(int(e.SourceScope))
	return
}

// The UL (Update Lease) EDNS0 (draft RFC) option is used to tell the server to set
// an expiration on an update RR. This is helpful for clients that cannot clean
// up after themselves. This is a draft RFC and more information can be found at
// http://files.dns-sd.org/draft-sekar-dns-ul.txt
//
//	o := new(dns.OPT)
//	o.Hdr.Name = "."
//	o.Hdr.Rrtype = dns.TypeOPT
//	e := new(dns.EDNS0_UL)
//	e.Code = dns.EDNS0UL
//	e.Lease = 120 // in seconds
//	o.Option = append(o.Option, e)
type EDNS0_UL struct {
	Code  uint16 // Always EDNS0UL
	Lease uint32
}

func (e *EDNS0_UL) Option() uint16 {
	return EDNS0UL
}

// Copied: http://golang.org/src/pkg/net/dnsmsg.go
func (e *EDNS0_UL) pack() ([]byte, error) {
	b := make([]byte, 4)
	b[0] = byte(e.Lease >> 24)
	b[1] = byte(e.Lease >> 16)
	b[2] = byte(e.Lease >> 8)
	b[3] = byte(e.Lease)
	return b, nil
}

func (e *EDNS0_UL) unpack(b []byte) {
	e.Lease = uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
}

func (e *EDNS0_UL) String() string {
	return strconv.Itoa(int(e.Lease))
}

// Long Lived Queries: ....
//
//
type EDNS0_LLQ struct {
	Code uint16 // Always EDNS0LLQ
}

func (e *EDNS0_LLQ) Option() uint16 {
	return EDNS0LLQ
}

func (e *EDNS0_LLQ) pack() ([]byte, error) {
	return nil, nil
}

func (e *EDNS0_LLQ) unpack(b []byte) {

}

func (e *EDNS0_LLQ) String() string {
	return ""
}

<|endoftext|>


<commit_before>package fluentd_forwarder

import (
	"bytes"
	"compress/gzip"
	"crypto/x509"
	"errors"
	logging "github.com/op/go-logging"
	td_client "github.com/treasure-data/td-client-go"
	"github.com/ugorji/go/codec"
	"net"
	"os"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type tdOutputSpooler struct {
	daemon         *tdOutputSpoolerDaemon
	ticker         *time.Ticker
	tag            string
	databaseName   string
	tableName      string
	key            string
	journal        Journal
	shutdownChan   chan struct{}
	isShuttingDown uintptr
	client         *td_client.TDClient
}

type tdOutputSpoolerDaemon struct {
	output       *TDOutput
	shutdownChan chan struct{}
	spoolersMtx  sync.Mutex
	spoolers     map[string]*tdOutputSpooler
	tempFactory  TempFileRandomAccessStoreFactory
	wg           sync.WaitGroup
	endNotify    func(*tdOutputSpoolerDaemon)
}

type TDOutput struct {
	logger         *logging.Logger
	codec          *codec.MsgpackHandle
	databaseName   string
	tableName      string
	tempDir        string
	enc            *codec.Encoder
	conn           net.Conn
	flushInterval  time.Duration
	wg             sync.WaitGroup
	journalGroup   JournalGroup
	emitterChan    chan FluentRecordSet
	spoolerDaemon  *tdOutputSpoolerDaemon
	isShuttingDown uintptr
	client         *td_client.TDClient
	sem            chan struct{}
	gcChan         chan *os.File
	completion     sync.Cond
}

func encodeRecords(encoder *codec.Encoder, records []TinyFluentRecord) error {
	for _, record := range records {
		e := map[string]interface{}{"time": record.Timestamp}
		for k, v := range record.Data {
			e[k] = v
		}
		err := encoder.Encode(e)
		if err != nil {
			return err
		}
	}
	return nil
}

func (spooler *tdOutputSpooler) cleanup() {
	spooler.ticker.Stop()
	spooler.journal.Dispose()
	spooler.daemon.wg.Done()
}

func (spooler *tdOutputSpooler) handle() {
	defer spooler.cleanup()
	spooler.daemon.output.logger.Notice("Spooler started")
outer:
	for {
		select {
		case <-spooler.ticker.C:
			spooler.daemon.output.logger.Notice("Flushing...")
			err := spooler.journal.Flush(func(chunk JournalChunk) interface{} {
				defer chunk.Dispose()
				if atomic.LoadUintptr(&spooler.isShuttingDown) != 0 {
					return errors.New("Flush aborted")
				}
				spooler.daemon.output.logger.Info("Flushing chunk %s", chunk.String())
				size, err := chunk.Size()
				if err != nil {
					return err
				}
				if size == 0 {
					return nil
				}
				futureErr := make(chan error, 1)
				sem := spooler.daemon.output.sem
				sem <- struct{}{}
				go func(size int64, chunk JournalChunk, futureErr chan error) {
					defer func() {
						<-sem
						chunk.Dispose()
					}()
					err := func() error {
						compressingBlob := NewCompressingBlob(
							chunk,
							maxInt(4096, int(size/4)),
							gzip.BestSpeed,
							&spooler.daemon.tempFactory,
						)
						defer compressingBlob.Dispose()
						_, err := spooler.client.Import(
							spooler.databaseName,
							spooler.tableName,
							"msgpack.gz",
							td_client.NewBufferingBlobSize(
								compressingBlob,
								maxInt(4096, int(size/16)),
							),
							chunk.Id(),
						)
						return err
					}()
					futureErr <- err
					if err != nil {
						spooler.daemon.output.logger.Info("Failed to flush chunk %s (reason: %s)", chunk.String(), err.Error())
					} else {
						spooler.daemon.output.logger.Info("Completed flushing chunk %s", chunk.String())
					}
				}(size, chunk.Dup(), futureErr)
				return (<-chan error)(futureErr)
			})
			if err != nil {
				spooler.daemon.output.logger.Error("Error during reading from the journal: %s", err.Error())
			}
		case <-spooler.shutdownChan:
			break outer
		}
	}
	spooler.daemon.output.logger.Notice("Spooler ended")
}

func normalizeDatabaseName(name string) (string, error) {
	name_ := ([]byte)(name)
	if len(name_) == 0 {
		return "", errors.New("Empty name is not allowed")
	}
	if len(name_) < 3 {
		name_ = append(name_, ("___"[0 : 3-len(name)])...)
	}
	if 255 < len(name_) {
		name_ = append(name_[0:253], "__"...)
	}
	name_ = bytes.ToLower(name_)
	for i, c := range name_ {
		if !((c >= 'a' || c <= 'z') || (c >= '0' && c <= '9') || c == '_') {
			c = '_'
		}
		name_[i] = c
	}
	return (string)(name_), nil
}

func normalizeTableName(name string) (string, error) {
	return normalizeDatabaseName(name)
}

func newTDOutputSpooler(daemon *tdOutputSpoolerDaemon, databaseName, tableName, key string) *tdOutputSpooler {
	journal := daemon.output.journalGroup.GetJournal(key)
	return &tdOutputSpooler{
		daemon:       daemon,
		ticker:       time.NewTicker(daemon.output.flushInterval),
		databaseName: databaseName,
		tableName:    tableName,
		key:          key,
		journal:      journal,
		shutdownChan: make(chan struct{}, 1),
		isShuttingDown: 0,
		client:       daemon.output.client,
	}
}

func (daemon *tdOutputSpoolerDaemon) spawnSpooler(databaseName, tableName, key string) *tdOutputSpooler {
	daemon.spoolersMtx.Lock()
	defer daemon.spoolersMtx.Unlock()
	spooler, exists := daemon.spoolers[key]
	if exists {
		return spooler
	}
	spooler = newTDOutputSpooler(daemon, databaseName, tableName, key)
	daemon.output.logger.Notice("Spawning spooler " + spooler.key)
	daemon.spoolers[spooler.key] = spooler
	daemon.wg.Add(1)
	go spooler.handle()
	return spooler
}

func (daemon *tdOutputSpoolerDaemon) cleanup() {
	func() {
		daemon.spoolersMtx.Lock()
		defer daemon.spoolersMtx.Unlock()
		for _, spooler := range daemon.spoolers {
			if atomic.CompareAndSwapUintptr(&spooler.isShuttingDown, 0, 1) {
				spooler.shutdownChan <- struct{}{}
			}
		}
	}()
	daemon.wg.Wait()
	if daemon.endNotify != nil {
		daemon.endNotify(daemon)
	}
	daemon.output.wg.Done()
}

func (daemon *tdOutputSpoolerDaemon) handle() {
	defer daemon.cleanup()
	daemon.output.logger.Notice("Spooler daemon started")
	// spawn Spooler according to the existing journals
	for _, key := range daemon.output.journalGroup.GetJournalKeys() {
		pair := strings.SplitN(key, ".", 2)
		if len(pair) != 2 {
			daemon.output.logger.Warning("Journal %s ignored", key)
			continue
		}
		daemon.spawnSpooler(pair[0], pair[1], key)
	}
outer:
	for {
		select {
		case <-daemon.shutdownChan:
			break outer
		}
	}
	daemon.output.logger.Notice("Spooler daemon ended")
}

func newTDOutputSpoolerDaemon(output *TDOutput) *tdOutputSpoolerDaemon {
	return &tdOutputSpoolerDaemon{
		output:       output,
		shutdownChan: make(chan struct{}, 1),
		spoolers:     make(map[string]*tdOutputSpooler),
		tempFactory:  TempFileRandomAccessStoreFactory{output.tempDir, "", output.gcChan},
		wg:           sync.WaitGroup{},
		endNotify: func(*tdOutputSpoolerDaemon) {
			close(output.gcChan)
		},
	}
}

func (output *TDOutput) spawnSpoolerDaemon() {
	output.logger.Notice("Spawning spooler daemon")
	output.spoolerDaemon = newTDOutputSpoolerDaemon(output)
	output.wg.Add(1)
	go output.spoolerDaemon.handle()
}

func (daemon *tdOutputSpoolerDaemon) getSpooler(tag string) (*tdOutputSpooler, error) {
	databaseName := daemon.output.databaseName
	tableName := daemon.output.tableName
	if databaseName == "*" {
		if tableName == "*" {
			c := strings.SplitN(tag, ".", 2)
			if len(c) == 1 {
				tableName = c[0]
			} else if len(c) == 2 {
				databaseName = c[0]
				tableName = c[1]
			}
		} else {
			databaseName = tag
		}
	} else {
		if tableName == "*" {
			tableName = tag
		}
	}
	databaseName, err := normalizeDatabaseName(databaseName)
	if err != nil {
		return nil, err
	}
	tableName, err = normalizeTableName(tableName)
	if err != nil {
		return nil, err
	}
	key := databaseName + "." + tableName
	return daemon.spawnSpooler(databaseName, tableName, key), nil
}

func (output *TDOutput) spawnEmitter() {
	output.logger.Notice("Spawning emitter")
	output.wg.Add(1)
	go func() {
		defer func() {
			output.spoolerDaemon.shutdownChan <- struct{}{}
			output.wg.Done()
		}()
		output.logger.Notice("Emitter started")
		buffer := bytes.Buffer{}
		for recordSet := range output.emitterChan {
			buffer.Reset()
			encoder := codec.NewEncoder(&buffer, output.codec)
			err := func() error {
				spooler, err := output.spoolerDaemon.getSpooler(recordSet.Tag)
				if err != nil {
					return err
				}
				err = encodeRecords(encoder, recordSet.Records)
				if err != nil {
					return err
				}
				output.logger.Debug("Emitter processed %d entries", len(recordSet.Records))
				return spooler.journal.Write(buffer.Bytes())
			}()
			if err != nil {
				output.logger.Error("%s", err.Error())
				continue
			}
		}
		output.logger.Notice("Emitter ended")
	}()
}

func (output *TDOutput) spawnTempFileCollector() {
	output.logger.Notice("Spawning temporary file collector")
	output.wg.Add(1)
	go func() {
		defer func() {
			output.wg.Done()
		}()
		for f := range output.gcChan {
			output.logger.Debug("Deleting %s...", f.Name())
			err := os.Remove(f.Name())
			if err != nil {
				output.logger.Warning("Failed to delete %s: %s", f.Name(), err.Error())
			}
		}
		output.logger.Debug("temporary file collector ended")
	}()
}

func (output *TDOutput) Emit(recordSets []FluentRecordSet) error {
	defer func() {
		recover()
	}()
	for _, recordSet := range recordSets {
		output.emitterChan <- recordSet
	}
	return nil
}

func (output *TDOutput) String() string {
	return "output"
}

func (output *TDOutput) Stop() {
	if atomic.CompareAndSwapUintptr(&output.isShuttingDown, 0, 1) {
		close(output.emitterChan)
	}
}

func (output *TDOutput) WaitForShutdown() {
	output.completion.L.Lock()
	output.completion.Wait()
	output.completion.L.Unlock()
}

func (output *TDOutput) Start() {
	syncCh := make(chan struct{})
	go func() {
		<-syncCh
		output.wg.Wait()
		err := output.journalGroup.Dispose()
		if err != nil {
			output.logger.Error("%s", err.Error())
		}
		output.completion.Broadcast()
	}()
	output.spawnTempFileCollector()
	output.spawnEmitter()
	output.spawnSpoolerDaemon()
	syncCh <- struct{}{}
}

func NewTDOutput(
	logger *logging.Logger,
	endpoint string,
	connectionTimeout time.Duration,
	writeTimeout time.Duration,
	flushInterval time.Duration,
	parallelism int,
	journalGroupPath string,
	maxJournalChunkSize int64,
	apiKey string,
	databaseName string,
	tableName string,
	tempDir string,
	useSsl bool,
	rootCAs *x509.CertPool,
	httpProxy string,
) (*TDOutput, error) {
	_codec := codec.MsgpackHandle{}
	_codec.MapType = reflect.TypeOf(map[string]interface{}(nil))
	_codec.RawToString = false
	_codec.StructToArray = true

	journalFactory := NewFileJournalGroupFactory(
		logger,
		randSource,
		time.Now,
		".log",
		os.FileMode(0600),
		maxJournalChunkSize,
	)
	router := (td_client.EndpointRouter)(nil)
	if endpoint != "" {
		router = &td_client.FixedEndpointRouter{endpoint}
	}
	httpProxy_ := (interface{})(nil)
	if httpProxy != "" {
		httpProxy_ = httpProxy
	}
	client, err := td_client.NewTDClient(td_client.Settings{
		ApiKey:            apiKey,
		Router:            router,
		ConnectionTimeout: connectionTimeout,
		// ReadTimeout: readTimeout, // TODO
		SendTimeout: writeTimeout,
		Ssl:         useSsl,
		RootCAs:     rootCAs,
		Proxy:       httpProxy_,
	})
	if err != nil {
		return nil, err
	}
	output := &TDOutput{
		logger:         logger,
		codec:          &_codec,
		wg:             sync.WaitGroup{},
		flushInterval:  flushInterval,
		emitterChan:    make(chan FluentRecordSet),
		isShuttingDown: 0,
		client:         client,
		databaseName:   databaseName,
		tableName:      tableName,
		tempDir:        tempDir,
		sem:            make(chan struct{}, parallelism),
		gcChan:         make(chan *os.File, 10),
		completion:     sync.Cond{L: &sync.Mutex{}},
	}
	journalGroup, err := journalFactory.GetJournalGroup(journalGroupPath, output)
	if err != nil {
		return nil, err
	}
	output.journalGroup = journalGroup
	return output, nil
}
<commit_msg>Fix race that occationally happens during flushing.<commit_after>package fluentd_forwarder

import (
	"bytes"
	"compress/gzip"
	"crypto/x509"
	"errors"
	logging "github.com/op/go-logging"
	td_client "github.com/treasure-data/td-client-go"
	"github.com/ugorji/go/codec"
	"net"
	"os"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type tdOutputSpooler struct {
	daemon         *tdOutputSpoolerDaemon
	ticker         *time.Ticker
	tag            string
	databaseName   string
	tableName      string
	key            string
	journal        Journal
	shutdownChan   chan struct{}
	isShuttingDown uintptr
	client         *td_client.TDClient
}

type tdOutputSpoolerDaemon struct {
	output       *TDOutput
	shutdownChan chan struct{}
	spoolersMtx  sync.Mutex
	spoolers     map[string]*tdOutputSpooler
	tempFactory  TempFileRandomAccessStoreFactory
	wg           sync.WaitGroup
	endNotify    func(*tdOutputSpoolerDaemon)
}

type TDOutput struct {
	logger         *logging.Logger
	codec          *codec.MsgpackHandle
	databaseName   string
	tableName      string
	tempDir        string
	enc            *codec.Encoder
	conn           net.Conn
	flushInterval  time.Duration
	wg             sync.WaitGroup
	journalGroup   JournalGroup
	emitterChan    chan FluentRecordSet
	spoolerDaemon  *tdOutputSpoolerDaemon
	isShuttingDown uintptr
	client         *td_client.TDClient
	sem            chan struct{}
	gcChan         chan *os.File
	completion     sync.Cond
}

func encodeRecords(encoder *codec.Encoder, records []TinyFluentRecord) error {
	for _, record := range records {
		e := map[string]interface{}{"time": record.Timestamp}
		for k, v := range record.Data {
			e[k] = v
		}
		err := encoder.Encode(e)
		if err != nil {
			return err
		}
	}
	return nil
}

func (spooler *tdOutputSpooler) cleanup() {
	spooler.ticker.Stop()
	spooler.journal.Dispose()
	spooler.daemon.wg.Done()
}

func (spooler *tdOutputSpooler) handle() {
	defer spooler.cleanup()
	spooler.daemon.output.logger.Notice("Spooler started")
outer:
	for {
		select {
		case <-spooler.ticker.C:
			spooler.daemon.output.logger.Notice("Flushing...")
			err := spooler.journal.Flush(func(chunk JournalChunk) interface{} {
				defer chunk.Dispose()
				if atomic.LoadUintptr(&spooler.isShuttingDown) != 0 {
					return errors.New("Flush aborted")
				}
				spooler.daemon.output.logger.Info("Flushing chunk %s", chunk.String())
				size, err := chunk.Size()
				if err != nil {
					return err
				}
				if size == 0 {
					return nil
				}
				futureErr := make(chan error, 1)
				sem := spooler.daemon.output.sem
				sem <- struct{}{}
				go func(size int64, chunk JournalChunk, futureErr chan error) {
					err := (error)(nil)
					defer func() {
						if err != nil {
							spooler.daemon.output.logger.Info("Failed to flush chunk %s (reason: %s)", chunk.String(), err.Error())
						} else {
							spooler.daemon.output.logger.Info("Completed flushing chunk %s", chunk.String())
						}
						<-sem
						// disposal must be done before notifying the initiator
						chunk.Dispose()
						futureErr <- err
					}()
					err = func() error {
						compressingBlob := NewCompressingBlob(
							chunk,
							maxInt(4096, int(size/4)),
							gzip.BestSpeed,
							&spooler.daemon.tempFactory,
						)
						defer compressingBlob.Dispose()
						_, err := spooler.client.Import(
							spooler.databaseName,
							spooler.tableName,
							"msgpack.gz",
							td_client.NewBufferingBlobSize(
								compressingBlob,
								maxInt(4096, int(size/16)),
							),
							chunk.Id(),
						)
						return err
					}()
				}(size, chunk.Dup(), futureErr)
				return (<-chan error)(futureErr)
			})
			if err != nil {
				spooler.daemon.output.logger.Error("Error during reading from the journal: %s", err.Error())
			}
		case <-spooler.shutdownChan:
			break outer
		}
	}
	spooler.daemon.output.logger.Notice("Spooler ended")
}

func normalizeDatabaseName(name string) (string, error) {
	name_ := ([]byte)(name)
	if len(name_) == 0 {
		return "", errors.New("Empty name is not allowed")
	}
	if len(name_) < 3 {
		name_ = append(name_, ("___"[0 : 3-len(name)])...)
	}
	if 255 < len(name_) {
		name_ = append(name_[0:253], "__"...)
	}
	name_ = bytes.ToLower(name_)
	for i, c := range name_ {
		if !((c >= 'a' || c <= 'z') || (c >= '0' && c <= '9') || c == '_') {
			c = '_'
		}
		name_[i] = c
	}
	return (string)(name_), nil
}

func normalizeTableName(name string) (string, error) {
	return normalizeDatabaseName(name)
}

func newTDOutputSpooler(daemon *tdOutputSpoolerDaemon, databaseName, tableName, key string) *tdOutputSpooler {
	journal := daemon.output.journalGroup.GetJournal(key)
	return &tdOutputSpooler{
		daemon:       daemon,
		ticker:       time.NewTicker(daemon.output.flushInterval),
		databaseName: databaseName,
		tableName:    tableName,
		key:          key,
		journal:      journal,
		shutdownChan: make(chan struct{}, 1),
		isShuttingDown: 0,
		client:       daemon.output.client,
	}
}

func (daemon *tdOutputSpoolerDaemon) spawnSpooler(databaseName, tableName, key string) *tdOutputSpooler {
	daemon.spoolersMtx.Lock()
	defer daemon.spoolersMtx.Unlock()
	spooler, exists := daemon.spoolers[key]
	if exists {
		return spooler
	}
	spooler = newTDOutputSpooler(daemon, databaseName, tableName, key)
	daemon.output.logger.Notice("Spawning spooler " + spooler.key)
	daemon.spoolers[spooler.key] = spooler
	daemon.wg.Add(1)
	go spooler.handle()
	return spooler
}

func (daemon *tdOutputSpoolerDaemon) cleanup() {
	func() {
		daemon.spoolersMtx.Lock()
		defer daemon.spoolersMtx.Unlock()
		for _, spooler := range daemon.spoolers {
			if atomic.CompareAndSwapUintptr(&spooler.isShuttingDown, 0, 1) {
				spooler.shutdownChan <- struct{}{}
			}
		}
	}()
	daemon.wg.Wait()
	if daemon.endNotify != nil {
		daemon.endNotify(daemon)
	}
	daemon.output.wg.Done()
}

func (daemon *tdOutputSpoolerDaemon) handle() {
	defer daemon.cleanup()
	daemon.output.logger.Notice("Spooler daemon started")
	// spawn Spooler according to the existing journals
	for _, key := range daemon.output.journalGroup.GetJournalKeys() {
		pair := strings.SplitN(key, ".", 2)
		if len(pair) != 2 {
			daemon.output.logger.Warning("Journal %s ignored", key)
			continue
		}
		daemon.spawnSpooler(pair[0], pair[1], key)
	}
outer:
	for {
		select {
		case <-daemon.shutdownChan:
			break outer
		}
	}
	daemon.output.logger.Notice("Spooler daemon ended")
}

func newTDOutputSpoolerDaemon(output *TDOutput) *tdOutputSpoolerDaemon {
	return &tdOutputSpoolerDaemon{
		output:       output,
		shutdownChan: make(chan struct{}, 1),
		spoolers:     make(map[string]*tdOutputSpooler),
		tempFactory:  TempFileRandomAccessStoreFactory{output.tempDir, "", output.gcChan},
		wg:           sync.WaitGroup{},
		endNotify: func(*tdOutputSpoolerDaemon) {
			close(output.gcChan)
		},
	}
}

func (output *TDOutput) spawnSpoolerDaemon() {
	output.logger.Notice("Spawning spooler daemon")
	output.spoolerDaemon = newTDOutputSpoolerDaemon(output)
	output.wg.Add(1)
	go output.spoolerDaemon.handle()
}

func (daemon *tdOutputSpoolerDaemon) getSpooler(tag string) (*tdOutputSpooler, error) {
	databaseName := daemon.output.databaseName
	tableName := daemon.output.tableName
	if databaseName == "*" {
		if tableName == "*" {
			c := strings.SplitN(tag, ".", 2)
			if len(c) == 1 {
				tableName = c[0]
			} else if len(c) == 2 {
				databaseName = c[0]
				tableName = c[1]
			}
		} else {
			databaseName = tag
		}
	} else {
		if tableName == "*" {
			tableName = tag
		}
	}
	databaseName, err := normalizeDatabaseName(databaseName)
	if err != nil {
		return nil, err
	}
	tableName, err = normalizeTableName(tableName)
	if err != nil {
		return nil, err
	}
	key := databaseName + "." + tableName
	return daemon.spawnSpooler(databaseName, tableName, key), nil
}

func (output *TDOutput) spawnEmitter() {
	output.logger.Notice("Spawning emitter")
	output.wg.Add(1)
	go func() {
		defer func() {
			output.spoolerDaemon.shutdownChan <- struct{}{}
			output.wg.Done()
		}()
		output.logger.Notice("Emitter started")
		buffer := bytes.Buffer{}
		for recordSet := range output.emitterChan {
			buffer.Reset()
			encoder := codec.NewEncoder(&buffer, output.codec)
			err := func() error {
				spooler, err := output.spoolerDaemon.getSpooler(recordSet.Tag)
				if err != nil {
					return err
				}
				err = encodeRecords(encoder, recordSet.Records)
				if err != nil {
					return err
				}
				output.logger.Debug("Emitter processed %d entries", len(recordSet.Records))
				return spooler.journal.Write(buffer.Bytes())
			}()
			if err != nil {
				output.logger.Error("%s", err.Error())
				continue
			}
		}
		output.logger.Notice("Emitter ended")
	}()
}

func (output *TDOutput) spawnTempFileCollector() {
	output.logger.Notice("Spawning temporary file collector")
	output.wg.Add(1)
	go func() {
		defer func() {
			output.wg.Done()
		}()
		for f := range output.gcChan {
			output.logger.Debug("Deleting %s...", f.Name())
			err := os.Remove(f.Name())
			if err != nil {
				output.logger.Warning("Failed to delete %s: %s", f.Name(), err.Error())
			}
		}
		output.logger.Debug("temporary file collector ended")
	}()
}

func (output *TDOutput) Emit(recordSets []FluentRecordSet) error {
	defer func() {
		recover()
	}()
	for _, recordSet := range recordSets {
		output.emitterChan <- recordSet
	}
	return nil
}

func (output *TDOutput) String() string {
	return "output"
}

func (output *TDOutput) Stop() {
	if atomic.CompareAndSwapUintptr(&output.isShuttingDown, 0, 1) {
		close(output.emitterChan)
	}
}

func (output *TDOutput) WaitForShutdown() {
	output.completion.L.Lock()
	output.completion.Wait()
	output.completion.L.Unlock()
}

func (output *TDOutput) Start() {
	syncCh := make(chan struct{})
	go func() {
		<-syncCh
		output.wg.Wait()
		err := output.journalGroup.Dispose()
		if err != nil {
			output.logger.Error("%s", err.Error())
		}
		output.completion.Broadcast()
	}()
	output.spawnTempFileCollector()
	output.spawnEmitter()
	output.spawnSpoolerDaemon()
	syncCh <- struct{}{}
}

func NewTDOutput(
	logger *logging.Logger,
	endpoint string,
	connectionTimeout time.Duration,
	writeTimeout time.Duration,
	flushInterval time.Duration,
	parallelism int,
	journalGroupPath string,
	maxJournalChunkSize int64,
	apiKey string,
	databaseName string,
	tableName string,
	tempDir string,
	useSsl bool,
	rootCAs *x509.CertPool,
	httpProxy string,
) (*TDOutput, error) {
	_codec := codec.MsgpackHandle{}
	_codec.MapType = reflect.TypeOf(map[string]interface{}(nil))
	_codec.RawToString = false
	_codec.StructToArray = true

	journalFactory := NewFileJournalGroupFactory(
		logger,
		randSource,
		time.Now,
		".log",
		os.FileMode(0600),
		maxJournalChunkSize,
	)
	router := (td_client.EndpointRouter)(nil)
	if endpoint != "" {
		router = &td_client.FixedEndpointRouter{endpoint}
	}
	httpProxy_ := (interface{})(nil)
	if httpProxy != "" {
		httpProxy_ = httpProxy
	}
	client, err := td_client.NewTDClient(td_client.Settings{
		ApiKey:            apiKey,
		Router:            router,
		ConnectionTimeout: connectionTimeout,
		// ReadTimeout: readTimeout, // TODO
		SendTimeout: writeTimeout,
		Ssl:         useSsl,
		RootCAs:     rootCAs,
		Proxy:       httpProxy_,
	})
	if err != nil {
		return nil, err
	}
	output := &TDOutput{
		logger:         logger,
		codec:          &_codec,
		wg:             sync.WaitGroup{},
		flushInterval:  flushInterval,
		emitterChan:    make(chan FluentRecordSet),
		isShuttingDown: 0,
		client:         client,
		databaseName:   databaseName,
		tableName:      tableName,
		tempDir:        tempDir,
		sem:            make(chan struct{}, parallelism),
		gcChan:         make(chan *os.File, 10),
		completion:     sync.Cond{L: &sync.Mutex{}},
	}
	journalGroup, err := journalFactory.GetJournalGroup(journalGroupPath, output)
	if err != nil {
		return nil, err
	}
	output.journalGroup = journalGroup
	return output, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Package productsearch contains samples for Google Cloud Vision API Product Search.
package productsearch

// [START vision_product_search_import_product_images]

import (
	"context"
	"fmt"
	"io"

	vision "cloud.google.com/go/vision/apiv1"
	visionpb "google.golang.org/genproto/googleapis/cloud/vision/v1"
)

// importProductSets creates a product set using information in a csv file on GCS.
func importProductSets(w io.Writer, projectID string, location string, gcsURI string) error {
	ctx := context.Background()
	c, err := vision.NewProductSearchClient(ctx)
	if err != nil {
		return fmt.Errorf("NewProductSearchClient: %v", err)
	}

	req := &visionpb.ImportProductSetsRequest{
		Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location),
		InputConfig: &visionpb.ImportProductSetsInputConfig{
			Source: &visionpb.ImportProductSetsInputConfig_GcsSource{
				GcsSource: &visionpb.ImportProductSetsGcsSource{
					CsvFileUri: gcsURI,
				},
			},
		},
	}

	op, err := c.ImportProductSets(ctx, req)
	if err != nil {
		return fmt.Errorf("ImportProductSets: %v", err)
	}

	fmt.Fprintf(w, "Processing operation name: %s\n", op.Name())

	resp, err := op.Wait(ctx)
	if err != nil {
		return fmt.Errorf("Wait: %v", err)
	}

	fmt.Fprintf(w, "processing done.\n")

	for i, status := range resp.Statuses {
		// `0` is the coee for OK in google.rpc.code
		fmt.Fprintf(w, "Status of processing line %d of the csv: %d\n", i, status.Code)

		if status.Code == 0 {
			fmt.Fprintf(w, "Reference image name: %s\n", resp.ReferenceImages[i].Name)
		} else {
			fmt.Fprintf(w, "Status code not OK: %s\n", status.Message)
		}
	}

	return nil
}

// [END vision_product_search_import_product_images]
<commit_msg>vision: add region tag vision_product_search_tutorial_import (#718)<commit_after>// Copyright 2018 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Package productsearch contains samples for Google Cloud Vision API Product Search.
package productsearch

// [START vision_product_search_import_product_images]

// [START vision_product_search_tutorial_import]

import (
	"context"
	"fmt"
	"io"

	vision "cloud.google.com/go/vision/apiv1"
	visionpb "google.golang.org/genproto/googleapis/cloud/vision/v1"
)

// [END vision_product_search_tutorial_import]

// importProductSets creates a product set using information in a csv file on GCS.
func importProductSets(w io.Writer, projectID string, location string, gcsURI string) error {
	ctx := context.Background()
	c, err := vision.NewProductSearchClient(ctx)
	if err != nil {
		return fmt.Errorf("NewProductSearchClient: %v", err)
	}

	req := &visionpb.ImportProductSetsRequest{
		Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location),
		InputConfig: &visionpb.ImportProductSetsInputConfig{
			Source: &visionpb.ImportProductSetsInputConfig_GcsSource{
				GcsSource: &visionpb.ImportProductSetsGcsSource{
					CsvFileUri: gcsURI,
				},
			},
		},
	}

	op, err := c.ImportProductSets(ctx, req)
	if err != nil {
		return fmt.Errorf("ImportProductSets: %v", err)
	}

	fmt.Fprintf(w, "Processing operation name: %s\n", op.Name())

	resp, err := op.Wait(ctx)
	if err != nil {
		return fmt.Errorf("Wait: %v", err)
	}

	fmt.Fprintf(w, "processing done.\n")

	for i, status := range resp.Statuses {
		// `0` is the coee for OK in google.rpc.code
		fmt.Fprintf(w, "Status of processing line %d of the csv: %d\n", i, status.Code)

		if status.Code == 0 {
			fmt.Fprintf(w, "Reference image name: %s\n", resp.ReferenceImages[i].Name)
		} else {
			fmt.Fprintf(w, "Status code not OK: %s\n", status.Message)
		}
	}

	return nil
}

// [END vision_product_search_import_product_images]
<|endoftext|>


<commit_before>package radial

import (
	"testing"
)

const MAX_DIFF float64 = 0.0000001

func diff(a, b float64) bool {
	if a > b {
		return (a - b) > MAX_DIFF
	}
	return (b - a) > MAX_DIFF
}

func TestDeg2rad(t *testing.T) {
	if diff(DegToRad(1), 0.0174533) {
		t.Error("! DegToRad is not accurately converting between degrees and radians.")
	}

	if diff(DegToRad(5), 0.0872665) {
		t.Error("! DegToRad is not accurately converting between degrees and radians.")
	}
}

func TestRad2deg(t *testing.T) {
	if diff(RadToDeg(1), 57.2957795) {
		t.Error("! RadToDeg is not accurately converting between degrees and radians.")
	}

	if diff(RadToDeg(5), 286.4788975) {
		t.Error("! RadToDeg is not accurately converting between degrees and radians.")
	}
}
<commit_msg>fixed constant name to make it compliant with Go coding style<commit_after>package radial

import (
	"testing"
)

const MaxDiff float64 = 0.0000001

func diff(a, b float64) bool {
	if a > b {
		return (a - b) > MaxDiff
	}
	return (b - a) > MaxDiff
}

func TestDeg2rad(t *testing.T) {
	if diff(DegToRad(1), 0.0174533) {
		t.Error("! DegToRad is not accurately converting between degrees and radians.")
	}

	if diff(DegToRad(5), 0.0872665) {
		t.Error("! DegToRad is not accurately converting between degrees and radians.")
	}
}

func TestRad2deg(t *testing.T) {
	if diff(RadToDeg(1), 57.2957795) {
		t.Error("! RadToDeg is not accurately converting between degrees and radians.")
	}

	if diff(RadToDeg(5), 286.4788975) {
		t.Error("! RadToDeg is not accurately converting between degrees and radians.")
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"

	log "github.com/golang/glog"
	"github.com/google/kne/deploy"
	cpb "github.com/google/kne/proto/controller"
	tpb "github.com/google/kne/proto/topo"
	"github.com/google/kne/topo"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/prototext"
	"k8s.io/client-go/util/homedir"
)

var (
	defaultKubeCfg            = ""
	defaultTopoBasePath       = ""
	defaultMetallbManifestDir = ""
	defaultMeshnetManifestDir = ""
	// Flags.
	port = flag.Int("port", 50051, "Controller server port")
)

func init() {
	if home := homedir.HomeDir(); home != "" {
		defaultKubeCfg = filepath.Join(home, ".kube", "config")
		defaultTopoBasePath = filepath.Join(home, "kne", "examples")
		defaultMeshnetManifestDir = filepath.Join(home, "kne", "manifests", "meshnet", "base")
		defaultMetallbManifestDir = filepath.Join(home, "kne", "manifests", "metallb")
	}
}

type server struct {
	cpb.UnimplementedTopologyManagerServer

	muDeploy    sync.Mutex // guards deployements map
	deployments map[string]*deploy.Deployment
	muTopo      sync.Mutex        // guards topos map
	topos       map[string][]byte // stores the topology protobuf from the initial topology creation request
}

func newServer() *server {
	return &server{
		deployments: map[string]*deploy.Deployment{},
		topos:       map[string][]byte{},
	}
}

func newDeployment(req *cpb.CreateClusterRequest) (*deploy.Deployment, error) {
	d := &deploy.Deployment{}
	switch kind := req.ClusterSpec.(type) {
	case *cpb.CreateClusterRequest_Kind:
		d.Cluster = &deploy.KindSpec{
			Name:                     req.GetKind().Name,
			Recycle:                  req.GetKind().Recycle,
			Version:                  req.GetKind().Version,
			Image:                    req.GetKind().Image,
			Retain:                   req.GetKind().Retain,
			GoogleArtifactRegistries: req.GetKind().GoogleArtifactRegistries,
			ContainerImages:          req.GetKind().ContainerImages,
		}
	default:
		return nil, fmt.Errorf("cluster type not supported: %T", kind)
	}
	switch metallb := req.IngressSpec.(type) {
	case *cpb.CreateClusterRequest_Metallb:
		l := &deploy.MetalLBSpec{}
		path := defaultMetallbManifestDir
		if req.GetMetallb().ManifestDir != "" {
			path = req.GetMetallb().ManifestDir
		}
		p, err := validatePath(path)
		if err != nil {
			return nil, fmt.Errorf("failed to validate path %q", path)
		}
		l.ManifestDir = p
		l.IPCount = int(req.GetMetallb().IpCount)
		l.Version = req.GetMetallb().Version
		d.Ingress = l
	default:
		return nil, fmt.Errorf("ingress spec not supported: %T", metallb)
	}
	switch meshnet := req.CniSpec.(type) {
	case *cpb.CreateClusterRequest_Meshnet:
		m := &deploy.MeshnetSpec{}
		path := defaultMeshnetManifestDir
		if req.GetMeshnet().ManifestDir != "" {
			path = req.GetMeshnet().ManifestDir
		}
		p, err := validatePath(path)
		if err != nil {
			return nil, fmt.Errorf("failed to validate path %q", path)
		}
		m.Image = req.GetMeshnet().Image
		m.ManifestDir = p
		d.CNI = m
	default:
		return nil, fmt.Errorf("cni type not supported: %T", meshnet)
	}
	return d, nil
}

func (s *server) CreateCluster(ctx context.Context, req *cpb.CreateClusterRequest) (*cpb.CreateClusterResponse, error) {
	log.Infof("Received CreateCluster request: %v", req)
	d, err := newDeployment(req)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "unable to parse request: %v", err)
	}
	log.Infof("Parsed request into deployment: %v", d)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	if _, ok := s.deployments[d.Cluster.GetName()]; ok { // if OK
		return nil, status.Errorf(codes.AlreadyExists, "cluster %q already exists", d.Cluster.GetName())
	}
	if err := d.Deploy(ctx, defaultKubeCfg); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to deploy cluster: %v", err)
	}
	s.deployments[d.Cluster.GetName()] = d
	log.Infof("Cluster %q deployed and ready for topology", d.Cluster.GetName())
	resp := &cpb.CreateClusterResponse{
		Name:  d.Cluster.GetName(),
		State: cpb.ClusterState_CLUSTER_STATE_RUNNING,
	}
	return resp, nil
}

func (s *server) DeleteCluster(ctx context.Context, req *cpb.DeleteClusterRequest) (*cpb.DeleteClusterResponse, error) {
	log.Infof("Received DeleteCluster request: %v", req)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	d, ok := s.deployments[req.GetName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "cluster %q not found, can only delete clusters created using TopologyManager", req.GetName())
	}
	if err := d.Delete(); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete cluster: %v", err)
	}
	delete(s.deployments, req.GetName())
	log.Infof("Deleted cluster %q", d.Cluster.GetName())
	return &cpb.DeleteClusterResponse{}, nil
}

func (s *server) ShowCluster(ctx context.Context, req *cpb.ShowClusterRequest) (*cpb.ShowClusterResponse, error) {
	log.Infof("Received ShowCluster request: %v", req)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	d, ok := s.deployments[req.GetName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "cluster %q not found, can only show clusters created using TopologyManager", req.GetName())
	}
	if err := d.Healthy(ctx); err != nil {
		return &cpb.ShowClusterResponse{State: cpb.ClusterState_CLUSTER_STATE_ERROR}, nil
	}
	return &cpb.ShowClusterResponse{State: cpb.ClusterState_CLUSTER_STATE_RUNNING}, nil
}

func (s *server) CreateTopology(ctx context.Context, req *cpb.CreateTopologyRequest) (*cpb.CreateTopologyResponse, error) {
	log.Infof("Received CreateTopology request: %v", req)
	topoPb := req.GetTopology()
	if topoPb == nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid request: missing topology protobuf")
	}
	if topoPb.GetName() == "" {
		return nil, status.Errorf(codes.InvalidArgument, "missing topology name")
	}

	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	if _, ok := s.topos[topoPb.GetName()]; ok {
		return nil, status.Errorf(codes.AlreadyExists, "topology %q already exists", req.Topology.GetName())
	}

	for _, node := range topoPb.Nodes {
		if node.GetConfig() == nil || node.GetConfig().GetFile() == "" {
			// A config section is not required: you are allowed to bring up a
			// topology with no initial config.
			continue
		}
		path := node.GetConfig().GetFile()
		if !filepath.IsAbs(path) {
			path = filepath.Join(defaultTopoBasePath, path)
		}
		node.GetConfig().ConfigData = &tpb.Config_File{File: path}
		log.Infof("Checking config path: %q", node.GetConfig().GetFile())
		if _, err := validatePath(node.GetConfig().GetFile()); err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "config file not found for node %q: %v", node.GetName(), err)
		}
	}
	// Saves the original topology protobuf.
	txtPb, err := prototext.Marshal(topoPb)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	path := defaultKubeCfg
	if req.Kubecfg != "" {
		path = req.Kubecfg
	}
	kcfg, err := validatePath(path)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "kubecfg %q does not exist: %v", path, err)
	}
	if err := topo.CreateTopology(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	}); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create topology: %v", err)
	}

	s.topos[topoPb.GetName()] = txtPb
	return &cpb.CreateTopologyResponse{
		TopologyName: req.Topology.GetName(),
		State:        cpb.TopologyState_TOPOLOGY_STATE_RUNNING,
	}, nil
}

func (s *server) DeleteTopology(ctx context.Context, req *cpb.DeleteTopologyRequest) (*cpb.DeleteTopologyResponse, error) {
	log.Infof("Received DeleteTopology request: %v", req)
	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	txtPb, ok := s.topos[req.GetTopologyName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "topology %q not found", req.GetTopologyName())
	}
	topoPb := &tpb.Topology{}
	if err := prototext.Unmarshal(txtPb, topoPb); err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	kcfg, err := validatePath(defaultKubeCfg)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "default kubecfg %q does not exist: %v", defaultKubeCfg, err)
	}
	if err := topo.DeleteTopology(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	}); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete topology: %v", err)
	}
	return &cpb.DeleteTopologyResponse{}, nil
}

func (s *server) ShowTopology(ctx context.Context, req *cpb.ShowTopologyRequest) (*cpb.ShowTopologyResponse, error) {
	log.Infof("Received ShowTopology request: %v", req)
	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	txtPb, ok := s.topos[req.GetTopologyName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "topology %q not found", req.GetTopologyName())
	}
	topoPb := &tpb.Topology{}
	if err := prototext.Unmarshal(txtPb, topoPb); err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	kcfg, err := validatePath(defaultKubeCfg)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "default kubecfg %q does not exist: %v", defaultKubeCfg, err)
	}
	resp, err := topo.GetTopologyServices(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to show topology: %v", err)
	}
	return resp, nil
}

func validatePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", fmt.Errorf("failed to evaluate absolute path: %q", path)
	}
	if _, err := os.Stat(path); err != nil {
		return "", fmt.Errorf("path %q does not exist", path)
	}
	return path, nil
}

func main() {
	flag.Parse()
	addr := fmt.Sprintf(":%d", *port)
	lis, err := net.Listen("tcp6", addr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	creds := alts.NewServerCreds(alts.DefaultServerOptions())
	s := grpc.NewServer(grpc.Creds(creds))
	cpb.RegisterTopologyManagerServer(s, newServer())
	log.Infof("Controller server listening at %v", lis.Addr())
	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}
<commit_msg>Reorder logic for assignment<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"

	log "github.com/golang/glog"
	"github.com/google/kne/deploy"
	cpb "github.com/google/kne/proto/controller"
	tpb "github.com/google/kne/proto/topo"
	"github.com/google/kne/topo"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/prototext"
	"k8s.io/client-go/util/homedir"
)

var (
	defaultKubeCfg            = ""
	defaultTopoBasePath       = ""
	defaultMetallbManifestDir = ""
	defaultMeshnetManifestDir = ""
	// Flags.
	port = flag.Int("port", 50051, "Controller server port")
)

func init() {
	if home := homedir.HomeDir(); home != "" {
		defaultKubeCfg = filepath.Join(home, ".kube", "config")
		defaultTopoBasePath = filepath.Join(home, "kne", "examples")
		defaultMeshnetManifestDir = filepath.Join(home, "kne", "manifests", "meshnet", "base")
		defaultMetallbManifestDir = filepath.Join(home, "kne", "manifests", "metallb")
	}
}

type server struct {
	cpb.UnimplementedTopologyManagerServer

	muDeploy    sync.Mutex // guards deployements map
	deployments map[string]*deploy.Deployment
	muTopo      sync.Mutex        // guards topos map
	topos       map[string][]byte // stores the topology protobuf from the initial topology creation request
}

func newServer() *server {
	return &server{
		deployments: map[string]*deploy.Deployment{},
		topos:       map[string][]byte{},
	}
}

func newDeployment(req *cpb.CreateClusterRequest) (*deploy.Deployment, error) {
	d := &deploy.Deployment{}
	switch kind := req.ClusterSpec.(type) {
	case *cpb.CreateClusterRequest_Kind:
		d.Cluster = &deploy.KindSpec{
			Name:                     req.GetKind().Name,
			Recycle:                  req.GetKind().Recycle,
			Version:                  req.GetKind().Version,
			Image:                    req.GetKind().Image,
			Retain:                   req.GetKind().Retain,
			GoogleArtifactRegistries: req.GetKind().GoogleArtifactRegistries,
			ContainerImages:          req.GetKind().ContainerImages,
		}
	default:
		return nil, fmt.Errorf("cluster type not supported: %T", kind)
	}
	switch metallb := req.IngressSpec.(type) {
	case *cpb.CreateClusterRequest_Metallb:
		l := &deploy.MetalLBSpec{}
		path := defaultMetallbManifestDir
		if req.GetMetallb().ManifestDir != "" {
			path = req.GetMetallb().ManifestDir
		}
		p, err := validatePath(path)
		if err != nil {
			return nil, fmt.Errorf("failed to validate path %q", path)
		}
		l.ManifestDir = p
		l.IPCount = int(req.GetMetallb().IpCount)
		l.Version = req.GetMetallb().Version
		d.Ingress = l
	default:
		return nil, fmt.Errorf("ingress spec not supported: %T", metallb)
	}
	switch meshnet := req.CniSpec.(type) {
	case *cpb.CreateClusterRequest_Meshnet:
		m := &deploy.MeshnetSpec{}
		path := defaultMeshnetManifestDir
		if req.GetMeshnet().ManifestDir != "" {
			path = req.GetMeshnet().ManifestDir
		}
		p, err := validatePath(path)
		if err != nil {
			return nil, fmt.Errorf("failed to validate path %q", path)
		}
		m.Image = req.GetMeshnet().Image
		m.ManifestDir = p
		d.CNI = m
	default:
		return nil, fmt.Errorf("cni type not supported: %T", meshnet)
	}
	return d, nil
}

func (s *server) CreateCluster(ctx context.Context, req *cpb.CreateClusterRequest) (*cpb.CreateClusterResponse, error) {
	log.Infof("Received CreateCluster request: %v", req)
	d, err := newDeployment(req)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "unable to parse request: %v", err)
	}
	log.Infof("Parsed request into deployment: %v", d)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	if _, ok := s.deployments[d.Cluster.GetName()]; ok { // if OK
		return nil, status.Errorf(codes.AlreadyExists, "cluster %q already exists", d.Cluster.GetName())
	}
	if err := d.Deploy(ctx, defaultKubeCfg); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to deploy cluster: %v", err)
	}
	s.deployments[d.Cluster.GetName()] = d
	log.Infof("Cluster %q deployed and ready for topology", d.Cluster.GetName())
	resp := &cpb.CreateClusterResponse{
		Name:  d.Cluster.GetName(),
		State: cpb.ClusterState_CLUSTER_STATE_RUNNING,
	}
	return resp, nil
}

func (s *server) DeleteCluster(ctx context.Context, req *cpb.DeleteClusterRequest) (*cpb.DeleteClusterResponse, error) {
	log.Infof("Received DeleteCluster request: %v", req)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	d, ok := s.deployments[req.GetName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "cluster %q not found, can only delete clusters created using TopologyManager", req.GetName())
	}
	if err := d.Delete(); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete cluster: %v", err)
	}
	delete(s.deployments, req.GetName())
	log.Infof("Deleted cluster %q", d.Cluster.GetName())
	return &cpb.DeleteClusterResponse{}, nil
}

func (s *server) ShowCluster(ctx context.Context, req *cpb.ShowClusterRequest) (*cpb.ShowClusterResponse, error) {
	log.Infof("Received ShowCluster request: %v", req)
	s.muDeploy.Lock()
	defer s.muDeploy.Unlock()
	d, ok := s.deployments[req.GetName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "cluster %q not found, can only show clusters created using TopologyManager", req.GetName())
	}
	if err := d.Healthy(ctx); err != nil {
		return &cpb.ShowClusterResponse{State: cpb.ClusterState_CLUSTER_STATE_ERROR}, nil
	}
	return &cpb.ShowClusterResponse{State: cpb.ClusterState_CLUSTER_STATE_RUNNING}, nil
}

func (s *server) CreateTopology(ctx context.Context, req *cpb.CreateTopologyRequest) (*cpb.CreateTopologyResponse, error) {
	log.Infof("Received CreateTopology request: %v", req)
	topoPb := req.GetTopology()
	if topoPb == nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid request: missing topology protobuf")
	}
	if topoPb.GetName() == "" {
		return nil, status.Errorf(codes.InvalidArgument, "missing topology name")
	}

	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	if _, ok := s.topos[topoPb.GetName()]; ok {
		return nil, status.Errorf(codes.AlreadyExists, "topology %q already exists", req.Topology.GetName())
	}

	for _, node := range topoPb.Nodes {
		if node.GetConfig() == nil || node.GetConfig().GetFile() == "" {
			// A config section is not required: you are allowed to bring up a
			// topology with no initial config.
			continue
		}
		path := node.GetConfig().GetFile()
		if !filepath.IsAbs(path) {
			path = filepath.Join(defaultTopoBasePath, path)
		}
		log.Infof("Checking config path: %q", path)
		if _, err := validatePath(path); err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "config file not found for node %q: %v", node.GetName(), err)
		}
		node.GetConfig().ConfigData = &tpb.Config_File{File: path}
	}
	// Saves the original topology protobuf.
	txtPb, err := prototext.Marshal(topoPb)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	path := defaultKubeCfg
	if req.Kubecfg != "" {
		path = req.Kubecfg
	}
	kcfg, err := validatePath(path)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "kubecfg %q does not exist: %v", path, err)
	}
	if err := topo.CreateTopology(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	}); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create topology: %v", err)
	}

	s.topos[topoPb.GetName()] = txtPb
	return &cpb.CreateTopologyResponse{
		TopologyName: req.Topology.GetName(),
		State:        cpb.TopologyState_TOPOLOGY_STATE_RUNNING,
	}, nil
}

func (s *server) DeleteTopology(ctx context.Context, req *cpb.DeleteTopologyRequest) (*cpb.DeleteTopologyResponse, error) {
	log.Infof("Received DeleteTopology request: %v", req)
	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	txtPb, ok := s.topos[req.GetTopologyName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "topology %q not found", req.GetTopologyName())
	}
	topoPb := &tpb.Topology{}
	if err := prototext.Unmarshal(txtPb, topoPb); err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	kcfg, err := validatePath(defaultKubeCfg)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "default kubecfg %q does not exist: %v", defaultKubeCfg, err)
	}
	if err := topo.DeleteTopology(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	}); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete topology: %v", err)
	}
	return &cpb.DeleteTopologyResponse{}, nil
}

func (s *server) ShowTopology(ctx context.Context, req *cpb.ShowTopologyRequest) (*cpb.ShowTopologyResponse, error) {
	log.Infof("Received ShowTopology request: %v", req)
	s.muTopo.Lock()
	defer s.muTopo.Unlock()
	txtPb, ok := s.topos[req.GetTopologyName()]
	if !ok {
		return nil, status.Errorf(codes.NotFound, "topology %q not found", req.GetTopologyName())
	}
	topoPb := &tpb.Topology{}
	if err := prototext.Unmarshal(txtPb, topoPb); err != nil {
		return nil, status.Errorf(codes.Internal, "invalid topology protobuf: %v", err)
	}
	kcfg, err := validatePath(defaultKubeCfg)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "default kubecfg %q does not exist: %v", defaultKubeCfg, err)
	}
	resp, err := topo.GetTopologyServices(ctx, topo.TopologyParams{
		TopoNewOptions: []topo.Option{topo.WithTopology(topoPb)},
		Kubecfg:        kcfg,
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to show topology: %v", err)
	}
	return resp, nil
}

func validatePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", fmt.Errorf("failed to evaluate absolute path: %q", path)
	}
	if _, err := os.Stat(path); err != nil {
		return "", fmt.Errorf("path %q does not exist", path)
	}
	return path, nil
}

func main() {
	flag.Parse()
	addr := fmt.Sprintf(":%d", *port)
	lis, err := net.Listen("tcp6", addr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	creds := alts.NewServerCreds(alts.DefaultServerOptions())
	s := grpc.NewServer(grpc.Creds(creds))
	cpb.RegisterTopologyManagerServer(s, newServer())
	log.Infof("Controller server listening at %v", lis.Addr())
	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}
<|endoftext|>


<commit_before>package actions

import (
	"encoding/json"
	"os"
	"testing"
)

func TestLink(t *testing.T) {
	f, err := os.Open("infoUrl.json")
	if err != nil {
		t.Error(err.Error())
	}

	data := make(map[string]string)
	dec := json.NewDecoder(f)
	dec.Decode(&data)

	if data["bter"] != "https://bter.com/" {
		t.Log(data)
		t.Fail()
	}
}
<commit_msg>test multi lines<commit_after>package actions

import (
	"encoding/json"
	"os"
	"testing"
)

func TestLink(t *testing.T) {
	f, err := os.Open("../infoUrl.json")
	if err != nil {
		t.Error(err.Error())
	}

	data := make(map[string][]string)
	dec := json.NewDecoder(f)
	dec.Decode(&data)

	if data["bter"][0] != "https://bter.com/" {
		t.Log(data)
		t.Fail()
	}
}
<|endoftext|>