<commit_before>package main

import (
  "log"
  "github.com/gorilla/mux"
  "net/http"
  "net/http/httputil"
  "net/url"
  "github.com/fhs/gompd/mpd"
  "github.com/ascherkus/go-id3/src/id3"
  "os"
  "fmt"
  "encoding/json"
  "strconv"
  "time"
  "io"
)

// TODO: consider if global is really the best idea, or if we should
//       make some classes, or something...
var mpd_conn *mpd.Client

func main() {
  // setup our global MPD connection
  mpd_conn = mpdConnect("localhost:6600")
  defer mpd_conn.Close()

  if mpd_conn == nil {
    log.Fatal("MPD Connection is nil!")
  }

  // create a new mux router for our server.
  r := mux.NewRouter()

  // requests to `/stream` are proxied to the MPD httpd.
  r.HandleFunc("/stream",
    httputil.NewSingleHostReverseProxy(
      &url.URL{
        Scheme:"http",
        Host: "localhost:8000",
        Path: "/",
      }).ServeHTTP)

  r.HandleFunc("/songs", listSongs)
  r.HandleFunc("/current", getCurrentSong)
  r.HandleFunc("/upcoming", getUpcomingSongs)

  // This MUST go last! It takes precidence over any after it, meaning
  // the server will try to serve a file, which most likely doesn't exist,
  // and will 404.
  //
  // serve up the frontend files.
  r.PathPrefix("/").Handler(http.FileServer(http.Dir("../frontend/turbo_wookie/web")))

  // sit, waiting, like a hunter, spying on its prey.
  log.Println("Starting server on port 9000")
  http.ListenAndServe(":9000", r)
}

/********************
  Handler Functions
 ********************/

// return all songs known to MPD to the client.
func listSongs(w http.ResponseWriter, r *http.Request) {
  // get all files from MPD
  mpdfiles, err := mpd_conn.GetFiles()
  if err != nil {
    count := 0

    for err != nil && count < 10 {
      time.Sleep(10)

      mpdfiles, err := mpd_conn.GetFiles()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get a list of files...", err)
      return
    }
  }

  // create a slice of id3.File s
  //files := make([]*id3.File, 0)
  files := make([]*TBFile, 0)

  for _, song := range mpdfiles {
    // grab the file on the filesystem
    file, err := os.Open("mpd/music/" + song)
    if err != nil {
      error(w, "Couldn't open file: " + song, err)
      return
    }

    // add the current file to our slice
    id3_file := id3Read(file, song)
    files = append(files, id3_file)
  }

  // send the json to the client.
  fmt.Fprintf(w, jsoniffy(files))
}

// Return a JSON representation of the currently playing song.
func getCurrentSong(w http.ResponseWriter, r *http.Request) {
  currentSong, err := mpd_conn.CurrentSong()
  if err != nil {

    count := 0;
    for err != nil && count < 10 {
      time.Sleep(10)

      currentSong, err = mpd_conn.CurrentSong()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get current song info for upcoming list", err)
      return
    }
  }

  fmt.Fprintf(w, jsoniffy(currentSong))
}

func getUpcomingSongs(w http.ResponseWriter, r *http.Request) {
  currentSong, err := mpd_conn.CurrentSong()
  if err != nil {

    count := 0;
    for err != nil && count < 10 {
      time.Sleep(10)

      currentSong, err = mpd_conn.CurrentSong()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get current song info for upcoming list", err)
      return
    }
  }

  pos, err := strconv.Atoi(currentSong["Pos"])
  if err != nil {
    error(w, "Couldn't turn current song's position to int", err)
    return
  }

  playlist, err := mpd_conn.PlaylistInfo(-1, -1)
  if err != nil {
    count := 0
    for err != nil && count < 10 {
      time.Sleep(10)

      playlist, err = mpd_conn.PlaylistInfo(-1, -1)
      count ++
    }

    if err != nil {
      error(w, "Couldn't get the current playlist", err)
      return
    }
  }

  upcoming := playlist[pos + 1:]

  fmt.Fprintf(w, jsoniffy(upcoming))
}

/*******************
  Helper Functions
 *******************/

// Connect to MPD's control channel, and set the global mpd_conn to it.
func mpdConnect(url string) *mpd.Client {
  conn, err := mpd.Dial("tcp", url)

  // if we can't connect to MPD everything's fucked, nothing's going to work
  // kill all humans, and die, respectfully, after explaining what the issue
  // is.
  if err != nil {
    log.Println("\n\nServer quiting because it can't connect to MPD");
    log.Println(err)

    return nil
  }

  return conn
}

// helper struct; used to hold some ID3 info, plus an MPD file path.
type TBFile struct {id3.File; FilePath string;}

// helper method, returns a pointer to one of our helper structs (see above).
func id3Read(reader io.Reader, filePath string) *TBFile {
  id3File := id3.Read(reader)

  file := new(TBFile)
  file.Header = id3File.Header
  file.Name = id3File.Name
  file.Artist = id3File.Artist
  file.Album = id3File.Album
  file.Year = id3File.Year
  file.FilePath = filePath

  return file
}

// turn anything into JSON.
func jsoniffy(v interface {}) string {
  obj, err := json.MarshalIndent(v, "", "  ")
  if err != nil {
    log.Print("Couldn't turn something into JSON: ", v)
    log.Fatal(err)
  }

  return string(obj)
}

func error(w http.ResponseWriter, message string, err interface{Error() string;}) {
  log.Println("An error occured; telling the client.")
  log.Println("Message:", message)
  log.Println("Error:", err, "\n\n")

  fmt.Fprintf(w, message + "\n")
}

func jsonError(w http.ResponseWriter, message string, err interface{Error() string;}) {
  message = "{error:\"" + message + "\"}"
  error(w, message, err)
}

<commit_msg>that was ... interesting<commit_after>package main

import (
  "log"
  "github.com/gorilla/mux"
  "net/http"
  "net/http/httputil"
  "net/url"
  "github.com/fhs/gompd/mpd"
  "github.com/ascherkus/go-id3/src/id3"
  "os"
  "fmt"
  "encoding/json"
  "strconv"
  "time"
  "io"
)

// TODO: consider if global is really the best idea, or if we should
//       make some classes, or something...
var mpd_conn *mpd.Client

func main() {
  // setup our global MPD connection
  mpd_conn = mpdConnect("localhost:6600")
  defer mpd_conn.Close()

  if mpd_conn == nil {
    log.Fatal("MPD Connection is nil!")
  }

  // create a new mux router for our server.
  r := mux.NewRouter()

  // requests to `/stream` are proxied to the MPD httpd.
  r.HandleFunc("/stream",
    httputil.NewSingleHostReverseProxy(
      &url.URL{
        Scheme:"http",
        Host: "localhost:8000",
        Path: "/",
      }).ServeHTTP)

  r.HandleFunc("/songs", listSongs)
  r.HandleFunc("/current", getCurrentSong)
  r.HandleFunc("/upcoming", getUpcomingSongs)

  // This MUST go last! It takes precidence over any after it, meaning
  // the server will try to serve a file, which most likely doesn't exist,
  // and will 404.
  //
  // serve up the frontend files.
  r.PathPrefix("/").Handler(http.FileServer(http.Dir("../frontend/turbo_wookie/web")))

  // sit, waiting, like a hunter, spying on its prey.
  log.Println("Starting server on port 9000")
  http.ListenAndServe(":9000", r)
}

/********************
  Handler Functions
 ********************/

// return all songs known to MPD to the client.
func listSongs(w http.ResponseWriter, r *http.Request) {
  // get all files from MPD
  mpdfiles, err := mpd_conn.GetFiles()
  if err != nil {
    count := 0

    for err != nil && count < 10 {
      time.Sleep(10)

      mpdfiles, err = mpd_conn.GetFiles()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get a list of files...", err)
      return
    }
  }

  // create a slice of id3.File s
  //files := make([]*id3.File, 0)
  files := make([]*TBFile, 0)

  for _, song := range mpdfiles {
    // grab the file on the filesystem
    file, err := os.Open("mpd/music/" + song)
    if err != nil {
      error(w, "Couldn't open file: " + song, err)
      return
    }

    // add the current file to our slice
    id3_file := id3Read(file, song)
    files = append(files, id3_file)
  }

  // send the json to the client.
  fmt.Fprintf(w, jsoniffy(files))
}

// Return a JSON representation of the currently playing song.
func getCurrentSong(w http.ResponseWriter, r *http.Request) {
  currentSong, err := mpd_conn.CurrentSong()
  if err != nil {

    count := 0;
    for err != nil && count < 10 {
      time.Sleep(10)

      currentSong, err = mpd_conn.CurrentSong()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get current song info for upcoming list", err)
      return
    }
  }

  fmt.Fprintf(w, jsoniffy(currentSong))
}

func getUpcomingSongs(w http.ResponseWriter, r *http.Request) {
  currentSong, err := mpd_conn.CurrentSong()
  if err != nil {

    count := 0;
    for err != nil && count < 10 {
      time.Sleep(10)

      currentSong, err = mpd_conn.CurrentSong()
      count ++
    }

    if err != nil {
      error(w, "Couldn't get current song info for upcoming list", err)
      return
    }
  }

  pos, err := strconv.Atoi(currentSong["Pos"])
  if err != nil {
    error(w, "Couldn't turn current song's position to int", err)
    return
  }

  playlist, err := mpd_conn.PlaylistInfo(-1, -1)
  if err != nil {
    count := 0
    for err != nil && count < 10 {
      time.Sleep(10)

      playlist, err = mpd_conn.PlaylistInfo(-1, -1)
      count ++
    }

    if err != nil {
      error(w, "Couldn't get the current playlist", err)
      return
    }
  }

  upcoming := playlist[pos + 1:]

  fmt.Fprintf(w, jsoniffy(upcoming))
}

/*******************
  Helper Functions
 *******************/

// Connect to MPD's control channel, and set the global mpd_conn to it.
func mpdConnect(url string) *mpd.Client {
  conn, err := mpd.Dial("tcp", url)

  // if we can't connect to MPD everything's fucked, nothing's going to work
  // kill all humans, and die, respectfully, after explaining what the issue
  // is.
  if err != nil {
    log.Println("\n\nServer quiting because it can't connect to MPD");
    log.Println(err)

    return nil
  }

  return conn
}

// helper struct; used to hold some ID3 info, plus an MPD file path.
type TBFile struct {id3.File; FilePath string;}

// helper method, returns a pointer to one of our helper structs (see above).
func id3Read(reader io.Reader, filePath string) *TBFile {
  id3File := id3.Read(reader)

  file := new(TBFile)
  file.Header = id3File.Header
  file.Name = id3File.Name
  file.Artist = id3File.Artist
  file.Album = id3File.Album
  file.Year = id3File.Year
  file.FilePath = filePath

  return file
}

// turn anything into JSON.
func jsoniffy(v interface {}) string {
  obj, err := json.MarshalIndent(v, "", "  ")
  if err != nil {
    log.Print("Couldn't turn something into JSON: ", v)
    log.Fatal(err)
  }

  return string(obj)
}

