<commit_before>// Copyright 2020 The Chromium 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 main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_driver/go/td"
)

func main() {
	var (
		projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
		taskId    = flag.String("task_id", "", "ID of this task.")
		bot       = flag.String("bot", "", "Name of the task.")
		output    = flag.String("o", "", "Dump JSON step data to the given file, or stdout if -.")
		local     = flag.Bool("local", true, "Running locally (else on the bots)?")

		resources = flag.String("resources", "resources", "Passed to fm -i.")
		script    = flag.String("script", "", "File (or - for stdin) with one job per line.")
	)
	ctx := td.StartRun(projectId, taskId, bot, output, local)
	defer td.EndRun(ctx)

	actualStdout := os.Stdout
	actualStderr := os.Stderr
	verbosity := exec.Info
	if *local {
		// Task Driver echoes every exec.Run() stdout and stderr to the console,
		// which makes it hard to find failures (especially stdout).  Send them to /dev/null.
		devnull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
		if err != nil {
			td.Fatal(ctx, err)
		}
		os.Stdout = devnull
		os.Stderr = devnull
		// Having stifled stderr/stdout, changing Command.Verbose won't have any visible effect,
		// but setting it to Silent will bypass a fair chunk of wasted formatting work.
		verbosity = exec.Silent
	}

	if flag.NArg() < 1 {
		td.Fatalf(ctx, "Please pass an fm binary.")
	}
	fm := flag.Arg(0)

	// Run `fm <flag>` to find the names of all linked GMs or tests.
	query := func(flag string) []string {
		stdout := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout, Verbose: verbosity}
		cmd.Args = append(cmd.Args, "-i", *resources)
		cmd.Args = append(cmd.Args, flag)
		if err := exec.Run(ctx, cmd); err != nil {
			td.Fatal(ctx, err)
		}

		lines := []string{}
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
		return lines
	}
	gms := query("--listGMs")
	tests := query("--listTests")

	// Query Gold for all known hashes when running as a bot.
	known := map[string]bool{
		"0832f708a97acc6da385446384647a8f": true, // MD5 of passing unit test.
	}
	if *bot != "" {
		func() {
			url := "https://storage.googleapis.com/skia-infra-gm/hash_files/gold-prod-hashes.txt"
			resp, err := http.Get(url)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer resp.Body.Close()

			scanner := bufio.NewScanner(resp.Body)
			for scanner.Scan() {
				known[scanner.Text()] = true
			}
			if err := scanner.Err(); err != nil {
				td.Fatal(ctx, err)
			}

			fmt.Fprintf(actualStdout, "Gold knew %v unique hashes.\n", len(known))
		}()
	}

	type Work struct {
		Sources []string // Passed to FM -s: names of gms/tests, paths to image files, .skps, etc.
		Flags   []string // Other flags to pass to FM: --ct 565, --msaa 16, etc.
	}

	queue := make(chan Work, 1<<20) // Arbitrarily huge buffer to avoid ever blocking.
	wg := &sync.WaitGroup{}
	var failures int32 = 0

	worker := func(sources, flags []string) {
		defer wg.Done()

		stdout := &bytes.Buffer{}
		stderr := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout, Stderr: stderr, Verbose: verbosity}
		cmd.Args = append(cmd.Args, "-i", *resources, "-s")
		cmd.Args = append(cmd.Args, sources...)
		cmd.Args = append(cmd.Args, flags...)

		// Run our FM command.
		err := exec.Run(ctx, cmd)

		// On success, scan stdout for any unknown hashes.
		unknownHash := func() string {
			if err == nil && *bot != "" { // We only fetch known hashes when using -bot.
				scanner := bufio.NewScanner(stdout)
				for scanner.Scan() {
					if parts := strings.Fields(scanner.Text()); len(parts) == 3 {
						md5 := parts[1]
						if !known[md5] {
							return md5
						}
					}
				}
				if err := scanner.Err(); err != nil {
					td.Fatal(ctx, err)
				}
			}
			return ""
		}()

		// If a batch failed or produced an unknown hash, isolate with individual reruns.
		if len(sources) > 1 && (err != nil || unknownHash != "") {
			wg.Add(len(sources))
			for i := range sources {
				queue <- Work{sources[i : i+1], flags}
			}
			return
		}

		// If an individual run failed, nothing more to do but fail.
		if err != nil {
			atomic.AddInt32(&failures, 1)
			td.FailStep(ctx, err)
			if *local {
				lines := []string{}
				scanner := bufio.NewScanner(stderr)
				for scanner.Scan() {
					lines = append(lines, scanner.Text())
				}
				if err := scanner.Err(); err != nil {
					td.Fatal(ctx, err)
				}
				fmt.Fprintf(actualStderr, "%v %v #failed:\n\t%v\n",
					cmd.Name,
					strings.Join(cmd.Args, " "),
					strings.Join(lines, "\n\t"))
			}
			return
		}

		// If an individual run succeeded but produced an unknown hash, TODO upload .png to Gold.
		// For now just print out the command and the hash it produced.
		if unknownHash != "" {
			fmt.Fprintf(actualStdout, "%v %v #%v\n",
				cmd.Name,
				strings.Join(cmd.Args, " "),
				unknownHash)
		}
	}

	for i := 0; i < runtime.NumCPU(); i++ {
		go func() {
			for w := range queue {
				worker(w.Sources, w.Flags)
			}
		}()
	}

	// Get some work going, first breaking it into batches to increase our parallelism.
	kickoff := func(sources, flags []string) {
		if len(sources) == 0 {
			return // A blank or commented job line from -script or the command line.
		}

		// Shuffle the sources randomly as a cheap way to approximate evenly expensive batches.
		// (Intentionally not rand.Seed()'d to stay deterministically reproducible.)
		rand.Shuffle(len(sources), func(i, j int) {
			sources[i], sources[j] = sources[j], sources[i]
		})

		nbatches := runtime.NumCPU()                      // Arbitrary, nice to scale ~= cores.
		batch := (len(sources) + nbatches - 1) / nbatches // Round up to avoid empty batches.
		util.ChunkIter(len(sources), batch, func(start, end int) error {
			wg.Add(1)
			queue <- Work{sources[start:end], flags}
			return nil
		})
	}

	// Parse a job like "gms b=cpu ct=8888" into sources and flags for kickoff().
	parse := func(job []string) (sources, flags []string) {
		for _, token := range job {
			// Everything after # is a comment.
			if strings.HasPrefix(token, "#") {
				break
			}

			// Treat "gm" or "gms" as a shortcut for all known GMs.
			if token == "gm" || token == "gms" {
				sources = append(sources, gms...)
				continue
			}
			// Same for tests.
			if token == "test" || token == "tests" {
				sources = append(sources, tests...)
				continue
			}

			// Is this a flag to pass through to FM?
			if parts := strings.Split(token, "="); len(parts) == 2 {
				f := "-"
				if len(parts[0]) > 1 {
					f += "-"
				}
				f += parts[0]

				flags = append(flags, f, parts[1])
				continue
			}

			// Anything else must be the name of a source for FM to run.
			sources = append(sources, token)
		}
		return
	}

	// Parse one job from the command line, handy for ad hoc local runs.
	kickoff(parse(flag.Args()[1:]))

	// Any number of jobs can come from -script.
	if *script != "" {
		file := os.Stdin
		if *script != "-" {
			file, err := os.Open(*script)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer file.Close()
		}
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			kickoff(parse(strings.Fields(scanner.Text())))
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
	}

	// If we're a bot (or acting as if we are one), kick off its work.
	if *bot != "" {
		parts := strings.Split(*bot, "-")
		OS := parts[1]

		// For no reason but as a demo, skip GM aarectmodes and test GoodHash.
		filter := func(in []string, test func(string) bool) (out []string) {
			for _, s := range in {
				if test(s) {
					out = append(out, s)
				}
			}
			return
		}
		if OS == "Debian10" {
			gms = filter(gms, func(s string) bool { return s != "aarectmodes" })
			tests = filter(tests, func(s string) bool { return s != "GoodHash" })
		}

		kickoff(tests, strings.Fields("-b cpu"))
		kickoff(gms, strings.Fields("-b cpu"))
		kickoff(gms, strings.Fields("-b cpu --skvm"))
	}

	wg.Wait()
	if failures > 0 {
		if *local {
			// td.Fatalf() would work fine, but barfs up a panic that we don't need to see.
			fmt.Fprintf(actualStderr, "%v runs of %v failed after retries.\n", failures, fm)
			os.Exit(1)
		} else {
			td.Fatalf(ctx, "%v runs of %v failed after retries.", failures, fm)
		}
	}
}
<commit_msg>test more with FM<commit_after>// Copyright 2020 The Chromium 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 main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_driver/go/td"
)

