<commit_before>package headlessChrome

import (
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/integrii/interactive"
)

// Debug enables debug output for this package to console
var Debug bool

// BrowserStartupTime is how long chrome has to startup the console
// before we consider it a failure
var BrowserStartupTime = time.Minute

// ChromePath is the command to execute chrome
var ChromePath = `/Applications/Google Chrome.app/Contents/MacOS/Google Chrome`

// Args are the args that will be used to start chrome
var Args = []string{
	"--headless",
	"--disable-gpu",
	"--repl",
	// "--dump-dom",
	// "--window-size=1024,768",
	// "--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
	// "--verbose",
}

const expectedFirstLine = `Type a Javascript expression to evaluate or "quit" to exit.`
const promptPrefix = `>>>`

// OutputSanitizer puts output coming from the consolw that
// does not begin with the input prompt into the session
// output channel
func (cs *ChromeSession) OutputSanitizer() {
	for text := range cs.Session.Output {
		if !strings.HasPrefix(text, promptPrefix) {
			cs.Output <- text
		}
	}
}

// ChromeSession is an interactive console Session with a Chrome
// instance.
type ChromeSession struct {
	Session *interactive.Session
	Output  chan string
	Input   chan string
}

// Exit exits the running command out by ossuing a 'quit'
// to the chrome console
func (cs *ChromeSession) Exit() {
	cs.Session.Write(`quit`)
	cs.Session.Exit()
}

// Write writes to the Session
func (cs *ChromeSession) Write(s string) {
	debug("Writing to console:")
	debug(s)
	cs.Session.Write(s)
}

// OutputPrinter prints all outputs from the output channel to the cli
func (cs *ChromeSession) OutputPrinter() {
	for l := range cs.Session.Output {
		fmt.Println(l)
	}
}

// forceClose issues a force kill to the command
func (cs *ChromeSession) forceClose() {
	cs.Session.ForceClose()
}

// ClickSelector calls a click() on the supplied selector
func (cs *ChromeSession) ClickSelector(s string) {
	cs.Write(`document.querySelector("` + s + `").click()`)
}

// ClickItemWithInnerHTML clicks an item that has the matching inner html
func (cs *ChromeSession) ClickItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `].click()`)
}

// GetItemWithInnerHTML fetches the item with the specified innerHTML content
func (cs *ChromeSession) GetItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML.indexOf("` + s + `") == 0; });x[` + strconv.Itoa(itemIndex) + `]`)
}

// GetContentOfItemWithClasses fetches the content of the element with the specified classes
func (cs *ChromeSession) GetContentOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML`)
}

// GetValueOfItemWithClasses returns the form value of the specified item
func (cs *ChromeSession) GetValueOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value`)
}

// GetContentOfItemWithSelector gets the content of an element with the specified selector
func (cs *ChromeSession) GetContentOfItemWithSelector(selector string) {
	cs.Write(`document.querySelector("` + selector + `").innerHTML()`)
}

// ClickItemWithClasses clicks on the first item it finds with the provided classes.
// Multiple classes are separated by spaces
func (cs *ChromeSession) ClickItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].click()`)
}

// SetTextByID sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByID(id string, text string) {
	cs.Write(`document.getElementById("` + id + `").innerHTML = "` + text + `"`)
}

// ClickItemWithID clicks an item with the specified id
func (cs *ChromeSession) ClickItemWithID(id string) {
	cs.Write(`document.getElementById("` + id + `").click()`)
}

// SetTextByClasses sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML = "` + text + `"`)
}

// SetInputTextByClasses sets the input text for an input field
func (cs *ChromeSession) SetInputTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value = "` + text + `"`)
}

// NewBrowserWithTimeout starts a new chrome headless session
// but limits how long it can run before its killed forcefully.
// A time limit of 0 means there is not a time limit
func NewBrowserWithTimeout(url string, timeout time.Duration) (*ChromeSession, error) {
	var err error

	chromeSession := ChromeSession{}
	chromeSession.Output = make(chan string, 5000)

	// add url as last arg and create new Session
	args := append(Args, url)
	chromeSession.Session, err = interactive.NewSessionWithTimeout(ChromePath, args, timeout)
	if err != nil {
		return &chromeSession, err
	}

	// map output and input channels for easy use
	chromeSession.Input = chromeSession.Session.Input

	go chromeSession.OutputSanitizer()

	// wait for the console ready line from the browser
	// and if it does not start in time, move on
	select {
	case firstLine := <-chromeSession.Output:
		if !strings.Contains(firstLine, expectedFirstLine) {
			log.Println("ERROR: Unespected first line when initializing headless Chrome console:", firstLine)
		}
	case <-time.After(BrowserStartupTime):
		log.Println("ERROR: Browser failed to start before browser startup time cutoff")
		chromeSession.forceClose() // force cloe the session because it failed
		err = errors.New("Chrome console failed to init in the alotted time")
	}

	return &chromeSession, err
}

// NewBrowser starts a new chrome headless Session.
func NewBrowser(url string) (*ChromeSession, error) {
	return NewBrowserWithTimeout(url, 0)
}

func debug(s ...interface{}) {
	if Debug {
		fmt.Println(s...)
	}
}
<commit_msg>improving the way divs with specific content are handled<commit_after>package headlessChrome

import (
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/integrii/interactive"
)

// Debug enables debug output for this package to console
var Debug bool

// BrowserStartupTime is how long chrome has to startup the console
// before we consider it a failure
var BrowserStartupTime = time.Minute

// ChromePath is the command to execute chrome
var ChromePath = `/Applications/Google Chrome.app/Contents/MacOS/Google Chrome`

// Args are the args that will be used to start chrome
var Args = []string{
	"--headless",
	"--disable-gpu",
	"--repl",
	// "--dump-dom",
	// "--window-size=1024,768",
	// "--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/61.0.3163.100 Safari/537.36",
	// "--verbose",
}

const expectedFirstLine = `Type a Javascript expression to evaluate or "quit" to exit.`
const promptPrefix = `>>>`

// OutputSanitizer puts output coming from the consolw that
// does not begin with the input prompt into the session
// output channel
func (cs *ChromeSession) OutputSanitizer() {
	for text := range cs.Session.Output {
		if !strings.HasPrefix(text, promptPrefix) {
			cs.Output <- text
		}
	}
}

// ChromeSession is an interactive console Session with a Chrome
// instance.
type ChromeSession struct {
	Session *interactive.Session
	Output  chan string
	Input   chan string
}

// Exit exits the running command out by ossuing a 'quit'
// to the chrome console
func (cs *ChromeSession) Exit() {
	cs.Session.Write(`quit`)
	cs.Session.Exit()
}

// Write writes to the Session
func (cs *ChromeSession) Write(s string) {
	debug(s)
	cs.Session.Write(s)
}

// OutputPrinter prints all outputs from the output channel to the cli
func (cs *ChromeSession) OutputPrinter() {
	for l := range cs.Session.Output {
		fmt.Println(l)
	}
}

// forceClose issues a force kill to the command
func (cs *ChromeSession) forceClose() {
	cs.Session.ForceClose()
}

// ClickSelector calls a click() on the supplied selector
func (cs *ChromeSession) ClickSelector(s string) {
	cs.Write(`document.querySelector("` + s + `").click()`)
}

// ClickItemWithInnerHTML clicks an item that has the matching inner html
func (cs *ChromeSession) ClickItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML == ` + s + `")});x[` + strconv.Itoa(itemIndex) + `].click()`)
}

// GetItemWithInnerHTML fetches the item with the specified innerHTML content
func (cs *ChromeSession) GetItemWithInnerHTML(elementType string, s string, itemIndex int) {
	cs.Write(`var x = $("` + elementType + `").filter(function(idx) { return this.innerHTML == ` + s + `")});x[` + strconv.Itoa(itemIndex) + `]`)
}

// GetContentOfItemWithClasses fetches the content of the element with the specified classes
func (cs *ChromeSession) GetContentOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML`)
}

// GetValueOfItemWithClasses returns the form value of the specified item
func (cs *ChromeSession) GetValueOfItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value`)
}

// GetContentOfItemWithSelector gets the content of an element with the specified selector
func (cs *ChromeSession) GetContentOfItemWithSelector(selector string) {
	cs.Write(`document.querySelector("` + selector + `").innerHTML()`)
}

// ClickItemWithClasses clicks on the first item it finds with the provided classes.
// Multiple classes are separated by spaces
func (cs *ChromeSession) ClickItemWithClasses(classes string, itemIndex int) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].click()`)
}

// SetTextByID sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByID(id string, text string) {
	cs.Write(`document.getElementById("` + id + `").innerHTML = "` + text + `"`)
}

// ClickItemWithID clicks an item with the specified id
func (cs *ChromeSession) ClickItemWithID(id string) {
	cs.Write(`document.getElementById("` + id + `").click()`)
}

// SetTextByClasses sets the text on the div with the specified id
func (cs *ChromeSession) SetTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].innerHTML = "` + text + `"`)
}

// SetInputTextByClasses sets the input text for an input field
func (cs *ChromeSession) SetInputTextByClasses(classes string, itemIndex int, text string) {
	cs.Write(`document.getElementsByClassName("` + classes + `")[` + strconv.Itoa(itemIndex) + `].value = "` + text + `"`)
}

// NewBrowserWithTimeout starts a new chrome headless session
// but limits how long it can run before its killed forcefully.
// A time limit of 0 means there is not a time limit
func NewBrowserWithTimeout(url string, timeout time.Duration) (*ChromeSession, error) {
	var err error

	chromeSession := ChromeSession{}
	chromeSession.Output = make(chan string, 5000)

	// add url as last arg and create new Session
	args := append(Args, url)
	chromeSession.Session, err = interactive.NewSessionWithTimeout(ChromePath, args, timeout)
	if err != nil {
		return &chromeSession, err
	}

	// map output and input channels for easy use
	chromeSession.Input = chromeSession.Session.Input

	go chromeSession.OutputSanitizer()

	// wait for the console ready line from the browser
	// and if it does not start in time, move on
	select {
	case firstLine := <-chromeSession.Output:
		if !strings.Contains(firstLine, expectedFirstLine) {
			log.Println("WARNING: Unespected first line when initializing headless Chrome console:", firstLine)
		}
	case <-time.After(BrowserStartupTime):
		log.Println("ERROR: Browser failed to start before browser startup time cutoff")
		chromeSession.forceClose() // force cloe the session because it failed
		err = errors.New("Chrome console failed to init in the alotted time")
	}

	return &chromeSession, err
}

// NewBrowser starts a new chrome headless Session.
func NewBrowser(url string) (*ChromeSession, error) {
	return NewBrowserWithTimeout(url, 0)
}

func debug(s ...interface{}) {
	if Debug {
		fmt.Println(s...)
	}
}
<|endoftext|>


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

// Package chrometracing writes per-process Chrome trace_event files that can be
// loaded into chrome://tracing.
package chrometracing

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/google/chrometracing/traceinternal"
)

var trace = struct {
	start time.Time
	pid   uint64

	fileMu sync.Mutex
	file   *os.File
}{
	pid: uint64(os.Getpid()),
}

var out = setup(false)

// Path returns the full path of the chrome://tracing trace_event file for
// display in log messages.
func Path() string { return out }

// EnableTracing turns on tracing, regardless of running in a test or
// not. Tracing is enabled by default if the CHROMETRACING_DIR environment
// variable is present and non-empty.
func EnableTracing() {
	trace.fileMu.Lock()
	alreadyEnabled := trace.file != nil
	trace.fileMu.Unlock()
	if alreadyEnabled {
		return
	}
	out = setup(true)
}