func error(w http.ResponseWriter, message string, err interface{Error() string;}) {
  log.Println("An error occured; telling the client.")
  log.Println("Message:", message)
  log.Println("Error:", err, "\n\n")

  fmt.Fprintf(w, message + "\n")
}

func jsonError(w http.ResponseWriter, message string, err interface{Error() string;}) {
  message = "{error:\"" + message + "\"}"
  error(w, message, err)
}

<|endoftext|>


<commit_before>package mantle

import (
	"github.com/garyburd/redigo/redis"
	"github.com/youtube/vitess/go/pools"
	"strings"
	"time"
)

//cant make these guys const as []string is not allowed in consts
var PoolSize = 10
var DefaultIpAndHost = []string{"localhost:6379"}

type Redis struct {
	Settings PoolSettings
	pool     *ResourcePool
}

func (r *Redis) SetDefaults() {
	if len(r.Settings.HostAndPorts) == 0 {
		r.Settings.HostAndPorts = DefaultIpAndHost
	}
	if r.Settings.Capacity == 0 {
		r.Settings.Capacity = PoolSize
	}
	if r.Settings.MaxCapacity == 0 {
		r.Settings.MaxCapacity = PoolSize
	}
	r.Settings.Timeout = time.Minute
	r.pool = NewPool(Connect, r.Settings)
}

//Alias to SetDefaults
func (r *Redis) Configure() {
	r.SetDefaults()
}

//Wrapping redis connection
type RedisConn struct {
	redis.Conn
}

//Close a redis connection
func (r *RedisConn) Close() {
	_ = r.Conn.Close()
}

func (r *Redis) GetClient() (*RedisConn, error) {
	connection, err := r.pool.GetConn()
	if err != nil {
		return nil, err
	}
	return connection.(*RedisConn), nil
}

func (r *Redis) PutClient(c *RedisConn) {
	r.pool.PutConn(c)
}

//This method creates a redis connection
func Connect(IpAndHost []string) (pools.Resource, error) {
	if len(IpAndHost) > 1 {
		panic("we can only connect to 1 server at the moment")
	}
	hostNPort := strings.Split(IpAndHost[0], ":")
	cli, err := redis.Dial("tcp", hostNPort[0]+":"+hostNPort[1])
	if err != nil {
		panic(err)
	}
	return &RedisConn{cli}, nil
}

//Generic method to execute any redis call
func (r *Redis) Execute(cmd string, args ...interface{}) (interface{}, error) {
	client, err := r.GetClient()
	if err != nil {
		return nil, err
	}
	defer r.PutClient(client)
	return client.Do(cmd, args...)
}

func (r *Redis) Delete(keys ...interface{}) int {
	value, err := redis.Int(r.Execute("DEL", keys...))
	if err != nil {
		return -1
	}
	return value
}

func (r *Redis) Get(key string) string {
	value, err := redis.String(r.Execute("GET", key))
	if err != nil {
		return ""
	}
	return value
}

func (r *Redis) Set(key string, value interface{}) bool {
	_, err := r.Execute("SET", key, value)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) MGet(keys ...interface{}) []string {
	values, err := redis.Strings(r.Execute("MGET", keys...))
	if err != nil {
		return []string{}
	}
	return values
}

func (r *Redis) MSet(mapOfKeyVal map[string]interface{}) bool {
	_, err := r.Execute("MSET", redis.Args{}.AddFlat(mapOfKeyVal)...)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) Expire(key string, duration int) bool {
	_, err := r.Execute("EXPIRE", key, duration)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) Setex(key string, duration int, val interface{}) bool {
	_, err := r.Execute("SETEX", key, duration, val)
	if err != nil {
		return false
	}
	return true
}
<commit_msg>using settings in a much better way<commit_after>package mantle

import (
	"github.com/garyburd/redigo/redis"
	"github.com/youtube/vitess/go/pools"
	"strings"
	"time"
)

//cant make these guys const as []string is not allowed in consts
var PoolSize = 10
var DefaultIpAndHost = []string{"localhost:6379"}

type Redis struct {
	Settings PoolSettings
	pool     *ResourcePool
}

func (r *Redis) SetDefaults() {
	if len(r.Settings.HostAndPorts) == 0 {
		r.Settings.HostAndPorts = DefaultIpAndHost
	}
	if r.Settings.Capacity == 0 {
		r.Settings.Capacity = PoolSize
	}
	if r.Settings.MaxCapacity == 0 {
		r.Settings.MaxCapacity = PoolSize
	}
	r.Settings.Timeout = time.Minute
	r.pool = NewPool(Connect, r.Settings)
}

//Alias to SetDefaults
func (r *Redis) Configure(settings PoolSettings) {
	r.Settings = settings
	r.SetDefaults()
}

//Wrapping redis connection
type RedisConn struct {
	redis.Conn
}

//Close a redis connection
func (r *RedisConn) Close() {
	_ = r.Conn.Close()
}

func (r *Redis) GetClient() (*RedisConn, error) {
	connection, err := r.pool.GetConn()
	if err != nil {
		return nil, err
	}
	return connection.(*RedisConn), nil
}

func (r *Redis) PutClient(c *RedisConn) {
	r.pool.PutConn(c)
}

//This method creates a redis connection
func Connect(IpAndHost []string) (pools.Resource, error) {
	if len(IpAndHost) > 1 {
		panic("we can only connect to 1 server at the moment")
	}
	hostNPort := strings.Split(IpAndHost[0], ":")
	cli, err := redis.Dial("tcp", hostNPort[0]+":"+hostNPort[1])
	if err != nil {
		panic(err)
	}
	return &RedisConn{cli}, nil
}

//Generic method to execute any redis call
func (r *Redis) Execute(cmd string, args ...interface{}) (interface{}, error) {
	client, err := r.GetClient()
	if err != nil {
		return nil, err
	}
	defer r.PutClient(client)
	return client.Do(cmd, args...)
}

func (r *Redis) Delete(keys ...interface{}) int {
	value, err := redis.Int(r.Execute("DEL", keys...))
	if err != nil {
		return -1
	}
	return value
}

func (r *Redis) Get(key string) string {
	value, err := redis.String(r.Execute("GET", key))
	if err != nil {
		return ""
	}
	return value
}

func (r *Redis) Set(key string, value interface{}) bool {
	_, err := r.Execute("SET", key, value)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) MGet(keys ...interface{}) []string {
	values, err := redis.Strings(r.Execute("MGET", keys...))
	if err != nil {
		return []string{}
	}
	return values
}

func (r *Redis) MSet(mapOfKeyVal map[string]interface{}) bool {
	_, err := r.Execute("MSET", redis.Args{}.AddFlat(mapOfKeyVal)...)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) Expire(key string, duration int) bool {
	_, err := r.Execute("EXPIRE", key, duration)
	if err != nil {
		return false
	}
	return true
}

