<commit_before>package appui

import (
	"sort"
	"strconv"
	"strings"
	"sync"

	gizaktermui "github.com/gizak/termui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/moncho/dry/ui/termui"
)

var defaultNetworkTableHeader = networkTableHeader()

var networkTableHeaders = []SortableColumnHeader{
	{`NETWORK ID`, docker.SortNetworksByID},
	{`NAME`, docker.SortNetworksByName},
	{`DRIVER`, docker.SortNetworksByDriver},
	{`CONTAINERS`, docker.SortNetworksByContainerCount},
	{`SERVICES`, docker.SortNetworksByServiceCount},
	{`SCOPE`, docker.NoSortNetworks},
	{`SUBNET`, docker.SortNetworksBySubnet},
	{`GATEWAY`, docker.NoSortNetworks},
}

//DockerNetworksWidget knows how render a container list
type DockerNetworksWidget struct {
	dockerDaemon         docker.NetworkAPI
	header               *termui.TableHeader
	filteredRows         []*NetworkRow
	totalRows            []*NetworkRow
	filterPattern        string
	height, width        int
	selectedIndex        int
	startIndex, endIndex int
	x, y                 int
	sortMode             docker.SortMode
	mounted              bool
	sync.RWMutex
}

//NewDockerNetworksWidget creates a renderer for a network list
func NewDockerNetworksWidget(dockerDaemon docker.NetworkAPI, y int) *DockerNetworksWidget {
	return &DockerNetworksWidget{
		dockerDaemon: dockerDaemon,
		y:            y,
		header:       defaultNetworkTableHeader,
		height:       MainScreenAvailableHeight(),
		sortMode:     docker.SortNetworksByID,
		width:        ui.ActiveScreen.Dimensions.Width}
}

//Buffer returns the content of this widget as a termui.Buffer
func (s *DockerNetworksWidget) Buffer() gizaktermui.Buffer {
	s.Lock()
	defer s.Unlock()
	y := s.y
	buf := gizaktermui.NewBuffer()
	if s.mounted {
		s.prepareForRendering()
		widgetHeader := NewWidgetHeader()
		widgetHeader.HeaderEntry("Networks", strconv.Itoa(s.RowCount()))
		if s.filterPattern != "" {
			widgetHeader.HeaderEntry("Active filter", s.filterPattern)
		}

		widgetHeader.Y = y
		buf.Merge(widgetHeader.Buffer())
		y += widgetHeader.GetHeight()
		//Empty line between the header and the rest of the content
		y++
		s.updateHeader()
		s.header.SetY(y)
		buf.Merge(s.header.Buffer())
		y += s.header.GetHeight()

		selected := s.selectedIndex - s.startIndex

		for i, imageRow := range s.visibleRows() {
			imageRow.SetY(y)
			y += imageRow.GetHeight()
			if i != selected {
				imageRow.NotHighlighted()
			} else {
				imageRow.Highlighted()
			}
			buf.Merge(imageRow.Buffer())
		}
	}
	return buf
}

//Filter filters the network list by the given filter
func (s *DockerNetworksWidget) Filter(filter string) {
	s.Lock()
	defer s.Unlock()
	s.filterPattern = filter
}

//Mount tells this widget to be ready for rendering
func (s *DockerNetworksWidget) Mount() error {
	s.Lock()
	defer s.Unlock()
	if !s.mounted {
		networks, err := s.dockerDaemon.Networks()
		if err != nil {
			return err
		}

		networkRows := make([]*NetworkRow, len(networks))
		for i, network := range networks {
			networkRows[i] = NewNetworkRow(network, s.header)
		}
		s.totalRows = networkRows
		s.mounted = true
		s.align()
	}
	return nil
}

//Name returns this widget name
func (s *DockerNetworksWidget) Name() string {
	return "DockerNetworksWidget"
}

//OnEvent runs the given command
func (s *DockerNetworksWidget) OnEvent(event EventCommand) error {
	if s.RowCount() > 0 {
		return event(s.filteredRows[s.selectedIndex].network.ID)
	}
	return nil
}

//RowCount returns the number of rows of this widget.
func (s *DockerNetworksWidget) RowCount() int {
	return len(s.filteredRows)

}

//Sort rotates to the next sort mode.
//SortNetworksByID -> SortNetworksByName -> SortNetworksByDriver
func (s *DockerNetworksWidget) Sort() {
	s.RLock()
	defer s.RUnlock()
	switch s.sortMode {
	case docker.SortNetworksByID:
		s.sortMode = docker.SortNetworksByName
	case docker.SortNetworksByName:
		s.sortMode = docker.SortNetworksByDriver
	case docker.SortNetworksByDriver:
		s.sortMode = docker.SortNetworksByContainerCount
	case docker.SortNetworksByContainerCount:
		s.sortMode = docker.SortNetworksByServiceCount
	case docker.SortNetworksByServiceCount:
		s.sortMode = docker.SortNetworksBySubnet
	case docker.SortNetworksBySubnet:
		s.sortMode = docker.SortNetworksByID
	}
}

//Unmount tells this widget that it will not be rendering anymore
func (s *DockerNetworksWidget) Unmount() error {
	s.RLock()
	defer s.RUnlock()
	s.mounted = false
	return nil
}

//Align aligns rows
func (s *DockerNetworksWidget) align() {
	x := s.x
	width := s.width

	s.header.SetWidth(width)
	s.header.SetX(x)

	for _, network := range s.totalRows {
		network.SetX(x)
		network.SetWidth(width)
	}

}

func (s *DockerNetworksWidget) filterRows() {

	if s.filterPattern != "" {
		var rows []*NetworkRow

		for _, row := range s.totalRows {
			if RowFilters.ByPattern(s.filterPattern)(row) {
				rows = append(rows, row)
			}
		}
		s.filteredRows = rows
	} else {
		s.filteredRows = s.totalRows
	}
}

func (s *DockerNetworksWidget) calculateVisibleRows() {

	count := s.RowCount()

	//no screen
	if s.height < 0 || count == 0 {
		s.startIndex = 0
		s.endIndex = 0
		return
	}
	selected := s.selectedIndex
	//everything fits
	if count <= s.height {
		s.startIndex = 0
		s.endIndex = count
		return
	}
	//at the the start
	if selected == 0 {
		s.startIndex = 0
		s.endIndex = s.height
	} else if selected >= count-1 { //at the end
		s.startIndex = count - s.height
		s.endIndex = count
	} else if selected == s.endIndex { //scroll down by one
		s.startIndex++
		s.endIndex++
	} else if selected <= s.startIndex { //scroll up by one
		s.startIndex--
		s.endIndex--
	} else if selected > s.endIndex { // scroll
		s.startIndex = selected - s.height
		s.endIndex = selected
	}
}

//prepareForRendering sets the internal state of this widget so it is ready for
//rendering (i.e. Buffer()).
func (s *DockerNetworksWidget) prepareForRendering() {
	s.sortRows()
	s.filterRows()
	index := ui.ActiveScreen.Cursor.Position()
	if index < 0 {
		index = 0
	} else if index > s.RowCount() {
		index = s.RowCount() - 1
	}
	s.selectedIndex = index
	s.calculateVisibleRows()
}
func (s *DockerNetworksWidget) updateHeader() {
	sortMode := s.sortMode

	for _, c := range s.header.Columns {
		colTitle := c.Text
		var header SortableColumnHeader
		if strings.Contains(colTitle, DownArrow) {
			colTitle = colTitle[DownArrowLength:]
		}
		for _, h := range networkTableHeaders {
			if colTitle == h.Title {
				header = h
				break
			}
		}
		if header.Mode == sortMode {
			c.Text = DownArrow + colTitle
		} else {
			c.Text = colTitle
		}

	}
}

func (s *DockerNetworksWidget) sortRows() {
	rows := s.totalRows
	mode := s.sortMode
	if s.sortMode == docker.NoSortNetworks {
		return
	}
	var sortAlg func(i, j int) bool

	switch mode {
	case docker.SortNetworksByID:
		sortAlg = func(i, j int) bool {
			return rows[i].ID.Text < rows[j].ID.Text
		}
	case docker.SortNetworksByName:
		sortAlg = func(i, j int) bool {
			return rows[i].Name.Text < rows[j].Name.Text
		}
	case docker.SortNetworksByDriver:
		sortAlg = func(i, j int) bool {
			return rows[i].Driver.Text < rows[j].Driver.Text
		}
	case docker.SortNetworksByContainerCount:
		sortAlg = func(i, j int) bool {
			return rows[i].Containers.Text < rows[j].Containers.Text
		}
	case docker.SortNetworksByServiceCount:
		sortAlg = func(i, j int) bool {
			return rows[i].Services.Text < rows[j].Services.Text
		}
	case docker.SortNetworksBySubnet:
		sortAlg = func(i, j int) bool {
			return rows[i].Subnet.Text < rows[j].Subnet.Text
		}

	}
	sort.SliceStable(rows, sortAlg)
}

func (s *DockerNetworksWidget) visibleRows() []*NetworkRow {
	return s.filteredRows[s.startIndex:s.endIndex]
}

func networkTableHeader() *termui.TableHeader {
	header := termui.NewHeader(DryTheme)
	header.ColumnSpacing = DefaultColumnSpacing
	header.AddColumn(networkTableHeaders[0].Title)
	header.AddColumn(networkTableHeaders[1].Title)
	header.AddFixedWidthColumn(networkTableHeaders[2].Title, 12)
	header.AddFixedWidthColumn(networkTableHeaders[3].Title, 12)
	header.AddFixedWidthColumn(networkTableHeaders[4].Title, 12)
	header.AddColumn(networkTableHeaders[5].Title)
	header.AddColumn(networkTableHeaders[6].Title)
	header.AddColumn(networkTableHeaders[7].Title)

	return header
}
<commit_msg>Minor improvements<commit_after>package appui

import (
	"sort"
	"strconv"
	"strings"
	"sync"

	gizaktermui "github.com/gizak/termui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/moncho/dry/ui/termui"
)

var defaultNetworkTableHeader = networkTableHeader()

var networkTableHeaders = []SortableColumnHeader{
	{`NETWORK ID`, docker.SortNetworksByID},
	{`NAME`, docker.SortNetworksByName},
	{`DRIVER`, docker.SortNetworksByDriver},
	{`CONTAINERS`, docker.SortNetworksByContainerCount},
	{`SERVICES`, docker.SortNetworksByServiceCount},
	{`SCOPE`, docker.NoSortNetworks},
	{`SUBNET`, docker.SortNetworksBySubnet},
	{`GATEWAY`, docker.NoSortNetworks},
}

//DockerNetworksWidget knows how render a container list
type DockerNetworksWidget struct {
	dockerDaemon         docker.NetworkAPI
	header               *termui.TableHeader
	filteredRows         []*NetworkRow
	totalRows            []*NetworkRow
	filterPattern        string
	height, width        int
	selectedIndex        int
	startIndex, endIndex int
	x, y                 int
	sortMode             docker.SortMode

	sync.RWMutex
	mounted bool
}

//NewDockerNetworksWidget creates a renderer for a network list
func NewDockerNetworksWidget(dockerDaemon docker.NetworkAPI, y int) *DockerNetworksWidget {
	return &DockerNetworksWidget{
		dockerDaemon: dockerDaemon,
		y:            y,
		header:       defaultNetworkTableHeader,
		height:       MainScreenAvailableHeight(),
		sortMode:     docker.SortNetworksByID,
		width:        ui.ActiveScreen.Dimensions.Width}
}

//Buffer returns the content of this widget as a termui.Buffer
func (s *DockerNetworksWidget) Buffer() gizaktermui.Buffer {
	s.RLock()
	defer s.RUnlock()
	buf := gizaktermui.NewBuffer()
	if !s.mounted {
		return buf
	}
	y := s.y

	s.prepareForRendering()
	widgetHeader := NewWidgetHeader()
	widgetHeader.HeaderEntry("Networks", strconv.Itoa(s.RowCount()))
	if s.filterPattern != "" {
		widgetHeader.HeaderEntry("Active filter", s.filterPattern)
	}

	widgetHeader.Y = y
	buf.Merge(widgetHeader.Buffer())
	y += widgetHeader.GetHeight()
	//Empty line between the header and the rest of the content
	y++
	s.updateHeader()
	s.header.SetY(y)
	buf.Merge(s.header.Buffer())
	y += s.header.GetHeight()

	selected := s.selectedIndex - s.startIndex

	for i, imageRow := range s.visibleRows() {
		imageRow.SetY(y)
		y += imageRow.GetHeight()
		if i != selected {
			imageRow.NotHighlighted()
		} else {
			imageRow.Highlighted()
		}
		buf.Merge(imageRow.Buffer())
	}

	return buf
}

//Filter filters the network list by the given filter
func (s *DockerNetworksWidget) Filter(filter string) {
	s.Lock()
	defer s.Unlock()
	s.filterPattern = filter
}