func setup(overrideEnable bool) string {
	inTest := os.Getenv("TEST_TMPDIR") != ""
	explicitlyEnabled := os.Getenv("CHROMETRACING_DIR") != ""
	enableTracing := inTest || explicitlyEnabled || overrideEnable
	if !enableTracing {
		return ""
	}

	var err error
	dir := os.Getenv("TEST_UNDECLARED_OUTPUTS_DIR")
	if dir == "" {
		dir = os.Getenv("CHROMETRACING_DIR")
	}
	if dir == "" {
		dir = os.TempDir()
	}
	fn := filepath.Join(dir, fmt.Sprintf("%s.%d.trace", filepath.Base(os.Args[0]), trace.pid))
	trace.file, err = os.OpenFile(fn, os.O_WRONLY|os.O_CREATE|os.O_TRUNC|os.O_EXCL, 0644)
	if err != nil {
		// Using the log package from func init results in an error message
		// being printed.
		fmt.Fprintf(os.Stderr, "continuing without tracing: %v\n", err)
		return ""
	}

	// We only ever open a JSON array. Ending the array is optional as per
	// go/trace_event so that not cleanly finished traces can still be read.
	trace.file.Write([]byte{'['})
	trace.start = time.Now()

	writeEvent(&traceinternal.ViewerEvent{
		Name:  "process_name",
		Phase: "M", // Metadata Event
		Pid:   trace.pid,
		Tid:   trace.pid,
		Arg: struct {
			Name string `json:"name"`
		}{
			Name: strings.Join(os.Args, " "),
		},
	})
	return fn
}

func writeEvent(ev *traceinternal.ViewerEvent) {
	b, err := json.Marshal(&ev)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
	trace.fileMu.Lock()
	defer trace.fileMu.Unlock()
	if _, err = trace.file.Write(b); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
	if _, err = trace.file.Write([]byte{',', '\n'}); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
}

const (
	begin = "B"
	end   = "E"
)

// A PendingEvent represents an ongoing unit of work. The begin trace event has
// already been written, and calling Done will write the end trace event.
type PendingEvent struct {
	name string
	tid  uint64
}

// Done writes the end trace event for this unit of work.
func (pe *PendingEvent) Done() {
	if pe == nil || pe.name == "" || trace.file == nil {
		return
	}
	writeEvent(&traceinternal.ViewerEvent{
		Name:  pe.name,
		Phase: end,
		Pid:   trace.pid,
		Tid:   pe.tid,
		Time:  float64(time.Since(trace.start).Microseconds()),
	})
	releaseTid(pe.tid)
}

// Event logs a unit of work. To instrument a Go function, use e.g.:
//
//   func calcPi() {
//     defer chrometracing.Event("calculate pi").Done()
//     // …
//   }
//
// For more finely-granular traces, use e.g.:
//
//   for _, cmd := range commands {
//     ev := chrometracing.Event("initialize " + cmd.Name)
//     cmd.Init()
//     ev.Done()
//   }
func Event(name string) *PendingEvent {
	if trace.file == nil {
		return &PendingEvent{}
	}
	tid := tid()
	writeEvent(&traceinternal.ViewerEvent{
		Name:  name,
		Phase: begin,
		Pid:   trace.pid,
		Tid:   tid,
		Time:  float64(time.Since(trace.start).Microseconds()),
	})
	return &PendingEvent{
		name: name,
		tid:  tid,
	}
}

// tids is a chrome://tracing thread id pool. Go does not permit accessing the
// goroutine id, so we need to maintain our own identifier. The chrome://tracing
// file format requires a numeric thread id, so we just increment whenever we
// need a thread id, and reuse the ones no longer in use.
//
// In practice, parallelized sections of the code (many goroutines) end up using
// only as few thread ids as are concurrently in use, and the rest of the events
// mirror the code call stack nicely. See e.g. http://screen/7MPcAcvXQNUE3JZ
var tids struct {
	sync.Mutex

	// We allocate chrome://tracing thread ids based on the index of the
	// corresponding entry in the used slice.
	used []bool

	// next points to the earliest unused tid to consider for the next tid to
	// hand out. This is purely a performance optimization to avoid O(n) slice
	// iteration.
	next int
}

func tid() uint64 {
	tids.Lock()
	defer tids.Unlock()
	// re-use released tids if any
	for t := tids.next; t < len(tids.used); t++ {
		if !tids.used[t] {
			tids.used[t] = true
			tids.next = t + 1
			return uint64(t)
		}
	}
	// allocate a new tid
	t := len(tids.used)
	tids.used = append(tids.used, true)
	tids.next = t + 1
	return uint64(t)
}

func releaseTid(t uint64) {
	tids.Lock()
	defer tids.Unlock()
	tids.used[int(t)] = false
	if tids.next > int(t) {
		tids.next = int(t)
	}
}
<commit_msg>format with Go 1.19 gofmt<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package chrometracing writes per-process Chrome trace_event files that can be
// loaded into chrome://tracing.
package chrometracing

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/google/chrometracing/traceinternal"
)

var trace = struct {
	start time.Time
	pid   uint64

	fileMu sync.Mutex
	file   *os.File
}{
	pid: uint64(os.Getpid()),
}

var out = setup(false)

// Path returns the full path of the chrome://tracing trace_event file for
// display in log messages.
func Path() string { return out }

// EnableTracing turns on tracing, regardless of running in a test or
// not. Tracing is enabled by default if the CHROMETRACING_DIR environment
// variable is present and non-empty.
func EnableTracing() {
	trace.fileMu.Lock()
	alreadyEnabled := trace.file != nil
	trace.fileMu.Unlock()
	if alreadyEnabled {
		return
	}
	out = setup(true)
}

func setup(overrideEnable bool) string {
	inTest := os.Getenv("TEST_TMPDIR") != ""
	explicitlyEnabled := os.Getenv("CHROMETRACING_DIR") != ""
	enableTracing := inTest || explicitlyEnabled || overrideEnable
	if !enableTracing {
		return ""
	}

	var err error
	dir := os.Getenv("TEST_UNDECLARED_OUTPUTS_DIR")
	if dir == "" {
		dir = os.Getenv("CHROMETRACING_DIR")
	}
	if dir == "" {
		dir = os.TempDir()
	}
	fn := filepath.Join(dir, fmt.Sprintf("%s.%d.trace", filepath.Base(os.Args[0]), trace.pid))
	trace.file, err = os.OpenFile(fn, os.O_WRONLY|os.O_CREATE|os.O_TRUNC|os.O_EXCL, 0644)
	if err != nil {
		// Using the log package from func init results in an error message
		// being printed.
		fmt.Fprintf(os.Stderr, "continuing without tracing: %v\n", err)
		return ""
	}

	// We only ever open a JSON array. Ending the array is optional as per
	// go/trace_event so that not cleanly finished traces can still be read.
	trace.file.Write([]byte{'['})
	trace.start = time.Now()

	writeEvent(&traceinternal.ViewerEvent{
		Name:  "process_name",
		Phase: "M", // Metadata Event
		Pid:   trace.pid,
		Tid:   trace.pid,
		Arg: struct {
			Name string `json:"name"`
		}{
			Name: strings.Join(os.Args, " "),
		},
	})
	return fn
}

func writeEvent(ev *traceinternal.ViewerEvent) {
	b, err := json.Marshal(&ev)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
	trace.fileMu.Lock()
	defer trace.fileMu.Unlock()
	if _, err = trace.file.Write(b); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
	if _, err = trace.file.Write([]byte{',', '\n'}); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return
	}
}

const (
	begin = "B"
	end   = "E"
)

// A PendingEvent represents an ongoing unit of work. The begin trace event has
// already been written, and calling Done will write the end trace event.
type PendingEvent struct {
	name string
	tid  uint64
}

// Done writes the end trace event for this unit of work.
func (pe *PendingEvent) Done() {
	if pe == nil || pe.name == "" || trace.file == nil {
		return
	}
	writeEvent(&traceinternal.ViewerEvent{
		Name:  pe.name,
		Phase: end,
		Pid:   trace.pid,
		Tid:   pe.tid,
		Time:  float64(time.Since(trace.start).Microseconds()),
	})
	releaseTid(pe.tid)
}

// Event logs a unit of work. To instrument a Go function, use e.g.:
//
//	func calcPi() {
//	  defer chrometracing.Event("calculate pi").Done()
//	  // …
//	}
//
// For more finely-granular traces, use e.g.:
//
//	for _, cmd := range commands {
//	  ev := chrometracing.Event("initialize " + cmd.Name)
//	  cmd.Init()
//	  ev.Done()
//	}
func Event(name string) *PendingEvent {
	if trace.file == nil {
		return &PendingEvent{}
	}
	tid := tid()
	writeEvent(&traceinternal.ViewerEvent{
		Name:  name,
		Phase: begin,
		Pid:   trace.pid,
		Tid:   tid,
		Time:  float64(time.Since(trace.start).Microseconds()),
	})
	return &PendingEvent{
		name: name,
		tid:  tid,
	}
}

// tids is a chrome://tracing thread id pool. Go does not permit accessing the
// goroutine id, so we need to maintain our own identifier. The chrome://tracing
// file format requires a numeric thread id, so we just increment whenever we
// need a thread id, and reuse the ones no longer in use.
//
// In practice, parallelized sections of the code (many goroutines) end up using
// only as few thread ids as are concurrently in use, and the rest of the events
// mirror the code call stack nicely. See e.g. http://screen/7MPcAcvXQNUE3JZ
var tids struct {
	sync.Mutex

	// We allocate chrome://tracing thread ids based on the index of the
	// corresponding entry in the used slice.
	used []bool

	// next points to the earliest unused tid to consider for the next tid to
	// hand out. This is purely a performance optimization to avoid O(n) slice
	// iteration.
	next int
}

func tid() uint64 {
	tids.Lock()
	defer tids.Unlock()
	// re-use released tids if any
	for t := tids.next; t < len(tids.used); t++ {
		if !tids.used[t] {
			tids.used[t] = true
			tids.next = t + 1
			return uint64(t)
		}
	}
	// allocate a new tid
	t := len(tids.used)
	tids.used = append(tids.used, true)
	tids.next = t + 1
	return uint64(t)
}

func releaseTid(t uint64) {
	tids.Lock()
	defer tids.Unlock()
	tids.used[int(t)] = false
	if tids.next > int(t) {
		tids.next = int(t)
	}
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	"log"

	"github.com/docker/machine/libmachine"
	"github.com/kubernetes/minikube/cli/cluster"
	"github.com/kubernetes/minikube/cli/constants"
	"github.com/spf13/cobra"
)

// startCmd represents the start command
var startCmd = &cobra.Command{
	Use:   "start",
	Short: "Starts a local kubernetes cluster.",
	Long: `Starts a local kubernetes cluster using Virtualbox. This command
assumes you already have Virtualbox installed.`,
	Run: runStart,
}

func runStart(cmd *cobra.Command, args []string) {

	fmt.Println("Starting local Kubernetes cluster...")
	api := libmachine.NewClient(constants.Minipath, constants.MakeMiniPath("certs"))
	defer api.Close()
	host, err := cluster.StartHost(api)
	if err != nil {
		fmt.Println("Error starting host: ", err)
	}
	kubeHost, err := cluster.StartCluster(host)
	if err != nil {
		fmt.Println("Error starting cluster: ", err)
	}
	log.Printf("Kubernetes is available at %s.\n", kubeHost)
	log.Println("Run this command to use the cluster: ")
	log.Printf("kubectl config set-cluster minikube --insecure-skip-tls-verify=true --server=%s\n", kubeHost)
}

func init() {
	RootCmd.AddCommand(startCmd)
}
<commit_msg>Exit with a non-zero code if we error.<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"log"
	"os"

	"github.com/docker/machine/libmachine"
	"github.com/kubernetes/minikube/cli/cluster"
	"github.com/kubernetes/minikube/cli/constants"
	"github.com/spf13/cobra"
)