func (r *Redis) Setex(key string, duration int, val interface{}) bool {
	_, err := r.Execute("SETEX", key, duration, val)
	if err != nil {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package sendgrid

import (
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/sendgrid/rest"
	"github.com/sendgrid/sendgrid-go/helpers/mail"
)

// Version is this client library's current version
const (
	Version        = "3.8.0"
	rateLimitRetry = 5
	rateLimitSleep = 1100
)

type options struct {
	Auth     string
	Endpoint string
	Host     string
	Subuser  string
}

// Client is the Twilio SendGrid Go client
type Client struct {
	rest.Request
}

func (o *options) baseURL() string {
	return o.Host + o.Endpoint
}

// requestNew create Request
// @return [Request] a default request object
func requestNew(options options) rest.Request {
	requestHeaders := map[string]string{
		"Authorization": options.Auth,
		"User-Agent":    "sendgrid/" + Version + ";go",
		"Accept":        "application/json",
	}

	if len(options.Subuser) != 0 {
		requestHeaders["On-Behalf-Of"] = options.Subuser
	}

	return rest.Request{
		BaseURL: options.baseURL(),
		Headers: requestHeaders,
	}
}

// Send sends an email through Twilio SendGrid
func (cl *Client) Send(email *mail.SGMailV3) (*rest.Response, error) {
	cl.Body = mail.GetRequestBody(email)
	return MakeRequest(cl.Request)
}

// DefaultClient is used if no custom HTTP client is defined
var DefaultClient = rest.DefaultClient

// API sets up the request to the Twilio SendGrid API, this is main interface.
// Please use the MakeRequest or MakeRequestAsync functions instead.
// (deprecated)
func API(request rest.Request) (*rest.Response, error) {
	return MakeRequest(request)
}

// MakeRequest attempts a Twilio SendGrid request synchronously.
func MakeRequest(request rest.Request) (*rest.Response, error) {
	return DefaultClient.Send(request)
}

// MakeRequestRetry a synchronous request, but retry in the event of a rate
// limited response.
func MakeRequestRetry(request rest.Request) (*rest.Response, error) {
	retry := 0
	var response *rest.Response
	var err error

	for {
		response, err = MakeRequest(request)
		if err != nil {
			return nil, err
		}

		if response.StatusCode != http.StatusTooManyRequests {
			return response, nil
		}

		if retry > rateLimitRetry {
			return nil, errors.New("rate limit retry exceeded")
		}
		retry++

		resetTime := time.Now().Add(rateLimitSleep * time.Millisecond)

		reset, ok := response.Headers["X-RateLimit-Reset"]
		if ok && len(reset) > 0 {
			t, err := strconv.Atoi(reset[0])
			if err == nil {
				resetTime = time.Unix(int64(t), 0)
			}
		}
		time.Sleep(resetTime.Sub(time.Now()))
	}
}

// MakeRequestAsync attempts a request asynchronously in a new go
// routine. This function returns two channels: responses
// and errors. This function will retry in the case of a
// rate limit.
func MakeRequestAsync(request rest.Request) (chan *rest.Response, chan error) {
	r := make(chan *rest.Response)
	e := make(chan error)

	go func() {
		response, err := MakeRequestRetry(request)
		if err != nil {
			e <- err
		}
		if response != nil {
			r <- response
		}
	}()

	return r, e
}
<commit_msg>Release v3.9.0<commit_after>package sendgrid

import (
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/sendgrid/rest"
	"github.com/sendgrid/sendgrid-go/helpers/mail"
)

// Version is this client library's current version
const (
	Version        = "3.9.0"
	rateLimitRetry = 5
	rateLimitSleep = 1100
)

type options struct {
	Auth     string
	Endpoint string
	Host     string
	Subuser  string
}

// Client is the Twilio SendGrid Go client
type Client struct {
	rest.Request
}

func (o *options) baseURL() string {
	return o.Host + o.Endpoint
}

// requestNew create Request
// @return [Request] a default request object
func requestNew(options options) rest.Request {
	requestHeaders := map[string]string{
		"Authorization": options.Auth,
		"User-Agent":    "sendgrid/" + Version + ";go",
		"Accept":        "application/json",
	}

	if len(options.Subuser) != 0 {
		requestHeaders["On-Behalf-Of"] = options.Subuser
	}

	return rest.Request{
		BaseURL: options.baseURL(),
		Headers: requestHeaders,
	}
}

// Send sends an email through Twilio SendGrid
func (cl *Client) Send(email *mail.SGMailV3) (*rest.Response, error) {
	cl.Body = mail.GetRequestBody(email)
	return MakeRequest(cl.Request)
}

// DefaultClient is used if no custom HTTP client is defined
var DefaultClient = rest.DefaultClient

// API sets up the request to the Twilio SendGrid API, this is main interface.
// Please use the MakeRequest or MakeRequestAsync functions instead.
// (deprecated)
func API(request rest.Request) (*rest.Response, error) {
	return MakeRequest(request)
}

// MakeRequest attempts a Twilio SendGrid request synchronously.
func MakeRequest(request rest.Request) (*rest.Response, error) {
	return DefaultClient.Send(request)
}

// MakeRequestRetry a synchronous request, but retry in the event of a rate
// limited response.
func MakeRequestRetry(request rest.Request) (*rest.Response, error) {
	retry := 0
	var response *rest.Response
	var err error

	for {
		response, err = MakeRequest(request)
		if err != nil {
			return nil, err
		}

		if response.StatusCode != http.StatusTooManyRequests {
			return response, nil
		}

		if retry > rateLimitRetry {
			return nil, errors.New("rate limit retry exceeded")
		}
		retry++

		resetTime := time.Now().Add(rateLimitSleep * time.Millisecond)

		reset, ok := response.Headers["X-RateLimit-Reset"]
		if ok && len(reset) > 0 {
			t, err := strconv.Atoi(reset[0])
			if err == nil {
				resetTime = time.Unix(int64(t), 0)
			}
		}
		time.Sleep(resetTime.Sub(time.Now()))
	}
}

// MakeRequestAsync attempts a request asynchronously in a new go
// routine. This function returns two channels: responses
// and errors. This function will retry in the case of a
// rate limit.
func MakeRequestAsync(request rest.Request) (chan *rest.Response, chan error) {
	r := make(chan *rest.Response)
	e := make(chan error)

	go func() {
		response, err := MakeRequestRetry(request)
		if err != nil {
			e <- err
		}
		if response != nil {
			r <- response
		}
	}()

	return r, e
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package installer // import "k8s.io/helm/pkg/plugin/installer"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"k8s.io/helm/pkg/helm/helmpath"

	"github.com/Masterminds/vcs"
)

var _ Installer = new(VCSInstaller)

type testRepo struct {
	local, remote, current string
	tags, branches         []string
	err                    error
	vcs.Repo
}

func (r *testRepo) LocalPath() string           { return r.local }
func (r *testRepo) Remote() string              { return r.remote }
func (r *testRepo) Update() error               { return r.err }
func (r *testRepo) Get() error                  { return r.err }
func (r *testRepo) IsReference(string) bool     { return false }
func (r *testRepo) Tags() ([]string, error)     { return r.tags, r.err }
func (r *testRepo) Branches() ([]string, error) { return r.branches, r.err }
func (r *testRepo) UpdateVersion(version string) error {
	r.current = version
	return r.err
}

func TestVCSInstaller(t *testing.T) {
	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"
	testRepoPath, _ := filepath.Abs("../testdata/plugdir/echo")
	repo := &testRepo{
		local: testRepoPath,
		tags:  []string{"0.1.0", "0.1.1"},
	}

	i, err := NewForSource(source, "~0.1.0", home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	vcsInstaller, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	// set the testRepo in the VCSInstaller
	vcsInstaller.Repo = repo

	if err := Install(i); err != nil {
		t.Fatal(err)
	}
	if repo.current != "0.1.1" {
		t.Errorf("expected version '0.1.1', got %q", repo.current)
	}
	if i.Path() != home.Path("plugins", "helm-env") {
		t.Errorf("expected path '$HELM_HOME/plugins/helm-env', got %q", i.Path())
	}

	// Install again to test plugin exists error
	if err := Install(i); err == nil {
		t.Error("expected error for plugin exists, got none")
	} else if err.Error() != "plugin already exists" {
		t.Errorf("expected error for plugin exists, got (%v)", err)
	}

	//Testing FindSource method, expect error because plugin code is not a cloned repository
	if _, err := FindSource(i.Path(), home); err == nil {
		t.Error("expected error for inability to find plugin source, got none")
	} else if err.Error() != "cannot get information about plugin source" {
		t.Errorf("expected error for inability to find plugin source, got (%v)", err)
	}
}

func TestVCSInstallerNonExistentVersion(t *testing.T) {
	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"
	version := "0.2.0"

	i, err := NewForSource(source, version, home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	_, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	if err := Install(i); err == nil {
		t.Error("expected error for version does not exists, got none")
	} else if err.Error() != fmt.Sprintf("requested version %q does not exist for plugin %q", version, source) {
		t.Errorf("expected error for version does not exists, got (%v)", err)
	}
}
func TestVCSInstallerUpdate(t *testing.T) {

	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"

	i, err := NewForSource(source, "", home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	_, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	if err := Update(i); err == nil {
		t.Fatal("expected error for plugin does not exist, got none")
	} else if err.Error() != "plugin does not exist" {
		t.Fatalf("expected error for plugin does not exist, got (%v)", err)
	}

	// Install plugin before update
	if err := Install(i); err != nil {
		t.Fatal(err)
	}

	// Test FindSource method for positive result
	pluginInfo, err := FindSource(i.Path(), home)
	if err != nil {
		t.Fatal(err)
	}

	repoRemote := pluginInfo.(*VCSInstaller).Repo.Remote()
	if repoRemote != source {
		t.Fatalf("invalid source found, expected %q got %q", source, repoRemote)
	}

	// Update plugin
	if err := Update(i); err != nil {
		t.Fatal(err)
	}

	// Test update failure
	os.Remove(filepath.Join(i.Path(), "plugin.yaml"))
	// Testing update for error
	if err := Update(i); err == nil {
		t.Error("expected error for plugin modified, got none")
	} else if err.Error() != "plugin repo was modified" {
		t.Errorf("expected error for plugin modified, got (%v)", err)
	}

}
<commit_msg>Correct the returned message<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package installer // import "k8s.io/helm/pkg/plugin/installer"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"k8s.io/helm/pkg/helm/helmpath"

	"github.com/Masterminds/vcs"
)

var _ Installer = new(VCSInstaller)

type testRepo struct {
	local, remote, current string
	tags, branches         []string
	err                    error
	vcs.Repo
}

func (r *testRepo) LocalPath() string           { return r.local }
func (r *testRepo) Remote() string              { return r.remote }
func (r *testRepo) Update() error               { return r.err }
func (r *testRepo) Get() error                  { return r.err }
func (r *testRepo) IsReference(string) bool     { return false }
func (r *testRepo) Tags() ([]string, error)     { return r.tags, r.err }
func (r *testRepo) Branches() ([]string, error) { return r.branches, r.err }
func (r *testRepo) UpdateVersion(version string) error {
	r.current = version
	return r.err
}

func TestVCSInstaller(t *testing.T) {
	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"
	testRepoPath, _ := filepath.Abs("../testdata/plugdir/echo")
	repo := &testRepo{
		local: testRepoPath,
		tags:  []string{"0.1.0", "0.1.1"},
	}

	i, err := NewForSource(source, "~0.1.0", home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	vcsInstaller, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	// set the testRepo in the VCSInstaller
	vcsInstaller.Repo = repo

	if err := Install(i); err != nil {
		t.Fatal(err)
	}
	if repo.current != "0.1.1" {
		t.Errorf("expected version '0.1.1', got %q", repo.current)
	}
	if i.Path() != home.Path("plugins", "helm-env") {
		t.Errorf("expected path '$HELM_HOME/plugins/helm-env', got %q", i.Path())
	}

	// Install again to test plugin exists error
	if err := Install(i); err == nil {
		t.Error("expected error for plugin exists, got none")
	} else if err.Error() != "plugin already exists" {
		t.Errorf("expected error for plugin exists, got (%v)", err)
	}

	//Testing FindSource method, expect error because plugin code is not a cloned repository
	if _, err := FindSource(i.Path(), home); err == nil {
		t.Error("expected error for inability to find plugin source, got none")
	} else if err.Error() != "cannot get information about plugin source" {
		t.Errorf("expected error for inability to find plugin source, got (%v)", err)
	}
}

func TestVCSInstallerNonExistentVersion(t *testing.T) {
	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"
	version := "0.2.0"

	i, err := NewForSource(source, version, home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	_, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	if err := Install(i); err == nil {
		t.Error("expected error for version does not exist, got none")
	} else if err.Error() != fmt.Sprintf("requested version %q does not exist for plugin %q", version, source) {
		t.Errorf("expected error for version does not exist, got (%v)", err)
	}
}
func TestVCSInstallerUpdate(t *testing.T) {

	hh, err := ioutil.TempDir("", "helm-home-")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(hh)

	home := helmpath.Home(hh)
	if err := os.MkdirAll(home.Plugins(), 0755); err != nil {
		t.Fatalf("Could not create %s: %s", home.Plugins(), err)
	}

	source := "https://github.com/adamreese/helm-env"

	i, err := NewForSource(source, "", home)
	if err != nil {
		t.Fatalf("unexpected error: %s", err)
	}

	// ensure a VCSInstaller was returned
	_, ok := i.(*VCSInstaller)
	if !ok {
		t.Fatal("expected a VCSInstaller")
	}

	if err := Update(i); err == nil {
		t.Fatal("expected error for plugin does not exist, got none")
	} else if err.Error() != "plugin does not exist" {
		t.Fatalf("expected error for plugin does not exist, got (%v)", err)
	}

	// Install plugin before update
	if err := Install(i); err != nil {
		t.Fatal(err)
	}

	// Test FindSource method for positive result
	pluginInfo, err := FindSource(i.Path(), home)
	if err != nil {
		t.Fatal(err)
	}

	repoRemote := pluginInfo.(*VCSInstaller).Repo.Remote()
	if repoRemote != source {
		t.Fatalf("invalid source found, expected %q got %q", source, repoRemote)
	}

	// Update plugin
	if err := Update(i); err != nil {
		t.Fatal(err)
	}

	// Test update failure
	os.Remove(filepath.Join(i.Path(), "plugin.yaml"))
	// Testing update for error
	if err := Update(i); err == nil {
		t.Error("expected error for plugin modified, got none")
	} else if err.Error() != "plugin repo was modified" {
		t.Errorf("expected error for plugin modified, got (%v)", err)
	}

}
<|endoftext|>


<commit_before>// Copyright Istio Authors
//
// 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.

// An example implementation of a client.

package forwarder

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/common/scheme"
	"istio.io/istio/pkg/test/echo/proto"
)

type request struct {
	URL         string
	Header      http.Header
	RequestID   int
	Message     string
	Timeout     time.Duration
	ServerFirst bool
	Method      string
}

type protocol interface {
	makeRequest(ctx context.Context, req *request) (string, error)
	Close() error
}

func newProtocol(cfg Config) (protocol, error) {
	var httpDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
	var wsDialContext func(network, addr string) (net.Conn, error)
	if len(cfg.UDS) > 0 {
		httpDialContext = func(_ context.Context, _, _ string) (net.Conn, error) {
			return net.Dial("unix", cfg.UDS)
		}

		wsDialContext = func(_, _ string) (net.Conn, error) {
			return net.Dial("unix", cfg.UDS)
		}
	}

	rawURL := cfg.Request.Url
	u, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("failed parsing request URL %s: %v", cfg.Request.Url, err)
	}

	timeout := common.GetTimeout(cfg.Request)
	headers := common.GetHeaders(cfg.Request)

	var getClientCertificate func(info *tls.CertificateRequestInfo) (*tls.Certificate, error)
	if cfg.Request.Cert != "" && cfg.Request.Key != "" {
		cert, err := tls.X509KeyPair([]byte(cfg.Request.Cert), []byte(cfg.Request.Key))
		if err != nil {
			return nil, fmt.Errorf("failed to parse x509 key pair: %v", err)
		}

		for _, c := range cert.Certificate {
			cert, err := x509.ParseCertificate(c)
			if err != nil {
				fwLog.Errorf("Failed to parse client certificate: %v", err)
			}
			fwLog.Debugf("Using client certificate [%s] issued by %s", cert.SerialNumber, cert.Issuer)
			for _, uri := range cert.URIs {
				fwLog.Debugf("  URI SAN: %s", uri)
			}
		}
		// nolint: unparam
		getClientCertificate = func(info *tls.CertificateRequestInfo) (*tls.Certificate, error) {
			fwLog.Debugf("Peer asking for client certificate")
			for i, ca := range info.AcceptableCAs {
				x := &pkix.RDNSequence{}
				if _, err := asn1.Unmarshal(ca, x); err != nil {
					fwLog.Errorf("Failed to decode AcceptableCA[%d]: %v", i, err)
				} else {
					name := &pkix.Name{}
					name.FillFromRDNSequence(x)
					fwLog.Debugf("  AcceptableCA[%d]: %s", i, name)
				}
			}

			return &cert, nil
		}
	}
	tlsConfig := &tls.Config{
		GetClientCertificate: getClientCertificate,
	}
	if cfg.Request.CaCert != "" {
		certPool := x509.NewCertPool()
		if !certPool.AppendCertsFromPEM([]byte(cfg.Request.CaCert)) {
			return nil, fmt.Errorf("failed to create cert pool")
		}
		tlsConfig.RootCAs = certPool
	} else {
		tlsConfig.InsecureSkipVerify = true
	}

	switch scheme.Instance(u.Scheme) {
	case scheme.HTTP, scheme.HTTPS:
		tlsConfig.NextProtos = []string{"http/1.1"}
		proto := &httpProtocol{
			client: &http.Client{
				Transport: &http.Transport{
					// We are creating a Transport on each ForwardEcho request. Transport is what holds connections,
					// so this means every ForwardEcho request will create a new connection. Without setting an idle timeout,
					// we would never close these connections.
					IdleConnTimeout: time.Second,
					TLSClientConfig: tlsConfig,
					DialContext:     httpDialContext,
				},
				Timeout: timeout,
			},
			do: cfg.Dialer.HTTP,
		}
		if cfg.Request.Http2 && scheme.Instance(u.Scheme) == scheme.HTTPS {
			tlsConfig.NextProtos = []string{"http/2"}
			proto.client.Transport = &http2.Transport{
				TLSClientConfig: tlsConfig,
				DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
					return tls.Dial(network, addr, cfg)
				},
			}
		} else if cfg.Request.Http2 {
			proto.client.Transport = &http2.Transport{
				// Golang doesn't have first class support for h2c, so we provide some workarounds
				// See https://www.mailgun.com/blog/http-2-cleartext-h2c-client-example-go/
				// So http2.Transport doesn't complain the URL scheme isn't 'https'
				AllowHTTP: true,
				// Pretend we are dialing a TLS endpoint. (Note, we ignore the passed tls.Config)
				DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
					return net.Dial(network, addr)
				},
			}
		}
		return proto, nil
	case scheme.GRPC:
		// grpc-go sets incorrect authority header
		authority := headers.Get(hostHeader)

		// transport security
		security := grpc.WithInsecure()
		if getClientCertificate != nil {
			security = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
		}

		// Strip off the scheme from the address.
		address := rawURL[len(u.Scheme+"://"):]

		// Connect to the GRPC server.
		ctx, cancel := context.WithTimeout(context.Background(), common.ConnectionTimeout)
		defer cancel()
		grpcConn, err := cfg.Dialer.GRPC(ctx,
			address,
			security,
			grpc.WithAuthority(authority))
		if err != nil {
			return nil, err
		}
		return &grpcProtocol{
			conn:   grpcConn,
			client: proto.NewEchoTestServiceClient(grpcConn),
		}, nil
	case scheme.WebSocket:
		dialer := &websocket.Dialer{
			TLSClientConfig:  tlsConfig,
			NetDial:          wsDialContext,
			HandshakeTimeout: timeout,
		}
		return &websocketProtocol{
			dialer: dialer,
		}, nil
	case scheme.TCP:
		return &tcpProtocol{
			conn: func() (net.Conn, error) {
				dialer := net.Dialer{
					Timeout: timeout,
				}
				address := rawURL[len(u.Scheme+"://"):]

				ctx, cancel := context.WithTimeout(context.Background(), common.ConnectionTimeout)
				defer cancel()

				if getClientCertificate == nil {
					return cfg.Dialer.TCP(dialer, ctx, address)
				}
				return tls.Dial("tcp", address, tlsConfig)

			},
		}, nil
	}

	return nil, fmt.Errorf("unrecognized protocol %q", u.String())
}
<commit_msg>Fix missing part of ALPN client code (#29568)<commit_after>// Copyright Istio Authors
//
// 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.

// An example implementation of a client.

package forwarder

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/common/scheme"
	"istio.io/istio/pkg/test/echo/proto"
)

type request struct {
	URL         string
	Header      http.Header
	RequestID   int
	Message     string
	Timeout     time.Duration
	ServerFirst bool
	Method      string
}

type protocol interface {
	makeRequest(ctx context.Context, req *request) (string, error)
	Close() error
}

func newProtocol(cfg Config) (protocol, error) {
	var httpDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
	var wsDialContext func(network, addr string) (net.Conn, error)
	if len(cfg.UDS) > 0 {
		httpDialContext = func(_ context.Context, _, _ string) (net.Conn, error) {
			return net.Dial("unix", cfg.UDS)
		}

		wsDialContext = func(_, _ string) (net.Conn, error) {
			return net.Dial("unix", cfg.UDS)
		}
	}

	rawURL := cfg.Request.Url
	u, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("failed parsing request URL %s: %v", cfg.Request.Url, err)
	}

	timeout := common.GetTimeout(cfg.Request)
	headers := common.GetHeaders(cfg.Request)

	var getClientCertificate func(info *tls.CertificateRequestInfo) (*tls.Certificate, error)
	if cfg.Request.Cert != "" && cfg.Request.Key != "" {
		cert, err := tls.X509KeyPair([]byte(cfg.Request.Cert), []byte(cfg.Request.Key))
		if err != nil {
			return nil, fmt.Errorf("failed to parse x509 key pair: %v", err)
		}

		for _, c := range cert.Certificate {
			cert, err := x509.ParseCertificate(c)
			if err != nil {
				fwLog.Errorf("Failed to parse client certificate: %v", err)
			}
			fwLog.Debugf("Using client certificate [%s] issued by %s", cert.SerialNumber, cert.Issuer)
			for _, uri := range cert.URIs {
				fwLog.Debugf("  URI SAN: %s", uri)
			}
		}
		// nolint: unparam
		getClientCertificate = func(info *tls.CertificateRequestInfo) (*tls.Certificate, error) {
			fwLog.Debugf("Peer asking for client certificate")
			for i, ca := range info.AcceptableCAs {
				x := &pkix.RDNSequence{}
				if _, err := asn1.Unmarshal(ca, x); err != nil {
					fwLog.Errorf("Failed to decode AcceptableCA[%d]: %v", i, err)
				} else {
					name := &pkix.Name{}
					name.FillFromRDNSequence(x)
					fwLog.Debugf("  AcceptableCA[%d]: %s", i, name)
				}
			}

			return &cert, nil
		}
	}
	tlsConfig := &tls.Config{
		GetClientCertificate: getClientCertificate,
		NextProtos:           cfg.Request.GetAlpn().GetValue(),
	}
	if cfg.Request.CaCert != "" {
		certPool := x509.NewCertPool()
		if !certPool.AppendCertsFromPEM([]byte(cfg.Request.CaCert)) {
			return nil, fmt.Errorf("failed to create cert pool")
		}
		tlsConfig.RootCAs = certPool
	} else {
		tlsConfig.InsecureSkipVerify = true
	}

	switch scheme.Instance(u.Scheme) {
	case scheme.HTTP, scheme.HTTPS:
		if cfg.Request.Alpn == nil {
			tlsConfig.NextProtos = []string{"http/1.1"}
		}
		proto := &httpProtocol{
			client: &http.Client{
				Transport: &http.Transport{
					// We are creating a Transport on each ForwardEcho request. Transport is what holds connections,
					// so this means every ForwardEcho request will create a new connection. Without setting an idle timeout,
					// we would never close these connections.
					IdleConnTimeout: time.Second,
					TLSClientConfig: tlsConfig,
					DialContext:     httpDialContext,
				},
				Timeout: timeout,
			},
			do: cfg.Dialer.HTTP,
		}
		if cfg.Request.Http2 && scheme.Instance(u.Scheme) == scheme.HTTPS {
			if cfg.Request.Alpn == nil {
				tlsConfig.NextProtos = []string{"h2"}
			}
			proto.client.Transport = &http2.Transport{
				TLSClientConfig: tlsConfig,
				DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
					return tls.Dial(network, addr, cfg)
				},
			}
		} else if cfg.Request.Http2 {
			proto.client.Transport = &http2.Transport{
				// Golang doesn't have first class support for h2c, so we provide some workarounds
				// See https://www.mailgun.com/blog/http-2-cleartext-h2c-client-example-go/
				// So http2.Transport doesn't complain the URL scheme isn't 'https'
				AllowHTTP: true,
				// Pretend we are dialing a TLS endpoint. (Note, we ignore the passed tls.Config)
				DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
					return net.Dial(network, addr)
				},
			}
		}
		return proto, nil
	case scheme.GRPC:
		// grpc-go sets incorrect authority header
		authority := headers.Get(hostHeader)

		// transport security
		security := grpc.WithInsecure()
		if getClientCertificate != nil {
			security = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
		}

		// Strip off the scheme from the address.
		address := rawURL[len(u.Scheme+"://"):]

		// Connect to the GRPC server.
		ctx, cancel := context.WithTimeout(context.Background(), common.ConnectionTimeout)
		defer cancel()
		grpcConn, err := cfg.Dialer.GRPC(ctx,
			address,
			security,
			grpc.WithAuthority(authority))
		if err != nil {
			return nil, err
		}
		return &grpcProtocol{
			conn:   grpcConn,
			client: proto.NewEchoTestServiceClient(grpcConn),
		}, nil
	case scheme.WebSocket:
		dialer := &websocket.Dialer{
			TLSClientConfig:  tlsConfig,
			NetDial:          wsDialContext,
			HandshakeTimeout: timeout,
		}
		return &websocketProtocol{
			dialer: dialer,
		}, nil
	case scheme.TCP:
		return &tcpProtocol{
			conn: func() (net.Conn, error) {
				dialer := net.Dialer{
					Timeout: timeout,
				}
				address := rawURL[len(u.Scheme+"://"):]

				ctx, cancel := context.WithTimeout(context.Background(), common.ConnectionTimeout)
				defer cancel()

				if getClientCertificate == nil {
					return cfg.Dialer.TCP(dialer, ctx, address)
				}
				return tls.Dial("tcp", address, tlsConfig)

			},
		}, nil
	}

	return nil, fmt.Errorf("unrecognized protocol %q", u.String())
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright 2018-2019 Authors of Cilium