func main() {
	var (
		projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
		taskId    = flag.String("task_id", "", "ID of this task.")
		bot       = flag.String("bot", "", "Name of the task.")
		output    = flag.String("o", "", "Dump JSON step data to the given file, or stdout if -.")
		local     = flag.Bool("local", true, "Running locally (else on the bots)?")

		resources = flag.String("resources", "resources", "Passed to fm -i.")
		script    = flag.String("script", "", "File (or - for stdin) with one job per line.")
	)
	ctx := td.StartRun(projectId, taskId, bot, output, local)
	defer td.EndRun(ctx)

	actualStdout := os.Stdout
	actualStderr := os.Stderr
	verbosity := exec.Info
	if *local {
		// Task Driver echoes every exec.Run() stdout and stderr to the console,
		// which makes it hard to find failures (especially stdout).  Send them to /dev/null.
		devnull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
		if err != nil {
			td.Fatal(ctx, err)
		}
		os.Stdout = devnull
		os.Stderr = devnull
		// Having stifled stderr/stdout, changing Command.Verbose won't have any visible effect,
		// but setting it to Silent will bypass a fair chunk of wasted formatting work.
		verbosity = exec.Silent
	}

	if flag.NArg() < 1 {
		td.Fatalf(ctx, "Please pass an fm binary.")
	}
	fm := flag.Arg(0)

	// Run `fm <flag>` to find the names of all linked GMs or tests.
	query := func(flag string) []string {
		stdout := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout, Verbose: verbosity}
		cmd.Args = append(cmd.Args, "-i", *resources)
		cmd.Args = append(cmd.Args, flag)
		if err := exec.Run(ctx, cmd); err != nil {
			td.Fatal(ctx, err)
		}

		lines := []string{}
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
		return lines
	}
	gms := query("--listGMs")
	tests := query("--listTests")

	// Query Gold for all known hashes when running as a bot.
	known := map[string]bool{
		"0832f708a97acc6da385446384647a8f": true, // MD5 of passing unit test.
	}
	if *bot != "" {
		func() {
			url := "https://storage.googleapis.com/skia-infra-gm/hash_files/gold-prod-hashes.txt"
			resp, err := http.Get(url)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer resp.Body.Close()

			scanner := bufio.NewScanner(resp.Body)
			for scanner.Scan() {
				known[scanner.Text()] = true
			}
			if err := scanner.Err(); err != nil {
				td.Fatal(ctx, err)
			}

			fmt.Fprintf(actualStdout, "Gold knew %v unique hashes.\n", len(known))
		}()
	}

	type Work struct {
		Sources []string // Passed to FM -s: names of gms/tests, paths to image files, .skps, etc.
		Flags   []string // Other flags to pass to FM: --ct 565, --msaa 16, etc.
	}

	queue := make(chan Work, 1<<20) // Arbitrarily huge buffer to avoid ever blocking.
	wg := &sync.WaitGroup{}
	var failures int32 = 0

	worker := func(sources, flags []string) {
		defer wg.Done()

		stdout := &bytes.Buffer{}
		stderr := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout, Stderr: stderr, Verbose: verbosity}
		cmd.Args = append(cmd.Args, "-i", *resources, "-s")
		cmd.Args = append(cmd.Args, sources...)
		cmd.Args = append(cmd.Args, flags...)

		// Run our FM command.
		err := exec.Run(ctx, cmd)

		// On success, scan stdout for any unknown hashes.
		unknownHash := func() string {
			if err == nil && *bot != "" { // We only fetch known hashes when using -bot.
				scanner := bufio.NewScanner(stdout)
				for scanner.Scan() {
					if parts := strings.Fields(scanner.Text()); len(parts) == 3 {
						md5 := parts[1]
						if !known[md5] {
							return md5
						}
					}
				}
				if err := scanner.Err(); err != nil {
					td.Fatal(ctx, err)
				}
			}
			return ""
		}()

		// If a batch failed or produced an unknown hash, isolate with individual reruns.
		if len(sources) > 1 && (err != nil || unknownHash != "") {
			wg.Add(len(sources))
			for i := range sources {
				queue <- Work{sources[i : i+1], flags}
			}
			return
		}

		// If an individual run failed, nothing more to do but fail.
		if err != nil {
			atomic.AddInt32(&failures, 1)
			td.FailStep(ctx, err)
			if *local {
				lines := []string{}
				scanner := bufio.NewScanner(stderr)
				for scanner.Scan() {
					lines = append(lines, scanner.Text())
				}
				if err := scanner.Err(); err != nil {
					td.Fatal(ctx, err)
				}
				fmt.Fprintf(actualStderr, "%v %v #failed:\n\t%v\n",
					cmd.Name,
					strings.Join(cmd.Args, " "),
					strings.Join(lines, "\n\t"))
			}
			return
		}

		// If an individual run succeeded but produced an unknown hash, TODO upload .png to Gold.
		// For now just print out the command and the hash it produced.
		if unknownHash != "" {
			fmt.Fprintf(actualStdout, "%v %v #%v\n",
				cmd.Name,
				strings.Join(cmd.Args, " "),
				unknownHash)
		}
	}

	for i := 0; i < runtime.NumCPU(); i++ {
		go func() {
			for w := range queue {
				worker(w.Sources, w.Flags)
			}
		}()
	}

	// Get some work going, first breaking it into batches to increase our parallelism.
	kickoff := func(sources, flags []string) {
		if len(sources) == 0 {
			return // A blank or commented job line from -script or the command line.
		}

		// Shuffle the sources randomly as a cheap way to approximate evenly expensive batches.
		// (Intentionally not rand.Seed()'d to stay deterministically reproducible.)
		rand.Shuffle(len(sources), func(i, j int) {
			sources[i], sources[j] = sources[j], sources[i]
		})

		nbatches := runtime.NumCPU()                      // Arbitrary, nice to scale ~= cores.
		batch := (len(sources) + nbatches - 1) / nbatches // Round up to avoid empty batches.
		util.ChunkIter(len(sources), batch, func(start, end int) error {
			wg.Add(1)
			queue <- Work{sources[start:end], flags}
			return nil
		})
	}

	// Parse a job like "gms b=cpu ct=8888" into sources and flags for kickoff().
	parse := func(job []string) (sources, flags []string) {
		for _, token := range job {
			// Everything after # is a comment.
			if strings.HasPrefix(token, "#") {
				break
			}

			// Treat "gm" or "gms" as a shortcut for all known GMs.
			if token == "gm" || token == "gms" {
				sources = append(sources, gms...)
				continue
			}
			// Same for tests.
			if token == "test" || token == "tests" {
				sources = append(sources, tests...)
				continue
			}

			// Is this a flag to pass through to FM?
			if parts := strings.Split(token, "="); len(parts) == 2 {
				f := "-"
				if len(parts[0]) > 1 {
					f += "-"
				}
				f += parts[0]

				flags = append(flags, f, parts[1])
				continue
			}

			// Anything else must be the name of a source for FM to run.
			sources = append(sources, token)
		}
		return
	}

	// Parse one job from the command line, handy for ad hoc local runs.
	kickoff(parse(flag.Args()[1:]))

	// Any number of jobs can come from -script.
	if *script != "" {
		file := os.Stdin
		if *script != "-" {
			file, err := os.Open(*script)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer file.Close()
		}
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			kickoff(parse(strings.Fields(scanner.Text())))
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
	}

	// If we're a bot (or acting as if we are one), kick off its work.
	if *bot != "" {
		parts := strings.Split(*bot, "-")
		model, CPU_or_GPU := parts[3], parts[4]

		commonFlags := []string{
			"--nativeFonts",
			strconv.FormatBool(strings.Contains(*bot, "NativeFonts")),
		}

		run := func(sources []string, extraFlags string) {
			kickoff(sources, append(strings.Fields(extraFlags), commonFlags...))
		}

		if CPU_or_GPU == "CPU" {
			commonFlags = append(commonFlags, "-b", "cpu")

			run(tests, "")
			run(gms, "--ct 8888 --legacy")  // Equivalent to DM --config 8888.

			if model == "GCE" {
				run(gms, "--ct g8 --legacy")                      // --config g8
				run(gms, "--ct 565 --legacy")                     // --config 565
				run(gms, "--ct 8888")                             // --config srgb
				run(gms, "--ct f16")                              // --config esrgb
				run(gms, "--ct f16 --tf linear")                  // --config f16
				run(gms, "--ct 8888 --gamut p3")                  // --config p3
				run(gms, "--ct 8888 --gamut narrow --tf 2.2")     // --config narrow
				run(gms, "--ct f16 --gamut rec2020 --tf rec2020") // --config erec2020

				run(gms, "--skvm")
				run(gms, "--skvm --ct f16")
			}

			// TODO: image/colorImage/svg tests
			// TODO: pic-8888 equivalent?
			// TODO: serialize-8888 equivalent?
		}
	}

	wg.Wait()
	if failures > 0 {
		if *local {
			// td.Fatalf() would work fine, but barfs up a panic that we don't need to see.
			fmt.Fprintf(actualStderr, "%v runs of %v failed after retries.\n", failures, fm)
			os.Exit(1)
		} else {
			td.Fatalf(ctx, "%v runs of %v failed after retries.", failures, fm)
		}
	}
}
<|endoftext|>


<commit_before>package support

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	neturl "net/url"
)

type HijackHttpOptions struct {
	Method             string
	Url                string
	Success            chan struct{}
	DockerTermProtocol bool
	InputStream        io.Reader
	ErrorStream        io.Writer
	OutputStream       io.Writer
	Data               interface{}
	Header             http.Header
}

// HijackHttpRequest performs an HTTP  request with given method, url and data and hijacks the request (after a successful connection) to stream
// data from/to the given input, output and error streams.
func HijackHttpRequest(options HijackHttpOptions) error {
	var params io.Reader
	if options.Data != nil {
		buf, err := json.Marshal(options.Data)
		if err != nil {
			return err
		}
		params = bytes.NewBuffer(buf)
	}

	stdout := options.OutputStream
	if stdout == nil {
		stdout = ioutil.Discard
	}
	stderr := options.ErrorStream
	if stderr == nil {
		stderr = ioutil.Discard
	}
	req, err := http.NewRequest(options.Method, options.Url, params)
	if err != nil {
		return err
	}
	if options.Header != nil {
		for k, values := range options.Header {
			req.Header.Del(k)
			for _, v := range values {
				req.Header.Set(k, v)
			}
		}
	}
	req.Header.Set("Content-Type", "text/plain")
	req.Header.Set("Connection", "Upgrade")
	req.Header.Set("Upgrade", "tcp")

	ep, err := neturl.Parse(options.Url)
	if err != nil {
		return err
	}

	protocol := ep.Scheme
	address := ep.Path
	if protocol != "unix" {
		protocol = "tcp"
		address = ep.Host
	}

	var dial net.Conn
	dial, err = net.Dial(protocol, address)
	if err != nil {
		return err
	}

	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()

	clientconn.Do(req)
	success := options.Success
	if success != nil {
		success <- struct{}{}
		<-success
	}

	rwc, br := clientconn.Hijack()
	defer rwc.Close()
	errs := make(chan error, 2)
	exit := make(chan bool)

	go func() {
		defer close(exit)
		var err error
		if !options.DockerTermProtocol {
			// When TTY is ON, use regular copy
			_, err = io.Copy(stdout, br)
		} else {
			_, err = dockerCopy(stdout, stderr, br)
		}
		errs <- err
	}()
	go func() {
		var err error
		in := options.InputStream
		if in != nil {
			_, err = io.Copy(rwc, in)
		}
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
		errs <- err
	}()
	<-exit
	return <-errs
}
<commit_msg>Refactored client for better structure<commit_after>package support

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httputil"
	neturl "net/url"
)

type HijackHttpOptions struct {
	Method             string
	Url                string
	Success            chan struct{}
	DockerTermProtocol bool
	InputStream        io.Reader
	ErrorStream        io.Writer
	OutputStream       io.Writer
	Data               interface{}
	Header             http.Header
}