// startCmd represents the start command
var startCmd = &cobra.Command{
	Use:   "start",
	Short: "Starts a local kubernetes cluster.",
	Long: `Starts a local kubernetes cluster using Virtualbox. This command
assumes you already have Virtualbox installed.`,
	Run: runStart,
}

func runStart(cmd *cobra.Command, args []string) {

	fmt.Println("Starting local Kubernetes cluster...")
	api := libmachine.NewClient(constants.Minipath, constants.MakeMiniPath("certs"))
	defer api.Close()
	host, err := cluster.StartHost(api)
	if err != nil {
		log.Println("Error starting host: ", err)
		os.Exit(1)
	}
	kubeHost, err := cluster.StartCluster(host)
	if err != nil {
		log.Println("Error starting cluster: ", err)
		os.Exit(1)
	}
	log.Printf("Kubernetes is available at %s.\n", kubeHost)
	log.Println("Run this command to use the cluster: ")
	log.Printf("kubectl config set-cluster minikube --insecure-skip-tls-verify=true --server=%s\n", kubeHost)
}

func init() {
	RootCmd.AddCommand(startCmd)
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package binaryprefix

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix MB to number of megabytes", t, func() {
		Convey("Given the string 1MB, the number of megabytes is 1", func() {
			str := "1MB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1)
		})
		Convey("Given the string 1mb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1mb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}
<commit_msg>Added tests for GB suffix<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package binaryprefix

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix MB to number of megabytes", t, func() {
		Convey("Given the string 1MB, the number of megabytes is 1", func() {
			str := "1MB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1)
		})
		Convey("Given the string 1mb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1mb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestGB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix GB to number of megabytes", t, func() {
		Convey("Given the string 1GB, the number of megabytes is 1024", func() {
			str := "1GB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1024)
		})
		Convey("Given the string 1gb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1gb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}
<|endoftext|>


<commit_before>package marathon

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Event struct {
	EventType  string    `json:"eventType"`
	Timestamp  time.Time `json:"timestamp"`
	SlaveID    string    `json:"slaveId"`
	TaskID     string    `json:"taskId"`
	TaskStatus string    `json:"taskStatus"`
	AppID      string    `json:"appId"`
	Host       string    `json:"host"`
	Ports      []int     `json:"ports"`
	Version    string    `json:"version"`
}

type Listener struct {
	events       chan Event
	host         string
	internalPort string // Internal/external ports are relative
	externalPort string // to the container this process runs in.
}

func NewListener(host string, internalPort, externalPort string) *Listener {
	listener := &Listener{
		events:       make(chan Event),
		host:         host,
		internalPort: internalPort,
		externalPort: externalPort,
	}
	http.HandleFunc("/push-listener", listener.handler)
	go http.ListenAndServe(":"+internalPort, nil)

	return listener
}

func (l *Listener) handler(res http.ResponseWriter, req *http.Request) {
	decoder := json.NewDecoder(req.Body)

	var event Event
	if err := decoder.Decode(&event); err != nil {
		log.Fatal(err)
	}

	if event.EventType == "status_update_event" { // We only care about container change events
		l.events <- event
	}

	res.Write([]byte("Thanks.")) // Marathon ignores replies.  Just being polite.
}

func (l *Listener) Events() <-chan Event {
	return l.events
}

func (l *Listener) Subscribe(marathonHost string) error {
	marathonURL := url.URL{Scheme: "http", Host: marathonHost, Path: "/v2/eventSubscriptions"}
	q := marathonURL.Query()
	q.Set("callbackUrl", fmt.Sprintf("http://%s:%s/push-listener", l.host, l.externalPort))
	marathonURL.RawQuery = q.Encode()

	res, err := http.Post(marathonURL.String(), "application/json", strings.NewReader(""))
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("Bad status code while subscribing to marathon events: " + res.Status)
	}

	var data map[string]interface{}
	decoder := json.NewDecoder(res.Body)
	if err := decoder.Decode(&data); err != nil {
		return err
	}

	return nil
}
<commit_msg>Used nil instead of dummy object.<commit_after>package marathon

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"time"
)

type Event struct {
	EventType  string    `json:"eventType"`
	Timestamp  time.Time `json:"timestamp"`
	SlaveID    string    `json:"slaveId"`
	TaskID     string    `json:"taskId"`
	TaskStatus string    `json:"taskStatus"`
	AppID      string    `json:"appId"`
	Host       string    `json:"host"`
	Ports      []int     `json:"ports"`
	Version    string    `json:"version"`
}

type Listener struct {
	events       chan Event
	host         string
	internalPort string // Internal/external ports are relative
	externalPort string // to the container this process runs in.
}

func NewListener(host string, internalPort, externalPort string) *Listener {
	listener := &Listener{
		events:       make(chan Event),
		host:         host,
		internalPort: internalPort,
		externalPort: externalPort,
	}
	http.HandleFunc("/push-listener", listener.handler)
	go http.ListenAndServe(":"+internalPort, nil)

	return listener
}

func (l *Listener) handler(res http.ResponseWriter, req *http.Request) {
	decoder := json.NewDecoder(req.Body)

	var event Event
	if err := decoder.Decode(&event); err != nil {
		log.Fatal(err)
	}

	if event.EventType == "status_update_event" { // We only care about container change events
		l.events <- event
	}

	res.Write([]byte("Thanks.")) // Marathon ignores replies.  Just being polite.
}

func (l *Listener) Events() <-chan Event {
	return l.events
}

func (l *Listener) Subscribe(marathonHost string) error {
	marathonURL := url.URL{Scheme: "http", Host: marathonHost, Path: "/v2/eventSubscriptions"}
	q := marathonURL.Query()
	q.Set("callbackUrl", fmt.Sprintf("http://%s:%s/push-listener", l.host, l.externalPort))
	marathonURL.RawQuery = q.Encode()

	res, err := http.Post(marathonURL.String(), "application/json", nil)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("Bad status code while subscribing to marathon events: " + res.Status)
	}

	var data map[string]interface{}
	decoder := json.NewDecoder(res.Body)
	if err := decoder.Decode(&data); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
	"os"
)

func main() {
	k := kite.New("client", "1.0.0")
        k.Config.DisableAuthentication = true

	fmt.Println(k.Config)

	client = k.NewClient("http://square.openshiftapps.com:6001/kite")
	client.Dial()

	response, err := fetchedKite.Tell("square", 4)
	if err != nil {
		panic(err)
	}
	fmt.Println(response.MustFloat64())
}
<commit_msg>hardcoded square service url<commit_after>package main

import (
	"fmt"
	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
	"os"
)

func main() {
	k := kite.New("client", "1.0.0")
        k.Config.DisableAuthentication = true

	fmt.Println(k.Config)

	client = k.NewClient("http://square.openshiftapps.com:6001/kite")
	client.Dial()

	response, err := client.Tell("square", 4)
	if err != nil {
		panic(err)
	}
	fmt.Println(response.MustFloat64())
}
<|endoftext|>


<commit_before>/*
Package client is a Go client for the Docker Engine API.

For more information about the Engine API, see the documentation:
https://docs.docker.com/engine/reference/api/

Usage

You use the library by creating a client object and calling methods on it. The
client can be created either from environment variables with NewEnvClient, or
configured manually with NewClient.

For example, to list running containers (the equivalent of "docker ps"):

	package main

	import (
		"context"
		"fmt"

		"github.com/docker/docker/api/types"
		"github.com/docker/docker/client"
	)

	func main() {
		cli, err := client.NewEnvClient()
		if err != nil {
			panic(err)
		}

		containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{})
		if err != nil {
			panic(err)
		}

		for _, container := range containers {
			fmt.Printf("%s %s\n", container.ID[:10], container.Image)
		}
	}

*/
package client // import "github.com/docker/docker/client"

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/docker/docker/api"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/versions"
	"github.com/docker/go-connections/sockets"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/pkg/errors"
)

// ErrRedirect is the error returned by checkRedirect when the request is non-GET.
var ErrRedirect = errors.New("unexpected redirect in response")

// Client is the API client that performs all operations
// against a docker server.
type Client struct {
	// scheme sets the scheme for the client
	scheme string
	// host holds the server address to connect to
	host string
	// proto holds the client protocol i.e. unix.
	proto string
	// addr holds the client address.
	addr string
	// basePath holds the path to prepend to the requests.
	basePath string
	// client used to send and receive http requests.
	client *http.Client
	// version of the server to talk to.
	version string
	// custom http headers configured by users.
	customHTTPHeaders map[string]string
	// manualOverride is set to true when the version was set by users.
	manualOverride bool
}

// CheckRedirect specifies the policy for dealing with redirect responses:
// If the request is non-GET return `ErrRedirect`. Otherwise use the last response.
//
// Go 1.8 changes behavior for HTTP redirects (specifically 301, 307, and 308) in the client .
// The Docker client (and by extension docker API client) can be made to to send a request
// like POST /containers//start where what would normally be in the name section of the URL is empty.
// This triggers an HTTP 301 from the daemon.
// In go 1.8 this 301 will be converted to a GET request, and ends up getting a 404 from the daemon.
// This behavior change manifests in the client in that before the 301 was not followed and
// the client did not generate an error, but now results in a message like Error response from daemon: page not found.
func CheckRedirect(req *http.Request, via []*http.Request) error {
	if via[0].Method == http.MethodGet {
		return http.ErrUseLastResponse
	}
	return ErrRedirect
}

// NewEnvClient initializes a new API client based on environment variables.
// See FromEnv for a list of support environment variables.
//
// Deprecated: use NewClientWithOpts(FromEnv)
func NewEnvClient() (*Client, error) {
	return NewClientWithOpts(FromEnv)
}

// FromEnv configures the client with values from environment variables.
//
// Supported environment variables:
// DOCKER_HOST to set the url to the docker server.
// DOCKER_API_VERSION to set the version of the API to reach, leave empty for latest.
// DOCKER_CERT_PATH to load the TLS certificates from.
// DOCKER_TLS_VERIFY to enable or disable TLS verification, off by default.
func FromEnv(c *Client) error {
	if dockerCertPath := os.Getenv("DOCKER_CERT_PATH"); dockerCertPath != "" {
		options := tlsconfig.Options{
			CAFile:             filepath.Join(dockerCertPath, "ca.pem"),
			CertFile:           filepath.Join(dockerCertPath, "cert.pem"),
			KeyFile:            filepath.Join(dockerCertPath, "key.pem"),
			InsecureSkipVerify: os.Getenv("DOCKER_TLS_VERIFY") == "",
		}
		tlsc, err := tlsconfig.Client(options)
		if err != nil {
			return err
		}

		c.client = &http.Client{
			Transport:     &http.Transport{TLSClientConfig: tlsc},
			CheckRedirect: CheckRedirect,
		}
	}

	if host := os.Getenv("DOCKER_HOST"); host != "" {
		if err := WithHost(host)(c); err != nil {
			return err
		}
	}

	if version := os.Getenv("DOCKER_API_VERSION"); version != "" {
		c.version = version
		c.manualOverride = true
	}
	return nil
}

// WithTLSClientConfig applies a tls config to the client transport.
func WithTLSClientConfig(cacertPath, certPath, keyPath string) func(*Client) error {
	return func(c *Client) error {
		opts := tlsconfig.Options{
			CAFile:             cacertPath,
			CertFile:           certPath,
			KeyFile:            keyPath,
			ExclusiveRootPools: true,
		}
		config, err := tlsconfig.Client(opts)
		if err != nil {
			return errors.Wrap(err, "failed to create tls config")
		}
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			transport.TLSClientConfig = config
			return nil
		}
		return errors.Errorf("cannot apply tls config to transport: %T", c.client.Transport)
	}
}