package model

import (
	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/labels"
)

func NewIdentityFromModel(base *models.Identity) *identity.Identity {
	if base == nil {
		return nil
	}

	id := &identity.Identity{
		ID:     identity.NumericIdentity(base.ID),
		Labels: make(labels.Labels),
	}
	for _, v := range base.Labels {
		lbl := labels.ParseLabel(v)
		id.Labels[lbl.Key] = lbl
	}
	id.Sanitize()

	return id
}

func CreateModel(id *identity.Identity) *models.Identity {
	if id == nil {
		return nil
	}

	ret := &models.Identity{
		ID:           int64(id.ID),
		Labels:       []string{},
		LabelsSHA256: "",
	}

	for _, v := range id.Labels {
		ret.Labels = append(ret.Labels, v.String())
	}
	ret.LabelsSHA256 = id.GetLabelsSHA256()
	return ret
}
<commit_msg>identity/model: preallocate slice and map with known size<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright 2018-2019 Authors of Cilium

package model

import (
	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/labels"
)

func NewIdentityFromModel(base *models.Identity) *identity.Identity {
	if base == nil {
		return nil
	}

	id := &identity.Identity{
		ID:     identity.NumericIdentity(base.ID),
		Labels: make(labels.Labels, len(base.Labels)),
	}
	for _, v := range base.Labels {
		lbl := labels.ParseLabel(v)
		id.Labels[lbl.Key] = lbl
	}
	id.Sanitize()

	return id
}