// HijackHttpRequest performs an HTTP  request with given method, url and data and hijacks the request (after a successful connection) to stream
// data from/to the given input, output and error streams.
func HijackHttpRequest(options HijackHttpOptions) error {
	req, err := createHijackHttpRequest(options)
	if err != nil {
		return err
	}

	// Parse URL for endpoint data
	ep, err := neturl.Parse(options.Url)
	if err != nil {
		return err
	}

	protocol := ep.Scheme
	address := ep.Path
	if protocol != "unix" {
		protocol = "tcp"
		address = ep.Host
	}

	// Dial the server
	var dial net.Conn
	dial, err = net.Dial(protocol, address)
	if err != nil {
		return err
	}

	// Start initial HTTP connection
	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()

	clientconn.Do(req)

	// Hijack HTTP connection
	success := options.Success
	if success != nil {
		success <- struct{}{}
		<-success
	}

	rwc, br := clientconn.Hijack()
	defer rwc.Close()

	// Stream data
	return streamData(rwc, br, options)
}

// createHijackHttpRequest creates an upgradable HTTP request according to the given options
func createHijackHttpRequest(options HijackHttpOptions) (*http.Request, error) {
	var params io.Reader
	if options.Data != nil {
		buf, err := json.Marshal(options.Data)
		if err != nil {
			return nil, err
		}
		params = bytes.NewBuffer(buf)
	}

	req, err := http.NewRequest(options.Method, options.Url, params)
	if err != nil {
		return nil, err
	}
	if options.Header != nil {
		for k, values := range options.Header {
			req.Header.Del(k)
			for _, v := range values {
				req.Header.Set(k, v)
			}
		}
	}
	req.Header.Set("Content-Type", "text/plain")
	req.Header.Set("Connection", "Upgrade")
	req.Header.Set("Upgrade", "tcp")
	return req, nil
}

// streamData copies both input/output/error streams to/from the hijacked streams
func streamData(rwc io.Writer, br io.Reader, options HijackHttpOptions) error {
	errs := make(chan error, 2)
	exit := make(chan bool)

	go func() {
		defer close(exit)
		var err error
		stdout := options.OutputStream
		if stdout == nil {
			stdout = ioutil.Discard
		}
		stderr := options.ErrorStream
		if stderr == nil {
			stderr = ioutil.Discard
		}
		if !options.DockerTermProtocol {
			// When TTY is ON, use regular copy
			_, err = io.Copy(stdout, br)
		} else {
			_, err = dockerCopy(stdout, stderr, br)
		}
		errs <- err
	}()
	go func() {
		var err error
		in := options.InputStream
		if in != nil {
			_, err = io.Copy(rwc, in)
		}
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
		errs <- err
	}()
	<-exit
	return <-errs
}
<|endoftext|>


<commit_before>package vkapi

import (
	"errors"
	"log"
	"net/url"
	"strconv"
	"strings"
)

// Client allows you to transparently send requests to API server.
type Client struct {
	apiClient *APIClient
	User      Users
	LongPoll  *LongPoll
	//	Group
	//	Wall
	//	WallComment
	//  Message
	//  Chat
	//	Note
	//	Page
	//	Board
	//	BoardComment
}

// SetLanguage sets the language in which different data will be returned,
// for example, names of countries and cities.
func (client *Client) SetLanguage(lang string) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Language = lang
	return nil
}

// SetLogger sets logger.
func (client *Client) SetLogger(logger *log.Logger) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Logger = logger
	return nil
}

// Log allow write log.
func (client *Client) Log(flag bool) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Log = flag
	return nil
}

// NewClientFromToken creates a new *Client instance.
func NewClientFromToken(token string) (client *Client, err error) {
	client = new(Client)
	client.apiClient = NewApiClient()
	client.apiClient.SetAccessToken(token)
	return
}

// NewClientFromLogin creates a new *Client instance
// and allows you to pass a authentication.
func NewClientFromLogin(username string, password string, scope int64) (client *Client, err error) {
	client = new(Client)
	client.apiClient = NewApiClient()
	err = client.apiClient.Authenticate(NewApplication(username, password, scope))
	if err != nil {
		return nil, err
	}

	return
}

// Do makes a request to a specific endpoint with our request
// and returns response.
func (client *Client) Do(request Request) (response *Response, err *Error) {
	if client.apiClient == nil {
		return nil, NewError(ErrBadCode, ErrApiClientNotFound)
	}

	if request.Token == "" && client.apiClient.AccessToken != nil {
		request.Token = client.apiClient.AccessToken.AccessToken
	}

	return client.apiClient.Do(request)
}

// Destination describes the final destination.
type Destination struct {
	UserID      int64    `json:"user_id"`
	PeerID      int64    `json:"peer_id"`
	Domain      string   `json:"domain"`
	ChatID      int64    `json:"chat_id"`
	GroupID     int64    `json:"group_id"`
	UserIDs     []int64  `json:"user_ids"`
	ScreenName  string   `json:"user_id"`
	ScreenNames []string `json:"user_ids"`
}

func (dst Destination) Values() (values url.Values) {
	values = url.Values{}

	switch {
	case dst.UserID != 0:
		values.Add("user_id", strconv.FormatInt(dst.UserID, 10))
	case dst.PeerID != 0:
		values.Add("peer_id", strconv.FormatInt(dst.PeerID, 10))
	case dst.Domain != "":
		values.Add("domain", dst.Domain)
	case dst.ChatID != 0:
		values.Add("chat_id", strconv.FormatInt(dst.ChatID, 10))
	case dst.GroupID != 0:
		values.Add("group_id", strconv.FormatInt(dst.GroupID, 10))
	case len(dst.UserIDs) != 0:
		values.Add("user_ids", ConcatInt64ToString(dst.UserIDs...))
	case dst.ScreenName != "":
		values.Add("user_id", dst.ScreenName)
	case len(dst.ScreenNames) > 0:
		values.Add("user_ids", strings.Join(dst.ScreenNames, ","))
	}

	return
}

// NewDstFromUserID creates a new MessageConfig instance from userID.
func NewDstFromUserID(userIDs ...int64) (dst Destination) {
	if len(userIDs) == 1 {
		dst.UserID = userIDs[0]
	} else {
		dst.UserIDs = userIDs
	}
	return
}

// NewDstFromScreenName creates a new MessageConfig instance from userID.
func NewDstFromScreenName(screenNames ...string) (dst Destination) {
	if len(screenNames) == 1 {
		dst.ScreenName = screenNames[0]
	} else {
		dst.ScreenNames = screenNames
	}
	return
}

// NewDstFromPeerID creates a new MessageConfig instance from peerID.
func NewDstFromPeerID(peerID int64) (dst Destination) {
	dst.PeerID = peerID
	return
}

// NewDstFromChatID creates a new MessageConfig instance from chatID.
func NewDstFromChatID(chatID int64) (dst Destination) {
	dst.ChatID = chatID
	return
}

// NewDstFromGroupID creates a new MessageConfig instance from groupID.
func NewDstFromGroupID(groupID int64) (dst Destination) {
	dst.GroupID = groupID
	return
}

// NewDstFromDomain creates a new MessageConfig instance from domain.
func NewDstFromDomain(domain string) (dst Destination) {
	dst.Domain = domain
	return
}
<commit_msg>Add the vkapi.NewClientFromApplication() function<commit_after>package vkapi

import (
	"errors"
	"log"
	"net/url"
	"strconv"
	"strings"
)

// Client allows you to transparently send requests to API server.
type Client struct {
	apiClient *APIClient
	User      Users
	LongPoll  *LongPoll
	//	Group
	//	Wall
	//	WallComment
	//  Message
	//  Chat
	//	Note
	//	Page
	//	Board
	//	BoardComment
}

// SetLanguage sets the language in which different data will be returned,
// for example, names of countries and cities.
func (client *Client) SetLanguage(lang string) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Language = lang
	return nil
}

// SetLogger sets logger.
func (client *Client) SetLogger(logger *log.Logger) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Logger = logger
	return nil
}

// Log allow write log.
func (client *Client) Log(flag bool) error {
	if client.apiClient == nil {
		return errors.New(ErrApiClientNotFound)
	}

	client.apiClient.Log = flag
	return nil
}

// NewClientFromToken creates a new *Client instance.
func NewClientFromToken(token string) (client *Client, err error) {
	client = new(Client)
	client.apiClient = NewApiClient()
	client.apiClient.SetAccessToken(token)
	return
}

// NewClientFromLogin creates a new *Client instance
// and allows you to pass a authentication.
func NewClientFromLogin(username string, password string, scope int64) (client *Client, err error) {
	client = new(Client)
	client.apiClient = NewApiClient()
	err = client.apiClient.Authenticate(NewApplication(username, password, scope))
	if err != nil {
		return nil, err
	}

	return
}

// NewClientFromApplication creates a new *Client instance
// and allows you to pass a custom application.
func NewClientFromApplication(app Application) (client *Client, err error) {
	client = new(Client)
	client.apiClient = NewApiClient()
	err = client.apiClient.Authenticate(app)
	if err != nil {
		return nil, err
	}

	return
}

// Do makes a request to a specific endpoint with our request
// and returns response.
func (client *Client) Do(request Request) (response *Response, err *Error) {
	if client.apiClient == nil {
		return nil, NewError(ErrBadCode, ErrApiClientNotFound)
	}

	if request.Token == "" && client.apiClient.AccessToken != nil {
		request.Token = client.apiClient.AccessToken.AccessToken
	}

	return client.apiClient.Do(request)
}

// Destination describes the final destination.
type Destination struct {
	UserID      int64    `json:"user_id"`
	PeerID      int64    `json:"peer_id"`
	Domain      string   `json:"domain"`
	ChatID      int64    `json:"chat_id"`
	GroupID     int64    `json:"group_id"`
	UserIDs     []int64  `json:"user_ids"`
	ScreenName  string   `json:"user_id"`
	ScreenNames []string `json:"user_ids"`
}