// WithDialer applies the dialer.DialContext to the client transport. This can be
// used to set the Timeout and KeepAlive settings of the client.
// Deprecated: use WithDialContext
func WithDialer(dialer *net.Dialer) func(*Client) error {
	return WithDialContext(dialer.DialContext)
}

// WithDialContext applies the dialer to the client transport. This can be
// used to set the Timeout and KeepAlive settings of the client.
func WithDialContext(dialContext func(ctx context.Context, network, addr string) (net.Conn, error)) func(*Client) error {
	return func(c *Client) error {
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			transport.DialContext = dialContext
			return nil
		}
		return errors.Errorf("cannot apply dialer to transport: %T", c.client.Transport)
	}
}

// WithVersion overrides the client version with the specified one
func WithVersion(version string) func(*Client) error {
	return func(c *Client) error {
		c.version = version
		return nil
	}
}

// WithHost overrides the client host with the specified one.
func WithHost(host string) func(*Client) error {
	return func(c *Client) error {
		hostURL, err := ParseHostURL(host)
		if err != nil {
			return err
		}
		c.host = host
		c.proto = hostURL.Scheme
		c.addr = hostURL.Host
		c.basePath = hostURL.Path
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			return sockets.ConfigureTransport(transport, c.proto, c.addr)
		}
		return errors.Errorf("cannot apply host to transport: %T", c.client.Transport)
	}
}

// WithHTTPClient overrides the client http client with the specified one
func WithHTTPClient(client *http.Client) func(*Client) error {
	return func(c *Client) error {
		if client != nil {
			c.client = client
		}
		return nil
	}
}

// WithHTTPHeaders overrides the client default http headers
func WithHTTPHeaders(headers map[string]string) func(*Client) error {
	return func(c *Client) error {
		c.customHTTPHeaders = headers
		return nil
	}
}

// NewClientWithOpts initializes a new API client with default values. It takes functors
// to modify values when creating it, like `NewClientWithOpts(WithVersion(…))`
// It also initializes the custom http headers to add to each request.
//
// It won't send any version information if the version number is empty. It is
// highly recommended that you set a version or your client may break if the
// server is upgraded.
func NewClientWithOpts(ops ...func(*Client) error) (*Client, error) {
	client, err := defaultHTTPClient(DefaultDockerHost)
	if err != nil {
		return nil, err
	}
	c := &Client{
		host:    DefaultDockerHost,
		version: api.DefaultVersion,
		scheme:  "http",
		client:  client,
		proto:   defaultProto,
		addr:    defaultAddr,
	}

	for _, op := range ops {
		if err := op(c); err != nil {
			return nil, err
		}
	}

	if _, ok := c.client.Transport.(http.RoundTripper); !ok {
		return nil, fmt.Errorf("unable to verify TLS configuration, invalid transport %v", c.client.Transport)
	}
	tlsConfig := resolveTLSConfig(c.client.Transport)
	if tlsConfig != nil {
		// TODO(stevvooe): This isn't really the right way to write clients in Go.
		// `NewClient` should probably only take an `*http.Client` and work from there.
		// Unfortunately, the model of having a host-ish/url-thingy as the connection
		// string has us confusing protocol and transport layers. We continue doing
		// this to avoid breaking existing clients but this should be addressed.
		c.scheme = "https"
	}

	return c, nil
}

func defaultHTTPClient(host string) (*http.Client, error) {
	url, err := ParseHostURL(host)
	if err != nil {
		return nil, err
	}
	transport := new(http.Transport)
	sockets.ConfigureTransport(transport, url.Scheme, url.Host)
	return &http.Client{
		Transport:     transport,
		CheckRedirect: CheckRedirect,
	}, nil
}

// NewClient initializes a new API client for the given host and API version.
// It uses the given http client as transport.
// It also initializes the custom http headers to add to each request.
//
// It won't send any version information if the version number is empty. It is
// highly recommended that you set a version or your client may break if the
// server is upgraded.
// Deprecated: use NewClientWithOpts
func NewClient(host string, version string, client *http.Client, httpHeaders map[string]string) (*Client, error) {
	return NewClientWithOpts(WithHost(host), WithVersion(version), WithHTTPClient(client), WithHTTPHeaders(httpHeaders))
}

// Close the transport used by the client
func (cli *Client) Close() error {
	if t, ok := cli.client.Transport.(*http.Transport); ok {
		t.CloseIdleConnections()
	}
	return nil
}

// getAPIPath returns the versioned request path to call the api.
// It appends the query parameters to the path if they are not empty.
func (cli *Client) getAPIPath(p string, query url.Values) string {
	var apiPath string
	if cli.version != "" {
		v := strings.TrimPrefix(cli.version, "v")
		apiPath = path.Join(cli.basePath, "/v"+v, p)
	} else {
		apiPath = path.Join(cli.basePath, p)
	}
	return (&url.URL{Path: apiPath, RawQuery: query.Encode()}).String()
}

// ClientVersion returns the API version used by this client.
func (cli *Client) ClientVersion() string {
	return cli.version
}

// NegotiateAPIVersion queries the API and updates the version to match the
// API version. Any errors are silently ignored.
func (cli *Client) NegotiateAPIVersion(ctx context.Context) {
	ping, _ := cli.Ping(ctx)
	cli.NegotiateAPIVersionPing(ping)
}

// NegotiateAPIVersionPing updates the client version to match the Ping.APIVersion
// if the ping version is less than the default version.
func (cli *Client) NegotiateAPIVersionPing(p types.Ping) {
	if cli.manualOverride {
		return
	}

	// try the latest version before versioning headers existed
	if p.APIVersion == "" {
		p.APIVersion = "1.24"
	}

	// if the client is not initialized with a version, start with the latest supported version
	if cli.version == "" {
		cli.version = api.DefaultVersion
	}

	// if server version is lower than the client version, downgrade
	if versions.LessThan(p.APIVersion, cli.version) {
		cli.version = p.APIVersion
	}
}

// DaemonHost returns the host address used by the client
func (cli *Client) DaemonHost() string {
	return cli.host
}

// HTTPClient returns a copy of the HTTP client bound to the server
func (cli *Client) HTTPClient() *http.Client {
	return &*cli.client
}

// ParseHostURL parses a url string, validates the string is a host url, and
// returns the parsed URL
func ParseHostURL(host string) (*url.URL, error) {
	protoAddrParts := strings.SplitN(host, "://", 2)
	if len(protoAddrParts) == 1 {
		return nil, fmt.Errorf("unable to parse docker host `%s`", host)
	}

	var basePath string
	proto, addr := protoAddrParts[0], protoAddrParts[1]
	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return nil, err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}
	return &url.URL{
		Scheme: proto,
		Host:   addr,
		Path:   basePath,
	}, nil
}

// CustomHTTPHeaders returns the custom http headers stored by the client.
func (cli *Client) CustomHTTPHeaders() map[string]string {
	m := make(map[string]string)
	for k, v := range cli.customHTTPHeaders {
		m[k] = v
	}
	return m
}

// SetCustomHTTPHeaders that will be set on every HTTP request made by the client.
// Deprecated: use WithHTTPHeaders when creating the client.
func (cli *Client) SetCustomHTTPHeaders(headers map[string]string) {
	cli.customHTTPHeaders = headers
}

// Dialer returns a dialer for a raw stream connection, with HTTP/1.1 header, that can be used for proxying the daemon connection.
// Used by `docker dial-stdio` (docker/cli#889).
func (cli *Client) Dialer() func(context.Context) (net.Conn, error) {
	return func(ctx context.Context) (net.Conn, error) {
		if transport, ok := cli.client.Transport.(*http.Transport); ok {
			if transport.DialContext != nil {
				return transport.DialContext(ctx, cli.proto, cli.addr)
			}
		}
		return fallbackDial(cli.proto, cli.addr, resolveTLSConfig(cli.client.Transport))
	}
}
<commit_msg>client: dial tls on Dialer if tls config is set<commit_after>/*
Package client is a Go client for the Docker Engine API.

For more information about the Engine API, see the documentation:
https://docs.docker.com/engine/reference/api/

Usage

You use the library by creating a client object and calling methods on it. The
client can be created either from environment variables with NewEnvClient, or
configured manually with NewClient.

For example, to list running containers (the equivalent of "docker ps"):

	package main

	import (
		"context"
		"fmt"

		"github.com/docker/docker/api/types"
		"github.com/docker/docker/client"
	)

	func main() {
		cli, err := client.NewEnvClient()
		if err != nil {
			panic(err)
		}

		containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{})
		if err != nil {
			panic(err)
		}

		for _, container := range containers {
			fmt.Printf("%s %s\n", container.ID[:10], container.Image)
		}
	}

*/
package client // import "github.com/docker/docker/client"

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/docker/docker/api"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/versions"
	"github.com/docker/go-connections/sockets"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/pkg/errors"
)

// ErrRedirect is the error returned by checkRedirect when the request is non-GET.
var ErrRedirect = errors.New("unexpected redirect in response")

// Client is the API client that performs all operations
// against a docker server.
type Client struct {
	// scheme sets the scheme for the client
	scheme string
	// host holds the server address to connect to
	host string
	// proto holds the client protocol i.e. unix.
	proto string
	// addr holds the client address.
	addr string
	// basePath holds the path to prepend to the requests.
	basePath string
	// client used to send and receive http requests.
	client *http.Client
	// version of the server to talk to.
	version string
	// custom http headers configured by users.
	customHTTPHeaders map[string]string
	// manualOverride is set to true when the version was set by users.
	manualOverride bool
}

// CheckRedirect specifies the policy for dealing with redirect responses:
// If the request is non-GET return `ErrRedirect`. Otherwise use the last response.
//
// Go 1.8 changes behavior for HTTP redirects (specifically 301, 307, and 308) in the client .
// The Docker client (and by extension docker API client) can be made to to send a request
// like POST /containers//start where what would normally be in the name section of the URL is empty.
// This triggers an HTTP 301 from the daemon.
// In go 1.8 this 301 will be converted to a GET request, and ends up getting a 404 from the daemon.
// This behavior change manifests in the client in that before the 301 was not followed and
// the client did not generate an error, but now results in a message like Error response from daemon: page not found.
func CheckRedirect(req *http.Request, via []*http.Request) error {
	if via[0].Method == http.MethodGet {
		return http.ErrUseLastResponse
	}
	return ErrRedirect
}

// NewEnvClient initializes a new API client based on environment variables.
// See FromEnv for a list of support environment variables.
//
// Deprecated: use NewClientWithOpts(FromEnv)
func NewEnvClient() (*Client, error) {
	return NewClientWithOpts(FromEnv)
}

// FromEnv configures the client with values from environment variables.
//
// Supported environment variables:
// DOCKER_HOST to set the url to the docker server.
// DOCKER_API_VERSION to set the version of the API to reach, leave empty for latest.
// DOCKER_CERT_PATH to load the TLS certificates from.
// DOCKER_TLS_VERIFY to enable or disable TLS verification, off by default.
func FromEnv(c *Client) error {
	if dockerCertPath := os.Getenv("DOCKER_CERT_PATH"); dockerCertPath != "" {
		options := tlsconfig.Options{
			CAFile:             filepath.Join(dockerCertPath, "ca.pem"),
			CertFile:           filepath.Join(dockerCertPath, "cert.pem"),
			KeyFile:            filepath.Join(dockerCertPath, "key.pem"),
			InsecureSkipVerify: os.Getenv("DOCKER_TLS_VERIFY") == "",
		}
		tlsc, err := tlsconfig.Client(options)
		if err != nil {
			return err
		}

		c.client = &http.Client{
			Transport:     &http.Transport{TLSClientConfig: tlsc},
			CheckRedirect: CheckRedirect,
		}
	}

	if host := os.Getenv("DOCKER_HOST"); host != "" {
		if err := WithHost(host)(c); err != nil {
			return err
		}
	}

	if version := os.Getenv("DOCKER_API_VERSION"); version != "" {
		c.version = version
		c.manualOverride = true
	}
	return nil
}