func CreateModel(id *identity.Identity) *models.Identity {
	if id == nil {
		return nil
	}

	ret := &models.Identity{
		ID:           int64(id.ID),
		Labels:       make([]string, 0, len(id.Labels)),
		LabelsSHA256: id.GetLabelsSHA256(),
	}

	for _, v := range id.Labels {
		ret.Labels = append(ret.Labels, v.String())
	}
	return ret
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package resource

import (
	"fmt"
	"reflect"

	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apimachinery/registered"
	"k8s.io/kubernetes/pkg/registry/thirdpartyresourcedata"
	"k8s.io/kubernetes/pkg/runtime"
)

// DisabledClientForMapping allows callers to avoid allowing remote calls when handling
// resources.
type DisabledClientForMapping struct {
	ClientMapper
}

func (f DisabledClientForMapping) ClientForMapping(mapping *meta.RESTMapping) (RESTClient, error) {
	return nil, nil
}

// Mapper is a convenience struct for holding references to the three interfaces
// needed to create Info for arbitrary objects.
type Mapper struct {
	runtime.ObjectTyper
	meta.RESTMapper
	ClientMapper
	runtime.Decoder
}

// InfoForData creates an Info object for the given data. An error is returned
// if any of the decoding or client lookup steps fail. Name and namespace will be
// set into Info if the mapping's MetadataAccessor can retrieve them.
func (m *Mapper) InfoForData(data []byte, source string) (*Info, error) {
	versions := &runtime.VersionedObjects{}
	_, gvk, err := m.Decode(data, nil, versions)
	var obj runtime.Object
	var versioned runtime.Object
	if registered.IsThirdPartyAPIGroupVersion(gvk.GroupVersion()) {
		obj, err = runtime.Decode(thirdpartyresourcedata.NewCodec(nil, gvk.Kind), data)
		versioned = obj
	} else {
		obj, versioned = versions.Last(), versions.First()
	}
	if err != nil {
		return nil, fmt.Errorf("unable to decode %q: %v [%v]", source, err, gvk)
	}
	mapping, err := m.RESTMapping(gvk.GroupKind(), gvk.Version)
	if err != nil {
		return nil, fmt.Errorf("unable to recognize %q: %v", source, err)
	}

	client, err := m.ClientForMapping(mapping)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to a server to handle %q: %v", mapping.Resource, err)
	}

	name, _ := mapping.MetadataAccessor.Name(obj)
	namespace, _ := mapping.MetadataAccessor.Namespace(obj)
	resourceVersion, _ := mapping.MetadataAccessor.ResourceVersion(obj)

	return &Info{
		Mapping:         mapping,
		Client:          client,
		Namespace:       namespace,
		Name:            name,
		Source:          source,
		VersionedObject: versioned,
		Object:          obj,
		ResourceVersion: resourceVersion,
	}, nil
}

// InfoForObject creates an Info object for the given Object. An error is returned
// if the object cannot be introspected. Name and namespace will be set into Info
// if the mapping's MetadataAccessor can retrieve them.
func (m *Mapper) InfoForObject(obj runtime.Object, preferredGVKs []unversioned.GroupVersionKind) (*Info, error) {
	groupVersionKinds, err := m.ObjectKinds(obj)
	if err != nil {
		return nil, fmt.Errorf("unable to get type info from the object %q: %v", reflect.TypeOf(obj), err)
	}

	groupVersionKind := groupVersionKinds[0]
	if len(groupVersionKinds) > 1 && len(preferredGVKs) > 0 {
		groupVersionKind = preferredObjectKind(groupVersionKinds, preferredGVKs)
	}

	mapping, err := m.RESTMapping(groupVersionKind.GroupKind(), groupVersionKind.Version)
	if err != nil {
		return nil, fmt.Errorf("unable to recognize %v: %v", groupVersionKind, err)
	}
	client, err := m.ClientForMapping(mapping)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to a server to handle %q: %v", mapping.Resource, err)
	}
	name, _ := mapping.MetadataAccessor.Name(obj)
	namespace, _ := mapping.MetadataAccessor.Namespace(obj)
	resourceVersion, _ := mapping.MetadataAccessor.ResourceVersion(obj)
	return &Info{
		Mapping:   mapping,
		Client:    client,
		Namespace: namespace,
		Name:      name,

		Object:          obj,
		ResourceVersion: resourceVersion,
	}, nil
}

// preferredObjectKind picks the possibility that most closely matches the priority list in this order:
// GroupVersionKind matches (exact match)
// GroupKind matches
// Group matches
func preferredObjectKind(possibilities []unversioned.GroupVersionKind, preferences []unversioned.GroupVersionKind) unversioned.GroupVersionKind {
	// Exact match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility == priority {
				return possibility
			}
		}
	}

	// GroupKind match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility.GroupKind() == priority.GroupKind() {
				return possibility
			}
		}
	}

	// Group match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility.Group == priority.Group {
				return possibility
			}
		}
	}

	// Just pick the first
	return possibilities[0]
}
<commit_msg>kubectl: avoid panic with nil gvk<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package resource

import (
	"fmt"
	"reflect"

	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apimachinery/registered"
	"k8s.io/kubernetes/pkg/registry/thirdpartyresourcedata"
	"k8s.io/kubernetes/pkg/runtime"
)

// DisabledClientForMapping allows callers to avoid allowing remote calls when handling
// resources.
type DisabledClientForMapping struct {
	ClientMapper
}

func (f DisabledClientForMapping) ClientForMapping(mapping *meta.RESTMapping) (RESTClient, error) {
	return nil, nil
}

// Mapper is a convenience struct for holding references to the three interfaces
// needed to create Info for arbitrary objects.
type Mapper struct {
	runtime.ObjectTyper
	meta.RESTMapper
	ClientMapper
	runtime.Decoder
}

// InfoForData creates an Info object for the given data. An error is returned
// if any of the decoding or client lookup steps fail. Name and namespace will be
// set into Info if the mapping's MetadataAccessor can retrieve them.
func (m *Mapper) InfoForData(data []byte, source string) (*Info, error) {
	versions := &runtime.VersionedObjects{}
	_, gvk, err := m.Decode(data, nil, versions)
	if err != nil {
		return nil, fmt.Errorf("unable to decode %q: %v", source, err)
	}
	var obj runtime.Object
	var versioned runtime.Object
	if registered.IsThirdPartyAPIGroupVersion(gvk.GroupVersion()) {
		obj, err = runtime.Decode(thirdpartyresourcedata.NewCodec(nil, gvk.Kind), data)
		versioned = obj
	} else {
		obj, versioned = versions.Last(), versions.First()
	}
	if err != nil {
		return nil, fmt.Errorf("unable to decode %q: %v [%v]", source, err, gvk)
	}
	mapping, err := m.RESTMapping(gvk.GroupKind(), gvk.Version)
	if err != nil {
		return nil, fmt.Errorf("unable to recognize %q: %v", source, err)
	}

	client, err := m.ClientForMapping(mapping)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to a server to handle %q: %v", mapping.Resource, err)
	}

	name, _ := mapping.MetadataAccessor.Name(obj)
	namespace, _ := mapping.MetadataAccessor.Namespace(obj)
	resourceVersion, _ := mapping.MetadataAccessor.ResourceVersion(obj)

	return &Info{
		Mapping:         mapping,
		Client:          client,
		Namespace:       namespace,
		Name:            name,
		Source:          source,
		VersionedObject: versioned,
		Object:          obj,
		ResourceVersion: resourceVersion,
	}, nil
}

// InfoForObject creates an Info object for the given Object. An error is returned
// if the object cannot be introspected. Name and namespace will be set into Info
// if the mapping's MetadataAccessor can retrieve them.
func (m *Mapper) InfoForObject(obj runtime.Object, preferredGVKs []unversioned.GroupVersionKind) (*Info, error) {
	groupVersionKinds, err := m.ObjectKinds(obj)
	if err != nil {
		return nil, fmt.Errorf("unable to get type info from the object %q: %v", reflect.TypeOf(obj), err)
	}

	groupVersionKind := groupVersionKinds[0]
	if len(groupVersionKinds) > 1 && len(preferredGVKs) > 0 {
		groupVersionKind = preferredObjectKind(groupVersionKinds, preferredGVKs)
	}

	mapping, err := m.RESTMapping(groupVersionKind.GroupKind(), groupVersionKind.Version)
	if err != nil {
		return nil, fmt.Errorf("unable to recognize %v: %v", groupVersionKind, err)
	}
	client, err := m.ClientForMapping(mapping)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to a server to handle %q: %v", mapping.Resource, err)
	}
	name, _ := mapping.MetadataAccessor.Name(obj)
	namespace, _ := mapping.MetadataAccessor.Namespace(obj)
	resourceVersion, _ := mapping.MetadataAccessor.ResourceVersion(obj)
	return &Info{
		Mapping:   mapping,
		Client:    client,
		Namespace: namespace,
		Name:      name,

		Object:          obj,
		ResourceVersion: resourceVersion,
	}, nil
}