//Mount tells this widget to be ready for rendering
func (s *DockerNetworksWidget) Mount() error {
	s.Lock()
	defer s.Unlock()
	if s.mounted {
		return nil
	}
	networks, err := s.dockerDaemon.Networks()
	if err != nil {
		return err
	}

	networkRows := make([]*NetworkRow, len(networks))
	for i, network := range networks {
		networkRows[i] = NewNetworkRow(network, s.header)
	}
	s.totalRows = networkRows
	s.mounted = true
	s.align()

	return nil
}

//Name returns this widget name
func (s *DockerNetworksWidget) Name() string {
	return "DockerNetworksWidget"
}

//OnEvent runs the given command
func (s *DockerNetworksWidget) OnEvent(event EventCommand) error {
	if s.RowCount() > 0 {
		return event(s.filteredRows[s.selectedIndex].network.ID)
	}
	return nil
}

//RowCount returns the number of rows of this widget.
func (s *DockerNetworksWidget) RowCount() int {
	return len(s.filteredRows)

}

//Sort rotates to the next sort mode.
//SortNetworksByID -> SortNetworksByName -> SortNetworksByDriver
func (s *DockerNetworksWidget) Sort() {
	s.RLock()
	defer s.RUnlock()
	switch s.sortMode {
	case docker.SortNetworksByID:
		s.sortMode = docker.SortNetworksByName
	case docker.SortNetworksByName:
		s.sortMode = docker.SortNetworksByDriver
	case docker.SortNetworksByDriver:
		s.sortMode = docker.SortNetworksByContainerCount
	case docker.SortNetworksByContainerCount:
		s.sortMode = docker.SortNetworksByServiceCount
	case docker.SortNetworksByServiceCount:
		s.sortMode = docker.SortNetworksBySubnet
	case docker.SortNetworksBySubnet:
		s.sortMode = docker.SortNetworksByID
	}
}

//Unmount tells this widget that it will not be rendering anymore
func (s *DockerNetworksWidget) Unmount() error {
	s.Lock()
	defer s.Unlock()
	s.mounted = false
	return nil
}

//Align aligns rows
func (s *DockerNetworksWidget) align() {
	x := s.x
	width := s.width

	s.header.SetWidth(width)
	s.header.SetX(x)

	for _, network := range s.totalRows {
		network.SetX(x)
		network.SetWidth(width)
	}

}

func (s *DockerNetworksWidget) filterRows() {

	if s.filterPattern != "" {
		var rows []*NetworkRow

		for _, row := range s.totalRows {
			if RowFilters.ByPattern(s.filterPattern)(row) {
				rows = append(rows, row)
			}
		}
		s.filteredRows = rows
	} else {
		s.filteredRows = s.totalRows
	}
}

func (s *DockerNetworksWidget) calculateVisibleRows() {

	count := s.RowCount()

	//no screen
	if s.height < 0 || count == 0 {
		s.startIndex = 0
		s.endIndex = 0
		return
	}
	selected := s.selectedIndex
	//everything fits
	if count <= s.height {
		s.startIndex = 0
		s.endIndex = count
		return
	}
	//at the the start
	if selected == 0 {
		s.startIndex = 0
		s.endIndex = s.height
	} else if selected >= count-1 { //at the end
		s.startIndex = count - s.height
		s.endIndex = count
	} else if selected == s.endIndex { //scroll down by one
		s.startIndex++
		s.endIndex++
	} else if selected <= s.startIndex { //scroll up by one
		s.startIndex--
		s.endIndex--
	} else if selected > s.endIndex { // scroll
		s.startIndex = selected - s.height
		s.endIndex = selected
	}
}

//prepareForRendering sets the internal state of this widget so it is ready for
//rendering (i.e. Buffer()).
func (s *DockerNetworksWidget) prepareForRendering() {
	s.sortRows()
	s.filterRows()
	index := ui.ActiveScreen.Cursor.Position()
	if index < 0 {
		index = 0
	} else if index > s.RowCount() {
		index = s.RowCount() - 1
	}
	s.selectedIndex = index
	s.calculateVisibleRows()
}
func (s *DockerNetworksWidget) updateHeader() {
	sortMode := s.sortMode

	for _, c := range s.header.Columns {
		colTitle := c.Text
		var header SortableColumnHeader
		if strings.Contains(colTitle, DownArrow) {
			colTitle = colTitle[DownArrowLength:]
		}
		for _, h := range networkTableHeaders {
			if colTitle == h.Title {
				header = h
				break
			}
		}
		if header.Mode == sortMode {
			c.Text = DownArrow + colTitle
		} else {
			c.Text = colTitle
		}

	}
}

func (s *DockerNetworksWidget) sortRows() {
	rows := s.totalRows
	mode := s.sortMode
	if s.sortMode == docker.NoSortNetworks {
		return
	}
	var sortAlg func(i, j int) bool

	switch mode {
	case docker.SortNetworksByID:
		sortAlg = func(i, j int) bool {
			return rows[i].ID.Text < rows[j].ID.Text
		}
	case docker.SortNetworksByName:
		sortAlg = func(i, j int) bool {
			return rows[i].Name.Text < rows[j].Name.Text
		}
	case docker.SortNetworksByDriver:
		sortAlg = func(i, j int) bool {
			return rows[i].Driver.Text < rows[j].Driver.Text
		}
	case docker.SortNetworksByContainerCount:
		sortAlg = func(i, j int) bool {
			return rows[i].Containers.Text < rows[j].Containers.Text
		}
	case docker.SortNetworksByServiceCount:
		sortAlg = func(i, j int) bool {
			return rows[i].Services.Text < rows[j].Services.Text
		}
	case docker.SortNetworksBySubnet:
		sortAlg = func(i, j int) bool {
			return rows[i].Subnet.Text < rows[j].Subnet.Text
		}

	}
	sort.SliceStable(rows, sortAlg)
}

func (s *DockerNetworksWidget) visibleRows() []*NetworkRow {
	return s.filteredRows[s.startIndex:s.endIndex]
}

func networkTableHeader() *termui.TableHeader {
	header := termui.NewHeader(DryTheme)
	header.ColumnSpacing = DefaultColumnSpacing
	header.AddColumn(networkTableHeaders[0].Title)
	header.AddColumn(networkTableHeaders[1].Title)
	header.AddFixedWidthColumn(networkTableHeaders[2].Title, 12)
	header.AddFixedWidthColumn(networkTableHeaders[3].Title, 12)
	header.AddFixedWidthColumn(networkTableHeaders[4].Title, 12)
	header.AddColumn(networkTableHeaders[5].Title)
	header.AddColumn(networkTableHeaders[6].Title)
	header.AddColumn(networkTableHeaders[7].Title)

	return header
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build windows

package libkb

import (
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"github.com/kardianos/osext"

	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/registry"
)

type GUID struct {
	Data1 uint32
	Data2 uint16
	Data3 uint16
	Data4 [8]byte
}

// 3EB685DB-65F9-4CF6-A03A-E3EF65729F3D
var (
	FOLDERID_RoamingAppData = GUID{0x3EB685DB, 0x65F9, 0x4CF6, [8]byte{0xA0, 0x3A, 0xE3, 0xEF, 0x65, 0x72, 0x9F, 0x3D}}
)

// F1B32785-6FBA-4FCF-9D55-7B8E7F157091
var (
	FOLDERID_LocalAppData = GUID{0xF1B32785, 0x6FBA, 0x4FCF, [8]byte{0x9D, 0x55, 0x7B, 0x8E, 0x7F, 0x15, 0x70, 0x91}}
)

var (
	modShell32               = windows.NewLazySystemDLL("Shell32.dll")
	modOle32                 = windows.NewLazySystemDLL("Ole32.dll")
	procSHGetKnownFolderPath = modShell32.NewProc("SHGetKnownFolderPath")
	procCoTaskMemFree        = modOle32.NewProc("CoTaskMemFree")
	shChangeNotifyProc       = modShell32.NewProc("SHChangeNotify")
)

// LookPath searches for an executable binary named file
// in the directories named by the PATH environment variable.
// If file contains a slash, it is tried directly and the PATH is not consulted.

func canExec(s string) error {
	if strings.IndexAny(s, `:\/`) == -1 {
		s = s + "/"
	}
	_, err := exec.LookPath(s)
	return err
}

func PosixLineEndings(arg string) string {
	return strings.Replace(arg, "\r", "", -1)
}

func coTaskMemFree(pv uintptr) {
	syscall.Syscall(procCoTaskMemFree.Addr(), 1, uintptr(pv), 0, 0)
	return
}

func GetDataDir(id GUID) (string, error) {

	var pszPath uintptr
	r0, _, _ := procSHGetKnownFolderPath.Call(uintptr(unsafe.Pointer(&id)), uintptr(0), uintptr(0), uintptr(unsafe.Pointer(&pszPath)))
	if r0 != 0 {
		return "", errors.New("can't get FOLDERID_RoamingAppData")
	}

	defer coTaskMemFree(pszPath)

	// go vet: "possible misuse of unsafe.Pointer"
	folder := syscall.UTF16ToString((*[1 << 16]uint16)(unsafe.Pointer(pszPath))[:])

	if len(folder) == 0 {
		return "", errors.New("can't get AppData directory")
	}

	return folder, nil
}

func AppDataDir() (string, error) {
	return GetDataDir(FOLDERID_RoamingAppData)
}

func LocalDataDir() (string, error) {
	return GetDataDir(FOLDERID_LocalAppData)
}

// SafeWriteToFile retries safeWriteToFileOnce a few times on Windows,
// in case AV programs interfere with 2 writes in quick succession.
func SafeWriteToFile(g SafeWriteLogger, t SafeWriter, mode os.FileMode) error {

	var err error
	for i := 0; i < 5; i++ {
		if err != nil {
			g.Debug("Retrying failed safeWriteToFileOnce - %s", err)
			time.Sleep(10 * time.Millisecond)
		}
		err = safeWriteToFileOnce(g, t, mode)
		if err == nil {
			break
		}
	}
	return err
}

// helper for RemoteSettingsRepairman
func moveNonChromiumFiles(g *GlobalContext, oldHome string, currentHome string) error {
	g.Log.Info("RemoteSettingsRepairman moving from %s to %s", oldHome, currentHome)
	files, _ := filepath.Glob(filepath.Join(oldHome, "*"))
	for _, oldPathName := range files {
		_, name := filepath.Split(oldPathName)
		// Chromium seems stubborn about these - TBD
		switch name {
		case "GPUCache":
			continue
		case "lockfile":
			continue
		case "app-state.json":
			continue
		case "Cache":
			continue
		case "Cookies":
			continue
		case "Cookies-journal":
			continue
		case "Local Storage":
			continue
		}
		// explicitly skip logs
		if strings.HasSuffix(name, ".log") {
			continue
		}
		newPathName := filepath.Join(currentHome, name)
		var err error
		g.Log.Info("   moving %s", name)

		for i := 0; i < 5; i++ {
			if err != nil {
				time.Sleep(100 * time.Millisecond)
			}
			err = os.Rename(oldPathName, newPathName)
			if err == nil {
				break
			}
		}
		if err != nil {
			g.Log.Error("RemoteSettingsRepairman error moving %s to %s - %s", oldPathName, newPathName, err)
			return err
		}
	}
	return nil
}

// RemoteSettingsRepairman does a one-time move of everyting from the roaming
// target directory to local. We depend on the .exe files having been uninstalled from
// there first.
// Note that Chromium still insists on keeping some stuff in roaming,
// exceptions for which are hardcoded.
func RemoteSettingsRepairman(g *GlobalContext) error {
	w := Win32{Base{"keybase",
		func() string { return g.Env.getHomeFromCmdOrConfig() },
		func() RunMode { return g.Env.GetRunMode() }}}

	currentHome := w.Home(false)
	kbDir := filepath.Base(currentHome)
	currentConfig := g.Env.GetConfigFilename()
	oldDir, err := AppDataDir()
	if err != nil {
		return err
	}
	_, configName := filepath.Split(currentConfig)
	oldHome := filepath.Join(oldDir, kbDir)
	oldConfig := filepath.Join(oldHome, configName)
	if oldExists, _ := FileExists(oldConfig); oldExists {
		if currentExists, _ := FileExists(currentConfig); !currentExists {
			return moveNonChromiumFiles(g, oldHome, currentHome)
		}
	}
	return nil
}

// Notify the shell that the thing located at path has changed
func notifyShell(path string) {
	shChangeNotifyProc.Call(
		uintptr(0x00002000), // SHCNE_UPDATEITEM
		uintptr(0x0005),     // SHCNF_PATHW
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(path))),
		0)
}

// Manipulate registry entries to reflect the mount point icon in the shell
func ChangeMountIcon(oldMount string, newMount string) error {
	if oldMount != "" {
		// DeleteKey doesn't work if there are subkeys
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1]+`\DefaultIcon`)
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1]+`\DefaultLabel`)
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1])
		notifyShell(oldMount)
	}
	if newMount == "" {
		return nil
	}
	k, _, err := registry.CreateKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+newMount[:1]+`\DefaultIcon`, registry.SET_VALUE|registry.CREATE_SUB_KEY|registry.WRITE)
	defer k.Close()
	if err != nil {
		return err
	}
	keybaseExe, err := osext.Executable()
	if err != nil {
		return err
	}
	// Use the second icon bound into keybase.exe - hence the 1
	err = k.SetStringValue("", keybaseExe+",1")
	if err != nil {
		return err
	}

	// Also give a nice label
	k2, _, err := registry.CreateKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+newMount[:1]+`\DefaultLabel`, registry.SET_VALUE|registry.CREATE_SUB_KEY|registry.WRITE)
	defer k2.Close()
	err = k2.SetStringValue("", "Keybase")
	notifyShell(newMount)
	return err
}
<commit_msg>Zanderz/repairman 2017.02.13 (#5855)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build windows

package libkb

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"github.com/kardianos/osext"

	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/registry"
)