func (dst Destination) Values() (values url.Values) {
	values = url.Values{}

	switch {
	case dst.UserID != 0:
		values.Add("user_id", strconv.FormatInt(dst.UserID, 10))
	case dst.PeerID != 0:
		values.Add("peer_id", strconv.FormatInt(dst.PeerID, 10))
	case dst.Domain != "":
		values.Add("domain", dst.Domain)
	case dst.ChatID != 0:
		values.Add("chat_id", strconv.FormatInt(dst.ChatID, 10))
	case dst.GroupID != 0:
		values.Add("group_id", strconv.FormatInt(dst.GroupID, 10))
	case len(dst.UserIDs) != 0:
		values.Add("user_ids", ConcatInt64ToString(dst.UserIDs...))
	case dst.ScreenName != "":
		values.Add("user_id", dst.ScreenName)
	case len(dst.ScreenNames) > 0:
		values.Add("user_ids", strings.Join(dst.ScreenNames, ","))
	}

	return
}

// NewDstFromUserID creates a new MessageConfig instance from userID.
func NewDstFromUserID(userIDs ...int64) (dst Destination) {
	if len(userIDs) == 1 {
		dst.UserID = userIDs[0]
	} else {
		dst.UserIDs = userIDs
	}
	return
}

// NewDstFromScreenName creates a new MessageConfig instance from userID.
func NewDstFromScreenName(screenNames ...string) (dst Destination) {
	if len(screenNames) == 1 {
		dst.ScreenName = screenNames[0]
	} else {
		dst.ScreenNames = screenNames
	}
	return
}

// NewDstFromPeerID creates a new MessageConfig instance from peerID.
func NewDstFromPeerID(peerID int64) (dst Destination) {
	dst.PeerID = peerID
	return
}

// NewDstFromChatID creates a new MessageConfig instance from chatID.
func NewDstFromChatID(chatID int64) (dst Destination) {
	dst.ChatID = chatID
	return
}

// NewDstFromGroupID creates a new MessageConfig instance from groupID.
func NewDstFromGroupID(groupID int64) (dst Destination) {
	dst.GroupID = groupID
	return
}

// NewDstFromDomain creates a new MessageConfig instance from domain.
func NewDstFromDomain(domain string) (dst Destination) {
	dst.Domain = domain
	return
}
<|endoftext|>


<commit_before>package micro

import (
	"fmt"
	"log"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/shouyingo/consul"
)

type serviceEntry struct {
	state int32 // 0: alive, 1: down
	id    string
	name  string
	addr  string
}

type Client struct {
	deps   map[string]*conngroup // service => group
	mgr    *contextManager
	nreqid uint64

	svcmu sync.RWMutex
	svcs  map[string]*serviceEntry // service-id => entry

	registry *consul.Client
	OnError  ErrFunc
}

func (c *Client) handleServer(svc *serviceEntry) {
	g := c.deps[svc.name]
	if g == nil {
		return
	}
	for atomic.LoadInt32(&svc.state) == 0 {
		nc, err := net.Dial("tcp", svc.addr)
		if err != nil {
			if eh := c.OnError; eh != nil {
				eh("dial", err)
			}
			time.Sleep(time.Second)
			continue
		}
		server := goconn(nc)
		g.add(server)
	mainloop:
		for {
			select {
			case pack := <-server.chrq:
				if pack.Flag&FlagReply == 0 {
					continue
				}
				ctx := c.mgr.remove(pack.Id)
				if ctx != nil && ctx.done() {
					ctx.fn(ctx.method, int(pack.Code), pack.Body)
				}
			case <-server.chdown:
				if isDebug {
					log.Printf("server(%d) close: %s", server.id, server.err)
				}
				break mainloop
			}
		}
		g.remove(server)
	}
}

func (d *Client) onWatch(action int, id string, svc *consul.CatalogService) {
	switch action {
	case 0:
		s := &serviceEntry{
			id:   id,
			name: svc.ServiceName,
			addr: net.JoinHostPort(svc.ServiceAddress, strconv.Itoa(svc.ServicePort)),
		}
		d.svcmu.Lock()
		old := d.svcs[id]
		d.svcs[id] = s
		d.svcmu.Unlock()
		if old != nil {
			atomic.StoreInt32(&old.state, 1)
		}
		go d.handleServer(s)
	case 2:
		d.svcmu.Lock()
		s := d.svcs[id]
		if s != nil {
			atomic.StoreInt32(&s.state, 1)
			delete(d.svcs, id)
		}
		d.svcmu.Unlock()
	}
}

func (c *Client) watchService(service string) {
	for {
		err := c.registry.WatchCatalogService(service, "", c.onWatch)
		if err != nil {
			if eh := c.OnError; eh != nil {
				eh("watch", err)
			}
		}
		time.Sleep(time.Second)
	}
}

func (c *Client) Call(service string, method string, params []byte, fn Callback) error {
	g := c.deps[service]
	if g == nil {
		return fmt.Errorf("%s not dependent service", service)
	}

	reqid := atomic.AddUint64(&c.nreqid, 1)
	ctx := c.mgr.add(reqid, method, fn, rpcTimeout)
	mc := g.randone()
	if mc == nil {
		if ctx.done() {
			c.mgr.remove(reqid)
			return fmt.Errorf("service(%s) no alive session", service)
		}
		return nil
	}

	ok := mc.send(&Packet{
		Id:   reqid,
		Name: method,
		Body: params,
	})
	if ok || !ctx.done() {
		return nil
	}
	c.mgr.remove(reqid)
	return fmt.Errorf("service(%s) send request failed", service)
}

func NewClient(r *consul.Client, depends []string) *Client {
	deps := make(map[string]*conngroup, len(depends))
	for _, dep := range depends {
		deps[dep] = &conngroup{}
	}

	c := &Client{
		deps: deps,
		mgr: &contextManager{
			ctxs:  make(map[uint64]*rpccontext),
			timer: time.NewTimer(timerIdle),
		},
		svcs:     make(map[string]*serviceEntry),
		registry: r,
	}
	c.mgr.c = c

	for svc := range c.deps {
		go c.watchService(svc)
	}
	go c.mgr.cleanExpired()
	return c
}
<commit_msg>add stop duplicate id server<commit_after>package micro

import (
	"fmt"
	"log"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/shouyingo/consul"
)

type serviceEntry struct {
	state int32 // 0: alive, 1: down
	mu    sync.Mutex
	c     *conn
	id    string
	name  string
	addr  string
}

func (s *serviceEntry) stop() {
	s.mu.Lock()
	atomic.StoreInt32(&s.state, 1)
	c := s.c
	s.c = nil
	s.mu.Unlock()
	if c != nil {
		c.shutdown(nil)
	}
}

type Client struct {
	deps   map[string]*conngroup // service => group
	mgr    *contextManager
	nreqid uint64

	svcmu sync.RWMutex
	svcs  map[string]*serviceEntry // service-id => entry

	registry *consul.Client
	OnError  ErrFunc
}

func (c *Client) handleServer(svc *serviceEntry) {
	g := c.deps[svc.name]
	if g == nil {
		return
	}
	for atomic.LoadInt32(&svc.state) == 0 {
		nc, err := net.Dial("tcp", svc.addr)
		if err != nil {
			if eh := c.OnError; eh != nil {
				eh("dial", err)
			}
			time.Sleep(time.Second)
			continue
		}
		server := goconn(nc)

		svc.mu.Lock()
		if svc.state != 0 {
			svc.mu.Unlock()
			server.shutdown(nil)
			break
		}
		svc.c = server
		svc.mu.Unlock()

		g.add(server)
	mainloop:
		for {
			select {
			case pack := <-server.chrq:
				if pack.Flag&FlagReply == 0 {
					continue
				}
				ctx := c.mgr.remove(pack.Id)
				if ctx != nil && ctx.done() {
					ctx.fn(ctx.method, int(pack.Code), pack.Body)
				}
			case <-server.chdown:
				if isDebug {
					log.Printf("server(%d) close: %s", server.id, server.err)
				}
				break mainloop
			}
		}
		g.remove(server)
	}
}

func (d *Client) onWatch(action int, id string, svc *consul.CatalogService) {
	switch action {
	case 0:
		s := &serviceEntry{
			id:   id,
			name: svc.ServiceName,
			addr: net.JoinHostPort(svc.ServiceAddress, strconv.Itoa(svc.ServicePort)),
		}
		d.svcmu.Lock()
		old := d.svcs[id]
		d.svcs[id] = s
		d.svcmu.Unlock()
		if old != nil {
			old.stop()
		}
		go d.handleServer(s)
	case 2:
		d.svcmu.Lock()
		s := d.svcs[id]
		if s != nil {
			delete(d.svcs, id)
		}
		d.svcmu.Unlock()
		s.stop()
	}
}

func (c *Client) watchService(service string) {
	for {
		err := c.registry.WatchCatalogService(service, "", c.onWatch)
		if err != nil {
			if eh := c.OnError; eh != nil {
				eh("watch", err)
			}
		}
		time.Sleep(time.Second)
	}
}

func (c *Client) Call(service string, method string, params []byte, fn Callback) error {
	g := c.deps[service]
	if g == nil {
		return fmt.Errorf("%s not dependent service", service)
	}

	reqid := atomic.AddUint64(&c.nreqid, 1)
	ctx := c.mgr.add(reqid, method, fn, rpcTimeout)
	mc := g.randone()
	if mc == nil {
		if ctx.done() {
			c.mgr.remove(reqid)
			return fmt.Errorf("service(%s) no alive session", service)
		}
		return nil
	}

	ok := mc.send(&Packet{
		Id:   reqid,
		Name: method,
		Body: params,
	})
	if ok || !ctx.done() {
		return nil
	}
	c.mgr.remove(reqid)
	return fmt.Errorf("service(%s) send request failed", service)
}

func NewClient(r *consul.Client, depends []string) *Client {
	deps := make(map[string]*conngroup, len(depends))
	for _, dep := range depends {
		deps[dep] = &conngroup{}
	}

	c := &Client{
		deps: deps,
		mgr: &contextManager{
			ctxs:  make(map[uint64]*rpccontext),
			timer: time.NewTimer(timerIdle),
		},
		svcs:     make(map[string]*serviceEntry),
		registry: r,
	}
	c.mgr.c = c

	for svc := range c.deps {
		go c.watchService(svc)
	}
	go c.mgr.cleanExpired()
	return c
}
<|endoftext|>


<commit_before>// Package untappd provides an Untappd APIv4 client, written in Go.  MIT Licensed.
//
// To use this client with the Untappd APIv4, you must register for an API key
// here: https://untappd.com/api/register.
//
// This package is inspired by Google's go-github library, as well as
// Antoine Grondin's canlii library.  Both can be found on GitHub:
//  - https://github.com/google/go-github
//  - https://github.com/aybabtme/canlii
package untappd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