// preferredObjectKind picks the possibility that most closely matches the priority list in this order:
// GroupVersionKind matches (exact match)
// GroupKind matches
// Group matches
func preferredObjectKind(possibilities []unversioned.GroupVersionKind, preferences []unversioned.GroupVersionKind) unversioned.GroupVersionKind {
	// Exact match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility == priority {
				return possibility
			}
		}
	}

	// GroupKind match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility.GroupKind() == priority.GroupKind() {
				return possibility
			}
		}
	}

	// Group match
	for _, priority := range preferences {
		for _, possibility := range possibilities {
			if possibility.Group == priority.Group {
				return possibility
			}
		}
	}

	// Just pick the first
	return possibilities[0]
}
<|endoftext|>


<commit_before>package proxyquerier

import (
	"github.com/prometheus/prometheus/storage"
)

// NewSeriesSet returns a SeriesSet for the given series
func NewSeriesSet(series []storage.Series, warnings storage.Warnings, err error) *SeriesSet {
	return &SeriesSet{
		series: series,
	}
}

// SeriesSet implements prometheus' SeriesSet interface
type SeriesSet struct {
	offset int // 0 means we haven't seen anything
	series []storage.Series

	err      error
	warnings storage.Warnings
}

// Next will attempt to move the iterator up
func (s *SeriesSet) Next() bool {
	if s.offset < len(s.series) {
		s.offset++
		return true
	}
	return false
}

// At returns the current Series for this iterator
func (s *SeriesSet) At() storage.Series {
	return s.series[s.offset-1]
}

// Err returns any error found in this iterator
func (s *SeriesSet) Err() error {
	return s.err
}

// A collection of warnings for the whole set.
// Warnings could be return even iteration has not failed with error.
func (s *SeriesSet) Warnings() storage.Warnings {
	return s.warnings
}
<commit_msg>Correctly pass down warnings/errors to SeriesSet<commit_after>package proxyquerier

import (
	"github.com/prometheus/prometheus/storage"
)

// NewSeriesSet returns a SeriesSet for the given series
func NewSeriesSet(series []storage.Series, warnings storage.Warnings, err error) *SeriesSet {
	return &SeriesSet{
		series:   series,
		warnings: warnings,
		err:      err,
	}
}

// SeriesSet implements prometheus' SeriesSet interface
type SeriesSet struct {
	offset int // 0 means we haven't seen anything
	series []storage.Series

	err      error
	warnings storage.Warnings
}

// Next will attempt to move the iterator up
func (s *SeriesSet) Next() bool {
	if s.offset < len(s.series) {
		s.offset++
		return true
	}
	return false
}

// At returns the current Series for this iterator
func (s *SeriesSet) At() storage.Series {
	return s.series[s.offset-1]
}

// Err returns any error found in this iterator
func (s *SeriesSet) Err() error {
	return s.err
}

// A collection of warnings for the whole set.
// Warnings could be return even iteration has not failed with error.
func (s *SeriesSet) Warnings() storage.Warnings {
	return s.warnings
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"fmt"
	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	bus.AddHandler("sql", GetOrgQuotaByTarget)
	bus.AddHandler("sql", GetOrgQuotas)
	bus.AddHandler("sql", UpdateOrgQuota)
	bus.AddHandler("sql", GetUserQuotaByTarget)
	bus.AddHandler("sql", GetUserQuotas)
	bus.AddHandler("sql", UpdateUserQuota)
	bus.AddHandler("sql", GetGlobalQuotaByTarget)
}

type targetCount struct {
	Count int64
}

func GetOrgQuotaByTarget(query *m.GetOrgQuotaByTargetQuery) error {
	quota := m.Quota{
		Target: query.Target,
		OrgId:  query.OrgId,
	}
	has, err := x.Get(quota)
	if err != nil {
		return err
	} else if has == false {
		quota.Limit = query.Default
	}

	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql, query.OrgId).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.OrgQuotaDTO{
		Target: query.Target,
		Limit:  quota.Limit,
		OrgId:  query.OrgId,
		Used:   resp[0].Count,
	}

	return nil
}

func GetOrgQuotas(query *m.GetOrgQuotasQuery) error {
	quotas := make([]*m.Quota, 0)
	sess := x.Table("quota")
	if err := sess.Where("org_id=? AND user_id=0", query.OrgId).Find(&quotas); err != nil {
		return err
	}

	defaultQuotas := m.QuotaToMap(setting.Quota.Org)

	seenTargets := make(map[string]bool)
	for _, q := range quotas {
		seenTargets[q.Target] = true
	}

	for t, v := range defaultQuotas {
		if _, ok := seenTargets[t]; !ok {
			quotas = append(quotas, &m.Quota{
				OrgId:  query.OrgId,
				Target: t,
				Limit:  v,
			})
		}
	}

	result := make([]*m.OrgQuotaDTO, len(quotas))
	for i, q := range quotas {
		//get quota used.
		rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
		resp := make([]*targetCount, 0)
		if err := x.Sql(rawSql, q.OrgId).Find(&resp); err != nil {
			return err
		}
		result[i] = &m.OrgQuotaDTO{
			Target: q.Target,
			Limit:  q.Limit,
			OrgId:  q.OrgId,
			Used:   resp[0].Count,
		}
	}
	query.Result = result
	return nil
}

func UpdateOrgQuota(cmd *m.UpdateOrgQuotaCmd) error {
	return inTransaction2(func(sess *session) error {
		//Check if quota is already defined in the DB
		quota := m.Quota{
			Target: cmd.Target,
			OrgId:  cmd.OrgId,
		}
		has, err := sess.Get(quota)
		if err != nil {
			return err
		}
		quota.Limit = cmd.Limit
		if has == false {
			//No quota in the DB for this target, so create a new one.
			if _, err := sess.Insert(&quota); err != nil {
				return err
			}
		} else {
			//update existing quota entry in the DB.
			if _, err := sess.Id(quota.Id).Update(&quota); err != nil {
				return err
			}
		}

		return nil
	})
}

func GetUserQuotaByTarget(query *m.GetUserQuotaByTargetQuery) error {
	quota := m.Quota{
		Target: query.Target,
		UserId: query.UserId,
	}
	has, err := x.Get(quota)
	if err != nil {
		return err
	} else if has == false {
		quota.Limit = query.Default
	}

	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql, query.UserId).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.UserQuotaDTO{
		Target: query.Target,
		Limit:  quota.Limit,
		UserId: query.UserId,
		Used:   resp[0].Count,
	}

	return nil
}

func GetUserQuotas(query *m.GetUserQuotasQuery) error {
	quotas := make([]*m.Quota, 0)
	sess := x.Table("quota")
	if err := sess.Where("user_id=? AND org_id=0", query.UserId).Find(&quotas); err != nil {
		return err
	}

	defaultQuotas := m.QuotaToMap(setting.Quota.User)

	seenTargets := make(map[string]bool)
	for _, q := range quotas {
		seenTargets[q.Target] = true
	}

	for t, v := range defaultQuotas {
		if _, ok := seenTargets[t]; !ok {
			quotas = append(quotas, &m.Quota{
				UserId: query.UserId,
				Target: t,
				Limit:  v,
			})
		}
	}

	result := make([]*m.UserQuotaDTO, len(quotas))
	for i, q := range quotas {
		//get quota used.
		rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(q.Target))
		resp := make([]*targetCount, 0)
		if err := x.Sql(rawSql, q.UserId).Find(&resp); err != nil {
			return err
		}
		result[i] = &m.UserQuotaDTO{
			Target: q.Target,
			Limit:  q.Limit,
			UserId: q.UserId,
			Used:   resp[0].Count,
		}
	}
	query.Result = result
	return nil
}

func UpdateUserQuota(cmd *m.UpdateUserQuotaCmd) error {
	return inTransaction2(func(sess *session) error {
		//Check if quota is already defined in the DB
		quota := m.Quota{
			Target: cmd.Target,
			UserId: cmd.UserId,
		}
		has, err := sess.Get(quota)
		if err != nil {
			return err
		}
		quota.Limit = cmd.Limit
		if has == false {
			//No quota in the DB for this target, so create a new one.
			if _, err := sess.Insert(&quota); err != nil {
				return err
			}
		} else {
			//update existing quota entry in the DB.
			if _, err := sess.Id(quota.Id).Update(&quota); err != nil {
				return err
			}
		}

		return nil
	})
}

func GetGlobalQuotaByTarget(query *m.GetGlobalQuotaByTargetQuery) error {
	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.GlobalQuotaDTO{
		Target: query.Target,
		Limit:  query.Default,
		Used:   resp[0].Count,
	}

	return nil
}
<commit_msg>be sure to pass result obj by reference to xorm.<commit_after>package sqlstore

import (
	"fmt"
	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	bus.AddHandler("sql", GetOrgQuotaByTarget)
	bus.AddHandler("sql", GetOrgQuotas)
	bus.AddHandler("sql", UpdateOrgQuota)
	bus.AddHandler("sql", GetUserQuotaByTarget)
	bus.AddHandler("sql", GetUserQuotas)
	bus.AddHandler("sql", UpdateUserQuota)
	bus.AddHandler("sql", GetGlobalQuotaByTarget)
}

type targetCount struct {
	Count int64
}

func GetOrgQuotaByTarget(query *m.GetOrgQuotaByTargetQuery) error {
	quota := m.Quota{
		Target: query.Target,
		OrgId:  query.OrgId,
	}
	has, err := x.Get(&quota)
	if err != nil {
		return err
	} else if has == false {
		quota.Limit = query.Default
	}

	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql, query.OrgId).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.OrgQuotaDTO{
		Target: query.Target,
		Limit:  quota.Limit,
		OrgId:  query.OrgId,
		Used:   resp[0].Count,
	}

	return nil
}