// WithTLSClientConfig applies a tls config to the client transport.
func WithTLSClientConfig(cacertPath, certPath, keyPath string) func(*Client) error {
	return func(c *Client) error {
		opts := tlsconfig.Options{
			CAFile:             cacertPath,
			CertFile:           certPath,
			KeyFile:            keyPath,
			ExclusiveRootPools: true,
		}
		config, err := tlsconfig.Client(opts)
		if err != nil {
			return errors.Wrap(err, "failed to create tls config")
		}
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			transport.TLSClientConfig = config
			return nil
		}
		return errors.Errorf("cannot apply tls config to transport: %T", c.client.Transport)
	}
}

// WithDialer applies the dialer.DialContext to the client transport. This can be
// used to set the Timeout and KeepAlive settings of the client.
// Deprecated: use WithDialContext
func WithDialer(dialer *net.Dialer) func(*Client) error {
	return WithDialContext(dialer.DialContext)
}

// WithDialContext applies the dialer to the client transport. This can be
// used to set the Timeout and KeepAlive settings of the client.
func WithDialContext(dialContext func(ctx context.Context, network, addr string) (net.Conn, error)) func(*Client) error {
	return func(c *Client) error {
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			transport.DialContext = dialContext
			return nil
		}
		return errors.Errorf("cannot apply dialer to transport: %T", c.client.Transport)
	}
}

// WithVersion overrides the client version with the specified one
func WithVersion(version string) func(*Client) error {
	return func(c *Client) error {
		c.version = version
		return nil
	}
}

// WithHost overrides the client host with the specified one.
func WithHost(host string) func(*Client) error {
	return func(c *Client) error {
		hostURL, err := ParseHostURL(host)
		if err != nil {
			return err
		}
		c.host = host
		c.proto = hostURL.Scheme
		c.addr = hostURL.Host
		c.basePath = hostURL.Path
		if transport, ok := c.client.Transport.(*http.Transport); ok {
			return sockets.ConfigureTransport(transport, c.proto, c.addr)
		}
		return errors.Errorf("cannot apply host to transport: %T", c.client.Transport)
	}
}

// WithHTTPClient overrides the client http client with the specified one
func WithHTTPClient(client *http.Client) func(*Client) error {
	return func(c *Client) error {
		if client != nil {
			c.client = client
		}
		return nil
	}
}

// WithHTTPHeaders overrides the client default http headers
func WithHTTPHeaders(headers map[string]string) func(*Client) error {
	return func(c *Client) error {
		c.customHTTPHeaders = headers
		return nil
	}
}

// NewClientWithOpts initializes a new API client with default values. It takes functors
// to modify values when creating it, like `NewClientWithOpts(WithVersion(…))`
// It also initializes the custom http headers to add to each request.
//
// It won't send any version information if the version number is empty. It is
// highly recommended that you set a version or your client may break if the
// server is upgraded.
func NewClientWithOpts(ops ...func(*Client) error) (*Client, error) {
	client, err := defaultHTTPClient(DefaultDockerHost)
	if err != nil {
		return nil, err
	}
	c := &Client{
		host:    DefaultDockerHost,
		version: api.DefaultVersion,
		scheme:  "http",
		client:  client,
		proto:   defaultProto,
		addr:    defaultAddr,
	}

	for _, op := range ops {
		if err := op(c); err != nil {
			return nil, err
		}
	}

	if _, ok := c.client.Transport.(http.RoundTripper); !ok {
		return nil, fmt.Errorf("unable to verify TLS configuration, invalid transport %v", c.client.Transport)
	}
	tlsConfig := resolveTLSConfig(c.client.Transport)
	if tlsConfig != nil {
		// TODO(stevvooe): This isn't really the right way to write clients in Go.
		// `NewClient` should probably only take an `*http.Client` and work from there.
		// Unfortunately, the model of having a host-ish/url-thingy as the connection
		// string has us confusing protocol and transport layers. We continue doing
		// this to avoid breaking existing clients but this should be addressed.
		c.scheme = "https"
	}

	return c, nil
}

func defaultHTTPClient(host string) (*http.Client, error) {
	url, err := ParseHostURL(host)
	if err != nil {
		return nil, err
	}
	transport := new(http.Transport)
	sockets.ConfigureTransport(transport, url.Scheme, url.Host)
	return &http.Client{
		Transport:     transport,
		CheckRedirect: CheckRedirect,
	}, nil
}

// NewClient initializes a new API client for the given host and API version.
// It uses the given http client as transport.
// It also initializes the custom http headers to add to each request.
//
// It won't send any version information if the version number is empty. It is
// highly recommended that you set a version or your client may break if the
// server is upgraded.
// Deprecated: use NewClientWithOpts
func NewClient(host string, version string, client *http.Client, httpHeaders map[string]string) (*Client, error) {
	return NewClientWithOpts(WithHost(host), WithVersion(version), WithHTTPClient(client), WithHTTPHeaders(httpHeaders))
}

// Close the transport used by the client
func (cli *Client) Close() error {
	if t, ok := cli.client.Transport.(*http.Transport); ok {
		t.CloseIdleConnections()
	}
	return nil
}

// getAPIPath returns the versioned request path to call the api.
// It appends the query parameters to the path if they are not empty.
func (cli *Client) getAPIPath(p string, query url.Values) string {
	var apiPath string
	if cli.version != "" {
		v := strings.TrimPrefix(cli.version, "v")
		apiPath = path.Join(cli.basePath, "/v"+v, p)
	} else {
		apiPath = path.Join(cli.basePath, p)
	}
	return (&url.URL{Path: apiPath, RawQuery: query.Encode()}).String()
}

// ClientVersion returns the API version used by this client.
func (cli *Client) ClientVersion() string {
	return cli.version
}

// NegotiateAPIVersion queries the API and updates the version to match the
// API version. Any errors are silently ignored.
func (cli *Client) NegotiateAPIVersion(ctx context.Context) {
	ping, _ := cli.Ping(ctx)
	cli.NegotiateAPIVersionPing(ping)
}

// NegotiateAPIVersionPing updates the client version to match the Ping.APIVersion
// if the ping version is less than the default version.
func (cli *Client) NegotiateAPIVersionPing(p types.Ping) {
	if cli.manualOverride {
		return
	}

	// try the latest version before versioning headers existed
	if p.APIVersion == "" {
		p.APIVersion = "1.24"
	}

	// if the client is not initialized with a version, start with the latest supported version
	if cli.version == "" {
		cli.version = api.DefaultVersion
	}

	// if server version is lower than the client version, downgrade
	if versions.LessThan(p.APIVersion, cli.version) {
		cli.version = p.APIVersion
	}
}

// DaemonHost returns the host address used by the client
func (cli *Client) DaemonHost() string {
	return cli.host
}

// HTTPClient returns a copy of the HTTP client bound to the server
func (cli *Client) HTTPClient() *http.Client {
	return &*cli.client
}

// ParseHostURL parses a url string, validates the string is a host url, and
// returns the parsed URL
func ParseHostURL(host string) (*url.URL, error) {
	protoAddrParts := strings.SplitN(host, "://", 2)
	if len(protoAddrParts) == 1 {
		return nil, fmt.Errorf("unable to parse docker host `%s`", host)
	}

	var basePath string
	proto, addr := protoAddrParts[0], protoAddrParts[1]
	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return nil, err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}
	return &url.URL{
		Scheme: proto,
		Host:   addr,
		Path:   basePath,
	}, nil
}

// CustomHTTPHeaders returns the custom http headers stored by the client.
func (cli *Client) CustomHTTPHeaders() map[string]string {
	m := make(map[string]string)
	for k, v := range cli.customHTTPHeaders {
		m[k] = v
	}
	return m
}

// SetCustomHTTPHeaders that will be set on every HTTP request made by the client.
// Deprecated: use WithHTTPHeaders when creating the client.
func (cli *Client) SetCustomHTTPHeaders(headers map[string]string) {
	cli.customHTTPHeaders = headers
}

// Dialer returns a dialer for a raw stream connection, with HTTP/1.1 header, that can be used for proxying the daemon connection.
// Used by `docker dial-stdio` (docker/cli#889).
func (cli *Client) Dialer() func(context.Context) (net.Conn, error) {
	return func(ctx context.Context) (net.Conn, error) {
		if transport, ok := cli.client.Transport.(*http.Transport); ok {
			if transport.DialContext != nil && transport.TLSClientConfig == nil {
				return transport.DialContext(ctx, cli.proto, cli.addr)
			}
		}
		return fallbackDial(cli.proto, cli.addr, resolveTLSConfig(cli.client.Transport))
	}
}
<|endoftext|>


<commit_before>package hammy

import (
	"fmt"
	"log"
	"time"
	"strings"
)

type CmdBufferProcessorImpl struct {
	// Send buffer
	SBuffer SendBuffer
	// Data saver
	Saver SendBuffer
}

func (cbp *CmdBufferProcessorImpl) Process(key string, cmdb *CmdBuffer) error {
	for _, c := range *cmdb {
		switch c.Cmd {
			case "log":
				log.Printf("[%s] %s", key, c.Options["message"])
			case "send":
				cbp.processSend(key, c.Options)
			case "save":
				cbp.processSave(key, c.Options)
			default:
				log.Printf("[%s] %s %v", key, c.Cmd, c.Options)
		}
	}

	return nil
}

func (cbp *CmdBufferProcessorImpl) log(key string, message string) error {
	cmdb := make(CmdBuffer, 1)
	cmdb[0].Cmd = "log"
	cmdb[0].Options = make(map[string]interface{})
	cmdb[0].Options["message"] = message
	return cbp.Process(key, &cmdb)
}

func (cbp *CmdBufferProcessorImpl) processSend(key string, opts map[string]interface{}) {
	hostNameRaw := opts["host"]
	var hostName string
	if hostNameRaw == nil {
		hostName = key
	} else {
		var converted bool
		hostName, converted = hostNameRaw.(string)
		if !converted {
			cbp.log(key, fmt.Sprintf("Invalid send: invalid host name (command options: %v)", opts))
			return
		}
	}

	itemKey, itemKeyConverted := opts["key"].(string)
	if !itemKeyConverted  || itemKey == "" {
		cbp.log(key, fmt.Sprintf("Invalid send: key expected (command options: %v)", opts))
		return
	}

	value, valueFound := opts["value"]
	if !valueFound {
		cbp.log(key, fmt.Sprintf("Invalid send: value expected (command options: %v)", opts))
		return
	}

	data := make(IncomingData)
	hostData := make(IncomingHostData)
	hostValue := IncomingValueData{
		Timestamp: uint64(time.Now().Unix()),
		Value: value,
	}
	hostData[itemKey] = []IncomingValueData{hostValue}
	data[hostName] = hostData

	cbp.SBuffer.Push(&data)
}