const (
	// formEncodedContentType is the content type for key/value POST
	// body requests.
	formEncodedContentType = "application/x-www-form-urlencoded"

	// jsonContentType is the content type for JSON data.
	jsonContentType = "application/json"

	// untappdUserAgent is the default user agent this package will report to
	// the Untappd APIv4.
	untappdUserAgent = "github.com/mdlayher/untappd"
)

var (
	// ErrNoAccessToken is returned when an empty AccessToken is passed to
	// NewAuthenticatedClient.
	ErrNoAccessToken = errors.New("no client ID")

	// ErrNoClientID is returned when an empty Client ID is passed to NewClient.
	ErrNoClientID = errors.New("no client ID")

	// ErrNoClientSecret is returned when an empty Client Secret is passed
	// to NewClient.
	ErrNoClientSecret = errors.New("no client secret")
)

// Client is a HTTP client for the Untappd APIv4.  It enables access to various
// methods of the Untappd APIv4.
type Client struct {
	UserAgent string

	client *http.Client
	url    *url.URL

	clientID     string
	clientSecret string

	accessToken string

	// Methods which require authentication
	Auth interface {
		// https://untappd.com/api/docs#checkin
		Checkin(r CheckinRequest) (*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#activityfeed
		Checkins() ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)
	}

	// Methods involving a Beer
	Beer interface {
		// https://untappd.com/api/docs#beeractivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#beerinfo
		Info(id int, compact bool) (*Beer, *http.Response, error)

		// https://untappd.com/api/docs#beersearch
		Search(query string) ([]*Beer, *http.Response, error)
		SearchOffsetLimitSort(query string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)
	}

	// Methods involving a Brewery
	Brewery interface {
		// https://untappd.com/api/docs#breweryactivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#breweryinfo
		Info(id int, compact bool) (*Brewery, *http.Response, error)

		// https://untappd.com/api/docs#brewerysearch
		Search(query string) ([]*Brewery, *http.Response, error)
		SearchOffsetLimit(query string, offset int, limit int) ([]*Brewery, *http.Response, error)
	}

	// Methods involving a Local area
	Local interface {
		// https://untappd.com/api/docs#theppublocal
		Checkins(latitude float64, longitude float64) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimitRadius(r LocalCheckinsRequest) ([]*Checkin, *http.Response, error)
	}

	// Methods involving a User
	User interface {
		// https://untappd.com/api/docs#userbadges
		Badges(username string) ([]*Badge, *http.Response, error)
		BadgesOffsetLimit(username string, offset int, limit int) ([]*Badge, *http.Response, error)

		// https://untappd.com/api/docs#userbeers
		Beers(username string) ([]*Beer, *http.Response, error)
		BeersOffsetLimitSort(username string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)

		// https://untappd.com/api/docs#useractivityfeed
		Checkins(username string) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(username string, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#userfriends
		Friends(username string) ([]*User, *http.Response, error)
		FriendsOffsetLimit(username string, offset int, limit int) ([]*User, *http.Response, error)

		// https://untappd.com/api/docs#userinfo
		Info(username string, compact bool) (*User, *http.Response, error)

		// https://untappd.com/api/docs#userwishlist
		WishList(username string) ([]*Beer, *http.Response, error)
		WishListOffsetLimitSort(username string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)
	}

	// Methods involving a Venue
	Venue interface {
		// https://untappd.com/api/docs#venueactivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#venueinfo
		Info(id int, compact bool) (*Venue, *http.Response, error)
	}
}

// NewClient creates a properly initialized instance of Client, using the input
// client ID, client secret, and http.Client.
//
// To use a Client with the Untappd APIv4, you must register for an API key
// here: https://untappd.com/api/register.
func NewClient(clientID string, clientSecret string, client *http.Client) (*Client, error) {
	// Disallow empty ID and secret
	if clientID == "" {
		return nil, ErrNoClientID
	}
	if clientSecret == "" {
		return nil, ErrNoClientSecret
	}

	// Perform common client setup
	return newClient(clientID, clientSecret, "", client)
}

// NewAuthenticatedClient creates a properly initialized and authenticated instance
// of Client, using the input access token and http.Client.
//
// NewAuthenticatedClient must be called in order to create a Client which can
// access authenticated API actions, such as checking in beers, toasting other
// users' checkins, adding comments, etc.
//
// To use an authenticated Client with the Untappd APIv4, you must register
// for an API key here: https://untappd.com/api/register.  Next, you must follow
// the OAuth Authentication procedure documented here:
// https://untappd.com/api/docs#authentication.  Upon successful OAuth Authentication,
// you will receive an access token which can be used with NewAuthenticatedClient.
func NewAuthenticatedClient(accessToken string, client *http.Client) (*Client, error) {
	// Disallow empty access token
	if accessToken == "" {
		return nil, ErrNoAccessToken
	}

	// Perform common client setup
	return newClient("", "", accessToken, client)
}

// newClient handles common setup logic for a Client for NewClient and
// NewAuthenticatedClient.
func newClient(clientID string, clientSecret string, accessToken string, client *http.Client) (*Client, error) {
	// If input client is nil, use http.DefaultClient
	if client == nil {
		client = http.DefaultClient
	}

	// Set up basic client
	c := &Client{
		UserAgent: untappdUserAgent,

		client: client,
		url: &url.URL{
			Scheme: "https",
			Host:   "api.untappd.com",
			Path:   "v4",
		},

		clientID:     clientID,
		clientSecret: clientSecret,

		accessToken: accessToken,
	}

	// Add "services" which allow access to various API methods
	c.Auth = &AuthService{client: c}
	c.User = &UserService{client: c}
	c.Beer = &BeerService{client: c}
	c.Brewery = &BreweryService{client: c}
	c.Venue = &VenueService{client: c}
	c.Local = &LocalService{client: c}

	return c, nil
}

// Error represents an error returned from the Untappd APIv4.
type Error struct {
	Code              int
	Detail            string
	Type              string
	DeveloperFriendly string
	Duration          time.Duration
}

// Error returns the string representation of an Error.
func (e Error) Error() string {
	// Per APIv4 documentation, the "developer friendly" string should be used
	// in place of the regular "details" string wherever available
	details := e.Detail
	if e.DeveloperFriendly != "" {
		details = e.DeveloperFriendly
	}

	return fmt.Sprintf("%d [%s]: %s", e.Code, e.Type, details)
}

// request creates a new HTTP request, using the specified HTTP method and API endpoint.
// Additionally, it accepts POST body parameters, GET query parameters, and an
// optional struct which can be used to unmarshal result JSON.
func (c *Client) request(method string, endpoint string, body url.Values, query url.Values, v interface{}) (*http.Response, error) {
	// Generate relative URL using API root and endpoint
	rel, err := url.Parse(fmt.Sprintf("%s/%s/", c.url.Path, endpoint))
	if err != nil {
		return nil, err
	}

	// Resolve relative URL to base, using input host
	u := c.url.ResolveReference(rel)

	// Add any URL requested URL query parameters
	q := u.Query()
	for k, v := range query {
		for _, vv := range v {
			q.Add(k, vv)
		}
	}

	// Always prefer authenticated client access, using an access token.
	// If no token is found, fall back to unauthenticated client ID and
	// client secret.
	if c.accessToken != "" {
		q.Set("access_token", c.accessToken)
	} else {
		q.Set("client_id", c.clientID)
		q.Set("client_secret", c.clientSecret)
	}
	u.RawQuery = q.Encode()

	// Determine if request will contain a POST body
	hasBody := method == "POST" && len(body) > 0
	var length int

	// If performing a POST request and body parameters exist, encode
	// them now
	buf := bytes.NewBuffer(nil)
	if hasBody {
		// Encode and retrieve length to send to server
		buf = bytes.NewBufferString(body.Encode())
		length = buf.Len()
	}

	// Generate new HTTP request for appropriate URL
	req, err := http.NewRequest(method, u.String(), buf)
	if err != nil {
		return nil, err
	}

	// Set headers to indicate proper content type
	req.Header.Add("Accept", jsonContentType)

	// For POST requests, add proper headers
	if hasBody {
		req.Header.Add("Content-Type", formEncodedContentType)
		req.Header.Add("Content-Length", strconv.Itoa(length))
	}

	// Identify the client
	req.Header.Add("User-Agent", c.UserAgent)

	// Invoke request using underlying HTTP client
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	// Check response for errors
	if err := checkResponse(res); err != nil {
		return res, err
	}

	// If no second parameter was passed, do not attempt to handle response
	if v == nil {
		return res, nil
	}

	// Decode response body into v, returning response
	return res, json.NewDecoder(res.Body).Decode(v)
}

// getCheckins is the backing method for both any request which returns a
// list of checkins.  It handles performing the necessary HTTP request
// with the correct parameters, and returns a list of Checkins.
func (c *Client) getCheckins(endpoint string, q url.Values) ([]*Checkin, *http.Response, error) {
	// Temporary struct to unmarshal checkin JSON
	var v struct {
		Response struct {
			Checkins struct {
				Count int           `json:"count"`
				Items []*rawCheckin `json:"items"`
			} `json:"checkins"`
		} `json:"response"`
	}

	// Perform request for user checkins by ID
	res, err := c.request("GET", endpoint, nil, q, &v)
	if err != nil {
		return nil, res, err
	}

	// Build result slice from struct
	checkins := make([]*Checkin, v.Response.Checkins.Count)
	for i := range v.Response.Checkins.Items {
		checkins[i] = v.Response.Checkins.Items[i].export()
	}

	return checkins, res, nil
}

// checkResponse checks for a non-200 HTTP status code, and returns any errors
// encountered.
func checkResponse(res *http.Response) error {
	// Ensure correct content type
	if cType := res.Header.Get("Content-Type"); cType != jsonContentType {
		return fmt.Errorf("expected %s content type, but received %s", jsonContentType, cType)
	}

	// Check for 200-range status code
	if c := res.StatusCode; 200 <= c && c <= 299 {
		return nil
	}

	// Used as an intermediary form, but the contents are packed into
	// a more consumable form on error output
	var apiErr struct {
		Meta struct {
			Code              int              `json:"code"`
			ErrorDetail       string           `json:"error_detail"`
			ErrorType         string           `json:"error_type"`
			DeveloperFriendly string           `json:"developer_friendly"`
			ResponseTime      responseDuration `json:"response_time"`
		} `json:"meta"`
	}

	// Unmarshal error response
	if err := json.NewDecoder(res.Body).Decode(&apiErr); err != nil {
		return err
	}

	// Assemble Error struct from API response
	m := apiErr.Meta
	return &Error{
		Code:              m.Code,
		Detail:            m.ErrorDetail,
		Type:              m.ErrorType,
		DeveloperFriendly: m.DeveloperFriendly,
		Duration:          time.Duration(m.ResponseTime),
	}
}

// formatFloat converts a float64 to a string in a common way, to
// reduce inconsistencies with repeated calls to strconv.FormatFloat.
func formatFloat(f float64) string {
	return strconv.FormatFloat(f, 'f', -1, 64)
}
<commit_msg>Fix content type error (#11)<commit_after>// Package untappd provides an Untappd APIv4 client, written in Go.  MIT Licensed.
//
// To use this client with the Untappd APIv4, you must register for an API key
// here: https://untappd.com/api/register.
//
// This package is inspired by Google's go-github library, as well as
// Antoine Grondin's canlii library.  Both can be found on GitHub:
//  - https://github.com/google/go-github
//  - https://github.com/aybabtme/canlii
package untappd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	// formEncodedContentType is the content type for key/value POST
	// body requests.
	formEncodedContentType = "application/x-www-form-urlencoded"

	// jsonContentType is the content type for JSON data.
	jsonContentType = "application/json"

	// untappdUserAgent is the default user agent this package will report to
	// the Untappd APIv4.
	untappdUserAgent = "github.com/mdlayher/untappd"
)

var (
	// ErrNoAccessToken is returned when an empty AccessToken is passed to
	// NewAuthenticatedClient.
	ErrNoAccessToken = errors.New("no client ID")

	// ErrNoClientID is returned when an empty Client ID is passed to NewClient.
	ErrNoClientID = errors.New("no client ID")

	// ErrNoClientSecret is returned when an empty Client Secret is passed
	// to NewClient.
	ErrNoClientSecret = errors.New("no client secret")
)

// Client is a HTTP client for the Untappd APIv4.  It enables access to various
// methods of the Untappd APIv4.
type Client struct {
	UserAgent string

	client *http.Client
	url    *url.URL

	clientID     string
	clientSecret string

	accessToken string

	// Methods which require authentication
	Auth interface {
		// https://untappd.com/api/docs#checkin
		Checkin(r CheckinRequest) (*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#activityfeed
		Checkins() ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)
	}

	// Methods involving a Beer
	Beer interface {
		// https://untappd.com/api/docs#beeractivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#beerinfo
		Info(id int, compact bool) (*Beer, *http.Response, error)

		// https://untappd.com/api/docs#beersearch
		Search(query string) ([]*Beer, *http.Response, error)
		SearchOffsetLimitSort(query string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)
	}

	// Methods involving a Brewery
	Brewery interface {
		// https://untappd.com/api/docs#breweryactivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#breweryinfo
		Info(id int, compact bool) (*Brewery, *http.Response, error)

		// https://untappd.com/api/docs#brewerysearch
		Search(query string) ([]*Brewery, *http.Response, error)
		SearchOffsetLimit(query string, offset int, limit int) ([]*Brewery, *http.Response, error)
	}

	// Methods involving a Local area
	Local interface {
		// https://untappd.com/api/docs#theppublocal
		Checkins(latitude float64, longitude float64) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimitRadius(r LocalCheckinsRequest) ([]*Checkin, *http.Response, error)
	}

	// Methods involving a User
	User interface {
		// https://untappd.com/api/docs#userbadges
		Badges(username string) ([]*Badge, *http.Response, error)
		BadgesOffsetLimit(username string, offset int, limit int) ([]*Badge, *http.Response, error)

		// https://untappd.com/api/docs#userbeers
		Beers(username string) ([]*Beer, *http.Response, error)
		BeersOffsetLimitSort(username string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)

		// https://untappd.com/api/docs#useractivityfeed
		Checkins(username string) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(username string, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#userfriends
		Friends(username string) ([]*User, *http.Response, error)
		FriendsOffsetLimit(username string, offset int, limit int) ([]*User, *http.Response, error)

		// https://untappd.com/api/docs#userinfo
		Info(username string, compact bool) (*User, *http.Response, error)

		// https://untappd.com/api/docs#userwishlist
		WishList(username string) ([]*Beer, *http.Response, error)
		WishListOffsetLimitSort(username string, offset int, limit int, sort Sort) ([]*Beer, *http.Response, error)
	}

	// Methods involving a Venue
	Venue interface {
		// https://untappd.com/api/docs#venueactivityfeed
		Checkins(id int) ([]*Checkin, *http.Response, error)
		CheckinsMinMaxIDLimit(id int, minID int, maxID int, limit int) ([]*Checkin, *http.Response, error)

		// https://untappd.com/api/docs#venueinfo
		Info(id int, compact bool) (*Venue, *http.Response, error)
	}
}

// NewClient creates a properly initialized instance of Client, using the input
// client ID, client secret, and http.Client.
//
// To use a Client with the Untappd APIv4, you must register for an API key
// here: https://untappd.com/api/register.
func NewClient(clientID string, clientSecret string, client *http.Client) (*Client, error) {
	// Disallow empty ID and secret
	if clientID == "" {
		return nil, ErrNoClientID
	}
	if clientSecret == "" {
		return nil, ErrNoClientSecret
	}

	// Perform common client setup
	return newClient(clientID, clientSecret, "", client)
}

// NewAuthenticatedClient creates a properly initialized and authenticated instance
// of Client, using the input access token and http.Client.
//
// NewAuthenticatedClient must be called in order to create a Client which can
// access authenticated API actions, such as checking in beers, toasting other
// users' checkins, adding comments, etc.
//
// To use an authenticated Client with the Untappd APIv4, you must register
// for an API key here: https://untappd.com/api/register.  Next, you must follow
// the OAuth Authentication procedure documented here:
// https://untappd.com/api/docs#authentication.  Upon successful OAuth Authentication,
// you will receive an access token which can be used with NewAuthenticatedClient.
func NewAuthenticatedClient(accessToken string, client *http.Client) (*Client, error) {
	// Disallow empty access token
	if accessToken == "" {
		return nil, ErrNoAccessToken
	}

	// Perform common client setup
	return newClient("", "", accessToken, client)
}

// newClient handles common setup logic for a Client for NewClient and
// NewAuthenticatedClient.
func newClient(clientID string, clientSecret string, accessToken string, client *http.Client) (*Client, error) {
	// If input client is nil, use http.DefaultClient
	if client == nil {
		client = http.DefaultClient
	}

	// Set up basic client
	c := &Client{
		UserAgent: untappdUserAgent,

		client: client,
		url: &url.URL{
			Scheme: "https",
			Host:   "api.untappd.com",
			Path:   "v4",
		},

		clientID:     clientID,
		clientSecret: clientSecret,

		accessToken: accessToken,
	}

	// Add "services" which allow access to various API methods
	c.Auth = &AuthService{client: c}
	c.User = &UserService{client: c}
	c.Beer = &BeerService{client: c}
	c.Brewery = &BreweryService{client: c}
	c.Venue = &VenueService{client: c}
	c.Local = &LocalService{client: c}

	return c, nil
}

// Error represents an error returned from the Untappd APIv4.
type Error struct {
	Code              int
	Detail            string
	Type              string
	DeveloperFriendly string
	Duration          time.Duration
}

// Error returns the string representation of an Error.
func (e Error) Error() string {
	// Per APIv4 documentation, the "developer friendly" string should be used
	// in place of the regular "details" string wherever available
	details := e.Detail
	if e.DeveloperFriendly != "" {
		details = e.DeveloperFriendly
	}

	return fmt.Sprintf("%d [%s]: %s", e.Code, e.Type, details)
}

// request creates a new HTTP request, using the specified HTTP method and API endpoint.
// Additionally, it accepts POST body parameters, GET query parameters, and an
// optional struct which can be used to unmarshal result JSON.
func (c *Client) request(method string, endpoint string, body url.Values, query url.Values, v interface{}) (*http.Response, error) {
	// Generate relative URL using API root and endpoint
	rel, err := url.Parse(fmt.Sprintf("%s/%s/", c.url.Path, endpoint))
	if err != nil {
		return nil, err
	}

	// Resolve relative URL to base, using input host
	u := c.url.ResolveReference(rel)

	// Add any URL requested URL query parameters
	q := u.Query()
	for k, v := range query {
		for _, vv := range v {
			q.Add(k, vv)
		}
	}

	// Always prefer authenticated client access, using an access token.
	// If no token is found, fall back to unauthenticated client ID and
	// client secret.
	if c.accessToken != "" {
		q.Set("access_token", c.accessToken)
	} else {
		q.Set("client_id", c.clientID)
		q.Set("client_secret", c.clientSecret)
	}
	u.RawQuery = q.Encode()

	// Determine if request will contain a POST body
	hasBody := method == "POST" && len(body) > 0
	var length int

	// If performing a POST request and body parameters exist, encode
	// them now
	buf := bytes.NewBuffer(nil)
	if hasBody {
		// Encode and retrieve length to send to server
		buf = bytes.NewBufferString(body.Encode())
		length = buf.Len()
	}

	// Generate new HTTP request for appropriate URL
	req, err := http.NewRequest(method, u.String(), buf)
	if err != nil {
		return nil, err
	}

	// Set headers to indicate proper content type
	req.Header.Add("Accept", jsonContentType)

	// For POST requests, add proper headers
	if hasBody {
		req.Header.Add("Content-Type", formEncodedContentType)
		req.Header.Add("Content-Length", strconv.Itoa(length))
	}

	// Identify the client
	req.Header.Add("User-Agent", c.UserAgent)

	// Invoke request using underlying HTTP client
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	// Check response for errors
	if err := checkResponse(res); err != nil {
		return res, err
	}

	// If no second parameter was passed, do not attempt to handle response
	if v == nil {
		return res, nil
	}

	// Decode response body into v, returning response
	return res, json.NewDecoder(res.Body).Decode(v)
}

// getCheckins is the backing method for both any request which returns a
// list of checkins.  It handles performing the necessary HTTP request
// with the correct parameters, and returns a list of Checkins.
func (c *Client) getCheckins(endpoint string, q url.Values) ([]*Checkin, *http.Response, error) {
	// Temporary struct to unmarshal checkin JSON
	var v struct {
		Response struct {
			Checkins struct {
				Count int           `json:"count"`
				Items []*rawCheckin `json:"items"`
			} `json:"checkins"`
		} `json:"response"`
	}

	// Perform request for user checkins by ID
	res, err := c.request("GET", endpoint, nil, q, &v)
	if err != nil {
		return nil, res, err
	}

	// Build result slice from struct
	checkins := make([]*Checkin, v.Response.Checkins.Count)
	for i := range v.Response.Checkins.Items {
		checkins[i] = v.Response.Checkins.Items[i].export()
	}

	return checkins, res, nil
}

// checkResponse checks for a non-200 HTTP status code, and returns any errors
// encountered.
func checkResponse(res *http.Response) error {
	// Ensure correct content type
	if cType := res.Header.Get("Content-Type"); !strings.HasPrefix(cType, jsonContentType) {
		return fmt.Errorf("expected %s content type, but received %s", jsonContentType, cType)
	}

	// Check for 200-range status code
	if c := res.StatusCode; 200 <= c && c <= 299 {
		return nil
	}

	// Used as an intermediary form, but the contents are packed into
	// a more consumable form on error output
	var apiErr struct {
		Meta struct {
			Code              int              `json:"code"`
			ErrorDetail       string           `json:"error_detail"`
			ErrorType         string           `json:"error_type"`
			DeveloperFriendly string           `json:"developer_friendly"`
			ResponseTime      responseDuration `json:"response_time"`
		} `json:"meta"`
	}

	// Unmarshal error response
	if err := json.NewDecoder(res.Body).Decode(&apiErr); err != nil {
		return err
	}

	// Assemble Error struct from API response
	m := apiErr.Meta
	return &Error{
		Code:              m.Code,
		Detail:            m.ErrorDetail,
		Type:              m.ErrorType,
		DeveloperFriendly: m.DeveloperFriendly,
		Duration:          time.Duration(m.ResponseTime),
	}
}

// formatFloat converts a float64 to a string in a common way, to
// reduce inconsistencies with repeated calls to strconv.FormatFloat.
func formatFloat(f float64) string {
	return strconv.FormatFloat(f, 'f', -1, 64)
}
<|endoftext|>


<commit_before>package mdns

import (
	"code.google.com/p/go.net/ipv4"
	"code.google.com/p/go.net/ipv6"
	"fmt"
	"github.com/miekg/dns"
	"log"
	"net"
	"strings"
	"sync"
	"time"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name string
	Addr net.IP
	Port int
	Info string

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return s.Addr != nil && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service   string               // Service to lookup
	Domain    string               // Lookup domain, default "local"
	Timeout   time.Duration        // Lookup timeout, default 1 second
	Interface *net.Interface       // Multicast interface to use
	Entries   chan<- *ServiceEntry // Entries Channel
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service: service,
		Domain:  "local",
		Timeout: time.Second,
		Entries: make(chan *ServiceEntry),
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Set the multicast interface
	if params.Interface != nil {
		if err := client.setInterface(params.Interface); err != nil {
			return err
		}
	}

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4List *net.UDPConn
	ipv6List *net.UDPConn

	closed    bool
	closedCh  chan struct{}
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// Create a IPv4 listener
	ipv4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	ipv6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if ipv4 == nil && ipv6 == nil {
		return nil, fmt.Errorf("Failed to bind to any udp port!")
	}

	c := &client{
		ipv4List: ipv4,
		ipv6List: ipv6,
		closedCh: make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true
	close(c.closedCh)

	if c.ipv4List != nil {
		c.ipv4List.Close()
	}
	if c.ipv6List != nil {
		c.ipv6List.Close()
	}
	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4List)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6List)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4List, msgCh)
	go c.recv(c.ipv6List, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypeANY)
	if err := c.sendQuery(m); err != nil {
		return nil
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range resp.Answer {
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Name = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA
				}
			}

			// Check if this entry is complete
			if inp.complete() && !inp.sent {
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypeANY)
				if err := c.sendQuery(m); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
	return nil
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4List != nil {
		c.ipv4List.WriteTo(buf, ipv4Addr)
	}
	if c.ipv6List != nil {
		c.ipv6List.WriteTo(buf, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<commit_msg>Suppress recursive mDNS queries and query for PTR records only<commit_after>package mdns

import (
	"code.google.com/p/go.net/ipv4"
	"code.google.com/p/go.net/ipv6"
	"fmt"
	"github.com/miekg/dns"
	"log"
	"net"
	"strings"
	"sync"
	"time"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name string
	Addr net.IP
	Port int
	Info string

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return s.Addr != nil && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service   string               // Service to lookup
	Domain    string               // Lookup domain, default "local"
	Timeout   time.Duration        // Lookup timeout, default 1 second
	Interface *net.Interface       // Multicast interface to use
	Entries   chan<- *ServiceEntry // Entries Channel
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service: service,
		Domain:  "local",
		Timeout: time.Second,
		Entries: make(chan *ServiceEntry),
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Set the multicast interface
	if params.Interface != nil {
		if err := client.setInterface(params.Interface); err != nil {
			return err
		}
	}

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4List *net.UDPConn
	ipv6List *net.UDPConn

	closed    bool
	closedCh  chan struct{}
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// Create a IPv4 listener
	ipv4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	ipv6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if ipv4 == nil && ipv6 == nil {
		return nil, fmt.Errorf("Failed to bind to any udp port!")
	}

	c := &client{
		ipv4List: ipv4,
		ipv6List: ipv6,
		closedCh: make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true
	close(c.closedCh)

	if c.ipv4List != nil {
		c.ipv4List.Close()
	}
	if c.ipv6List != nil {
		c.ipv6List.Close()
	}
	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4List)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6List)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4List, msgCh)
	go c.recv(c.ipv6List, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypePTR)
	m.RecursionDesired = false
	if err := c.sendQuery(m); err != nil {
		return nil
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range resp.Answer {
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Name = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA
				}
			}

			// Check if this entry is complete
			if inp.complete() && !inp.sent {
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypePTR)
				m.RecursionDesired = false
				if err := c.sendQuery(m); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
	return nil
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4List != nil {
		c.ipv4List.WriteTo(buf, ipv4Addr)
	}
	if c.ipv6List != nil {
		c.ipv6List.WriteTo(buf, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<|endoftext|>


<commit_before>/*
 * Copyright 2014 Xuyuan Pang <xuyuanp # gmail dot com>
 *
 * 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 gots

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
)

const (
	DefaultEncoding      = "utf8"
	DefaultSocketTimeout = 50
	DefaultMaxConnection = 50
)

type Client struct {
	EndPoint      string
	AccessID      string
	AccessKey     string
	InstanceName  string
	Encoding      string
	SocketTimeout float32
	MaxConnection int
	Debug         bool
	Logger        *log.Logger
	protocol      *Protocol
	encoder       *Encoder
	decoder       *Decoder
}

func NewClient(endPoint, accessID, accessKey, instanceName string) *Client {
	return &Client{
		EndPoint:      endPoint,
		AccessID:      accessID,
		AccessKey:     accessKey,
		InstanceName:  instanceName,
		Encoding:      DefaultEncoding,
		SocketTimeout: DefaultSocketTimeout,
		MaxConnection: DefaultMaxConnection,
		Debug:         false,
	}
}

func (c *Client) Init() error {
	c.protocol = &Protocol{
		EndPoint:     c.EndPoint,
		AccessID:     c.AccessID,
		AccessKey:    c.AccessKey,
		InstanceName: c.InstanceName,
	}
	c.encoder = &Encoder{encoding: c.Encoding}
	c.decoder = &Decoder{encoding: c.Encoding}
	return nil
}

func (c *Client) Visit(apiName string, message proto.Message) (data []byte, err error) {
	body, err := proto.Marshal(message)
	if err != nil {
		return nil, &OTSClientError{Message: fmt.Sprintf("%s Marshal protocol buffer failed", err.Error())}
	}
	if c.Debug && c.Logger != nil {
		c.Logger.Printf(`Request: %s data: %s`, apiName, message.String())
	}
	req, err := c.protocol.MakeRequest(apiName, body)
	if err != nil {
		return nil, err
	}
	// TODO: Use connection pool
	response, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, &OTSClientError{Message: fmt.Sprintf("%s Send request failed", err.Error())}
	}
	defer response.Body.Close()
	data, err = ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, &OTSClientError{Message: "Read data faild in response"}
	}

	headers := response.Header
	newHeaders := make(map[string]string, len(headers))
	for k, _ := range headers {
		lk := strings.ToLower(k)
		newHeaders[lk] = headers.Get(k)
	}
	return data, c.protocol.ParseResponse(apiName, response.StatusCode, newHeaders, data)
}

func (c *Client) ListTable() (names []string, err error) {
	message, err := c.encoder.EncodeListTable()
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("ListTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeListTable(data)
}

func (c *Client) CreateTable(name string, primaryKey []*ColumnSchema, rt *ReservedThroughput) (*CreateTableResponse, error) {
	message, err := c.encoder.EncodeCreateTable(name, primaryKey, rt)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("CreateTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeCreateTable(data)
}

func (c *Client) DeleteTable(name string) (*DeleteTableResponse, error) {
	message, err := c.encoder.EncodeDeleteTable(name)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("DeleteTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeDeleteTable(data)
}

func (c *Client) DescribeTable(name string) (*TableMeta, *ReservedThoughputDetails, error) {
	message, err := c.encoder.EncodeDescribeTable(name)
	if err != nil {
		return nil, nil, err
	}
	data, err := c.Visit("DescribeTable", message)
	if err != nil {
		return nil, nil, err
	}
	return c.decoder.DecodeDescribeTable(data)
}

func (c *Client) UpdateTable(name string, reservedThroughput *ReservedThroughput) (*UpdateTableResponse, error) {
	message, err := c.encoder.EncodeUpdateTable(name, reservedThroughput)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("UpdateTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeUpdateTable(data)
}

func (c *Client) GetRow(name string, primaryKey map[string]interface{}, columnNames []string) (*GetRowResponse, error) {
	message, err := c.encoder.EncodeGetRow(name, primaryKey, columnNames)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("GetRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeGetRow(data)
}

func (c *Client) PutRow(name string, condition *Condition, primaryKey map[string]interface{}, columns map[string]interface{}) (response *PutRowResponse, err error) {
	message, err := c.encoder.EncodePutRow(name, condition, primaryKey, columns)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("PutRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodePutRow(data)
}

func (c *Client) UpdateRow(name string, condition *Condition, primaryKey map[string]interface{}, columnsPut map[string]interface{}, columnsDelete []string) (*UpdateRowResponse, error) {
	message, err := c.encoder.EncodeUpdateRow(name, condition, primaryKey, columnsPut, columnsDelete)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("UpdateRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeUpdateRow(data)
}

func (c *Client) DeleteRow(name string, condition *Condition, primaryKey map[string]interface{}) (*DeleteRowResponse, error) {
	message, err := c.encoder.EncodeDeleteRow(name, condition, primaryKey)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("DeleteRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeDeleteRow(data)
}

func (c *Client) BatchGetRow(items map[string]BatchGetRowItem) (*BatchGetRowResponse, error) {
	message, err := c.encoder.EncodeBatchGetRow(items)
	if err != nil {
		return nil, err
	}
	data, err := c.Visit("BatchGetRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeBatchGetRow(data)
}

// func (c *Client) BatchWriteRow(batchList []map[string]interface{}) (items []map[string]interface{}, err error) {
// 	return nil, nil
// }
//
// func (c *Client) GetRange(name string, direction Direction, incStartPrimaryKey *PrimaryKey, excEndPrimaryKey *PrimaryKey,
// 	colums []string, limit int) (consumed *CapacityUnit, next []*PrimaryKey, rows []interface{}, err error) {
// 	return nil, nil, nil, nil
// }
//
// func (c *Client) XGetRange(name string, direction Direction, incStartPrimaryKey *PrimaryKey, excEndPrimaryKey *PrimaryKey,
// 	consumedCounter *CapacityUnit, colums []string, limit int) (consumed *CapacityUnit, next []*PrimaryKey, rows []interface{}, err error) {
// 	return nil, nil, nil, nil
// }
<commit_msg>Add comments<commit_after>/*
 * Copyright 2014 Xuyuan Pang <xuyuanp # gmail dot com>
 *
 * 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 gots

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
)

const (
	// DefaultEncoding 默认编码
	DefaultEncoding = "utf8"
	// DefaultSocketTimeout 默认socket链接过期时间
	DefaultSocketTimeout = 50
	// DefaultMaxConnection 默认链接池最大链接数
	DefaultMaxConnection = 50
)

// Client 实现了OTS服务的所有接口。用户可以通过NewClient方法创建Client实例
type Client struct {
	EndPoint      string
	AccessID      string
	AccessKey     string
	InstanceName  string
	Encoding      string
	SocketTimeout float32
	MaxConnection int
	Debug         bool
	Logger        *log.Logger
	protocol      *Protocol
	encoder       *Encoder
	decoder       *Decoder
}

// NewClient 方法返回一个Client实例
// 示例:
//
// import "github.com/Xuyuanp/gots
// client := ots.NewClient("your_instance_endpoint", "your_user_id", "your_user_key", "your_instance_name")
// client.Init()
func NewClient(endPoint, accessID, accessKey, instanceName string) *Client {
	return &Client{
		EndPoint:      endPoint,
		AccessID:      accessID,
		AccessKey:     accessKey,
		InstanceName:  instanceName,
		Encoding:      DefaultEncoding,
		SocketTimeout: DefaultSocketTimeout,
		MaxConnection: DefaultMaxConnection,
		Debug:         false,
	}
}

// Init 方法初始化Client，在创建Client实例之后，调用任何访问接口之前必须调用此方法
func (c *Client) Init() error {
	c.protocol = &Protocol{
		EndPoint:     c.EndPoint,
		AccessID:     c.AccessID,
		AccessKey:    c.AccessKey,
		InstanceName: c.InstanceName,
	}
	c.encoder = &Encoder{encoding: c.Encoding}
	c.decoder = &Decoder{encoding: c.Encoding}
	return nil
}

func (c *Client) vist(apiName string, message proto.Message) (data []byte, err error) {
	body, err := proto.Marshal(message)
	if err != nil {
		return nil, &OTSClientError{Message: fmt.Sprintf("%s Marshal protocol buffer failed", err.Error())}
	}
	if c.Debug && c.Logger != nil {
		c.Logger.Printf(`Request: %s data: %s`, apiName, message.String())
	}
	req, err := c.protocol.MakeRequest(apiName, body)
	if err != nil {
		return nil, err
	}
	// TODO: Use connection pool
	response, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, &OTSClientError{Message: fmt.Sprintf("%s Send request failed", err.Error())}
	}
	defer response.Body.Close()
	data, err = ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, &OTSClientError{Message: "Read data faild in response"}
	}

	headers := response.Header
	newHeaders := make(map[string]string, len(headers))
	for k := range headers {
		lk := strings.ToLower(k)
		newHeaders[lk] = headers.Get(k)
	}
	return data, c.protocol.ParseResponse(apiName, response.StatusCode, newHeaders, data)
}

// ListTable 方法用于获取所有表名。
// 示例:
//
// names, err := client.ListTable()
func (c *Client) ListTable() (names []string, err error) {
	message, err := c.encoder.EncodeListTable()
	if err != nil {
		return nil, err
	}
	data, err := c.vist("ListTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeListTable(data)
}

// CreateTable 方法用于创建新表。
// name: 表名
// primaryKey: 主键列表
// rt: 预留读写吞吐量
// 示例：
//
// primaryKey := []*gots.ColumnSchema{
//      &gots.ColumnSchema{
//              Name: "gid",
//              Type: gots.ColumnTypeInteger,
//      },
//      &gots.ColumnSchema{
//              Name: "uid",
//              Type: gots.ColumnTypeInteger,
//      },
// }
// rt := &gots.ReservedThroughput{
//      CapacityUnit: &gots.CapacityUnit{
//              Read:  100,
//              Write: 100,
//      },
// }
//
// resp, err := client.CreateTable("sample_table", primaryKey, rt)
func (c *Client) CreateTable(name string, primaryKey []*ColumnSchema, rt *ReservedThroughput) (*CreateTableResponse, error) {
	message, err := c.encoder.EncodeCreateTable(name, primaryKey, rt)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("CreateTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeCreateTable(data)
}

// DeleteTable 方法用于删除表
// name: 表名
// 示例:
//
// resp, err := client.DeleteTable("sample_table")
func (c *Client) DeleteTable(name string) (*DeleteTableResponse, error) {
	message, err := c.encoder.EncodeDeleteTable(name)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("DeleteTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeDeleteTable(data)
}

// DescribeTable 方法用于获取表描述信息
// name: 表名
// 示例:
//
// resp, err := client.DescribeTable("sample_table")
func (c *Client) DescribeTable(name string) (*TableMeta, *ReservedThoughputDetails, error) {
	message, err := c.encoder.EncodeDescribeTable(name)
	if err != nil {
		return nil, nil, err
	}
	data, err := c.vist("DescribeTable", message)
	if err != nil {
		return nil, nil, err
	}
	return c.decoder.DecodeDescribeTable(data)
}

// UpdateTable 跟新表属性，目前只支持修改预留读写吞吐量
// name: 表名
// reservedThroughput: 预留读写吞吐量
// 示例:
//
// rt := &gots.ReservedThroughput{
//      CapacityUnit: &gots.CapacityUnit{
//              Read:  150,
//              Write: 150,
//      },
// }
// resp, err := client.UpdateTable("sample_table", rt)
func (c *Client) UpdateTable(name string, reservedThroughput *ReservedThroughput) (*UpdateTableResponse, error) {
	message, err := c.encoder.EncodeUpdateTable(name, reservedThroughput)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("UpdateTable", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeUpdateTable(data)
}

func (c *Client) GetRow(name string, primaryKey map[string]interface{}, columnNames []string) (*GetRowResponse, error) {
	message, err := c.encoder.EncodeGetRow(name, primaryKey, columnNames)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("GetRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeGetRow(data)
}

func (c *Client) PutRow(name string, condition *Condition, primaryKey map[string]interface{}, columns map[string]interface{}) (response *PutRowResponse, err error) {
	message, err := c.encoder.EncodePutRow(name, condition, primaryKey, columns)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("PutRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodePutRow(data)
}

func (c *Client) UpdateRow(name string, condition *Condition, primaryKey map[string]interface{}, columnsPut map[string]interface{}, columnsDelete []string) (*UpdateRowResponse, error) {
	message, err := c.encoder.EncodeUpdateRow(name, condition, primaryKey, columnsPut, columnsDelete)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("UpdateRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeUpdateRow(data)
}

func (c *Client) DeleteRow(name string, condition *Condition, primaryKey map[string]interface{}) (*DeleteRowResponse, error) {
	message, err := c.encoder.EncodeDeleteRow(name, condition, primaryKey)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("DeleteRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeDeleteRow(data)
}

func (c *Client) BatchGetRow(items map[string]BatchGetRowItem) (*BatchGetRowResponse, error) {
	message, err := c.encoder.EncodeBatchGetRow(items)
	if err != nil {
		return nil, err
	}
	data, err := c.vist("BatchGetRow", message)
	if err != nil {
		return nil, err
	}
	return c.decoder.DecodeBatchGetRow(data)
}

// func (c *Client) BatchWriteRow(batchList []map[string]interface{}) (items []map[string]interface{}, err error) {
// 	return nil, nil
// }
//
// func (c *Client) GetRange(name string, direction Direction, incStartPrimaryKey *PrimaryKey, excEndPrimaryKey *PrimaryKey,
// 	colums []string, limit int) (consumed *CapacityUnit, next []*PrimaryKey, rows []interface{}, err error) {
// 	return nil, nil, nil, nil
// }
//
// func (c *Client) XGetRange(name string, direction Direction, incStartPrimaryKey *PrimaryKey, excEndPrimaryKey *PrimaryKey,
// 	consumedCounter *CapacityUnit, colums []string, limit int) (consumed *CapacityUnit, next []*PrimaryKey, rows []interface{}, err error) {
// 	return nil, nil, nil, nil
// }
<|endoftext|>