func GetOrgQuotas(query *m.GetOrgQuotasQuery) error {
	quotas := make([]*m.Quota, 0)
	sess := x.Table("quota")
	if err := sess.Where("org_id=? AND user_id=0", query.OrgId).Find(&quotas); err != nil {
		return err
	}

	defaultQuotas := m.QuotaToMap(setting.Quota.Org)

	seenTargets := make(map[string]bool)
	for _, q := range quotas {
		seenTargets[q.Target] = true
	}

	for t, v := range defaultQuotas {
		if _, ok := seenTargets[t]; !ok {
			quotas = append(quotas, &m.Quota{
				OrgId:  query.OrgId,
				Target: t,
				Limit:  v,
			})
		}
	}

	result := make([]*m.OrgQuotaDTO, len(quotas))
	for i, q := range quotas {
		//get quota used.
		rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
		resp := make([]*targetCount, 0)
		if err := x.Sql(rawSql, q.OrgId).Find(&resp); err != nil {
			return err
		}
		result[i] = &m.OrgQuotaDTO{
			Target: q.Target,
			Limit:  q.Limit,
			OrgId:  q.OrgId,
			Used:   resp[0].Count,
		}
	}
	query.Result = result
	return nil
}

func UpdateOrgQuota(cmd *m.UpdateOrgQuotaCmd) error {
	return inTransaction2(func(sess *session) error {
		//Check if quota is already defined in the DB
		quota := m.Quota{
			Target: cmd.Target,
			OrgId:  cmd.OrgId,
		}
		has, err := sess.Get(&quota)
		if err != nil {
			return err
		}
		quota.Limit = cmd.Limit
		if has == false {
			//No quota in the DB for this target, so create a new one.
			if _, err := sess.Insert(&quota); err != nil {
				return err
			}
		} else {
			//update existing quota entry in the DB.
			if _, err := sess.Id(quota.Id).Update(&quota); err != nil {
				return err
			}
		}

		return nil
	})
}

func GetUserQuotaByTarget(query *m.GetUserQuotaByTargetQuery) error {
	quota := m.Quota{
		Target: query.Target,
		UserId: query.UserId,
	}
	has, err := x.Get(&quota)
	if err != nil {
		return err
	} else if has == false {
		quota.Limit = query.Default
	}

	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql, query.UserId).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.UserQuotaDTO{
		Target: query.Target,
		Limit:  quota.Limit,
		UserId: query.UserId,
		Used:   resp[0].Count,
	}

	return nil
}

func GetUserQuotas(query *m.GetUserQuotasQuery) error {
	quotas := make([]*m.Quota, 0)
	sess := x.Table("quota")
	if err := sess.Where("user_id=? AND org_id=0", query.UserId).Find(&quotas); err != nil {
		return err
	}

	defaultQuotas := m.QuotaToMap(setting.Quota.User)

	seenTargets := make(map[string]bool)
	for _, q := range quotas {
		seenTargets[q.Target] = true
	}

	for t, v := range defaultQuotas {
		if _, ok := seenTargets[t]; !ok {
			quotas = append(quotas, &m.Quota{
				UserId: query.UserId,
				Target: t,
				Limit:  v,
			})
		}
	}

	result := make([]*m.UserQuotaDTO, len(quotas))
	for i, q := range quotas {
		//get quota used.
		rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(q.Target))
		resp := make([]*targetCount, 0)
		if err := x.Sql(rawSql, q.UserId).Find(&resp); err != nil {
			return err
		}
		result[i] = &m.UserQuotaDTO{
			Target: q.Target,
			Limit:  q.Limit,
			UserId: q.UserId,
			Used:   resp[0].Count,
		}
	}
	query.Result = result
	return nil
}

func UpdateUserQuota(cmd *m.UpdateUserQuotaCmd) error {
	return inTransaction2(func(sess *session) error {
		//Check if quota is already defined in the DB
		quota := m.Quota{
			Target: cmd.Target,
			UserId: cmd.UserId,
		}
		has, err := sess.Get(&quota)
		if err != nil {
			return err
		}
		quota.Limit = cmd.Limit
		if has == false {
			//No quota in the DB for this target, so create a new one.
			if _, err := sess.Insert(&quota); err != nil {
				return err
			}
		} else {
			//update existing quota entry in the DB.
			if _, err := sess.Id(quota.Id).Update(&quota); err != nil {
				return err
			}
		}

		return nil
	})
}

func GetGlobalQuotaByTarget(query *m.GetGlobalQuotaByTargetQuery) error {
	//get quota used.
	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s", dialect.Quote(query.Target))
	resp := make([]*targetCount, 0)
	if err := x.Sql(rawSql).Find(&resp); err != nil {
		return err
	}

	query.Result = &m.GlobalQuotaDTO{
		Target: query.Target,
		Limit:  query.Default,
		Used:   resp[0].Count,
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/hashicorp/nomad/api"
)

func (n *NomadRegion) watchNodes() {
	q := &api.QueryOptions{WaitIndex: 1}
	for {
		nodes, meta, err := n.Client.Nodes().List(q)
		if err != nil {
			logger.Errorf("watch: unable to fetch nodes: %s", err)
			time.Sleep(10 * time.Second)
			continue
		}

		remoteWaitIndex := meta.LastIndex
		localWaitIndex := q.WaitIndex

		// only work if the WaitIndex have changed
		if remoteWaitIndex == localWaitIndex {
			logger.Debugf("Nodes wait-index is unchanged (%d <> %d)", localWaitIndex, remoteWaitIndex)
			continue
		}

		n.nodes = nodes
		n.broadcastChannels.nodes.Update(&Action{Type: fetchedNodes, Payload: nodes, Index: remoteWaitIndex})
		q = &api.QueryOptions{WaitIndex: remoteWaitIndex}
	}
}
<commit_msg>[nomad] sort clients by name for easier overview<commit_after>package main

import (
	"sort"
	"time"

	api "github.com/hashicorp/nomad/api"
)

// ClientNameSorter sorts planets by name
type ClientNameSorter []*api.NodeListStub

func (a ClientNameSorter) Len() int           { return len(a) }
func (a ClientNameSorter) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ClientNameSorter) Less(i, j int) bool { return a[i].Name < a[j].Name }

func (n *NomadRegion) watchNodes() {
	q := &api.QueryOptions{WaitIndex: 1}
	for {
		nodes, meta, err := n.Client.Nodes().List(q)
		if err != nil {
			logger.Errorf("watch: unable to fetch nodes: %s", err)
			time.Sleep(10 * time.Second)
			continue
		}

		remoteWaitIndex := meta.LastIndex
		localWaitIndex := q.WaitIndex

		// only work if the WaitIndex have changed
		if remoteWaitIndex == localWaitIndex {
			logger.Debugf("Nodes wait-index is unchanged (%d <> %d)", localWaitIndex, remoteWaitIndex)
			continue
		}

		// http://stackoverflow.com/a/28999886
		sort.Sort(ClientNameSorter(nodes))

		n.nodes = nodes
		n.broadcastChannels.nodes.Update(&Action{Type: fetchedNodes, Payload: nodes, Index: remoteWaitIndex})
		q = &api.QueryOptions{WaitIndex: remoteWaitIndex}
	}
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"bytes"
	"database/sql"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
)

type TB testing.B

func (tb *TB) check(err error) {
	if err != nil {
		tb.Fatal(err)
	}
}

func (tb *TB) checkDB(db *sql.DB, err error) *sql.DB {
	tb.check(err)
	return db
}

func (tb *TB) checkRows(rows *sql.Rows, err error) *sql.Rows {
	tb.check(err)
	return rows
}

func (tb *TB) checkStmt(stmt *sql.Stmt, err error) *sql.Stmt {
	tb.check(err)
	return stmt
}

func initDB(b *testing.B, queries ...string) *sql.DB {
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	for _, query := range queries {
		if _, err := db.Exec(query); err != nil {
			b.Fatalf("Error on %q: %v", query, err)
		}
	}
	return db
}

// by Brad Fitzpatrick
const concurrencyLevel = 10

func BenchmarkQuery(b *testing.B) {
	tb := (*TB)(b)
	b.StopTimer()
	b.ReportAllocs()
	db := initDB(b,
		"DROP TABLE IF EXISTS foo",
		"CREATE TABLE foo (id INT PRIMARY KEY, val CHAR(50))",
		`INSERT INTO foo VALUES (1, "one")`,
		`INSERT INTO foo VALUES (2, "two")`,
	)
	db.SetMaxIdleConns(concurrencyLevel)
	defer db.Close()

	stmt := tb.checkStmt(db.Prepare("SELECT val FROM foo WHERE id=?"))
	defer stmt.Close()
	b.StartTimer()

	remain := int64(b.N)
	var wg sync.WaitGroup
	wg.Add(concurrencyLevel)
	defer wg.Wait()
	for i := 0; i < concurrencyLevel; i++ {
		go func() {
			defer wg.Done()
			for {
				if atomic.AddInt64(&remain, -1) < 0 {
					return
				}
				var got string
				tb.check(stmt.QueryRow(1).Scan(&got))
				if got != "one" {
					b.Errorf("query = %q; want one", got)
					return
				}
			}
		}()
	}
}

// data, but no db writes
var roundtripSample []byte

func initRoundtripBenchmarks() ([]byte, int, int) {
	if roundtripSample == nil {
		roundtripSample = []byte(strings.Repeat("0123456789abcdef", 1024*1024))
	}
	return roundtripSample, 16, len(roundtripSample)
}

func BenchmarkRoundtripTxt(b *testing.B) {
	b.StopTimer()
	sample, min, max := initRoundtripBenchmarks()
	sampleString := string(sample)
	b.ReportAllocs()
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	defer db.Close()
	b.StartTimer()
	var result string
	for i := 0; i < b.N; i++ {
		length := min + i
		if length > max {
			length = max
		}
		test := sampleString[0:length]
		rows := tb.checkRows(db.Query(`SELECT "` + test + `"`))
		if !rows.Next() {
			rows.Close()
			b.Fatalf("crashed")
		}
		err := rows.Scan(&result)
		if err != nil {
			rows.Close()
			b.Fatalf("crashed")
		}
		if result != test {
			rows.Close()
			b.Errorf("mismatch")
		}
		rows.Close()
	}
}

func BenchmarkRoundtripBin(b *testing.B) {
	b.StopTimer()
	sample, min, max := initRoundtripBenchmarks()
	b.ReportAllocs()
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	defer db.Close()
	stmt := tb.checkStmt(db.Prepare("SELECT ?"))
	defer stmt.Close()
	b.StartTimer()
	var result sql.RawBytes
	for i := 0; i < b.N; i++ {
		length := min + i
		if length > max {
			length = max
		}
		test := sample[0:length]
		rows := tb.checkRows(stmt.Query(test))
		if !rows.Next() {
			rows.Close()
			b.Fatalf("crashed")
		}
		err := rows.Scan(&result)
		if err != nil {
			rows.Close()
			b.Fatalf("crashed")
		}
		if !bytes.Equal(result, test) {
			rows.Close()
			b.Errorf("mismatch")
		}
		rows.Close()
	}
}
<commit_msg>Add an Exec Benchmark<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"bytes"
	"database/sql"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
)

type TB testing.B

func (tb *TB) check(err error) {
	if err != nil {
		tb.Fatal(err)
	}
}

func (tb *TB) checkDB(db *sql.DB, err error) *sql.DB {
	tb.check(err)
	return db
}

func (tb *TB) checkRows(rows *sql.Rows, err error) *sql.Rows {
	tb.check(err)
	return rows
}