func (cbp *CmdBufferProcessorImpl) processSave(key string, opts map[string]interface{}) {
	itemKey, itemKeyConverted := opts["key"].(string)
	if !itemKeyConverted || itemKey == "" {
		cbp.log(key, fmt.Sprintf("Invalid save: key expected (command options: %v)", opts))
		return
	}

	value, valueFound := opts["value"]
	if !valueFound {
		cbp.log(key, fmt.Sprintf("Invalid save: value expected (command options: %v)", opts))
		return
	}

	switch value.(type) {
		case int:
			value = float64(value.(int))
		case int8:
			value = float64(value.(int8))
		case int16:
			value = float64(value.(int16))
		case int32:
			value = float64(value.(int32))
		case int64:
			value = float64(value.(int64))
		case uint:
			value = float64(value.(uint))
		case uint8:
			value = float64(value.(uint8))
		case uint16:
			value = float64(value.(uint16))
		case uint32:
			value = float64(value.(uint32))
		case uint64:
			value = float64(value.(uint64))
		case float32:
			value = float64(value.(float32))
		case float64:
			// Do nothing
		case string:
			// Do nothing
		default:
			value = fmt.Sprint(value)
	}

	if strings.HasSuffix(itemKey, "#log") {
		if _, converted := value.(string); !converted {
			value = fmt.Sprint(value)
		}
	} else {
		if _, converted := value.(float64); !converted {
			var val float64
			str, strConverted := value.(string)
			if !strConverted {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid value for non log key `%s` (command options: %v)", itemKey, opts))
				return
			}
			if n, _ := fmt.Sscan(str, &val); n != 1 {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid value for non log key `%s` (command options: %v)", itemKey, opts))
				return
			}
			value = val
		}
	}

	var ts uint64
	tsRaw := opts["timestamp"]
	switch tsRaw.(type) {
		case nil:
			ts = uint64(time.Now().Unix())
		case string:
			_, err := fmt.Sscan(tsRaw.(string), &ts)
			if err != nil {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid timestamp (command options: %v)", opts))
				return
			}
		default:
			cbp.log(key, fmt.Sprintf("Invalid save: invalid timestamp (command options: %v)", opts))
			return
	}

	data := make(IncomingData)
	hostData := make(IncomingHostData)
	hostValue := IncomingValueData{
		Timestamp: ts,
		Value: value,
	}
	hostData[itemKey] = []IncomingValueData{hostValue}
	data[key] = hostData

	cbp.Saver.Push(&data)
}
<commit_msg>fix type conversions in cmd_buffer_processor<commit_after>package hammy

import (
	"fmt"
	"log"
	"time"
	"strings"
	"reflect"
)

type CmdBufferProcessorImpl struct {
	// Send buffer
	SBuffer SendBuffer
	// Data saver
	Saver SendBuffer
}

func (cbp *CmdBufferProcessorImpl) Process(key string, cmdb *CmdBuffer) error {
	for _, c := range *cmdb {
		switch c.Cmd {
			case "log":
				log.Printf("[%s] %s", key, c.Options["message"])
			case "send":
				cbp.processSend(key, c.Options)
			case "save":
				cbp.processSave(key, c.Options)
			default:
				log.Printf("[%s] %s %v", key, c.Cmd, c.Options)
		}
	}

	return nil
}

func (cbp *CmdBufferProcessorImpl) log(key string, message string) error {
	cmdb := make(CmdBuffer, 1)
	cmdb[0].Cmd = "log"
	cmdb[0].Options = make(map[string]interface{})
	cmdb[0].Options["message"] = message
	return cbp.Process(key, &cmdb)
}

func (cbp *CmdBufferProcessorImpl) processSend(key string, opts map[string]interface{}) {
	hostNameRaw := opts["host"]
	var hostName string
	if hostNameRaw == nil {
		hostName = key
	} else {
		var converted bool
		hostName, converted = hostNameRaw.(string)
		if !converted {
			cbp.log(key, fmt.Sprintf("Invalid send: invalid host name (command options: %v)", opts))
			return
		}
	}

	itemKey, itemKeyConverted := opts["key"].(string)
	if !itemKeyConverted  || itemKey == "" {
		cbp.log(key, fmt.Sprintf("Invalid send: key expected (command options: %v)", opts))
		return
	}

	value, valueFound := opts["value"]
	if !valueFound {
		cbp.log(key, fmt.Sprintf("Invalid send: value expected (command options: %v)", opts))
		return
	}

	data := make(IncomingData)
	hostData := make(IncomingHostData)
	hostValue := IncomingValueData{
		Timestamp: uint64(time.Now().Unix()),
		Value: value,
	}
	hostData[itemKey] = []IncomingValueData{hostValue}
	data[hostName] = hostData

	cbp.SBuffer.Push(&data)
}

func (cbp *CmdBufferProcessorImpl) processSave(key string, opts map[string]interface{}) {
	itemKey, itemKeyConverted := opts["key"].(string)
	if !itemKeyConverted || itemKey == "" {
		cbp.log(key, fmt.Sprintf("Invalid save: key expected (command options: %v)", opts))
		return
	}

	value, valueFound := opts["value"]
	if !valueFound {
		cbp.log(key, fmt.Sprintf("Invalid save: value expected (command options: %v)", opts))
		return
	}

	switch value.(type) {
		case int, int8, int16, int32, int64:
			value = float64(reflect.ValueOf(value).Int())
		case uint, uint8, uint16, uint32, uint64:
			value = float64(reflect.ValueOf(value).Uint())
		case float32:
			value = float64(value.(float32))
		case float64:
			// Do nothing
		case string:
			// Do nothing
		default:
			value = fmt.Sprint(value)
	}

	if strings.HasSuffix(itemKey, "#log") {
		if _, converted := value.(string); !converted {
			value = fmt.Sprint(value)
		}
	} else {
		if _, converted := value.(float64); !converted {
			var val float64
			str, strConverted := value.(string)
			if !strConverted {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid value for non log key `%s` (command options: %v)", itemKey, opts))
				return
			}
			if n, _ := fmt.Sscan(str, &val); n != 1 {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid value for non log key `%s` (command options: %v)", itemKey, opts))
				return
			}
			value = val
		}
	}

	var ts uint64
	tsRaw := opts["timestamp"]
	switch tsRaw.(type) {
		case nil:
			ts = uint64(time.Now().Unix())
		case uint, uint8, uint16, uint32, uint64:
			ts = reflect.ValueOf(tsRaw).Uint()
		case int, int8, int16, int32, int64:
			ts = uint64(reflect.ValueOf(tsRaw).Int())
		case float32, float64:
			ts = uint64(reflect.ValueOf(tsRaw).Float())
		case string:
			_, err := fmt.Sscan(tsRaw.(string), &ts)
			if err != nil {
				cbp.log(key, fmt.Sprintf("Invalid save: invalid timestamp (command options: %v)", opts))
				return
			}
		default:
			cbp.log(key, fmt.Sprintf("Invalid save: invalid timestamp of type %T (command options: %v)", tsRaw, opts))
			return
	}

	data := make(IncomingData)
	hostData := make(IncomingHostData)
	hostValue := IncomingValueData{
		Timestamp: ts,
		Value: value,
	}
	hostData[itemKey] = []IncomingValueData{hostValue}
	data[key] = hostData

	cbp.Saver.Push(&data)
}
<|endoftext|>


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

package manager

import (
	"fmt"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/heapster/sinks"
	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	"github.com/GoogleCloudPlatform/heapster/sinks/cache"
	source_api "github.com/GoogleCloudPlatform/heapster/sources/api"
	"github.com/golang/glog"
)

// Manager provides an interface to control the core of heapster.
// Implementations are not required to be thread safe.
type Manager interface {
	// Housekeep collects data from all the configured sources and
	// stores the data to all the configured sinks.
	Housekeep()

	// Export the latest data point of all metrics.
	ExportMetrics() ([]*sink_api.Point, error)
}

type realManager struct {
	sources     []source_api.Source
	cache       cache.Cache
	sinkManager sinks.ExternalSinkManager
	lastSync    time.Time
	resolution  time.Duration
	decoder     sink_api.Decoder
}

type syncData struct {
	data  source_api.AggregateData
	mutex sync.Mutex
}

func NewManager(sources []source_api.Source, sinkManager sinks.ExternalSinkManager, res, bufferDuration time.Duration) (Manager, error) {
	return &realManager{
		sources:     sources,
		sinkManager: sinkManager,
		cache:       cache.NewCache(bufferDuration),
		lastSync:    time.Now(),
		resolution:  res,
		decoder:     sink_api.NewDecoder(),
	}, nil
}

func (rm *realManager) scrapeSource(s source_api.Source, start, end time.Time, sd *syncData, errChan chan<- error) {
	glog.V(2).Infof("attempting to get data from source %q", s.Name())
	data, err := s.GetInfo(start, end, rm.resolution)
	if err != nil {
		err = fmt.Errorf("failed to get information from source %q - %v", s.Name(), err)
	} else {
		sd.mutex.Lock()
		defer sd.mutex.Unlock()
		sd.data.Merge(&data)
	}
	errChan <- err
}

func (rm *realManager) Housekeep() {
	errChan := make(chan error, len(rm.sources))
	var sd syncData
	start := rm.lastSync
	end := time.Now()
	rm.lastSync = start
	glog.V(2).Infof("starting to scrape data from sources")
	for idx := range rm.sources {
		s := rm.sources[idx]
		go rm.scrapeSource(s, start, end, &sd, errChan)
	}
	var errors []string
	for i := 0; i < len(rm.sources); i++ {
		if err := <-errChan; err != nil {
			errors = append(errors, err.Error())
		}
	}
	glog.V(2).Infof("completed scraping data from sources. Errors: %v", errors)
	if err := rm.cache.StorePods(sd.data.Pods); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.cache.StoreContainers(sd.data.Machine); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.cache.StoreContainers(sd.data.Containers); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.sinkManager.Store(sd.data); err != nil {
		errors = append(errors, err.Error())
	}
	if len(errors) > 0 {
		glog.V(1).Infof("housekeeping resulted in following errors: %v", errors)
	}
}

func (rm *realManager) ExportMetrics() ([]*sink_api.Point, error) {
	var zero time.Time

	// Get all pods as points.
	pods := trimStatsForPods(rm.cache.GetPods(zero, zero))
	timeseries, err := rm.decoder.TimeseriesFromPods(pods)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points := make([]*sink_api.Point, 0, len(timeseries))
	points = appendPoints(points, timeseries)

	// Get all nodes as points.
	containers := trimStatsForContainers(rm.cache.GetNodes(zero, zero))
	timeseries, err = rm.decoder.TimeseriesFromContainers(containers)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points = appendPoints(points, timeseries)

	// Get all free containers as points.
	containers = trimStatsForContainers(rm.cache.GetFreeContainers(zero, zero))
	timeseries, err = rm.decoder.TimeseriesFromContainers(containers)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points = appendPoints(points, timeseries)

	return points, nil
}

// Extract the points from the specified timeseries and append them to output.
func appendPoints(output []*sink_api.Point, toExtract []sink_api.Timeseries) []*sink_api.Point {
	for i := range toExtract {
		output = append(output, toExtract[i].Point)
	}
	return output
}

// Only keep latest stats for the specified pods
func trimStatsForPods(pods []*cache.PodElement) []*cache.PodElement {
	for _, pod := range pods {
		trimStatsForContainers(pod.Containers)
	}
	return pods
}

// Only keep latest stats for the specified containers
func trimStatsForContainers(containers []*cache.ContainerElement) []*cache.ContainerElement {
	for _, cont := range containers {
		onlyKeepLatestStat(cont)
	}
	return containers
}

// Only keep the latest stats data point.
func onlyKeepLatestStat(cont *cache.ContainerElement) {
	if len(cont.Metrics) > 1 {
		cont.Metrics = cont.Metrics[len(cont.Metrics)-1:]
	}
}
<commit_msg>Fixed time range for collecting stats for every pass<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package manager

import (
	"fmt"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/heapster/sinks"
	sink_api "github.com/GoogleCloudPlatform/heapster/sinks/api/v1"
	"github.com/GoogleCloudPlatform/heapster/sinks/cache"
	source_api "github.com/GoogleCloudPlatform/heapster/sources/api"
	"github.com/golang/glog"
)

// Manager provides an interface to control the core of heapster.
// Implementations are not required to be thread safe.
type Manager interface {
	// Housekeep collects data from all the configured sources and
	// stores the data to all the configured sinks.
	Housekeep()

	// Export the latest data point of all metrics.
	ExportMetrics() ([]*sink_api.Point, error)
}