type GUID struct {
	Data1 uint32
	Data2 uint16
	Data3 uint16
	Data4 [8]byte
}

// 3EB685DB-65F9-4CF6-A03A-E3EF65729F3D
var (
	FOLDERID_RoamingAppData = GUID{0x3EB685DB, 0x65F9, 0x4CF6, [8]byte{0xA0, 0x3A, 0xE3, 0xEF, 0x65, 0x72, 0x9F, 0x3D}}
)

// F1B32785-6FBA-4FCF-9D55-7B8E7F157091
var (
	FOLDERID_LocalAppData = GUID{0xF1B32785, 0x6FBA, 0x4FCF, [8]byte{0x9D, 0x55, 0x7B, 0x8E, 0x7F, 0x15, 0x70, 0x91}}
)

var (
	modShell32               = windows.NewLazySystemDLL("Shell32.dll")
	modOle32                 = windows.NewLazySystemDLL("Ole32.dll")
	procSHGetKnownFolderPath = modShell32.NewProc("SHGetKnownFolderPath")
	procCoTaskMemFree        = modOle32.NewProc("CoTaskMemFree")
	shChangeNotifyProc       = modShell32.NewProc("SHChangeNotify")
)

// LookPath searches for an executable binary named file
// in the directories named by the PATH environment variable.
// If file contains a slash, it is tried directly and the PATH is not consulted.

func canExec(s string) error {
	if strings.IndexAny(s, `:\/`) == -1 {
		s = s + "/"
	}
	_, err := exec.LookPath(s)
	return err
}

func PosixLineEndings(arg string) string {
	return strings.Replace(arg, "\r", "", -1)
}

func coTaskMemFree(pv uintptr) {
	syscall.Syscall(procCoTaskMemFree.Addr(), 1, uintptr(pv), 0, 0)
	return
}

func GetDataDir(id GUID) (string, error) {

	var pszPath uintptr
	r0, _, _ := procSHGetKnownFolderPath.Call(uintptr(unsafe.Pointer(&id)), uintptr(0), uintptr(0), uintptr(unsafe.Pointer(&pszPath)))
	if r0 != 0 {
		return "", errors.New("can't get FOLDERID_RoamingAppData")
	}

	defer coTaskMemFree(pszPath)

	// go vet: "possible misuse of unsafe.Pointer"
	folder := syscall.UTF16ToString((*[1 << 16]uint16)(unsafe.Pointer(pszPath))[:])

	if len(folder) == 0 {
		return "", errors.New("can't get AppData directory")
	}

	return folder, nil
}

func AppDataDir() (string, error) {
	return GetDataDir(FOLDERID_RoamingAppData)
}

func LocalDataDir() (string, error) {
	return GetDataDir(FOLDERID_LocalAppData)
}

// SafeWriteToFile retries safeWriteToFileOnce a few times on Windows,
// in case AV programs interfere with 2 writes in quick succession.
func SafeWriteToFile(g SafeWriteLogger, t SafeWriter, mode os.FileMode) error {

	var err error
	for i := 0; i < 5; i++ {
		if err != nil {
			g.Debug("Retrying failed safeWriteToFileOnce - %s", err)
			time.Sleep(10 * time.Millisecond)
		}
		err = safeWriteToFileOnce(g, t, mode)
		if err == nil {
			break
		}
	}
	return err
}

func copyFile(src string, dest string) error {
	srcFile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	destFile, err := os.Create(dest) // creates if file doesn't exist
	if err != nil {
		return err
	}
	defer destFile.Close()

	_, err = io.Copy(destFile, srcFile) // check first var for number of bytes copied
	if err != nil {
		return err
	}

	err = destFile.Sync()
	return err
}

// These are the really important ones, so we'll copy first and then delete the old ones,
// undoing on failure.
func moveKeyFiles(g *GlobalContext, oldHome string, currentHome string) (bool, error) {
	var err error

	// See if any secret key files are in the new location. If so, don't repair.
	if newSecretKeyfiles, _ := filepath.Glob(filepath.Join(currentHome, "*.ss")); len(newSecretKeyfiles) > 0 {
		return false, nil
	}
	g.Log.Info("RemoteSettingsRepairman moving from %s to %s", oldHome, currentHome)

	files, _ := filepath.Glob(filepath.Join(oldHome, "*.mpack"))
	oldSecretKeyfiles, _ := filepath.Glob(filepath.Join(oldHome, "*.ss"))
	files = append(files, oldSecretKeyfiles...)
	var newFiles []string
	for _, oldPathName := range files {
		_, name := filepath.Split(oldPathName)
		newPathName := filepath.Join(currentHome, name)

		// If both copies exist, skip
		if exists, _ := FileExists(newPathName); !exists {
			err = copyFile(oldPathName, newPathName)
			if err != nil {
				g.Log.Error("RemoteSettingsRepairman fatal error copying %s to %s - %s", oldPathName, newPathName, err)
				break
			} else {
				newFiles = append(newFiles, newPathName)
			}
		}

	}
	if err != nil {
		// Undo any of the new copies and quit
		for _, newPathName := range newFiles {
			os.Remove(newPathName)
		}
		return false, err
	}
	// Now that we've successfully copied, delete the old ones - BUT don't bail out on error here
	for _, oldPathName := range files {
		os.Remove(oldPathName)
	}

	// Return true if we copied any
	if len(files) > 0 {
		return true, err
	}

	return false, err
}

// helper for RemoteSettingsRepairman
func moveNonChromiumFiles(g *GlobalContext, oldHome string, currentHome string) error {

	files, _ := filepath.Glob(filepath.Join(oldHome, "*"))
	for _, oldPathName := range files {
		_, name := filepath.Split(oldPathName)
		// Chromium seems stubborn about these - TBD
		switch name {
		case "GPUCache":
			continue
		case "lockfile":
			continue
		case "app-state.json":
			continue
		case "Cache":
			continue
		case "Cookies":
			continue
		case "Cookies-journal":
			continue
		case "Local Storage":
			continue
		}
		// explicitly skip logs
		if strings.HasSuffix(name, ".log") {
			continue
		}
		newPathName := filepath.Join(currentHome, name)
		var err error
		g.Log.Info("   moving %s", name)

		for i := 0; i < 5; i++ {
			if err != nil {
				time.Sleep(100 * time.Millisecond)
			}
			err = os.Rename(oldPathName, newPathName)
			if err == nil {
				break
			}
		}
		if err != nil {
			g.Log.Error("RemoteSettingsRepairman error moving %s to %s (continuing) - %s", oldPathName, newPathName, err)
		}
	}
	return nil
}

// RemoteSettingsRepairman does a one-time move of everyting from the roaming
// target directory to local. We depend on the .exe files having been uninstalled from
// there first.
// Note that Chromium still insists on keeping some stuff in roaming,
// exceptions for which are hardcoded.
func RemoteSettingsRepairman(g *GlobalContext) error {
	w := Win32{Base{"keybase",
		func() string { return g.Env.getHomeFromCmdOrConfig() },
		func() RunMode { return g.Env.GetRunMode() }}}

	currentHome := w.Home(false)
	kbDir := filepath.Base(currentHome)
	oldDir, err := AppDataDir()
	if err != nil {
		return err
	}
	oldHome := filepath.Join(oldDir, kbDir)

	// Only continue repairing if key files needed to be moved
	if moved, err := moveKeyFiles(g, oldHome, currentHome); !moved || err != nil {
		return err
	}
	// Don't fail the repairmain if these others can't be moved
	moveNonChromiumFiles(g, oldHome, currentHome)
	return nil
}

// Notify the shell that the thing located at path has changed
func notifyShell(path string) {
	shChangeNotifyProc.Call(
		uintptr(0x00002000), // SHCNE_UPDATEITEM
		uintptr(0x0005),     // SHCNF_PATHW
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(path))),
		0)
}