func (tb *TB) checkStmt(stmt *sql.Stmt, err error) *sql.Stmt {
	tb.check(err)
	return stmt
}

func initDB(b *testing.B, queries ...string) *sql.DB {
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	for _, query := range queries {
		if _, err := db.Exec(query); err != nil {
			b.Fatalf("Error on %q: %v", query, err)
		}
	}
	return db
}

const concurrencyLevel = 10

func BenchmarkQuery(b *testing.B) {
	tb := (*TB)(b)
	b.StopTimer()
	b.ReportAllocs()
	db := initDB(b,
		"DROP TABLE IF EXISTS foo",
		"CREATE TABLE foo (id INT PRIMARY KEY, val CHAR(50))",
		`INSERT INTO foo VALUES (1, "one")`,
		`INSERT INTO foo VALUES (2, "two")`,
	)
	db.SetMaxIdleConns(concurrencyLevel)
	defer db.Close()

	stmt := tb.checkStmt(db.Prepare("SELECT val FROM foo WHERE id=?"))
	defer stmt.Close()
	b.StartTimer()

	remain := int64(b.N)
	var wg sync.WaitGroup
	wg.Add(concurrencyLevel)
	defer wg.Wait()
	for i := 0; i < concurrencyLevel; i++ {
		go func() {
			defer wg.Done()
			for {
				if atomic.AddInt64(&remain, -1) < 0 {
					return
				}
				var got string
				tb.check(stmt.QueryRow(1).Scan(&got))
				if got != "one" {
					b.Errorf("query = %q; want one", got)
					return
				}
			}
		}()
	}
}

func BenchmarkExec(b *testing.B) {
	tb := (*TB)(b)
	b.StopTimer()
	b.ReportAllocs()
	db := tb.checkDB(sql.Open("mysql", dsn))
	db.SetMaxIdleConns(concurrencyLevel)
	defer db.Close()

	stmt := tb.checkStmt(db.Prepare("DO 1"))
	defer stmt.Close()

	remain := int64(b.N)
	var wg sync.WaitGroup
	wg.Add(concurrencyLevel)
	defer wg.Wait()
	b.StartTimer()

	for i := 0; i < concurrencyLevel; i++ {
		go func() {
			for {
				if atomic.AddInt64(&remain, -1) < 0 {
					wg.Done()
					return
				}

				if _, err := stmt.Exec(); err != nil {
					b.Fatal(err.Error())
				}
			}
		}()
	}
}

// data, but no db writes
var roundtripSample []byte

func initRoundtripBenchmarks() ([]byte, int, int) {
	if roundtripSample == nil {
		roundtripSample = []byte(strings.Repeat("0123456789abcdef", 1024*1024))
	}
	return roundtripSample, 16, len(roundtripSample)
}

func BenchmarkRoundtripTxt(b *testing.B) {
	b.StopTimer()
	sample, min, max := initRoundtripBenchmarks()
	sampleString := string(sample)
	b.ReportAllocs()
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	defer db.Close()
	b.StartTimer()
	var result string
	for i := 0; i < b.N; i++ {
		length := min + i
		if length > max {
			length = max
		}
		test := sampleString[0:length]
		rows := tb.checkRows(db.Query(`SELECT "` + test + `"`))
		if !rows.Next() {
			rows.Close()
			b.Fatalf("crashed")
		}
		err := rows.Scan(&result)
		if err != nil {
			rows.Close()
			b.Fatalf("crashed")
		}
		if result != test {
			rows.Close()
			b.Errorf("mismatch")
		}
		rows.Close()
	}
}

func BenchmarkRoundtripBin(b *testing.B) {
	b.StopTimer()
	sample, min, max := initRoundtripBenchmarks()
	b.ReportAllocs()
	tb := (*TB)(b)
	db := tb.checkDB(sql.Open("mysql", dsn))
	defer db.Close()
	stmt := tb.checkStmt(db.Prepare("SELECT ?"))
	defer stmt.Close()
	b.StartTimer()
	var result sql.RawBytes
	for i := 0; i < b.N; i++ {
		length := min + i
		if length > max {
			length = max
		}
		test := sample[0:length]
		rows := tb.checkRows(stmt.Query(test))
		if !rows.Next() {
			rows.Close()
			b.Fatalf("crashed")
		}
		err := rows.Scan(&result)
		if err != nil {
			rows.Close()
			b.Fatalf("crashed")
		}
		if !bytes.Equal(result, test) {
			rows.Close()
			b.Errorf("mismatch")
		}
		rows.Close()
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 PingCAP, Inc.
//
// 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.

//go:build !codes

package testkit

import (
	"flag"
	"testing"
	"time"

	"github.com/pingcap/tidb/ddl/schematracker"
	"github.com/pingcap/tidb/domain"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/store/driver"
	"github.com/pingcap/tidb/store/mockstore"
	"github.com/pingcap/tidb/util/gctuner"
	"github.com/stretchr/testify/require"
	"go.opencensus.io/stats/view"
)

// WithTiKV flag is only used for debugging locally with real tikv cluster.
var WithTiKV = flag.String("with-tikv", "", "address of tikv cluster, if set, running test with real tikv cluster")

// CreateMockStore return a new mock kv.Storage.
func CreateMockStore(t testing.TB, opts ...mockstore.MockTiKVStoreOption) kv.Storage {
	if *WithTiKV != "" {
		var d driver.TiKVDriver
		var err error
		store, err := d.Open("tikv://" + *WithTiKV)
		require.NoError(t, err)

		var dom *domain.Domain
		dom, err = session.BootstrapSession(store)
		t.Cleanup(func() {
			dom.Close()
			err := store.Close()
			require.NoError(t, err)
			view.Stop()
		})
		require.NoError(t, err)
		return store
	}
	t.Cleanup(func() {
		view.Stop()
	})
	gctuner.GlobalMemoryLimitTuner.Stop()
	store, _ := CreateMockStoreAndDomain(t, opts...)
	return store
}

// CreateMockStoreAndDomain return a new mock kv.Storage and *domain.Domain.
func CreateMockStoreAndDomain(t testing.TB, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
	store, err := mockstore.NewMockStore(opts...)
	require.NoError(t, err)
	dom := bootstrap(t, store, 500*time.Millisecond)
	sm := MockSessionManager{}
	dom.InfoSyncer().SetSessionManager(&sm)
	return schematracker.UnwrapStorage(store), dom
}

func bootstrap(t testing.TB, store kv.Storage, lease time.Duration) *domain.Domain {
	session.SetSchemaLease(lease)
	session.DisableStats4Test()
	dom, err := session.BootstrapSession(store)
	require.NoError(t, err)

	dom.SetStatsUpdating(true)

	t.Cleanup(func() {
		dom.Close()
		err := store.Close()
		require.NoError(t, err)
		view.Stop()
	})
	return dom
}

// CreateMockStoreWithSchemaLease return a new mock kv.Storage.
func CreateMockStoreWithSchemaLease(t testing.TB, lease time.Duration, opts ...mockstore.MockTiKVStoreOption) kv.Storage {
	store, _ := CreateMockStoreAndDomainWithSchemaLease(t, lease, opts...)
	return schematracker.UnwrapStorage(store)
}

// CreateMockStoreAndDomainWithSchemaLease return a new mock kv.Storage and *domain.Domain.
func CreateMockStoreAndDomainWithSchemaLease(t testing.TB, lease time.Duration, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
	store, err := mockstore.NewMockStore(opts...)
	require.NoError(t, err)
	dom := bootstrap(t, store, lease)
	sm := MockSessionManager{}
	dom.InfoSyncer().SetSessionManager(&sm)
	return schematracker.UnwrapStorage(store), dom
}
<commit_msg>*: fix goleak when to use CreateMockStoreAndDomain (#38734)<commit_after>// Copyright 2021 PingCAP, Inc.
//
// 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.

//go:build !codes

package testkit

import (
	"flag"
	"testing"
	"time"

	"github.com/pingcap/tidb/ddl/schematracker"
	"github.com/pingcap/tidb/domain"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/store/driver"
	"github.com/pingcap/tidb/store/mockstore"
	"github.com/pingcap/tidb/util/gctuner"
	"github.com/stretchr/testify/require"
	"go.opencensus.io/stats/view"
)

// WithTiKV flag is only used for debugging locally with real tikv cluster.
var WithTiKV = flag.String("with-tikv", "", "address of tikv cluster, if set, running test with real tikv cluster")

// CreateMockStore return a new mock kv.Storage.
func CreateMockStore(t testing.TB, opts ...mockstore.MockTiKVStoreOption) kv.Storage {
	if *WithTiKV != "" {
		var d driver.TiKVDriver
		var err error
		store, err := d.Open("tikv://" + *WithTiKV)
		require.NoError(t, err)

		var dom *domain.Domain
		dom, err = session.BootstrapSession(store)
		t.Cleanup(func() {
			dom.Close()
			err := store.Close()
			require.NoError(t, err)
			view.Stop()
		})
		require.NoError(t, err)
		return store
	}
	t.Cleanup(func() {
		view.Stop()
	})
	gctuner.GlobalMemoryLimitTuner.Stop()
	store, _ := CreateMockStoreAndDomain(t, opts...)
	return store
}

// CreateMockStoreAndDomain return a new mock kv.Storage and *domain.Domain.
func CreateMockStoreAndDomain(t testing.TB, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
	store, err := mockstore.NewMockStore(opts...)
	require.NoError(t, err)
	dom := bootstrap(t, store, 500*time.Millisecond)
	sm := MockSessionManager{}
	dom.InfoSyncer().SetSessionManager(&sm)
	t.Cleanup(func() {
		view.Stop()
		gctuner.GlobalMemoryLimitTuner.Stop()
	})
	return schematracker.UnwrapStorage(store), dom
}

func bootstrap(t testing.TB, store kv.Storage, lease time.Duration) *domain.Domain {
	session.SetSchemaLease(lease)
	session.DisableStats4Test()
	dom, err := session.BootstrapSession(store)
	require.NoError(t, err)

	dom.SetStatsUpdating(true)

	t.Cleanup(func() {
		dom.Close()
		err := store.Close()
		require.NoError(t, err)
		view.Stop()
	})
	return dom
}

// CreateMockStoreWithSchemaLease return a new mock kv.Storage.
func CreateMockStoreWithSchemaLease(t testing.TB, lease time.Duration, opts ...mockstore.MockTiKVStoreOption) kv.Storage {
	store, _ := CreateMockStoreAndDomainWithSchemaLease(t, lease, opts...)
	return schematracker.UnwrapStorage(store)
}

// CreateMockStoreAndDomainWithSchemaLease return a new mock kv.Storage and *domain.Domain.
func CreateMockStoreAndDomainWithSchemaLease(t testing.TB, lease time.Duration, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
	store, err := mockstore.NewMockStore(opts...)
	require.NoError(t, err)
	dom := bootstrap(t, store, lease)
	sm := MockSessionManager{}
	dom.InfoSyncer().SetSessionManager(&sm)
	return schematracker.UnwrapStorage(store), dom
}
<|endoftext|>