type realManager struct {
	sources     []source_api.Source
	cache       cache.Cache
	sinkManager sinks.ExternalSinkManager
	lastSync    time.Time
	resolution  time.Duration
	decoder     sink_api.Decoder
}

type syncData struct {
	data  source_api.AggregateData
	mutex sync.Mutex
}

func NewManager(sources []source_api.Source, sinkManager sinks.ExternalSinkManager, res, bufferDuration time.Duration) (Manager, error) {
	return &realManager{
		sources:     sources,
		sinkManager: sinkManager,
		cache:       cache.NewCache(bufferDuration),
		lastSync:    time.Now(),
		resolution:  res,
		decoder:     sink_api.NewDecoder(),
	}, nil
}

func (rm *realManager) scrapeSource(s source_api.Source, start, end time.Time, sd *syncData, errChan chan<- error) {
	glog.V(2).Infof("attempting to get data from source %q", s.Name())
	data, err := s.GetInfo(start, end, rm.resolution)
	if err != nil {
		err = fmt.Errorf("failed to get information from source %q - %v", s.Name(), err)
	} else {
		sd.mutex.Lock()
		defer sd.mutex.Unlock()
		sd.data.Merge(&data)
	}
	errChan <- err
}

func (rm *realManager) Housekeep() {
	errChan := make(chan error, len(rm.sources))
	var sd syncData
	start := rm.lastSync
	end := time.Now()
	rm.lastSync = end
	glog.V(2).Infof("starting to scrape data from sources")
	for idx := range rm.sources {
		s := rm.sources[idx]
		go rm.scrapeSource(s, start, end, &sd, errChan)
	}
	var errors []string
	for i := 0; i < len(rm.sources); i++ {
		if err := <-errChan; err != nil {
			errors = append(errors, err.Error())
		}
	}
	glog.V(2).Infof("completed scraping data from sources. Errors: %v", errors)
	if err := rm.cache.StorePods(sd.data.Pods); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.cache.StoreContainers(sd.data.Machine); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.cache.StoreContainers(sd.data.Containers); err != nil {
		errors = append(errors, err.Error())
	}
	if err := rm.sinkManager.Store(sd.data); err != nil {
		errors = append(errors, err.Error())
	}
	if len(errors) > 0 {
		glog.V(1).Infof("housekeeping resulted in following errors: %v", errors)
	}
}

func (rm *realManager) ExportMetrics() ([]*sink_api.Point, error) {
	var zero time.Time

	// Get all pods as points.
	pods := trimStatsForPods(rm.cache.GetPods(zero, zero))
	timeseries, err := rm.decoder.TimeseriesFromPods(pods)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points := make([]*sink_api.Point, 0, len(timeseries))
	points = appendPoints(points, timeseries)

	// Get all nodes as points.
	containers := trimStatsForContainers(rm.cache.GetNodes(zero, zero))
	timeseries, err = rm.decoder.TimeseriesFromContainers(containers)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points = appendPoints(points, timeseries)

	// Get all free containers as points.
	containers = trimStatsForContainers(rm.cache.GetFreeContainers(zero, zero))
	timeseries, err = rm.decoder.TimeseriesFromContainers(containers)
	if err != nil {
		return []*sink_api.Point{}, err
	}
	points = appendPoints(points, timeseries)

	return points, nil
}

// Extract the points from the specified timeseries and append them to output.
func appendPoints(output []*sink_api.Point, toExtract []sink_api.Timeseries) []*sink_api.Point {
	for i := range toExtract {
		output = append(output, toExtract[i].Point)
	}
	return output
}

// Only keep latest stats for the specified pods
func trimStatsForPods(pods []*cache.PodElement) []*cache.PodElement {
	for _, pod := range pods {
		trimStatsForContainers(pod.Containers)
	}
	return pods
}

// Only keep latest stats for the specified containers
func trimStatsForContainers(containers []*cache.ContainerElement) []*cache.ContainerElement {
	for _, cont := range containers {
		onlyKeepLatestStat(cont)
	}
	return containers
}

// Only keep the latest stats data point.
func onlyKeepLatestStat(cont *cache.ContainerElement) {
	if len(cont.Metrics) > 1 {
		cont.Metrics = cont.Metrics[len(cont.Metrics)-1:]
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The go-python 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 funcs

var (
	F1 func()
	F2 Func
	F3 S1
	F4 S2
	F5 []func()
	F6 []Func
	F7 [2]func()
	F8 [2]Func
)

type Func func()

type S1 struct {
	F1 Func
	F2 []Func
	F3 [4]Func
}

type S2 struct {
	F1 func()
	F2 []func()
	F3 [4]func()
}
<commit_msg>test: create different types of arrays of funcs<commit_after>// Copyright 2015 The go-python 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 funcs

var (
	F1 func()
	F2 Func
	F3 S1
	F4 S2
	F5 []func()
	F6 []Func
	F7 [2]func()
	F8 [3]Func
)

type Func func()

type S1 struct {
	F1 Func
	F2 []Func
	F3 [4]Func
}

type S2 struct {
	F1 func()
	F2 []func()
	F3 [5]func()
}
<|endoftext|>


<commit_before>package marshal

import (
	"errors"
	"reflect"
	"strings"

	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/parquet"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/types"
)

type Node struct {
	Val     reflect.Value
	PathMap *schema.PathMapType
	RL      int32
	DL      int32
}

//Improve Performance///////////////////////////
//NodeBuf
type NodeBufType struct {
	Index int
	Buf   []*Node
}

func NewNodeBuf(ln int) *NodeBufType {
	nodeBuf := new(NodeBufType)
	nodeBuf.Index = 0
	nodeBuf.Buf = make([]*Node, ln)
	for i := 0; i < ln; i++ {
		nodeBuf.Buf[i] = new(Node)
	}
	return nodeBuf
}

func (self *NodeBufType) GetNode() *Node {
	if self.Index >= len(self.Buf) {
		self.Buf = append(self.Buf, new(Node))
	}
	self.Index++
	return self.Buf[self.Index-1]
}

func (self *NodeBufType) Reset() {
	self.Index = 0
}

////////for improve performance///////////////////////////////////
type Marshaler interface {
	Marshal(node *Node, nodeBuf *NodeBufType) []*Node
}

type ParquetPtr struct{}

func (p *ParquetPtr) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	if node.Val.IsNil() {
		return nodes
	}
	node.Val = node.Val.Elem()
	node.DL++
	nodes = append(nodes, node)
	return nodes
}

type ParquetStruct struct{}

func (p *ParquetStruct) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	var ok bool

	numField := node.Val.Type().NumField()
	nodes := make([]*Node, 0, numField)
	for j := 0; j < numField; j++ {
		tf := node.Val.Type().Field(j)
		name := tf.Name
		newNode := nodeBuf.GetNode()

		//some ignored item
		if newNode.PathMap, ok = node.PathMap.Children[name]; !ok {
			continue
		}

		newNode.Val = node.Val.Field(j)
		newNode.RL = node.RL
		newNode.DL = node.DL
		nodes = append(nodes, newNode)
	}
	return nodes
}

type ParquetMapStruct struct{}

func (p *ParquetMapStruct) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	var ok bool

	nodes := make([]*Node, 0)
	keys := node.Val.MapKeys()
	if len(keys) <= 0 {
		return nodes
	}

	for j := len(keys) - 1; j >= 0; j-- {
		key := keys[j]
		newNode := nodeBuf.GetNode()

		//some ignored item
		if newNode.PathMap, ok = node.PathMap.Children[key.String()]; !ok {
			continue
		}

		newNode.Val = node.Val.MapIndex(key)
		newNode.RL = node.RL
		newNode.DL = node.DL
		nodes = append(nodes, newNode)
	}
	return nodes
}

type ParquetSlice struct {
	schemaHandler *schema.SchemaHandler
}

func (p *ParquetSlice) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	ln := node.Val.Len()
	pathMap := node.PathMap
	path := node.PathMap.Path
	if *p.schemaHandler.SchemaElements[p.schemaHandler.MapIndex[node.PathMap.Path]].RepetitionType != parquet.FieldRepetitionType_REPEATED {
		pathMap = pathMap.Children["List"].Children["Element"]
		path += ".List" + ".Element"
	}
	if ln <= 0 {
		return nodes
	}

	rlNow, _ := p.schemaHandler.MaxRepetitionLevel(common.StrToPath(path))
	for j := ln - 1; j >= 0; j-- {
		newNode := nodeBuf.GetNode()
		newNode.PathMap = pathMap
		newNode.Val = node.Val.Index(j)
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		newNode.DL = node.DL + 1
		nodes = append(nodes, newNode)
	}
	return nodes
}

type ParquetMap struct {
	schemaHandler *schema.SchemaHandler
}

func (p *ParquetMap) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	path := node.PathMap.Path + ".Key_value"
	keys := node.Val.MapKeys()
	if len(keys) <= 0 {
		return nodes
	}

	rlNow, _ := p.schemaHandler.MaxRepetitionLevel(common.StrToPath(path))
	for j := len(keys) - 1; j >= 0; j-- {
		key := keys[j]
		value := node.Val.MapIndex(key)
		newNode := nodeBuf.GetNode()
		newNode.PathMap = node.PathMap.Children["Key_value"].Children["Key"]
		newNode.Val = key
		newNode.DL = node.DL + 1
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		nodes = append(nodes, newNode)

		newNode = nodeBuf.GetNode()
		newNode.PathMap = node.PathMap.Children["Key_value"].Children["Value"]
		newNode.Val = value
		newNode.DL = node.DL + 1
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		nodes = append(nodes, newNode)
	}
	return nodes
}

//Convert the objects to table map. srcInterface is a slice of objects
func Marshal(srcInterface []interface{}, schemaHandler *schema.SchemaHandler) (tb *map[string]*layout.Table, err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("unkown error")
			}
		}
	}()

	src := reflect.ValueOf(srcInterface)
	res := make(map[string]*layout.Table)
	pathMap := schemaHandler.PathMap
	nodeBuf := NewNodeBuf(1)

	for i := 0; i < len(schemaHandler.SchemaElements); i++ {
		schema := schemaHandler.SchemaElements[i]
		pathStr := schemaHandler.IndexMap[int32(i)]
		numChildren := schema.GetNumChildren()
		if numChildren == 0 {
			res[pathStr] = layout.NewEmptyTable()
			res[pathStr].Path = common.StrToPath(pathStr)
			res[pathStr].MaxDefinitionLevel, _ = schemaHandler.MaxDefinitionLevel(res[pathStr].Path)
			res[pathStr].MaxRepetitionLevel, _ = schemaHandler.MaxRepetitionLevel(res[pathStr].Path)
			res[pathStr].RepetitionType = schema.GetRepetitionType()
			res[pathStr].Schema = schemaHandler.SchemaElements[schemaHandler.MapIndex[pathStr]]
			res[pathStr].Info = schemaHandler.Infos[i]
		}
	}

	stack := make([]*Node, 0, 100)
	for i := 0; i < len(srcInterface); i++ {
		stack = stack[:0]
		nodeBuf.Reset()

		node := nodeBuf.GetNode()
		node.Val = src.Index(i)
		if src.Index(i).Type().Kind() == reflect.Interface {
			node.Val = src.Index(i).Elem()
		}
		node.PathMap = pathMap
		stack = append(stack, node)

		for len(stack) > 0 {
			ln := len(stack)
			node := stack[ln-1]
			stack = stack[:ln-1]

			tk := node.Val.Type().Kind()
			var m Marshaler

			schemaIndex := schemaHandler.MapIndex[node.PathMap.Path]
			if tk == reflect.Interface && schemaHandler.SchemaElements[schemaIndex].GetNumChildren() > 0 {
				node.Val = node.Val.Elem()
				tk = node.Val.Type().Kind()
			}

			if tk == reflect.Ptr {
				m = &ParquetPtr{}
			} else if tk == reflect.Struct {
				m = &ParquetStruct{}
			} else if tk == reflect.Slice {
				m = &ParquetSlice{schemaHandler: schemaHandler}
			} else if tk == reflect.Map {
				schemaIndex := schemaHandler.MapIndex[node.PathMap.Path]
				sele := schemaHandler.SchemaElements[schemaIndex]
				if !sele.IsSetConvertedType() {
					m = &ParquetMapStruct{}
				} else {
					m = &ParquetMap{schemaHandler: schemaHandler}
				}
			} else {
				table := res[node.PathMap.Path]
				schema := schemaHandler.SchemaElements[schemaIndex]
				table.Values = append(table.Values, types.InterfaceToParquetType(node.Val.Interface(), schema.Type))
				table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
				table.RepetitionLevels = append(table.RepetitionLevels, node.RL)
				continue
			}

			nodes := m.Marshal(node, nodeBuf)
			if len(nodes) == 0 {
				path := node.PathMap.Path
				index := schemaHandler.MapIndex[path]
				numChildren := schemaHandler.SchemaElements[index].GetNumChildren()
				if numChildren > int32(0) {
					for key, table := range res {
						if strings.HasPrefix(key, path) &&
							(len(key) == len(path) || key[len(path)] == '.') {
							table.Values = append(table.Values, nil)
							table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
							table.RepetitionLevels = append(table.RepetitionLevels, node.RL)
						}
					}
				} else {
					table := res[path]
					table.Values = append(table.Values, nil)
					table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
					table.RepetitionLevels = append(table.RepetitionLevels, node.RL)

				}
			} else {
				for _, node := range nodes {
					stack = append(stack, node)
				}
			}
		}
	}

	return &res, nil
}
<commit_msg>Ability to use dynamically defined map[string]nterface{} instead of struct and define OPTIONAL columns.<commit_after>package marshal