// Manipulate registry entries to reflect the mount point icon in the shell
func ChangeMountIcon(oldMount string, newMount string) error {
	if oldMount != "" {
		// DeleteKey doesn't work if there are subkeys
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1]+`\DefaultIcon`)
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1]+`\DefaultLabel`)
		registry.DeleteKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+oldMount[:1])
		notifyShell(oldMount)
	}
	if newMount == "" {
		return nil
	}
	k, _, err := registry.CreateKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+newMount[:1]+`\DefaultIcon`, registry.SET_VALUE|registry.CREATE_SUB_KEY|registry.WRITE)
	defer k.Close()
	if err != nil {
		return err
	}
	keybaseExe, err := osext.Executable()
	if err != nil {
		return err
	}
	// Use the second icon bound into keybase.exe - hence the 1
	err = k.SetStringValue("", keybaseExe+",1")
	if err != nil {
		return err
	}

	// Also give a nice label
	k2, _, err := registry.CreateKey(registry.CURRENT_USER, `SOFTWARE\Classes\Applications\Explorer.exe\Drives\`+newMount[:1]+`\DefaultLabel`, registry.SET_VALUE|registry.CREATE_SUB_KEY|registry.WRITE)
	defer k2.Close()
	err = k2.SetStringValue("", "Keybase")
	notifyShell(newMount)
	return err
}
<|endoftext|>


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

package loader_test

// This file enumerates all packages beneath $GOROOT, loads them, plus
// their external tests if any, runs the type checker on them, and
// prints some summary information.

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"go/types"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/internal/testenv"
)

func TestStdlib(t *testing.T) {
	if runtime.GOOS == "android" {
		t.Skipf("incomplete std lib on %s", runtime.GOOS)
	}
	if testing.Short() {
		t.Skip("skipping in short mode; uses tons of memory (https://golang.org/issue/14113)")
	}
	testenv.NeedsTool(t, "go")

	runtime.GC()
	t0 := time.Now()
	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	alloc := memstats.Alloc

	// Load, parse and type-check the program.
	ctxt := build.Default // copy
	ctxt.GOPATH = ""      // disable GOPATH
	conf := loader.Config{Build: &ctxt}
	for _, path := range buildutil.AllPackages(conf.Build) {
		conf.ImportWithTests(path)
	}

	prog, err := conf.Load()
	if err != nil {
		t.Fatalf("Load failed: %v", err)
	}

	t1 := time.Now()
	runtime.GC()
	runtime.ReadMemStats(&memstats)

	numPkgs := len(prog.AllPackages)
	if want := 205; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Dump package members.
	if false {
		for pkg := range prog.AllPackages {
			fmt.Printf("Package %s:\n", pkg.Path())
			scope := pkg.Scope()
			qualifier := types.RelativeTo(pkg)
			for _, name := range scope.Names() {
				if ast.IsExported(name) {
					fmt.Printf("\t%s\n", types.ObjectString(scope.Lookup(name), qualifier))
				}
			}
			fmt.Println()
		}
	}

	// Check that Test functions for io/ioutil, regexp and
	// compress/bzip2 are all simultaneously present.
	// (The apparent cycle formed when augmenting all three of
	// these packages by their tests was the original motivation
	// for reporting b/7114.)
	//
	// compress/bzip2.TestBitReader in bzip2_test.go    imports io/ioutil
	// io/ioutil.TestTempFile       in tempfile_test.go imports regexp
	// regexp.TestRE2Search         in exec_test.go     imports compress/bzip2
	for _, test := range []struct{ pkg, fn string }{
		{"io/ioutil", "TestTempFile"},
		{"regexp", "TestRE2Search"},
		{"compress/bzip2", "TestBitReader"},
	} {
		info := prog.Imported[test.pkg]
		if info == nil {
			t.Errorf("failed to load package %q", test.pkg)
			continue
		}
		obj, _ := info.Pkg.Scope().Lookup(test.fn).(*types.Func)
		if obj == nil {
			t.Errorf("package %q has no func %q", test.pkg, test.fn)
			continue
		}
	}

	// Dump some statistics.

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("GOMAXPROCS:           ", runtime.GOMAXPROCS(0))
	t.Log("#Source lines:        ", lineCount)
	t.Log("Load/parse/typecheck: ", t1.Sub(t0))
	t.Log("#MB:                  ", int64(memstats.Alloc-alloc)/1000000)
}

func TestCgoOption(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in short mode; uses tons of memory (https://golang.org/issue/14113)")
	}
	switch runtime.GOOS {
	// On these systems, the net and os/user packages don't use cgo
	// or the std library is incomplete (Android).
	case "android", "plan9", "solaris", "windows":
		t.Skipf("no cgo or incomplete std lib on %s", runtime.GOOS)
	}
	// In nocgo builds (e.g. linux-amd64-nocgo),
	// there is no "runtime/cgo" package,
	// so cgo-generated Go files will have a failing import.
	if !build.Default.CgoEnabled {
		return
	}
	testenv.NeedsTool(t, "go")

	// Test that we can load cgo-using packages with
	// CGO_ENABLED=[01], which causes go/build to select pure
	// Go/native implementations, respectively, based on build
	// tags.
	//
	// Each entry specifies a package-level object and the generic
	// file expected to define it when cgo is disabled.
	// When cgo is enabled, the exact file is not specified (since
	// it varies by platform), but must differ from the generic one.
	//
	// The test also loads the actual file to verify that the
	// object is indeed defined at that location.
	for _, test := range []struct {
		pkg, name, genericFile string
	}{
		{"net", "cgoLookupHost", "cgo_stub.go"},
		{"os/user", "current", "lookup_stubs.go"},
	} {
		ctxt := build.Default
		for _, ctxt.CgoEnabled = range []bool{false, true} {
			conf := loader.Config{Build: &ctxt}
			conf.Import(test.pkg)
			prog, err := conf.Load()
			if err != nil {
				t.Errorf("Load failed: %v", err)
				continue
			}
			info := prog.Imported[test.pkg]
			if info == nil {
				t.Errorf("package %s not found", test.pkg)
				continue
			}
			obj := info.Pkg.Scope().Lookup(test.name)
			if obj == nil {
				t.Errorf("no object %s.%s", test.pkg, test.name)
				continue
			}
			posn := prog.Fset.Position(obj.Pos())
			t.Logf("%s: %s (CgoEnabled=%t)", posn, obj, ctxt.CgoEnabled)

			gotFile := filepath.Base(posn.Filename)
			filesMatch := gotFile == test.genericFile

			if ctxt.CgoEnabled && filesMatch {
				t.Errorf("CGO_ENABLED=1: %s found in %s, want native file",
					obj, gotFile)
			} else if !ctxt.CgoEnabled && !filesMatch {
				t.Errorf("CGO_ENABLED=0: %s found in %s, want %s",
					obj, gotFile, test.genericFile)
			}

			// Load the file and check the object is declared at the right place.
			b, err := ioutil.ReadFile(posn.Filename)
			if err != nil {
				t.Errorf("can't read %s: %s", posn.Filename, err)
				continue
			}
			line := string(bytes.Split(b, []byte("\n"))[posn.Line-1])
			ident := line[posn.Column-1:]
			if !strings.HasPrefix(ident, test.name) {
				t.Errorf("%s: %s not declared here (looking at %q)", posn, obj, ident)
			}
		}
	}
}
<commit_msg>go/loader: remove check on ioutil.TestTempFile in TestStdlib<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package loader_test

// This file enumerates all packages beneath $GOROOT, loads them, plus
// their external tests if any, runs the type checker on them, and
// prints some summary information.

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"go/types"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/internal/testenv"
)

func TestStdlib(t *testing.T) {
	if runtime.GOOS == "android" {
		t.Skipf("incomplete std lib on %s", runtime.GOOS)
	}
	if testing.Short() {
		t.Skip("skipping in short mode; uses tons of memory (https://golang.org/issue/14113)")
	}
	testenv.NeedsTool(t, "go")

	runtime.GC()
	t0 := time.Now()
	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	alloc := memstats.Alloc

	// Load, parse and type-check the program.
	ctxt := build.Default // copy
	ctxt.GOPATH = ""      // disable GOPATH
	conf := loader.Config{Build: &ctxt}
	for _, path := range buildutil.AllPackages(conf.Build) {
		conf.ImportWithTests(path)
	}

	prog, err := conf.Load()
	if err != nil {
		t.Fatalf("Load failed: %v", err)
	}

	t1 := time.Now()
	runtime.GC()
	runtime.ReadMemStats(&memstats)

	numPkgs := len(prog.AllPackages)
	if want := 205; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Dump package members.
	if false {
		for pkg := range prog.AllPackages {
			fmt.Printf("Package %s:\n", pkg.Path())
			scope := pkg.Scope()
			qualifier := types.RelativeTo(pkg)
			for _, name := range scope.Names() {
				if ast.IsExported(name) {
					fmt.Printf("\t%s\n", types.ObjectString(scope.Lookup(name), qualifier))
				}
			}
			fmt.Println()
		}
	}

	// Check that Test functions for regexp and compress/bzip2 are
	// simultaneously present. The apparent cycle formed when augmenting
	// these packages by their tests (together with io/ioutil's test, which is now
	// an xtest) was the original motivation or reporting golang.org/issue/7114.
	//
	// compress/bzip2.TestBitReader in bzip2_test.go    imports io/ioutil
	// io/ioutil.TestTempFile       in tempfile_test.go imports regexp (no longer exists)
	// regexp.TestRE2Search         in exec_test.go     imports compress/bzip2
	for _, test := range []struct{ pkg, fn string }{
		{"regexp", "TestRE2Search"},
		{"compress/bzip2", "TestBitReader"},
	} {
		info := prog.Imported[test.pkg]
		if info == nil {
			t.Errorf("failed to load package %q", test.pkg)
			continue
		}
		obj, _ := info.Pkg.Scope().Lookup(test.fn).(*types.Func)
		if obj == nil {
			t.Errorf("package %q has no func %q", test.pkg, test.fn)
			continue
		}
	}

	// Dump some statistics.

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("GOMAXPROCS:           ", runtime.GOMAXPROCS(0))
	t.Log("#Source lines:        ", lineCount)
	t.Log("Load/parse/typecheck: ", t1.Sub(t0))
	t.Log("#MB:                  ", int64(memstats.Alloc-alloc)/1000000)
}

func TestCgoOption(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in short mode; uses tons of memory (https://golang.org/issue/14113)")
	}
	switch runtime.GOOS {
	// On these systems, the net and os/user packages don't use cgo
	// or the std library is incomplete (Android).
	case "android", "plan9", "solaris", "windows":
		t.Skipf("no cgo or incomplete std lib on %s", runtime.GOOS)
	}
	// In nocgo builds (e.g. linux-amd64-nocgo),
	// there is no "runtime/cgo" package,
	// so cgo-generated Go files will have a failing import.
	if !build.Default.CgoEnabled {
		return
	}
	testenv.NeedsTool(t, "go")

	// Test that we can load cgo-using packages with
	// CGO_ENABLED=[01], which causes go/build to select pure
	// Go/native implementations, respectively, based on build
	// tags.
	//
	// Each entry specifies a package-level object and the generic
	// file expected to define it when cgo is disabled.
	// When cgo is enabled, the exact file is not specified (since
	// it varies by platform), but must differ from the generic one.
	//
	// The test also loads the actual file to verify that the
	// object is indeed defined at that location.
	for _, test := range []struct {
		pkg, name, genericFile string
	}{
		{"net", "cgoLookupHost", "cgo_stub.go"},
		{"os/user", "current", "lookup_stubs.go"},
	} {
		ctxt := build.Default
		for _, ctxt.CgoEnabled = range []bool{false, true} {
			conf := loader.Config{Build: &ctxt}
			conf.Import(test.pkg)
			prog, err := conf.Load()
			if err != nil {
				t.Errorf("Load failed: %v", err)
				continue
			}
			info := prog.Imported[test.pkg]
			if info == nil {
				t.Errorf("package %s not found", test.pkg)
				continue
			}
			obj := info.Pkg.Scope().Lookup(test.name)
			if obj == nil {
				t.Errorf("no object %s.%s", test.pkg, test.name)
				continue
			}
			posn := prog.Fset.Position(obj.Pos())
			t.Logf("%s: %s (CgoEnabled=%t)", posn, obj, ctxt.CgoEnabled)

			gotFile := filepath.Base(posn.Filename)
			filesMatch := gotFile == test.genericFile

			if ctxt.CgoEnabled && filesMatch {
				t.Errorf("CGO_ENABLED=1: %s found in %s, want native file",
					obj, gotFile)
			} else if !ctxt.CgoEnabled && !filesMatch {
				t.Errorf("CGO_ENABLED=0: %s found in %s, want %s",
					obj, gotFile, test.genericFile)
			}

			// Load the file and check the object is declared at the right place.
			b, err := ioutil.ReadFile(posn.Filename)
			if err != nil {
				t.Errorf("can't read %s: %s", posn.Filename, err)
				continue
			}
			line := string(bytes.Split(b, []byte("\n"))[posn.Line-1])
			ident := line[posn.Column-1:]
			if !strings.HasPrefix(ident, test.name) {
				t.Errorf("%s: %s not declared here (looking at %q)", posn, obj, ident)
			}
		}
	}
}
<|endoftext|>


<commit_before>package test

import (
	"strings"
	"testing"

	"gopkg.in/DATA-DOG/go-sqlmock.v1"
)

type ManagementDB struct {
	MockDB
}

func CreateManagementDB(context string, t *testing.T) (p ManagementDB, err error) {
	db, mock, err := createMockDB()
	if err != nil {
		t.Errorf("%s: Setup Project DB Failed with Error: %v", context, err)
		return p, err
	}
	p = ManagementDB{MockDB{db, mock, "management"}}

	return p, nil
}

// Metadata Helpers

var metadataColumns = []string{
	"mdid",
	"db",
	"property_id",
	"parent_id",
	"type",
	"name",
	"exists",
}

func (m *ManagementDB) MetadataGet(mid int, result DBRow, expectEmtpy bool) {

	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmtpy {
		query.Rows = []DBRow{result}
	}
	query.FormatQuery("SELECT * FROM `metadata` WHERE mdid=%d", mid)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataInsert(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `metadata` (`%s`)%s", strings.Join(metadataColumns, "`,`"), migrationValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

//'select * from migration WHERE status = 6'

//'select `mdid`,`db`,`property_id`,`parent_id`,`type`,`name`,`exists` from `metadata` where `mdid`=?;' with args [1] was not expected]

func (m *ManagementDB) MetadataSelectName(name string, result DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = append(query.Rows, result)
	}

	query.FormatQuery("SELECT * FROM metadata WHERE name=\"%s\"", name)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataSelectNameParent(name string, parentId string, result DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = append(query.Rows, result)
	}
	query.FormatQuery("SELECT * FROM metadata WHERE name=\"%s\" AND parent_id=\"%s\"", name, parentId)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataLoadAllTableMetadata(tblPropertyID string, dbID int64, results []DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = results
	}
	query.FormatQuery("select * from metadata WHERE name = \"%s\" OR parent_id = \"%s\" AND db=%d", tblPropertyID, tblPropertyID, dbID)

	m.ExpectQuery(query)
}

// Migration Helpers

var migrationColumns = []string{
	"mid",
	"db",
	"project",
	"version",
	"version_timestamp",
	"version_description",
	"status",
}

var migrationStepsColumns = []string{
	"sid",
	"mid",
	"op",
	"mdid",
	"name",
	"forward",
	"backward",
	"output",
	"status",
}

var migrationValuesTemplate = " values (null,?,?,?,?,?,?)"
var migrationStepsValuesTemplate = " values (null,?,?,?,?,?,?,?,?)"

func (m *ManagementDB) MigrationCount(result DBRow, expectEmpty bool) {

	query := DBQueryMock{
		Type:    QueryCmd,
		Query:   "select count(*) from migration",
		Columns: []string{"count"},
	}
	if !expectEmpty {
		query.Rows = []DBRow{result}
	}

	m.ExpectQuery(query)
}

func (m *ManagementDB) MigrationGetStatus(status int, results []DBRow, expectEmpty bool) {

	query := DBQueryMock{
		Columns: migrationColumns,
	}
	if !expectEmpty {
		query.Rows = results
	}
	query.FormatQuery("select * from migration WHERE status = %d", status)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MigrationSetStatus(mid int64, status int) {

	query := DBQueryMock{
		Columns: migrationColumns,
		Result:  sqlmock.NewResult(0, 1),
	}
	query.FormatQuery("update migration WHERE mid = %d SET status = %d", mid, status)

	m.ExpectExec(query)
}

func (m *ManagementDB) MigrationInsert(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `migration` (`%s`)%s", strings.Join(migrationColumns, "`,`"), migrationValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

func (m *ManagementDB) MigrationInsertStep(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `migration_steps` (`%s`)%s", strings.Join(migrationStepsColumns, "`,`"), migrationStepsValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

func (m *ManagementDB) StepSetStatus(sid int64, status int) {

	query := DBQueryMock{
		Columns: migrationStepsColumns,
		Result:  sqlmock.NewResult(0, 1),
	}
	query.FormatQuery("update migration_steps WHERE sid = %d SET status = %d", sid, status)

	m.ExpectExec(query)
}
<commit_msg>Added more utility functions to the management database test helper<commit_after>package test

import (
	"strings"
	"testing"

	"gopkg.in/DATA-DOG/go-sqlmock.v1"
)

type ManagementDB struct {
	MockDB
}

func CreateManagementDB(context string, t *testing.T) (p ManagementDB, err error) {
	db, mock, err := createMockDB()
	if err != nil {
		t.Errorf("%s: Setup Project DB Failed with Error: %v", context, err)
		return p, err
	}
	p = ManagementDB{MockDB{db, mock, "management"}}

	return p, nil
}

func (m *ManagementDB) ShowTables(results []DBRow, expectEmpty bool) {

	query := DBQueryMock{
		Query:   "SHOW TABLES IN management",
		Columns: []string{"table"},
	}

	if !expectEmpty {
		query.Rows = results
	}
	m.ExpectQuery(query)
}

// Database Helpers

var databaseColumns = []string{
	"dbid",
	"project",
	"name",
	"env",
}

func (m *ManagementDB) DatabaseGet(project string, name string, env string, result DBRow, expectEmtpy bool) {

	query := DBQueryMock{
		Columns: databaseColumns,
	}
	if !expectEmtpy {
		query.Rows = []DBRow{result}
	}
	query.FormatQuery("SELECT * FROM target_database WHERE project=\"%s\" AND name=\"%s\" AND env=\"%s\"", project, name, env)

	m.ExpectQuery(query)
}

// Metadata Helpers

var metadataColumns = []string{
	"mdid",
	"db",
	"property_id",
	"parent_id",
	"type",
	"name",
	"exists",
}

func (m *ManagementDB) MetadataGet(mid int, result DBRow, expectEmtpy bool) {

	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmtpy {
		query.Rows = []DBRow{result}
	}
	query.FormatQuery("SELECT * FROM `metadata` WHERE mdid=%d", mid)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataInsert(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `metadata` (`%s`)%s", strings.Join(metadataColumns, "`,`"), migrationValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

//'select * from migration WHERE status = 6'

//'select `mdid`,`db`,`property_id`,`parent_id`,`type`,`name`,`exists` from `metadata` where `mdid`=?;' with args [1] was not expected]

func (m *ManagementDB) MetadataSelectName(name string, result DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = append(query.Rows, result)
	}

	query.FormatQuery("SELECT * FROM metadata WHERE name=\"%s\"", name)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataSelectNameParent(name string, parentId string, result DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = append(query.Rows, result)
	}
	query.FormatQuery("SELECT * FROM metadata WHERE name=\"%s\" AND parent_id=\"%s\"", name, parentId)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MetadataLoadAllTableMetadata(tblPropertyID string, dbID int64, results []DBRow, expectEmpty bool) {
	query := DBQueryMock{
		Columns: metadataColumns,
	}
	if !expectEmpty {
		query.Rows = results
	}
	query.FormatQuery("select * from metadata WHERE name = \"%s\" OR parent_id = \"%s\" AND db=%d", tblPropertyID, tblPropertyID, dbID)

	m.ExpectQuery(query)
}

// Migration Helpers

var migrationColumns = []string{
	"mid",
	"db",
	"project",
	"version",
	"version_timestamp",
	"version_description",
	"status",
}

var migrationStepsColumns = []string{
	"sid",
	"mid",
	"op",
	"mdid",
	"name",
	"forward",
	"backward",
	"output",
	"status",
}

var migrationValuesTemplate = " values (null,?,?,?,?,?,?)"
var migrationStepsValuesTemplate = " values (null,?,?,?,?,?,?,?,?)"

func (m *ManagementDB) MigrationCount(result DBRow, expectEmpty bool) {

	query := DBQueryMock{
		Type:    QueryCmd,
		Query:   "select count(*) from migration",
		Columns: []string{"count"},
	}
	if !expectEmpty {
		query.Rows = []DBRow{result}
	}

	m.ExpectQuery(query)
}

func (m *ManagementDB) MigrationGetStatus(status int, results []DBRow, expectEmpty bool) {

	query := DBQueryMock{
		Columns: migrationColumns,
	}
	if !expectEmpty {
		query.Rows = results
	}
	query.FormatQuery("select * from migration WHERE status = %d", status)

	m.ExpectQuery(query)
}

func (m *ManagementDB) MigrationSetStatus(mid int64, status int) {

	query := DBQueryMock{
		Columns: migrationColumns,
		Result:  sqlmock.NewResult(0, 1),
	}
	query.FormatQuery("update migration WHERE mid = %d SET status = %d", mid, status)

	m.ExpectExec(query)
}

func (m *ManagementDB) MigrationInsert(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `migration` (`%s`)%s", strings.Join(migrationColumns, "`,`"), migrationValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

func (m *ManagementDB) MigrationInsertStep(args DBRow, lastInsert int64, rowsAffected int64) {

	query := DBQueryMock{
		Type:   ExecCmd,
		Result: sqlmock.NewResult(lastInsert, rowsAffected),
	}
	query.FormatQuery("insert into `migration_steps` (`%s`)%s", strings.Join(migrationStepsColumns, "`,`"), migrationStepsValuesTemplate)
	query.SetArgs(args...)

	m.ExpectExec(query)
}

func (m *ManagementDB) StepSetStatus(sid int64, status int) {

	query := DBQueryMock{
		Columns: migrationStepsColumns,
		Result:  sqlmock.NewResult(0, 1),
	}
	query.FormatQuery("update migration_steps WHERE sid = %d SET status = %d", sid, status)

	m.ExpectExec(query)
}
<|endoftext|>


<commit_before>package gocbcore

import (
	"encoding/binary"
	"time"
)

func (c *Agent) handleServerNmv(s *memdPipeline, req *memdQRequest, resp *memdResponse) {
	// Try to parse the value as a bucket configuration
	bk, err := parseConfig(resp.Value, s.Hostname())
	if err == nil {
		c.updateConfig(bk)
	}

	// Redirect it!  This may actually come back to this server, but I won't tell
	//   if you don't ;)
	c.redispatchDirect(req)
}

func (c *Agent) handleServerDeath(s *memdPipeline) {
	// Refresh the routing data with the existing configuration, this has
	//   the effect of attempting to rebuild the dead server.
	c.updateConfig(nil)

	// TODO(brett19): We probably should actually try other ways of resolving
	//  the issue, like requesting a new configuration.
}

func (c *Agent) tryHello(pipeline *memdPipeline, deadline time.Time) error {
	featuresBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(featuresBytes[0:], uint16(FeatureSeqNo))

	_, err := pipeline.ExecuteRequest(&memdQRequest{
		memdRequest: memdRequest{
			Magic:  ReqMagic,
			Opcode: CmdHello,
			Key:    []byte("Go SDK"),
			Value:  featuresBytes,
		},
	}, deadline)

	return err
}

// Attempt to connect a server, this function must be called
//  in its own goroutine and will ensure that offline servers
//  are not spammed with connection attempts.
func (agent *Agent) connectServer(server *memdPipeline) {
	for {
		agent.serverFailuresLock.Lock()
		failureTime := agent.serverFailures[server.address]
		agent.serverFailuresLock.Unlock()

		if !failureTime.IsZero() {
			waitedTime := time.Since(failureTime)
			if waitedTime < agent.serverWaitTimeout {
				time.Sleep(agent.serverWaitTimeout - waitedTime)

				if !agent.checkPendingServer(server) {
					// Server is no longer pending.  Stop trying.
					break
				}
			}
		}

		err := agent.connectPipeline(server, time.Now().Add(agent.serverConnectTimeout))
		if err != nil {
			agent.serverFailuresLock.Lock()
			agent.serverFailures[server.address] = time.Now()
			agent.serverFailuresLock.Unlock()

			// Try to connect again
			continue
		}

		if !agent.activatePendingServer(server) {
			// If this is no longer a valid pending server, we should shut
			//   it down!
			server.Close()
		}

		// Shut down this goroutine as we were successful
		break
	}
}

func (c *Agent) connectPipeline(pipeline *memdPipeline, deadline time.Time) error {
	logDebugf("Attempting to connect pipeline to %s", pipeline.address)
	memdConn, err := DialMemdConn(pipeline.address, c.tlsConfig, deadline)
	if err != nil {
		logDebugf("Failed to connect. %v", err)
		return err
	}

	logDebugf("Connected.")
	pipeline.conn = memdConn
	go pipeline.Run()

	logDebugf("Authenticating...")
	err = c.initFn(pipeline, deadline)
	if err != nil {
		logDebugf("Failed to authenticate. %v", err)
		memdConn.Close()
		return err
	}

	err = c.tryHello(pipeline, deadline)
	if err != nil {
		logDebugf("Failed to hello. %v", err)
		memdConn.Close()
		return err
	}

	return nil
}

// Drains all the requests out of the queue for this server.  This must be
//   invoked only once this server no longer exists in the routing data or an
//   infinite loop will likely occur.
func (c *Agent) shutdownPipeline(s *memdPipeline) {
	s.Drain(func(req *memdQRequest) {
		c.redispatchDirect(req)
	})
}

func (agent *Agent) checkPendingServer(server *memdPipeline) bool {
	oldRouting := agent.routingInfo.get()

	// Find the index of the pending server we want to swap
	var serverIdx int = -1
	for i, s := range oldRouting.pendingServers {
		if s == server {
			serverIdx = i
			break
		}
	}

	return serverIdx != -1
}

func (agent *Agent) activatePendingServer(server *memdPipeline) bool {
	logDebugf("Activating Server...")

	var oldRouting *routeData
	for {
		oldRouting = agent.routingInfo.get()

		// Find the index of the pending server we want to swap
		var serverIdx int = -1
		for i, s := range oldRouting.pendingServers {
			if s == server {
				serverIdx = i
				break
			}
		}

		if serverIdx == -1 {
			// This server is no longer in the list
			return false
		}

		var newRouting *routeData = &routeData{
			revId:      oldRouting.revId,
			capiEpList: oldRouting.capiEpList,
			mgmtEpList: oldRouting.mgmtEpList,
			n1qlEpList: oldRouting.n1qlEpList,
			vbMap:      oldRouting.vbMap,
			source:     oldRouting.source,
			deadQueue:  oldRouting.deadQueue,
		}

		// Copy the lists
		newRouting.queues = append(newRouting.queues, oldRouting.queues...)
		newRouting.servers = append(newRouting.servers, oldRouting.servers...)
		newRouting.pendingServers = append(newRouting.pendingServers, oldRouting.pendingServers...)

		// Swap around the pending server to being an active one
		newRouting.servers = append(newRouting.servers, server)
		newRouting.queues[serverIdx] = server.queue
		newRouting.pendingServers[serverIdx] = nil

		// Update to the new waitQueue
		for i, q := range newRouting.queues {
			if q == oldRouting.waitQueue {
				if newRouting.waitQueue == nil {
					newRouting.waitQueue = createMemdQueue()
				}
				newRouting.queues[i] = newRouting.waitQueue
			}
		}

		// Double-check the queue to make sure its still big enough.
		if len(newRouting.queues) != len(oldRouting.queues) {
			panic("Pending server swap corrupted the queue list.")
		}

		// Attempt to atomically update the routing data
		if !agent.routingInfo.update(oldRouting, newRouting) {
			// Someone preempted us, let's restart and try again...
			continue
		}

		server.SetHandlers(agent.handleServerNmv, agent.handleServerDeath)

		logDebugf("Switching routing data (server activation %d)...", serverIdx)
		oldRouting.logDebug()
		logDebugf("To new data...")
		newRouting.logDebug()

		// We've successfully swapped to the new config, lets finish building the
		//   new routing data's connections and destroy/draining old connections.
		break
	}

	oldWaitQueue := oldRouting.waitQueue
	if oldWaitQueue != nil {
		oldWaitQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}

	return true
}

// Accepts a cfgBucket object representing a cluster configuration and rebuilds the server list
//  along with any routing information for the Client.  Passing no config will refresh the existing one.
//  This method MUST NEVER BLOCK due to its use from various contention points.
func (agent *Agent) applyConfig(cfg *routeConfig) {
	// Check some basic things to ensure consistency!
	if len(cfg.vbMap) != agent.numVbuckets {
		panic("Received a configuration with a different number of vbuckets.")
	}

	var oldRouting *routeData
	var newRouting *routeData = &routeData{
		revId:      cfg.revId,
		capiEpList: cfg.capiEpList,
		mgmtEpList: cfg.mgmtEpList,
		n1qlEpList: cfg.n1qlEpList,
		vbMap:      cfg.vbMap,
		source:     cfg,
	}

	var needsDeadQueue bool
	for _, replicaList := range cfg.vbMap {
		for _, srvIdx := range replicaList {
			if srvIdx == -1 {
				needsDeadQueue = true
				break
			}
		}
	}
	if needsDeadQueue {
		newRouting.deadQueue = createMemdQueue()
	}

	var createdServers []*memdPipeline
	for {
		oldRouting = agent.routingInfo.get()

		if newRouting.revId < oldRouting.revId {
			logDebugf("Ignoring new configuration as it has an older revision id.")
			return
		}

		// Reset the lists before each iteration
		newRouting.queues = nil
		newRouting.servers = nil
		newRouting.pendingServers = nil

		for _, hostPort := range cfg.kvServerList {
			var thisServer *memdPipeline

			// See if this server exists in the old routing data and is still alive
			for _, oldServer := range oldRouting.servers {
				if oldServer.Address() == hostPort && !oldServer.IsClosed() {
					thisServer = oldServer
					break
				}
			}

			// If we found a still active connection in our old routing table,
			//   we just need to copy it over to the new routing table.
			if thisServer != nil {
				newRouting.pendingServers = append(newRouting.pendingServers, nil)
				newRouting.servers = append(newRouting.servers, thisServer)
				newRouting.queues = append(newRouting.queues, thisServer.queue)
				continue
			}

			// Search for any servers we are already trying to connect with.
			for _, oldServer := range oldRouting.pendingServers {
				if oldServer != nil && oldServer.Address() == hostPort {
					thisServer = oldServer
					break
				}
			}

			// If we do not already have a pending server we are trying to
			//   connect to, we should build one!
			if thisServer == nil {
				thisServer = CreateMemdPipeline(hostPort)
				createdServers = append(createdServers, thisServer)
			}

			if newRouting.waitQueue == nil {
				newRouting.waitQueue = createMemdQueue()
			}

			newRouting.pendingServers = append(newRouting.pendingServers, thisServer)
			newRouting.queues = append(newRouting.queues, newRouting.waitQueue)
		}

		// Check everything is sane
		if len(newRouting.queues) < len(cfg.kvServerList) {
			panic("Failed to generate a queues list that matches the config server list.")
		}

		// Attempt to atomically update the routing data
		if !agent.routingInfo.update(oldRouting, newRouting) {
			// Someone preempted us, let's restart and try again...
			continue
		}

		// We've successfully swapped to the new config, lets finish building the
		//   new routing data's connections and destroy/draining old connections.
		break
	}

	logDebugf("Switching routing data (update)...")
	oldRouting.logDebug()
	logDebugf("To new data...")
	newRouting.logDebug()

	// Launch all the new servers
	for _, newServer := range createdServers {
		go agent.connectServer(newServer)
	}

	// Identify all the dead servers and drain their requests
	for _, oldServer := range oldRouting.servers {
		found := false
		for _, newServer := range newRouting.servers {
			if newServer == oldServer {
				found = true
				break
			}
		}
		if !found {
			go agent.shutdownPipeline(oldServer)
		}
	}

	oldWaitQueue := oldRouting.waitQueue
	if oldWaitQueue != nil {
		oldWaitQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}

	oldDeadQueue := oldRouting.deadQueue
	if oldDeadQueue != nil {
		oldDeadQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}
}

func (agent *Agent) updateConfig(bk *cfgBucket) {
	if bk == nil {
		// Use the existing config if none was passed.
		oldRouting := agent.routingInfo.get()
		agent.applyConfig(oldRouting.source)
	} else {
		// Normalize the cfgBucket to a routeConfig and apply it.
		routeCfg := buildRouteConfig(bk, agent.IsSecure())
		agent.applyConfig(routeCfg)
	}
}

func (agent *Agent) routeRequest(req *memdQRequest) *memdQueue {
	routingInfo := agent.routingInfo.get()
	if routingInfo == nil {
		return nil
	}

	var srvIdx int

	repId := req.ReplicaIdx
	if repId < 0 {
		vbId := req.Vbucket
		srvIdx = routingInfo.vbMap[vbId][0]
	} else if req.Key == nil {
		vbId := req.Vbucket
		srvIdx = routingInfo.vbMap[vbId][repId]
	} else {
		vbId := cbCrc(req.Key) % uint32(len(routingInfo.vbMap))
		req.Vbucket = uint16(vbId)
		srvIdx = routingInfo.vbMap[vbId][repId]
	}

	if srvIdx == -1 {
		return routingInfo.deadQueue
	}

	return routingInfo.queues[srvIdx]
}

// This immediately dispatches a request to the appropriate server based on the
//  currently available routing data.
func (c *Agent) dispatchDirect(req *memdQRequest) error {
	for {
		pipeline := c.routeRequest(req)
		if pipeline == nil {
			// If no routing data exists this indicates that this Agent
			//   has been shut down!
			panic("Attempted to perform operation on closed agent.")
		}

		if !pipeline.QueueRequest(req) {
			continue
		}

		break
	}
	return nil
}

// This function is meant to be used when a memdRequest is internally shuffled
//   around.  It will fail to redispatch operations which are not allowed to be
//   moved between connections for whatever reason.
func (c *Agent) redispatchDirect(req *memdQRequest) {
	if req.ReplicaIdx >= 0 {
		// Reschedule the operation
		c.dispatchDirect(req)
	} else {
		// Callback advising that a network failure caused this operation to
		//   not be processed, nothing outside the agent should really see this.
		req.Callback(nil, networkError{})
	}
}
<commit_msg>Orphan: Ignore the result of the HELLO command.<commit_after>package gocbcore

import (
	"encoding/binary"
	"time"
)

func (c *Agent) handleServerNmv(s *memdPipeline, req *memdQRequest, resp *memdResponse) {
	// Try to parse the value as a bucket configuration
	bk, err := parseConfig(resp.Value, s.Hostname())
	if err == nil {
		c.updateConfig(bk)
	}

	// Redirect it!  This may actually come back to this server, but I won't tell
	//   if you don't ;)
	c.redispatchDirect(req)
}

func (c *Agent) handleServerDeath(s *memdPipeline) {
	// Refresh the routing data with the existing configuration, this has
	//   the effect of attempting to rebuild the dead server.
	c.updateConfig(nil)

	// TODO(brett19): We probably should actually try other ways of resolving
	//  the issue, like requesting a new configuration.
}

func (c *Agent) tryHello(pipeline *memdPipeline, deadline time.Time) error {
	featuresBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(featuresBytes[0:], uint16(FeatureSeqNo))

	_, err := pipeline.ExecuteRequest(&memdQRequest{
		memdRequest: memdRequest{
			Magic:  ReqMagic,
			Opcode: CmdHello,
			Key:    []byte("Go SDK"),
			Value:  featuresBytes,
		},
	}, deadline)

	return err
}

// Attempt to connect a server, this function must be called
//  in its own goroutine and will ensure that offline servers
//  are not spammed with connection attempts.
func (agent *Agent) connectServer(server *memdPipeline) {
	for {
		agent.serverFailuresLock.Lock()
		failureTime := agent.serverFailures[server.address]
		agent.serverFailuresLock.Unlock()

		if !failureTime.IsZero() {
			waitedTime := time.Since(failureTime)
			if waitedTime < agent.serverWaitTimeout {
				time.Sleep(agent.serverWaitTimeout - waitedTime)

				if !agent.checkPendingServer(server) {
					// Server is no longer pending.  Stop trying.
					break
				}
			}
		}

		err := agent.connectPipeline(server, time.Now().Add(agent.serverConnectTimeout))
		if err != nil {
			agent.serverFailuresLock.Lock()
			agent.serverFailures[server.address] = time.Now()
			agent.serverFailuresLock.Unlock()

			// Try to connect again
			continue
		}

		if !agent.activatePendingServer(server) {
			// If this is no longer a valid pending server, we should shut
			//   it down!
			server.Close()
		}

		// Shut down this goroutine as we were successful
		break
	}
}

func (c *Agent) connectPipeline(pipeline *memdPipeline, deadline time.Time) error {
	logDebugf("Attempting to connect pipeline to %s", pipeline.address)
	memdConn, err := DialMemdConn(pipeline.address, c.tlsConfig, deadline)
	if err != nil {
		logDebugf("Failed to connect. %v", err)
		return err
	}

	logDebugf("Connected.")
	pipeline.conn = memdConn
	go pipeline.Run()

	logDebugf("Authenticating...")
	err = c.initFn(pipeline, deadline)
	if err != nil {
		logDebugf("Failed to authenticate. %v", err)
		memdConn.Close()
		return err
	}

	c.tryHello(pipeline, deadline)

	return nil
}

// Drains all the requests out of the queue for this server.  This must be
//   invoked only once this server no longer exists in the routing data or an
//   infinite loop will likely occur.
func (c *Agent) shutdownPipeline(s *memdPipeline) {
	s.Drain(func(req *memdQRequest) {
		c.redispatchDirect(req)
	})
}

func (agent *Agent) checkPendingServer(server *memdPipeline) bool {
	oldRouting := agent.routingInfo.get()

	// Find the index of the pending server we want to swap
	var serverIdx int = -1
	for i, s := range oldRouting.pendingServers {
		if s == server {
			serverIdx = i
			break
		}
	}

	return serverIdx != -1
}

func (agent *Agent) activatePendingServer(server *memdPipeline) bool {
	logDebugf("Activating Server...")

	var oldRouting *routeData
	for {
		oldRouting = agent.routingInfo.get()

		// Find the index of the pending server we want to swap
		var serverIdx int = -1
		for i, s := range oldRouting.pendingServers {
			if s == server {
				serverIdx = i
				break
			}
		}

		if serverIdx == -1 {
			// This server is no longer in the list
			return false
		}

		var newRouting *routeData = &routeData{
			revId:      oldRouting.revId,
			capiEpList: oldRouting.capiEpList,
			mgmtEpList: oldRouting.mgmtEpList,
			n1qlEpList: oldRouting.n1qlEpList,
			vbMap:      oldRouting.vbMap,
			source:     oldRouting.source,
			deadQueue:  oldRouting.deadQueue,
		}

		// Copy the lists
		newRouting.queues = append(newRouting.queues, oldRouting.queues...)
		newRouting.servers = append(newRouting.servers, oldRouting.servers...)
		newRouting.pendingServers = append(newRouting.pendingServers, oldRouting.pendingServers...)

		// Swap around the pending server to being an active one
		newRouting.servers = append(newRouting.servers, server)
		newRouting.queues[serverIdx] = server.queue
		newRouting.pendingServers[serverIdx] = nil

		// Update to the new waitQueue
		for i, q := range newRouting.queues {
			if q == oldRouting.waitQueue {
				if newRouting.waitQueue == nil {
					newRouting.waitQueue = createMemdQueue()
				}
				newRouting.queues[i] = newRouting.waitQueue
			}
		}

		// Double-check the queue to make sure its still big enough.
		if len(newRouting.queues) != len(oldRouting.queues) {
			panic("Pending server swap corrupted the queue list.")
		}

		// Attempt to atomically update the routing data
		if !agent.routingInfo.update(oldRouting, newRouting) {
			// Someone preempted us, let's restart and try again...
			continue
		}

		server.SetHandlers(agent.handleServerNmv, agent.handleServerDeath)

		logDebugf("Switching routing data (server activation %d)...", serverIdx)
		oldRouting.logDebug()
		logDebugf("To new data...")
		newRouting.logDebug()

		// We've successfully swapped to the new config, lets finish building the
		//   new routing data's connections and destroy/draining old connections.
		break
	}

	oldWaitQueue := oldRouting.waitQueue
	if oldWaitQueue != nil {
		oldWaitQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}

	return true
}

// Accepts a cfgBucket object representing a cluster configuration and rebuilds the server list
//  along with any routing information for the Client.  Passing no config will refresh the existing one.
//  This method MUST NEVER BLOCK due to its use from various contention points.
func (agent *Agent) applyConfig(cfg *routeConfig) {
	// Check some basic things to ensure consistency!
	if len(cfg.vbMap) != agent.numVbuckets {
		panic("Received a configuration with a different number of vbuckets.")
	}

	var oldRouting *routeData
	var newRouting *routeData = &routeData{
		revId:      cfg.revId,
		capiEpList: cfg.capiEpList,
		mgmtEpList: cfg.mgmtEpList,
		n1qlEpList: cfg.n1qlEpList,
		vbMap:      cfg.vbMap,
		source:     cfg,
	}

	var needsDeadQueue bool
	for _, replicaList := range cfg.vbMap {
		for _, srvIdx := range replicaList {
			if srvIdx == -1 {
				needsDeadQueue = true
				break
			}
		}
	}
	if needsDeadQueue {
		newRouting.deadQueue = createMemdQueue()
	}

	var createdServers []*memdPipeline
	for {
		oldRouting = agent.routingInfo.get()

		if newRouting.revId < oldRouting.revId {
			logDebugf("Ignoring new configuration as it has an older revision id.")
			return
		}

		// Reset the lists before each iteration
		newRouting.queues = nil
		newRouting.servers = nil
		newRouting.pendingServers = nil

		for _, hostPort := range cfg.kvServerList {
			var thisServer *memdPipeline

			// See if this server exists in the old routing data and is still alive
			for _, oldServer := range oldRouting.servers {
				if oldServer.Address() == hostPort && !oldServer.IsClosed() {
					thisServer = oldServer
					break
				}
			}

			// If we found a still active connection in our old routing table,
			//   we just need to copy it over to the new routing table.
			if thisServer != nil {
				newRouting.pendingServers = append(newRouting.pendingServers, nil)
				newRouting.servers = append(newRouting.servers, thisServer)
				newRouting.queues = append(newRouting.queues, thisServer.queue)
				continue
			}

			// Search for any servers we are already trying to connect with.
			for _, oldServer := range oldRouting.pendingServers {
				if oldServer != nil && oldServer.Address() == hostPort {
					thisServer = oldServer
					break
				}
			}

			// If we do not already have a pending server we are trying to
			//   connect to, we should build one!
			if thisServer == nil {
				thisServer = CreateMemdPipeline(hostPort)
				createdServers = append(createdServers, thisServer)
			}

			if newRouting.waitQueue == nil {
				newRouting.waitQueue = createMemdQueue()
			}

			newRouting.pendingServers = append(newRouting.pendingServers, thisServer)
			newRouting.queues = append(newRouting.queues, newRouting.waitQueue)
		}

		// Check everything is sane
		if len(newRouting.queues) < len(cfg.kvServerList) {
			panic("Failed to generate a queues list that matches the config server list.")
		}

		// Attempt to atomically update the routing data
		if !agent.routingInfo.update(oldRouting, newRouting) {
			// Someone preempted us, let's restart and try again...
			continue
		}

		// We've successfully swapped to the new config, lets finish building the
		//   new routing data's connections and destroy/draining old connections.
		break
	}

	logDebugf("Switching routing data (update)...")
	oldRouting.logDebug()
	logDebugf("To new data...")
	newRouting.logDebug()

	// Launch all the new servers
	for _, newServer := range createdServers {
		go agent.connectServer(newServer)
	}

	// Identify all the dead servers and drain their requests
	for _, oldServer := range oldRouting.servers {
		found := false
		for _, newServer := range newRouting.servers {
			if newServer == oldServer {
				found = true
				break
			}
		}
		if !found {
			go agent.shutdownPipeline(oldServer)
		}
	}

	oldWaitQueue := oldRouting.waitQueue
	if oldWaitQueue != nil {
		oldWaitQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}

	oldDeadQueue := oldRouting.deadQueue
	if oldDeadQueue != nil {
		oldDeadQueue.Drain(func(req *memdQRequest) {
			agent.redispatchDirect(req)
		}, nil)
	}
}

func (agent *Agent) updateConfig(bk *cfgBucket) {
	if bk == nil {
		// Use the existing config if none was passed.
		oldRouting := agent.routingInfo.get()
		agent.applyConfig(oldRouting.source)
	} else {
		// Normalize the cfgBucket to a routeConfig and apply it.
		routeCfg := buildRouteConfig(bk, agent.IsSecure())
		agent.applyConfig(routeCfg)
	}
}

func (agent *Agent) routeRequest(req *memdQRequest) *memdQueue {
	routingInfo := agent.routingInfo.get()
	if routingInfo == nil {
		return nil
	}

	var srvIdx int

	repId := req.ReplicaIdx
	if repId < 0 {
		vbId := req.Vbucket
		srvIdx = routingInfo.vbMap[vbId][0]
	} else if req.Key == nil {
		vbId := req.Vbucket
		srvIdx = routingInfo.vbMap[vbId][repId]
	} else {
		vbId := cbCrc(req.Key) % uint32(len(routingInfo.vbMap))
		req.Vbucket = uint16(vbId)
		srvIdx = routingInfo.vbMap[vbId][repId]
	}

	if srvIdx == -1 {
		return routingInfo.deadQueue
	}

	return routingInfo.queues[srvIdx]
}

// This immediately dispatches a request to the appropriate server based on the
//  currently available routing data.
func (c *Agent) dispatchDirect(req *memdQRequest) error {
	for {
		pipeline := c.routeRequest(req)
		if pipeline == nil {
			// If no routing data exists this indicates that this Agent
			//   has been shut down!
			panic("Attempted to perform operation on closed agent.")
		}

		if !pipeline.QueueRequest(req) {
			continue
		}

		break
	}
	return nil
}

// This function is meant to be used when a memdRequest is internally shuffled
//   around.  It will fail to redispatch operations which are not allowed to be
//   moved between connections for whatever reason.
func (c *Agent) redispatchDirect(req *memdQRequest) {
	if req.ReplicaIdx >= 0 {
		// Reschedule the operation
		c.dispatchDirect(req)
	} else {
		// Callback advising that a network failure caused this operation to
		//   not be processed, nothing outside the agent should really see this.
		req.Callback(nil, networkError{})
	}
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "net/http"
  "encoding/json"
  "database/sql"
  "github.com/pilu/traffic"
)

type Artist struct {
  Gid       string  `json:"gid"`
  Name      string  `json:"name"`
  Comment   string  `json:"comment"`
  BeginDate *string `json:"beginDate"`
  EndDate   *string `json:"endDate"`
  SortName  string  `json:"sortName"`
  Type      string  `json:"type"`
}

const FindArtistByGidQuery = `
  SELECT
    A.gid, AN.name, A.comment,
    A.begin_date_year, A.begin_date_month, A.begin_date_day,
    A.end_date_year, A.end_date_month, A.end_date_day,
    ASN.name, AT.name
  FROM
    artist A
  LEFT JOIN artist_name AN
    ON A.name = AN.id
  LEFT JOIN artist_name ASN
    ON A.sort_name = ASN.id
  LEFT JOIN artist_type AT
    ON A.type = AT.id
  WHERE
    A.gid = $1 limit 1;`

func FindArtistByGid(gid string) (*Artist, error) {
  artist := &Artist{}

  if !isValidUUID(gid) {
    return artist, InvalidUUID{ gid }
  }

  var name     *sql.NullString
  var sortName *sql.NullString
  var _type    *sql.NullString

  var beginDateYear   *sql.NullInt64
  var beginDateMonth  *sql.NullInt64
  var beginDateDay    *sql.NullInt64
  var endDateYear     *sql.NullInt64
  var endDateMonth    *sql.NullInt64
  var endDateDay      *sql.NullInt64

  err := DB.QueryRow(FindArtistByGidQuery, gid).Scan(&artist.Gid, &name, &artist.Comment,
                                                      &beginDateYear, &beginDateMonth, &beginDateDay,
                                                      &endDateYear, &endDateMonth, &endDateDay,
                                                      &sortName, &_type)

  if beginDateYear != nil {
    date := fmt.Sprintf("%d", beginDateYear.Int64)

    if beginDateMonth != nil  {
      date = fmt.Sprintf("%s-%02d", date, beginDateMonth.Int64)
    }

    if beginDateDay != nil  {
      date = fmt.Sprintf("%s-%02d", date, beginDateDay.Int64)
    }

    artist.BeginDate = &date
  }

  if endDateYear != nil {
    date := fmt.Sprintf("%d", endDateYear.Int64)

    if endDateMonth != nil  {
      date = fmt.Sprintf("%s-%02d", date, endDateMonth.Int64)
    }

    if endDateDay != nil  {
      date = fmt.Sprintf("%s-%02d", date, endDateDay.Int64)
    }

    artist.EndDate = &date
  }

  if name != nil {
    artist.Name = name.String
  }

  if sortName != nil {
    artist.SortName = sortName.String
  }

  if _type != nil {
    artist.Type = _type.String
  }

  return artist, err
}

func ArtistHandler(w traffic.ResponseWriter, r *http.Request) {
  gid := r.URL.Query().Get("gid")
  artist, err := FindArtistByGid(gid)

  if err == nil {
    json.NewEncoder(w).Encode(artist)
  } else if err == sql.ErrNoRows {
    w.WriteHeader(http.StatusNotFound)
  } else if _, ok := err.(InvalidUUID); ok {
    w.WriteHeader(http.StatusBadRequest)
  } else {
    w.WriteHeader(http.StatusInternalServerError)
  }
}
<commit_msg>refactoring<commit_after>package main

import (
  "fmt"
  "net/http"
  "encoding/json"
  "database/sql"
  "github.com/pilu/traffic"
)

type Artist struct {
  Gid       string  `json:"gid"`
  Name      string  `json:"name"`
  Comment   string  `json:"comment"`
  BeginDate *string `json:"beginDate"`
  EndDate   *string `json:"endDate"`
  SortName  string  `json:"sortName"`
  Type      string  `json:"type"`
}

const FindArtistByGidQuery = `
  SELECT
    A.gid, AN.name, A.comment,
    A.begin_date_year, A.begin_date_month, A.begin_date_day,
    A.end_date_year, A.end_date_month, A.end_date_day,
    ASN.name, AT.name
  FROM
    artist A
  LEFT JOIN artist_name AN
    ON A.name = AN.id
  LEFT JOIN artist_name ASN
    ON A.sort_name = ASN.id
  LEFT JOIN artist_type AT
    ON A.type = AT.id
  WHERE
    A.gid = $1 limit 1;`

func FindArtistByGid(gid string) (*Artist, error) {
  artist := &Artist{}

  if !isValidUUID(gid) {
    return artist, InvalidUUID{ gid }
  }

  var name     *sql.NullString
  var sortName *sql.NullString
  var _type    *sql.NullString

  var beginDateYear   *sql.NullInt64
  var beginDateMonth  *sql.NullInt64
  var beginDateDay    *sql.NullInt64
  var endDateYear     *sql.NullInt64
  var endDateMonth    *sql.NullInt64
  var endDateDay      *sql.NullInt64

  row := DB.QueryRow(FindArtistByGidQuery, gid)
  err := row.Scan(&artist.Gid, &name, &artist.Comment,
                  &beginDateYear, &beginDateMonth, &beginDateDay,
                  &endDateYear, &endDateMonth, &endDateDay,
                  &sortName, &_type)

  if beginDateYear != nil {
    date := fmt.Sprintf("%d", beginDateYear.Int64)

    if beginDateMonth != nil  {
      date = fmt.Sprintf("%s-%02d", date, beginDateMonth.Int64)
    }

    if beginDateDay != nil  {
      date = fmt.Sprintf("%s-%02d", date, beginDateDay.Int64)
    }

    artist.BeginDate = &date
  }

  if endDateYear != nil {
    date := fmt.Sprintf("%d", endDateYear.Int64)

    if endDateMonth != nil  {
      date = fmt.Sprintf("%s-%02d", date, endDateMonth.Int64)
    }

    if endDateDay != nil  {
      date = fmt.Sprintf("%s-%02d", date, endDateDay.Int64)
    }

    artist.EndDate = &date
  }

  if name != nil {
    artist.Name = name.String
  }

  if sortName != nil {
    artist.SortName = sortName.String
  }

  if _type != nil {
    artist.Type = _type.String
  }

  return artist, err
}

func ArtistHandler(w traffic.ResponseWriter, r *http.Request) {
  gid := r.URL.Query().Get("gid")
  artist, err := FindArtistByGid(gid)

  if err == nil {
    json.NewEncoder(w).Encode(artist)
  } else if err == sql.ErrNoRows {
    w.WriteHeader(http.StatusNotFound)
  } else if _, ok := err.(InvalidUUID); ok {
    w.WriteHeader(http.StatusBadRequest)
  } else {
    w.WriteHeader(http.StatusInternalServerError)
  }
}
<|endoftext|>


<commit_before>package app_context

import (
	"database/sql"
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/comstud/go-rollbar/rollbar"
	_ "github.com/lib/pq"
	"github.com/tilteng/go-logger/logger"
	"github.com/tilteng/go-metrics/metrics"
)

type AppContext interface {
	APIPort() int
	AppName() string
	BaseExternalURL() string
	CodeVersion() string
	DB() *sql.DB
	Hostname() string
	JSONSchemaFilePath() string
	Logger() logger.CtxLogger
	MetricsClient() metrics.MetricsClient
	MetricsEnabled() bool
	RollbarClient() rollbar.Client
	RollbarEnabled() bool
	SetLogger(logger.CtxLogger) AppContext
}

type baseAppContext struct {
	apiPort            int
	appName            string
	baseExternalURL    string
	codeVersion        string
	db                 *sql.DB
	hostname           string
	jsonSchemaFilePath string
	logger             logger.CtxLogger
	metricsClient      metrics.MetricsClient
	metricsEnabled     bool
	rollbarClient      rollbar.Client
	rollbarEnabled     bool
}

func (self *baseAppContext) APIPort() int {
	return self.apiPort
}

func (self *baseAppContext) AppName() string {
	return self.appName
}

func (self *baseAppContext) BaseExternalURL() string {
	return self.baseExternalURL
}

func (self *baseAppContext) CodeVersion() string {
	return self.codeVersion
}

func (self *baseAppContext) DB() *sql.DB {
	return self.db
}

func (self *baseAppContext) Hostname() string {
	return self.hostname
}

func (self *baseAppContext) JSONSchemaFilePath() string {
	return self.jsonSchemaFilePath
}

func (self *baseAppContext) Logger() logger.CtxLogger {
	return self.logger
}

func (self *baseAppContext) MetricsClient() metrics.MetricsClient {
	return self.metricsClient
}

func (self *baseAppContext) MetricsEnabled() bool {
	return self.metricsEnabled
}

func (self *baseAppContext) RollbarClient() rollbar.Client {
	return self.rollbarClient
}

func (self *baseAppContext) RollbarEnabled() bool {
	return self.rollbarEnabled
}

func (self *baseAppContext) SetLogger(logger logger.CtxLogger) AppContext {
	self.logger = logger
	return self
}

func (self *baseAppContext) isDisabled(s string) (bool, error) {
	s += "_DISABLE"
	if disable, ok := os.LookupEnv(s); ok {
		if disable == "true" {
			return true, nil
		}
		if disable != "false" {
			return true, errors.New(s + " must be 'true' or 'false'")
		}
	}
	return false, nil
}

func (self *baseAppContext) setRollbarClientFromEnv() error {
	if disabled, err := self.isDisabled("ROLLBAR"); disabled {
		return err
	}

	api_key := os.Getenv("ROLLBAR_API_KEY")
	if api_key == "" {
		return nil
	}

	if rcli, err := rollbar.NewClient(api_key); err != nil {
		return fmt.Errorf("Error creating rollbar client: %s", err)
	} else {
		self.rollbarClient = rcli
		self.rollbarEnabled = true

		opts := self.rollbarClient.Options()

		env := os.Getenv("ROLLBAR_ENVIRONMENT")
		if env == "development" || env == "staging" || env == "production" {
			opts.Environment = env
		}

		if len(self.codeVersion) != 0 {
			opts.NotifierServer.CodeVersion = self.codeVersion
		}
	}

	return nil
}

func (self *baseAppContext) setMetricsClientFromEnv() error {
	if disabled, err := self.isDisabled("METRICS"); disabled {
		return err
	}

	metrics_addr := os.Getenv("METRICS_ADDR")
	metrics_tags := os.Getenv("METRICS_TAGS")

	metrics_namespace, ok := os.LookupEnv("METRICS_NAMESPACE")
	if !ok {
		metrics_namespace = self.appName + "."
	}

	metrics_hostname, ok := os.LookupEnv("METRICS_HOSTNAME")
	if !ok {
		metrics_hostname = self.hostname
	}

	tags_map := map[string]string{
		"application": self.appName,
	}

	if len(metrics_hostname) > 0 {
		tags_map["host"] = metrics_hostname
	}

	for _, kv := range strings.Split(metrics_tags, ",") {
		if len(kv) == 0 {
			continue
		}
		if !strings.Contains(kv, "=") {
			kv += "="
		}
		parts := strings.SplitN(kv, "=", 2)
		if len(parts[0]) > 0 {
			tags_map[parts[0]] = parts[1]
		}
	}

	if mcli, err := metrics.NewMetricsClient(metrics_addr); err != nil {
		return err
	} else {
		mcli.SetNamespace(metrics_namespace)
		mcli.SetTags(tags_map)

		if err := mcli.Init(); err != nil {
			return err
		}

		self.metricsClient = mcli
		self.metricsEnabled = true
	}

	return nil
}

func (self *baseAppContext) setDBFromEnv() error {
	db_string := os.Getenv("DB_DSN")
	if len(db_string) == 0 {
		return nil
	}

	db, err := sql.Open("postgres", db_string)
	if err != nil {
		return errors.New("Couldn't open the database. Check that DB_DSN is correct.")
	}

	self.db = db

	return nil
}

func getIntFromEnv(name string) (int, bool, error) {
	str, found := os.LookupEnv(name)
	if !found || str == "" {
		return 0, found, nil
	}

	num, err := strconv.Atoi(str)
	if err != nil {
		return 0, true, fmt.Errorf("Env '%s' is not a number: %s", name, err)
	}

	return num, true, nil
}

func NewAppContext(app_name string) (AppContext, error) {
	appctx := &baseAppContext{
		logger:         logger.DefaultStdoutCtxLogger(),
		appName:        app_name,
		rollbarEnabled: false,
		rollbarClient:  rollbar.NewNOOPClient(),
		metricsEnabled: false,
		metricsClient:  metrics.NewNOOPClient(),
	}

	if host, err := os.Hostname(); err != nil {
		return nil, fmt.Errorf("Couldn't figure out hostname: %s", err)
	} else {
		appctx.hostname = host
	}

	// Set this before we setup rollbarClient
	appctx.codeVersion = os.Getenv("CODE_VERSION")

	if port, _, err := getIntFromEnv("API_PORT"); err != nil {
		return nil, err
	} else if port < 0 {
		return nil, errors.New("API_PORT can't be a negative number")
	} else {
		appctx.apiPort = port
	}

	appctx.jsonSchemaFilePath = os.Getenv("JSON_SCHEMA_FILEPATH")
	appctx.baseExternalURL = os.Getenv("BASE_URL")

	if err := appctx.setMetricsClientFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting metrics client: %s", err)
	}

	if err := appctx.setRollbarClientFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting rollbar client: %s", err)
	}

	if err := appctx.setDBFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting DB object: %s", err)
	}

	return appctx, nil
}
<commit_msg>Rename API_PORT to SERVICE_PORT, add DB connection limits<commit_after>package app_context

import (
	"database/sql"
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/comstud/go-rollbar/rollbar"
	_ "github.com/lib/pq"
	"github.com/tilteng/go-logger/logger"
	"github.com/tilteng/go-metrics/metrics"
)

type AppContext interface {
	AppName() string
	BaseExternalURL() string
	CodeVersion() string
	DB() *sql.DB
	Hostname() string
	JSONSchemaFilePath() string
	Logger() logger.CtxLogger
	MetricsClient() metrics.MetricsClient
	MetricsEnabled() bool
	RollbarClient() rollbar.Client
	RollbarEnabled() bool
	ServicePort() int
	SetLogger(logger.CtxLogger) AppContext
}

type baseAppContext struct {
	appName            string
	baseExternalURL    string
	codeVersion        string
	db                 *sql.DB
	dbMaxIdleConns     int
	dbMaxOpenConns     int
	hostname           string
	jsonSchemaFilePath string
	logger             logger.CtxLogger
	metricsClient      metrics.MetricsClient
	metricsEnabled     bool
	rollbarClient      rollbar.Client
	rollbarEnabled     bool
	servicePort        int
}

func (self *baseAppContext) AppName() string {
	return self.appName
}

func (self *baseAppContext) BaseExternalURL() string {
	return self.baseExternalURL
}

func (self *baseAppContext) CodeVersion() string {
	return self.codeVersion
}

func (self *baseAppContext) DB() *sql.DB {
	return self.db
}

func (self *baseAppContext) DBMaxIdleConns() int {
	return self.dbMaxIdleConns
}

func (self *baseAppContext) DBMaxOpenConns() int {
	return self.dbMaxOpenConns
}

func (self *baseAppContext) Hostname() string {
	return self.hostname
}

func (self *baseAppContext) JSONSchemaFilePath() string {
	return self.jsonSchemaFilePath
}

func (self *baseAppContext) Logger() logger.CtxLogger {
	return self.logger
}

func (self *baseAppContext) MetricsClient() metrics.MetricsClient {
	return self.metricsClient
}

func (self *baseAppContext) MetricsEnabled() bool {
	return self.metricsEnabled
}

func (self *baseAppContext) RollbarClient() rollbar.Client {
	return self.rollbarClient
}

func (self *baseAppContext) RollbarEnabled() bool {
	return self.rollbarEnabled
}

func (self *baseAppContext) ServicePort() int {
	return self.servicePort
}

func (self *baseAppContext) SetLogger(logger logger.CtxLogger) AppContext {
	self.logger = logger
	return self
}

func (self *baseAppContext) isDisabled(s string) (bool, error) {
	s += "_DISABLE"
	if disable, ok := os.LookupEnv(s); ok {
		if disable == "true" {
			return true, nil
		}
		if disable != "false" {
			return true, errors.New(s + " must be 'true' or 'false'")
		}
	}
	return false, nil
}

func (self *baseAppContext) setRollbarClientFromEnv() error {
	if disabled, err := self.isDisabled("ROLLBAR"); disabled {
		return err
	}

	api_key := os.Getenv("ROLLBAR_API_KEY")
	if api_key == "" {
		return nil
	}

	if rcli, err := rollbar.NewClient(api_key); err != nil {
		return fmt.Errorf("Error creating rollbar client: %s", err)
	} else {
		self.rollbarClient = rcli
		self.rollbarEnabled = true

		opts := self.rollbarClient.Options()

		env := os.Getenv("ROLLBAR_ENVIRONMENT")
		if env == "development" || env == "staging" || env == "production" {
			opts.Environment = env
		}

		if len(self.codeVersion) != 0 {
			opts.NotifierServer.CodeVersion = self.codeVersion
		}
	}

	return nil
}

func (self *baseAppContext) setMetricsClientFromEnv() error {
	if disabled, err := self.isDisabled("METRICS"); disabled {
		return err
	}

	metrics_addr := os.Getenv("METRICS_ADDR")
	metrics_tags := os.Getenv("METRICS_TAGS")

	metrics_namespace, ok := os.LookupEnv("METRICS_NAMESPACE")
	if !ok {
		metrics_namespace = self.appName + "."
	}

	metrics_hostname, ok := os.LookupEnv("METRICS_HOSTNAME")
	if !ok {
		metrics_hostname = self.hostname
	}

	tags_map := map[string]string{
		"application": self.appName,
	}

	if len(metrics_hostname) > 0 {
		tags_map["host"] = metrics_hostname
	}

	for _, kv := range strings.Split(metrics_tags, ",") {
		if len(kv) == 0 {
			continue
		}
		if !strings.Contains(kv, "=") {
			kv += "="
		}
		parts := strings.SplitN(kv, "=", 2)
		if len(parts[0]) > 0 {
			tags_map[parts[0]] = parts[1]
		}
	}

	if mcli, err := metrics.NewMetricsClient(metrics_addr); err != nil {
		return err
	} else {
		mcli.SetNamespace(metrics_namespace)
		mcli.SetTags(tags_map)

		if err := mcli.Init(); err != nil {
			return err
		}

		self.metricsClient = mcli
		self.metricsEnabled = true
	}

	return nil
}

func (self *baseAppContext) setDBFromEnv() error {
	db_string := os.Getenv("DB_DSN")
	if len(db_string) == 0 {
		return nil
	}

	db, err := sql.Open("postgres", db_string)
	if err != nil {
		return errors.New("Couldn't open the database. Check that DB_DSN is correct.")
	}

	self.db = db

	return nil
}

func (self *baseAppContext) setDBMaxIdleConnsFromEnv() error {
	if max_conns, found, err := getIntFromEnv("DB_MAX_IDLE_CONNS"); !found {
		return nil
	} else if err != nil {
		return err
	} else if max_conns < 1 {
		return errors.New("DB_MAX_IDLE_CONNS must be > 0")
	} else {
		self.dbMaxIdleConns = max_conns
		return nil
	}
}

func (self *baseAppContext) setDBMaxOpenConnsFromEnv() error {
	if max_conns, found, err := getIntFromEnv("DB_MAX_OPEN_CONNS"); !found {
		return nil
	} else if err != nil {
		return err
	} else if max_conns < 1 {
		return errors.New("DB_MAX_OPEN_CONNS must be > 0")
	} else {
		self.dbMaxOpenConns = max_conns
		return nil
	}
}

func getIntFromEnv(name string) (int, bool, error) {
	str, found := os.LookupEnv(name)
	if !found || str == "" {
		return 0, found, nil
	}

	num, err := strconv.Atoi(str)
	if err != nil {
		return 0, true, fmt.Errorf("Env '%s' is not a number: %s", name, err)
	}

	return num, true, nil
}

func (self *baseAppContext) setServicePortFromEnv() error {
	if port, _, err := getIntFromEnv("SERVICE_PORT"); err != nil {
		return err
	} else if port < 0 {
		return errors.New("SERVICE_PORT must be >= 0")
	} else {
		self.servicePort = port
		return nil
	}
}

func NewAppContext(app_name string) (AppContext, error) {
	appctx := &baseAppContext{
		logger:         logger.DefaultStdoutCtxLogger(),
		appName:        app_name,
		rollbarEnabled: false,
		rollbarClient:  rollbar.NewNOOPClient(),
		metricsEnabled: false,
		metricsClient:  metrics.NewNOOPClient(),
	}

	if host, err := os.Hostname(); err != nil {
		return nil, fmt.Errorf("Couldn't figure out hostname: %s", err)
	} else {
		appctx.hostname = host
	}

	// Set this before we setup rollbarClient
	appctx.codeVersion = os.Getenv("CODE_VERSION")

	appctx.jsonSchemaFilePath = os.Getenv("JSON_SCHEMA_FILEPATH")
	appctx.baseExternalURL = os.Getenv("BASE_URL")

	if err := appctx.setServicePortFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting service port: %s", err)
	}

	if err := appctx.setMetricsClientFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting metrics client: %s", err)
	}

	if err := appctx.setRollbarClientFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting rollbar client: %s", err)
	}

	if err := appctx.setDBMaxIdleConnsFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting DB max idle connections: %s", err)
	}

	if err := appctx.setDBMaxOpenConnsFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting DB max open connections: %s", err)
	}

	if err := appctx.setDBFromEnv(); err != nil {
		return nil, fmt.Errorf("Error setting DB object: %s", err)
	}

	if appctx.db != nil {
		if appctx.dbMaxIdleConns > 0 {
			appctx.db.SetMaxIdleConns(appctx.dbMaxIdleConns)
		}
		if appctx.dbMaxOpenConns > 0 {
			appctx.db.SetMaxOpenConns(appctx.dbMaxOpenConns)
		}
	}

	return appctx, nil
}
<|endoftext|>