import (
	"errors"
	"reflect"
	"strings"

	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/parquet"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/types"
)

type Node struct {
	Val     reflect.Value
	PathMap *schema.PathMapType
	RL      int32
	DL      int32
}

//Improve Performance///////////////////////////
//NodeBuf
type NodeBufType struct {
	Index int
	Buf   []*Node
}

func NewNodeBuf(ln int) *NodeBufType {
	nodeBuf := new(NodeBufType)
	nodeBuf.Index = 0
	nodeBuf.Buf = make([]*Node, ln)
	for i := 0; i < ln; i++ {
		nodeBuf.Buf[i] = new(Node)
	}
	return nodeBuf
}

func (self *NodeBufType) GetNode() *Node {
	if self.Index >= len(self.Buf) {
		self.Buf = append(self.Buf, new(Node))
	}
	self.Index++
	return self.Buf[self.Index-1]
}

func (self *NodeBufType) Reset() {
	self.Index = 0
}

////////for improve performance///////////////////////////////////
type Marshaler interface {
	Marshal(node *Node, nodeBuf *NodeBufType) []*Node
}

type ParquetPtr struct{}

func (p *ParquetPtr) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	if node.Val.IsNil() {
		return nodes
	}
	node.Val = node.Val.Elem()
	node.DL++
	nodes = append(nodes, node)
	return nodes
}

type ParquetStruct struct{}

func (p *ParquetStruct) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	var ok bool

	numField := node.Val.Type().NumField()
	nodes := make([]*Node, 0, numField)
	for j := 0; j < numField; j++ {
		tf := node.Val.Type().Field(j)
		name := tf.Name
		newNode := nodeBuf.GetNode()

		//some ignored item
		if newNode.PathMap, ok = node.PathMap.Children[name]; !ok {
			continue
		}

		newNode.Val = node.Val.Field(j)
		newNode.RL = node.RL
		newNode.DL = node.DL
		nodes = append(nodes, newNode)
	}
	return nodes
}

type ParquetMapStruct struct {
	schemaHandler *schema.SchemaHandler
}

func (p *ParquetMapStruct) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	var ok bool

	nodes := make([]*Node, 0, len(node.PathMap.Children))
	keys := node.Val.MapKeys()
	if len(keys) <= 0 {
		return nodes
	}

	missingKeys := make(map[string]bool)
	for k := range node.PathMap.Children {
		missingKeys[k] = true
	}
	for j := len(keys) - 1; j >= 0; j-- {
		key := keys[j]
		newNode := nodeBuf.GetNode()

		//some ignored item
		k := key.String()
		if newNode.PathMap, ok = node.PathMap.Children[k]; !ok {
			continue
		}
		missingKeys[k] = false
		v := node.Val.MapIndex(key)
		newNode.RL = node.RL
		newNode.DL = node.DL
		if v.Type().Kind() == reflect.Interface {
			newNode.Val = v.Elem()
			if newNode.Val.IsValid() {
				if *p.schemaHandler.SchemaElements[p.schemaHandler.MapIndex[newNode.PathMap.Path]].RepetitionType != parquet.FieldRepetitionType_REQUIRED {
					newNode.DL++
				}
			}
		} else {
			newNode.Val = v
		}
		nodes = append(nodes, newNode)
	}

	var null interface{}
	for k, isMissing := range missingKeys {
		if isMissing {
			newNode := nodeBuf.GetNode()
			newNode.PathMap = node.PathMap.Children[k]
			newNode.Val = reflect.ValueOf(null)
			newNode.RL = node.RL
			newNode.DL = node.DL
			nodes = append(nodes, newNode)
		}
	}
	return nodes
}

type ParquetSlice struct {
	schemaHandler *schema.SchemaHandler
}

func (p *ParquetSlice) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	ln := node.Val.Len()
	pathMap := node.PathMap
	path := node.PathMap.Path
	if *p.schemaHandler.SchemaElements[p.schemaHandler.MapIndex[node.PathMap.Path]].RepetitionType != parquet.FieldRepetitionType_REPEATED {
		pathMap = pathMap.Children["List"].Children["Element"]
		path += ".List" + ".Element"
	}
	if ln <= 0 {
		return nodes
	}

	rlNow, _ := p.schemaHandler.MaxRepetitionLevel(common.StrToPath(path))
	for j := ln - 1; j >= 0; j-- {
		newNode := nodeBuf.GetNode()
		newNode.PathMap = pathMap
		v := node.Val.Index(j)
		if v.Type().Kind() == reflect.Interface {
			newNode.Val = v.Elem()
		} else {
			newNode.Val = v
		}
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		newNode.DL = node.DL + 1
		nodes = append(nodes, newNode)
	}
	return nodes
}

type ParquetMap struct {
	schemaHandler *schema.SchemaHandler
}

func (p *ParquetMap) Marshal(node *Node, nodeBuf *NodeBufType) []*Node {
	nodes := make([]*Node, 0)
	path := node.PathMap.Path + ".Key_value"
	keys := node.Val.MapKeys()
	if len(keys) <= 0 {
		return nodes
	}

	rlNow, _ := p.schemaHandler.MaxRepetitionLevel(common.StrToPath(path))
	for j := len(keys) - 1; j >= 0; j-- {
		key := keys[j]
		value := node.Val.MapIndex(key)
		newNode := nodeBuf.GetNode()
		newNode.PathMap = node.PathMap.Children["Key_value"].Children["Key"]
		newNode.Val = key
		newNode.DL = node.DL + 1
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		nodes = append(nodes, newNode)

		newNode = nodeBuf.GetNode()
		newNode.PathMap = node.PathMap.Children["Key_value"].Children["Value"]
		newNode.Val = value
		newNode.DL = node.DL + 1
		if j == 0 {
			newNode.RL = node.RL
		} else {
			newNode.RL = rlNow
		}
		nodes = append(nodes, newNode)
	}
	return nodes
}

//Convert the objects to table map. srcInterface is a slice of objects
func Marshal(srcInterface []interface{}, schemaHandler *schema.SchemaHandler) (tb *map[string]*layout.Table, err error) {
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("unkown error")
			}
		}
	}()

	src := reflect.ValueOf(srcInterface)
	res := make(map[string]*layout.Table)
	pathMap := schemaHandler.PathMap
	nodeBuf := NewNodeBuf(1)

	for i := 0; i < len(schemaHandler.SchemaElements); i++ {
		schema := schemaHandler.SchemaElements[i]
		pathStr := schemaHandler.IndexMap[int32(i)]
		numChildren := schema.GetNumChildren()
		if numChildren == 0 {
			res[pathStr] = layout.NewEmptyTable()
			res[pathStr].Path = common.StrToPath(pathStr)
			res[pathStr].MaxDefinitionLevel, _ = schemaHandler.MaxDefinitionLevel(res[pathStr].Path)
			res[pathStr].MaxRepetitionLevel, _ = schemaHandler.MaxRepetitionLevel(res[pathStr].Path)
			res[pathStr].RepetitionType = schema.GetRepetitionType()
			res[pathStr].Schema = schemaHandler.SchemaElements[schemaHandler.MapIndex[pathStr]]
			res[pathStr].Info = schemaHandler.Infos[i]
		}
	}

	stack := make([]*Node, 0, 100)
	for i := 0; i < len(srcInterface); i++ {
		stack = stack[:0]
		nodeBuf.Reset()

		node := nodeBuf.GetNode()
		node.Val = src.Index(i)
		if src.Index(i).Type().Kind() == reflect.Interface {
			node.Val = src.Index(i).Elem()
		}
		node.PathMap = pathMap
		stack = append(stack, node)

		for len(stack) > 0 {
			ln := len(stack)
			node := stack[ln-1]
			stack = stack[:ln-1]

			tk := reflect.Interface
			if node.Val.IsValid() {
				tk = node.Val.Type().Kind()
			}
			var m Marshaler

			schemaIndex := schemaHandler.MapIndex[node.PathMap.Path]

			if tk == reflect.Ptr {
				m = &ParquetPtr{}
			} else if tk == reflect.Struct {
				m = &ParquetStruct{}
			} else if tk == reflect.Slice {
				m = &ParquetSlice{schemaHandler: schemaHandler}
			} else if tk == reflect.Map {
				schemaIndex := schemaHandler.MapIndex[node.PathMap.Path]
				sele := schemaHandler.SchemaElements[schemaIndex]
				if !sele.IsSetConvertedType() {
					m = &ParquetMapStruct{schemaHandler: schemaHandler}
				} else {
					m = &ParquetMap{schemaHandler: schemaHandler}
				}
			} else {
				table := res[node.PathMap.Path]
				schema := schemaHandler.SchemaElements[schemaIndex]
				var v interface{}
				if node.Val.IsValid() {
					v = node.Val.Interface()
				}
				table.Values = append(table.Values, types.InterfaceToParquetType(v, schema.Type))
				table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
				table.RepetitionLevels = append(table.RepetitionLevels, node.RL)
				continue
			}

			nodes := m.Marshal(node, nodeBuf)
			if len(nodes) == 0 {
				path := node.PathMap.Path
				index := schemaHandler.MapIndex[path]
				numChildren := schemaHandler.SchemaElements[index].GetNumChildren()
				if numChildren > int32(0) {
					for key, table := range res {
						if strings.HasPrefix(key, path) &&
							(len(key) == len(path) || key[len(path)] == '.') {
							table.Values = append(table.Values, nil)
							table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
							table.RepetitionLevels = append(table.RepetitionLevels, node.RL)
						}
					}
				} else {
					table := res[path]
					table.Values = append(table.Values, nil)
					table.DefinitionLevels = append(table.DefinitionLevels, node.DL)
					table.RepetitionLevels = append(table.RepetitionLevels, node.RL)

				}
			} else {
				for _, node := range nodes {
					stack = append(stack, node)
				}
			}
		}
	}

	return &res, nil
}
<|endoftext|>