<commit_before>// Copyright 2015 The rkt 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.

package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/coreos/rkt/tools/common"
	"github.com/coreos/rkt/tools/common/filelist"
)

const (
	globCmd = "glob"
	// globMakeFunction is a template for generating all files for
	// given set of wildcards. See globMakeWildcard.
	globMakeFunction = `$(shell stat --format "%n: %F" !!!WILDCARDS!!! | grep -e 'regular file$$' | cut -f1 -d:)`
	// globMakeWildcard is a template for call wildcard function
	// for in a given directory with a given suffix. This wildcard
	// is for normal files.
	globMakeWildcard = "$(wildcard !!!DIR!!!/*!!!SUFFIX!!!)"
	// globMakeHiddenWildcard is a template for call wildcard
	// function for in a given directory with a given suffix. This
	// wildcard is for files beginning with a dot, which are
	// normally not taken into account by wildcard.
	globMakeHiddenWildcard = "$(wildcard !!!DIR!!!/.*!!!SUFFIX!!!)"
)

type globMode int

const (
	globNormal globMode = iota
	globDotFiles
	globAll
)

type globArgs struct {
	target   string
	suffix   string
	files    []string
	mode     globMode
	fMode    string
	filelist string
	mapTo    []string
}

func init() {
	cmds[globCmd] = globDeps
}

func globDeps(args []string) string {
	parsedArgs := globGetArgs(args)
	files := globGetFiles(parsedArgs)
	makeFunction := globGetMakeFunction(files, parsedArgs.suffix, parsedArgs.mode)
	return GenerateFileDeps(parsedArgs.target, makeFunction, files)
}

// globGetArgs parses given parameters and returns a target, a suffix
// and a list of files.
func globGetArgs(args []string) globArgs {
	f, target := standardFlags(globCmd)
	suffix := f.String("suffix", "", "File suffix (example: .go)")
	globbingMode := f.String("glob-mode", "all", "Which files to glob (normal, dot-files, all [default])")
	filesMode := f.String("mode", "args", "How to get files, either 'filelist' mode or 'args' [default]")
	filelist := f.String("filelist", "", "For filelist mode, read all the files from this file")
	mapTo := []string{}
	mapToWrapper := common.StringSliceWrapper{Slice: &mapTo}
	f.Var(&mapToWrapper, "map-to", "Map contents of filelist to this directory, can be used multiple times")

	f.Parse(args)
	if *target == "" {
		common.Die("--target parameter must be specified and cannot be empty")
	}
	mode := globModeFromString(*globbingMode)
	switch *filesMode {
	case "filelist":
		if *filelist == "" {
			common.Die("--filelist parameter must be specified and cannot be empty")
		}
		if len(mapTo) < 1 {
			common.Die("--map-to parameter must be specified at least once")
		}
	case "args":
		if *filelist != "" {
			common.Warn("--filelist parameter is ignored in args mode")
		}
		if len(mapTo) > 0 {
			common.Warn("--map-to parameter is ignored in args mode")
		}
	}
	return globArgs{
		target:   *target,
		suffix:   *suffix,
		files:    f.Args(),
		mode:     mode,
		fMode:    *filesMode,
		filelist: *filelist,
		mapTo:    mapTo,
	}
}

func globModeFromString(mode string) globMode {
	switch mode {
	case "normal":
		return globNormal
	case "dot-files":
		return globDotFiles
	case "all":
		return globAll
	default:
		common.Die("Unknown glob mode %q", mode)
	}
	panic("Should not happen")
}

func globGetFiles(args globArgs) []string {
	if args.fMode == "args" {
		return args.files
	}
	f, err := globGetFilesFromFilelist(args.filelist)
	if err != nil {
		common.Die("Failed to get files from filelist %q: %v", args.filelist, err)
	}
	return common.MapFilesToDirectories(f, args.mapTo)
}

func globGetFilesFromFilelist(filename string) ([]string, error) {
	fl, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("Failed to open filelist %q: %v", filename, err)
	}
	defer fl.Close()
	lists := filelist.Lists{}
	if err := lists.ParseFilelist(fl); err != nil {
		return nil, err
	}
	return lists.Files, nil
}

// globGetMakeFunction returns a make snippet which calls wildcard
// function in all directories where given files are and with a given
// suffix.
func globGetMakeFunction(files []string, suffix string, mode globMode) string {
	dirs := map[string]struct{}{}
	for _, file := range files {
		dirs[filepath.Dir(file)] = struct{}{}
	}
	makeWildcards := make([]string, 0, len(dirs))
	wildcard := globGetMakeSnippet(mode)
	for dir := range dirs {
		str := replacePlaceholders(wildcard, "SUFFIX", suffix, "DIR", dir)
		makeWildcards = append(makeWildcards, str)
	}
	return replacePlaceholders(globMakeFunction, "WILDCARDS", strings.Join(makeWildcards, " "))
}

func globGetMakeSnippet(mode globMode) string {
	switch mode {
	case globNormal:
		return globMakeWildcard
	case globDotFiles:
		return globMakeHiddenWildcard
	case globAll:
		return fmt.Sprintf("%s %s", globMakeWildcard, globMakeHiddenWildcard)
	}
	panic("Should not happen")
}
<commit_msg>build: Drop "args" mode in depsgen<commit_after>// Copyright 2015 The rkt 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.

package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/coreos/rkt/tools/common"
	"github.com/coreos/rkt/tools/common/filelist"
)

const (
	globCmd = "glob"
	// globMakeFunction is a template for generating all files for
	// given set of wildcards. See globMakeWildcard.
	globMakeFunction = `$(shell stat --format "%n: %F" !!!WILDCARDS!!! | grep -e 'regular file$$' | cut -f1 -d:)`
	// globMakeWildcard is a template for call wildcard function
	// for in a given directory with a given suffix. This wildcard
	// is for normal files.
	globMakeWildcard = "$(wildcard !!!DIR!!!/*!!!SUFFIX!!!)"
	// globMakeHiddenWildcard is a template for call wildcard
	// function for in a given directory with a given suffix. This
	// wildcard is for files beginning with a dot, which are
	// normally not taken into account by wildcard.
	globMakeHiddenWildcard = "$(wildcard !!!DIR!!!/.*!!!SUFFIX!!!)"
)

type globMode int

const (
	globNormal globMode = iota
	globDotFiles
	globAll
)

type globArgs struct {
	target   string
	suffix   string
	mode     globMode
	filelist string
	mapTo    []string
}

func init() {
	cmds[globCmd] = globDeps
}

func globDeps(args []string) string {
	parsedArgs := globGetArgs(args)
	files := globGetFiles(parsedArgs)
	makeFunction := globGetMakeFunction(files, parsedArgs.suffix, parsedArgs.mode)
	return GenerateFileDeps(parsedArgs.target, makeFunction, files)
}

// globGetArgs parses given parameters and returns a target, a suffix
// and a list of files.
func globGetArgs(args []string) globArgs {
	f, target := standardFlags(globCmd)
	suffix := f.String("suffix", "", "File suffix (example: .go)")
	globbingMode := f.String("glob-mode", "all", "Which files to glob (normal, dot-files, all [default])")
	filelist := f.String("filelist", "", "Read all the files from this file")
	mapTo := []string{}
	mapToWrapper := common.StringSliceWrapper{Slice: &mapTo}
	f.Var(&mapToWrapper, "map-to", "Map contents of filelist to this directory, can be used multiple times")

	f.Parse(args)
	if *target == "" {
		common.Die("--target parameter must be specified and cannot be empty")
	}
	mode := globModeFromString(*globbingMode)
	if *filelist == "" {
		common.Die("--filelist parameter must be specified and cannot be empty")
	}
	if len(mapTo) < 1 {
		common.Die("--map-to parameter must be specified at least once")
	}
	return globArgs{
		target:   *target,
		suffix:   *suffix,
		mode:     mode,
		filelist: *filelist,
		mapTo:    mapTo,
	}
}

func globModeFromString(mode string) globMode {
	switch mode {
	case "normal":
		return globNormal
	case "dot-files":
		return globDotFiles
	case "all":
		return globAll
	default:
		common.Die("Unknown glob mode %q", mode)
	}
	panic("Should not happen")
}

func globGetFiles(args globArgs) []string {
	f, err := globGetFilesFromFilelist(args.filelist)
	if err != nil {
		common.Die("Failed to get files from filelist %q: %v", args.filelist, err)
	}
	return common.MapFilesToDirectories(f, args.mapTo)
}

func globGetFilesFromFilelist(filename string) ([]string, error) {
	fl, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("Failed to open filelist %q: %v", filename, err)
	}
	defer fl.Close()
	lists := filelist.Lists{}
	if err := lists.ParseFilelist(fl); err != nil {
		return nil, err
	}
	return lists.Files, nil
}

// globGetMakeFunction returns a make snippet which calls wildcard
// function in all directories where given files are and with a given
// suffix.
func globGetMakeFunction(files []string, suffix string, mode globMode) string {
	dirs := map[string]struct{}{}
	for _, file := range files {
		dirs[filepath.Dir(file)] = struct{}{}
	}
	makeWildcards := make([]string, 0, len(dirs))
	wildcard := globGetMakeSnippet(mode)
	for dir := range dirs {
		str := replacePlaceholders(wildcard, "SUFFIX", suffix, "DIR", dir)
		makeWildcards = append(makeWildcards, str)
	}
	return replacePlaceholders(globMakeFunction, "WILDCARDS", strings.Join(makeWildcards, " "))
}

func globGetMakeSnippet(mode globMode) string {
	switch mode {
	case globNormal:
		return globMakeWildcard
	case globDotFiles:
		return globMakeHiddenWildcard
	case globAll:
		return fmt.Sprintf("%s %s", globMakeWildcard, globMakeHiddenWildcard)
	}
	panic("Should not happen")
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func doGetConfig(w http.ResponseWriter, req *http.Request) {
	err := updateConfig()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error updating config: %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(&globalConfig)
	if err != nil {
		log.Printf("Error sending config: %v", err)
	}
}

func putConfig(w http.ResponseWriter, req *http.Request) {
	d := json.NewDecoder(req.Body)
	conf := cbfsconfig.CBFSConfig{}

	err := d.Decode(&conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading config: %v", err)
		return
	}

	err = StoreConfig(conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error writing config: %v", err)
		return
	}

	err = updateConfig()
	if err != nil {
		log.Printf("Error fetching newly stored config: %v", err)
	}

	w.WriteHeader(204)
}

func doList(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(200)
	explen := getHash().Size() * 2
	filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {
			_, e := w.Write([]byte(info.Name() + "\n"))
			return e
		}
		return nil
	})
}

func doListTasks(w http.ResponseWriter, req *http.Request) {
	tasks, err := listRunningTasks()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error listing tasks:  %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	// Reformat for more APIish output.
	output := map[string]map[string]time.Time{}

	// Remove node prefix from local task names.
	npre := serverId + "/"

	for _, tl := range tasks {
		for k, v := range tl.Tasks {
			if strings.HasPrefix(k, npre) {
				delete(tl.Tasks, k)
				tl.Tasks[k[len(npre):]] = v
			}
		}
		output[tl.Node] = tl.Tasks
	}

	e := json.NewEncoder(w)
	err = e.Encode(output)
	if err != nil {
		log.Printf("Error encoding running tasks list: %v", err)
	}
}

func doGetMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	err := couchbase.Get(path, &got)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	if got.Userdata == nil {
		w.Write([]byte("{}"))
	} else {
		w.Write(*got.Userdata)
	}
}

func putMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	casid := uint64(0)
	err := couchbase.Gets(path, &got, &casid)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	r := json.RawMessage{}
	err = json.NewDecoder(req.Body).Decode(&r)
	if err != nil {
		w.WriteHeader(400)
		w.Write([]byte(err.Error()))
		return
	}

	got.Userdata = &r
	b := mustEncode(&got)

	err = couchbase.Do(path, func(mc *memcached.Client, vb uint16) error {
		req := &gomemcached.MCRequest{
			Opcode:  gomemcached.SET,
			VBucket: vb,
			Key:     []byte(path),
			Cas:     casid,
			Opaque:  0,
			Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
			Body:    b}
		resp, err := mc.Send(req)
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return resp
		}
		return nil
	})

	if err == nil {
		w.WriteHeader(201)
	} else {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}
}

func doListNodes(w http.ResponseWriter, req *http.Request) {

	nl, err := findAllNodes()
	if err != nil {
		log.Printf("Error executing nodes view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating node list: %v", err)
		return
	}

	respob := map[string]interface{}{}
	for _, node := range nl {
		age := time.Since(node.Time)
		respob[node.name] = map[string]interface{}{
			"size":      node.storageSize,
			"addr":      node.Address(),
			"hbtime":    node.Time,
			"hbage_ms":  age.Nanoseconds() / 1e6,
			"hbage_str": age.String(),
			"hash":      node.Hash,
			"used":      node.Used,
			"free":      node.Free,
			"addr_raw":  node.Addr,
			"bindaddr":  node.BindAddr,
		}
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(mustEncode(respob))
}

func proxyViewRequest(w http.ResponseWriter, req *http.Request,
	path string) {

	node := couchbase.Nodes[rand.Intn(len(couchbase.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}

	u.Path = "/" + path
	u.RawQuery = req.URL.RawQuery

	client := &http.Client{
		Transport: TimeoutTransport(*viewTimeout),
	}

	res, err := client.Get(u.String())
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}
	defer res.Body.Close()

	for k, vs := range res.Header {
		w.Header()[k] = vs
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	}
	w.WriteHeader(res.StatusCode)

	io.Copy(output, res.Body)
}

func doListDocs(w http.ResponseWriter, req *http.Request,
	path string) {

	// trim off trailing slash early so we handle them consistently
	if strings.HasSuffix(path, "/") {
		path = path[0 : len(path)-1]
	}

	includeMeta := req.FormValue("includeMeta")
	depthString := req.FormValue("depth")
	depth := 1
	if depthString != "" {
		i, err := strconv.Atoi(depthString)
		if err != nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "Error processing depth parameter: %v", err)
			return
		}
		depth = i
	}

	fl, err := listFiles(path, includeMeta == "true", depth)
	if err != nil {
		log.Printf("Error executing file browse view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating file list: %v", err)
		return
	}

	if len(fl.Dirs) == 0 && len(fl.Files) == 0 {
		w.WriteHeader(404)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(fl)
	if err != nil {
		log.Printf("Error writing json stream: %v", err)
	}
}
<commit_msg>Remove the individual node's prefix, not the API server.<commit_after>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func doGetConfig(w http.ResponseWriter, req *http.Request) {
	err := updateConfig()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error updating config: %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(&globalConfig)
	if err != nil {
		log.Printf("Error sending config: %v", err)
	}
}

func putConfig(w http.ResponseWriter, req *http.Request) {
	d := json.NewDecoder(req.Body)
	conf := cbfsconfig.CBFSConfig{}

	err := d.Decode(&conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading config: %v", err)
		return
	}

	err = StoreConfig(conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error writing config: %v", err)
		return
	}

	err = updateConfig()
	if err != nil {
		log.Printf("Error fetching newly stored config: %v", err)
	}

	w.WriteHeader(204)
}

func doList(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(200)
	explen := getHash().Size() * 2
	filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {
			_, e := w.Write([]byte(info.Name() + "\n"))
			return e
		}
		return nil
	})
}

func doListTasks(w http.ResponseWriter, req *http.Request) {
	tasks, err := listRunningTasks()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error listing tasks:  %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	// Reformat for more APIish output.
	output := map[string]map[string]time.Time{}

	for _, tl := range tasks {
		// Remove node prefix from local task names.
		npre := tl.Node + "/"

		for k, v := range tl.Tasks {
			if strings.HasPrefix(k, npre) {
				delete(tl.Tasks, k)
				tl.Tasks[k[len(npre):]] = v
			}
		}
		output[tl.Node] = tl.Tasks
	}

	e := json.NewEncoder(w)
	err = e.Encode(output)
	if err != nil {
		log.Printf("Error encoding running tasks list: %v", err)
	}
}

func doGetMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	err := couchbase.Get(path, &got)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	if got.Userdata == nil {
		w.Write([]byte("{}"))
	} else {
		w.Write(*got.Userdata)
	}
}

func putMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	casid := uint64(0)
	err := couchbase.Gets(path, &got, &casid)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	r := json.RawMessage{}
	err = json.NewDecoder(req.Body).Decode(&r)
	if err != nil {
		w.WriteHeader(400)
		w.Write([]byte(err.Error()))
		return
	}

	got.Userdata = &r
	b := mustEncode(&got)

	err = couchbase.Do(path, func(mc *memcached.Client, vb uint16) error {
		req := &gomemcached.MCRequest{
			Opcode:  gomemcached.SET,
			VBucket: vb,
			Key:     []byte(path),
			Cas:     casid,
			Opaque:  0,
			Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
			Body:    b}
		resp, err := mc.Send(req)
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return resp
		}
		return nil
	})

	if err == nil {
		w.WriteHeader(201)
	} else {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}
}

func doListNodes(w http.ResponseWriter, req *http.Request) {

	nl, err := findAllNodes()
	if err != nil {
		log.Printf("Error executing nodes view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating node list: %v", err)
		return
	}

	respob := map[string]interface{}{}
	for _, node := range nl {
		age := time.Since(node.Time)
		respob[node.name] = map[string]interface{}{
			"size":      node.storageSize,
			"addr":      node.Address(),
			"hbtime":    node.Time,
			"hbage_ms":  age.Nanoseconds() / 1e6,
			"hbage_str": age.String(),
			"hash":      node.Hash,
			"used":      node.Used,
			"free":      node.Free,
			"addr_raw":  node.Addr,
			"bindaddr":  node.BindAddr,
		}
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(mustEncode(respob))
}

func proxyViewRequest(w http.ResponseWriter, req *http.Request,
	path string) {

	node := couchbase.Nodes[rand.Intn(len(couchbase.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}

	u.Path = "/" + path
	u.RawQuery = req.URL.RawQuery

	client := &http.Client{
		Transport: TimeoutTransport(*viewTimeout),
	}

	res, err := client.Get(u.String())
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}
	defer res.Body.Close()

	for k, vs := range res.Header {
		w.Header()[k] = vs
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	}
	w.WriteHeader(res.StatusCode)

	io.Copy(output, res.Body)
}

func doListDocs(w http.ResponseWriter, req *http.Request,
	path string) {

	// trim off trailing slash early so we handle them consistently
	if strings.HasSuffix(path, "/") {
		path = path[0 : len(path)-1]
	}

	includeMeta := req.FormValue("includeMeta")
	depthString := req.FormValue("depth")
	depth := 1
	if depthString != "" {
		i, err := strconv.Atoi(depthString)
		if err != nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "Error processing depth parameter: %v", err)
			return
		}
		depth = i
	}

	fl, err := listFiles(path, includeMeta == "true", depth)
	if err != nil {
		log.Printf("Error executing file browse view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating file list: %v", err)
		return
	}

	if len(fl.Dirs) == 0 && len(fl.Files) == 0 {
		w.WriteHeader(404)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(fl)
	if err != nil {
		log.Printf("Error writing json stream: %v", err)
	}
}
<|endoftext|>


<commit_before>// Code generated by "stringer -type=lexItemType"; DO NOT EDIT.

package sqlread

import "fmt"

const _lexItemType_name = "TIllegalTEofTSemiTCommaTCommentTNullTStringTNumberTIdentifierTDropTableFullStmtTLockTableFullStmtTUnlockTablesFullStmtTSetFullStmtTLParenTRParenTCreateTableTCreateTableDetailTCreateTableExtraTColumnTypeTColumnSizeTColumnEnumValTColumnDetailsTInsertIntoTInsertValuesTInsertRow"

var _lexItemType_index = [...]uint16{0, 8, 12, 17, 23, 31, 36, 43, 50, 61, 79, 97, 118, 130, 137, 144, 156, 174, 191, 202, 213, 227, 241, 252, 265, 275}

func (i lexItemType) String() string {
	if i >= lexItemType(len(_lexItemType_index)-1) {
		return fmt.Sprintf("lexItemType(%d)", i)
	}
	return _lexItemType_name[_lexItemType_index[i]:_lexItemType_index[i+1]]
}
<commit_msg>and updates stringer<commit_after>// Code generated by "stringer -type=lexItemType"; DO NOT EDIT.

package sqlread

import "fmt"

const _lexItemType_name = "TIllegalTEofTSemiTCommaTCommentTNullTStringTNumberTIdentifierTDropTableFullStmtTLockTableFullStmtTUnlockTablesFullStmtTSetFullStmtTLParenTRParenTCreateTableTCreateTableDetailTColumnTypeTColumnSizeTColumnEnumValTColumnDetailsTInsertIntoTInsertValues"

var _lexItemType_index = [...]uint8{0, 8, 12, 17, 23, 31, 36, 43, 50, 61, 79, 97, 118, 130, 137, 144, 156, 174, 185, 196, 210, 224, 235, 248}

func (i lexItemType) String() string {
	if i >= lexItemType(len(_lexItemType_index)-1) {
		return fmt.Sprintf("lexItemType(%d)", i)
	}
	return _lexItemType_name[_lexItemType_index[i]:_lexItemType_index[i+1]]
}
<|endoftext|>


<commit_before>package hll

import (
	"fmt"
	"math"

	"testing"
	"time"

	"github.com/bmizerany/assert"
)

// adding values in the dense case
// determine if the maximum rho value for a determined index is correctly written to M.
func TestAddNormal(t *testing.T) {
	h := NewHll(14, 20)

	value := uint64(0xAABBCCDD00112210)
	value2 := uint64(0xAABBCCDD00112211)

	register := value >> (64 - h.p)
	register2 := value2 >> (64 - h.p)
	assert.Equal(t, register2, register)
	assert.T(t, rho(value) > rho(value2))

	h.addNormal(value)
	h.addNormal(value2)
	assert.Equal(t, h.M.Get(uint64(register)), rho(value))
}

// Check to make sure that the temp set gets merged when it's supposed to
// and if it changes to the dense represtation if it passes the sparse threshold
func TestAddSparse(t *testing.T) {
	h := NewHll(14, 20)

	assert.Equal(t, h.isSparse, true)
	// the maximum size of the sparseList is 6145: (2^18) * (6/4) / 64
	rands := randUint64s(t, 6145)

	for idx, randVal := range rands {
		h.Add(randVal)

		// tempSet should be reset after adding (2^p * (6 / 4) / 64 elements
		if uint64(idx*64)%h.mergeSizeBits == 1 {
			assert.Equal(t, len(h.tempSet), 1)
		}

		// Should no longer be using the sparse representation after (2^p * 6) / 64 elements have been added
		if h.sparseList.SizeInBits() > h.sparseThresholdBits {
			assert.Equal(t, h.isSparse, false)
		}
	}
}

// Tests cardinality accuracy with varying number of distinct uint64 inputs
func TestCardinality(t *testing.T) {
	// number of random values to estimate cardinalities for
	counts := []int{1000, 5000, 20000, 50000, 100000, 250000, 1000000, 10000000}

	for _, count := range counts {
		// Create new Hll struct with p = 14 & p' = 25
		h := NewHll(14, 25)

		rands := randUint64s(t, count)

		startTime := time.Now()
		for _, randomU64 := range rands {
			h.Add(randomU64)
		}
		card := h.Cardinality()
		endTime := time.Since(startTime)

		calculatedError := math.Abs(float64(card)-float64(count)) / float64(count)
		fmt.Printf("\nActual Cardinality: %d\n Estimated Cardinality: %d\nError: %v\nTime Elapsed: %v\n\n", count, card, calculatedError, endTime)
	}
}

// Test the weighted mean estimate for the bias for precision 4.
func TestEstimateBias(t *testing.T) {
	h_four := NewHll(4, 10)

	// according to empirical bias calculations, bias should be below 9.2 and above 8.78
	bias := h_four.estimateBias(12.5)
	assert.T(t, bias > 8.78 && bias < 9.20)

	// if estimate is not in the estimated range, return max bias
	max_bias := h_four.estimateBias(80.00)
	assert.Equal(t, max_bias, -1.7606)
}
<commit_msg>actually add the test<commit_after>package hll

import (
	"fmt"
	"math"

	"testing"
	"time"

	"github.com/bmizerany/assert"
)

// adding values in the dense case
// determine if the maximum rho value for a determined index is correctly written to M.
func TestAddNormal(t *testing.T) {
	h := NewHll(14, 20)

	value := uint64(0xAABBCCDD00112210)
	value2 := uint64(0xAABBCCDD00112211)

	register := value >> (64 - h.p)
	register2 := value2 >> (64 - h.p)
	assert.Equal(t, register2, register)
	assert.T(t, rho(value) > rho(value2))

	h.addNormal(value)
	h.addNormal(value2)
	assert.Equal(t, h.M.Get(uint64(register)), rho(value))
}

// Check to make sure that the temp set gets merged when it's supposed to
// and if it changes to the dense represtation if it passes the sparse threshold
func TestAddSparse(t *testing.T) {
	h := NewHll(14, 20)

	assert.Equal(t, h.isSparse, true)
	// the maximum size of the sparseList is 6145: (2^18) * (6/4) / 64
	rands := randUint64s(t, 6145)

	for idx, randVal := range rands {
		h.Add(randVal)

		// tempSet should be reset after adding (2^p * (6 / 4) / 64 elements
		if uint64(idx*64)%h.mergeSizeBits == 1 {
			assert.Equal(t, len(h.tempSet), 1)
		}

		// Should no longer be using the sparse representation after (2^p * 6) / 64 elements have been added
		if h.sparseList.SizeInBits() > h.sparseThresholdBits {
			assert.Equal(t, h.isSparse, false)
		}
	}
}

// Tests cardinality accuracy with varying number of distinct uint64 inputs
func TestCardinality(t *testing.T) {
	// number of random values to estimate cardinalities for
	counts := []int{1000, 5000, 20000, 50000, 100000, 250000, 1000000, 10000000}

	for _, count := range counts {
		// Create new Hll struct with p = 14 & p' = 25
		h := NewHll(14, 25)
		// Random uint64 values to test.
		rands := randUint64s(t, count)

		startTime := time.Now()
		for _, randomU64 := range rands {
			h.Add(randomU64)
		}
		card := h.Cardinality()
		endTime := time.Since(startTime)

		calculatedError := math.Abs(float64(card)-float64(count)) / float64(count)
		assert.T(t, calculatedError < 0.15)
		fmt.Printf("\nActual Cardinality: %d\n Estimated Cardinality: %d\nError: %v\nTime Elapsed: %v\n\n", count, card, calculatedError, endTime)
	}
}

// Test the weighted mean estimate for the bias for precision 4.
func TestEstimateBias(t *testing.T) {
	h_four := NewHll(4, 10)

	// according to empirical bias calculations, bias should be below 9.2 and above 8.78
	bias := h_four.estimateBias(12.5)
	assert.T(t, bias > 8.78 && bias < 9.20)

	// if estimate is not in the estimated range, return max bias
	max_bias := h_four.estimateBias(80.00)
	assert.Equal(t, max_bias, -1.7606)
}
<|endoftext|>


<commit_before>package hof

import "testing"
import "reflect"

func TestIntMap(t *testing.T) {
	var mapper func(func(int) int, []int) []int
	MakeMapFunc(&mapper)

	in := []int{1, 2, 3, 4, 5}
	f := func(x int) int { return x * 2 }
	exp := []int{2, 4, 6, 8, 10}

	out := mapper(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestStringIntMap(t *testing.T) {
	var mapper func(func(string) int, []string) []int
	MakeMapFunc(&mapper)

	in := []string{"try", "this", "thing"}
	f := func(x string) int { return len(x) }
	exp := []int{3, 4, 5}

	out := mapper(f, in)

	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestIntStringMap(t *testing.T) {
	var mapper func(func(int) string, []int) []string
	MakeMapFunc(&mapper)

	in := []int{1, 2, 3}
	exp := []string{"x", "xx", "xxx"}
	f := func(x int) string {
		out := ""
		for i := 0; i < x; i++ {
			out += "x"
		}
		return out
	}

	out := mapper(f, in)

	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestEmptyMap(t *testing.T) {
	var mapper func(func(int) int, []int) []int
	MakeMapFunc(&mapper)

	in := []int{}
	f := func(x int) int { return x * 2 }
	exp := []int{}

	out := mapper(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestFilter(t *testing.T) {
	var filter func(func(int) bool, []int) []int
	MakeFilterFunc(&filter)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x int) bool { return x%2 == 0 }
	exp := []int{2, 4, 6, 8, 10}

	out := filter(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduce(t *testing.T) {
	var reduce func(func(int, int) int, []int) int
	MakeReduceFunc(&reduce)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x, y int) int { return x + y }
	exp := 55

	out := reduce(f, in)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduceInit(t *testing.T) {
	var reduce func(func(int, int) int, []int, int) int
	MakeReduceFunc(&reduce)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x, y int) int { return x + y }
	init := 45
	exp := 100

	out := reduce(f, in, init)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduceTwoTypes(t *testing.T) {
	var reduce func(func(string, int) string, []int, string) string
	MakeReduceFunc(&reduce)

	in := []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
	f := func(x string, y int) string { return x + string('0'+y) }
	init := ""
	exp := "0123456789"

	out := reduce(f, in, init)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}
<commit_msg>add a little benchmarking.<commit_after>package hof

import "testing"
import "reflect"

func TestIntMap(t *testing.T) {
	var mapper func(func(int) int, []int) []int
	MakeMapFunc(&mapper)

	in := []int{1, 2, 3, 4, 5}
	f := func(x int) int { return x * 2 }
	exp := []int{2, 4, 6, 8, 10}

	out := mapper(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestStringIntMap(t *testing.T) {
	var mapper func(func(string) int, []string) []int
	MakeMapFunc(&mapper)

	in := []string{"try", "this", "thing"}
	f := func(x string) int { return len(x) }
	exp := []int{3, 4, 5}

	out := mapper(f, in)

	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestIntStringMap(t *testing.T) {
	var mapper func(func(int) string, []int) []string
	MakeMapFunc(&mapper)

	in := []int{1, 2, 3}
	exp := []string{"x", "xx", "xxx"}
	f := func(x int) string {
		out := ""
		for i := 0; i < x; i++ {
			out += "x"
		}
		return out
	}

	out := mapper(f, in)

	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestEmptyMap(t *testing.T) {
	var mapper func(func(int) int, []int) []int
	MakeMapFunc(&mapper)

	in := []int{}
	f := func(x int) int { return x * 2 }
	exp := []int{}

	out := mapper(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestFilter(t *testing.T) {
	var filter func(func(int) bool, []int) []int
	MakeFilterFunc(&filter)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x int) bool { return x%2 == 0 }
	exp := []int{2, 4, 6, 8, 10}

	out := filter(f, in)
	if !reflect.DeepEqual(exp, out) {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduce(t *testing.T) {
	var reduce func(func(int, int) int, []int) int
	MakeReduceFunc(&reduce)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x, y int) int { return x + y }
	exp := 55

	out := reduce(f, in)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduceInit(t *testing.T) {
	var reduce func(func(int, int) int, []int, int) int
	MakeReduceFunc(&reduce)

	in := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	f := func(x, y int) int { return x + y }
	init := 45
	exp := 100

	out := reduce(f, in, init)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}

func TestReduceTwoTypes(t *testing.T) {
	var reduce func(func(string, int) string, []int, string) string
	MakeReduceFunc(&reduce)

	in := []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
	f := func(x string, y int) string { return x + string('0'+y) }
	init := ""
	exp := "0123456789"

	out := reduce(f, in, init)
	if exp != out {
		t.Fatal("expected", exp, ", got", out)
	}
}

func sliceRange(n int) []int {
	out := make([]int, n, n)
	for i := 0; i < n; i++ {
		out[i] = i
	}
	return out
}

var benchmarkIn = sliceRange(100)

func BenchmarkForMap(b *testing.B) {
	for n := 0; n < b.N; n++ {
		l := len(benchmarkIn)
		out := make([]int, l, l)
		for i := 0; i < l; i++ {
			out[i] = benchmarkIn[i] * 2
		}
		if out[1] != 2 {
			panic("wrong result")
		}
	}
}

func BenchmarkMakeMapFunc(b *testing.B) {
	var mapper func(func(int) int, []int) []int
	MakeMapFunc(&mapper)

	for n := 0; n < b.N; n++ {
		f := func(x int) int { return x * 2 }
		out := mapper(f, benchmarkIn)
		if out[1] != 2 {
			panic("wrong result")
		}
	}
}

func interfaceMapper(f func(interface{}) interface{}, in []interface{}) []interface{} {
	out := make([]interface{}, len(in), len(in))
	for i, v := range in {
		out[i] = f(v)
	}
	return out
}

func BenchmarkInterfaceMapFunc(b *testing.B) {
	for n := 0; n < b.N; n++ {
		l := len(benchmarkIn)

		interfaceIn := make([]interface{}, l, l)
		for i, x := range benchmarkIn {
			interfaceIn[i] = x
		}

		f := func(x interface{}) interface{} {
			return x.(int) * 2
		}

		interfaceOut := interfaceMapper(f, interfaceIn)

		out := make([]int, l, l)
		for i, n := range interfaceOut {
			out[i] = n.(int)
		}
	}
}
<|endoftext|>


<commit_before>// Package httpdown provides http.ConnState enabled graceful termination of
// http.Server.
package httpdown

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/facebookgo/clock"
	"github.com/facebookgo/stats"
)

const (
	defaultStopTimeout = time.Minute
	defaultKillTimeout = time.Minute
)

// A Server allows encapsulates the process of accepting new connections and
// serving them, and gracefully shutting down the listener without dropping
// active connections.
type Server interface {
	// Wait waits for the serving loop to finish. This will happen when Stop is
	// called, at which point it returns no error, or if there is an error in the
	// serving loop. You must call Wait after calling Serve or ListenAndServe.
	Wait() error

	// Stop stops the listener. It will block until all connections have been
	// closed.
	Stop() error
}

// HTTP defines the configuration for serving a http.Server. Multiple calls to
// Serve or ListenAndServe can be made on the same HTTP instance. The default
// timeouts of 1 minute each result in a maximum of 2 minutes before a Stop()
// returns.
type HTTP struct {
	// StopTimeout is the duration before we begin force closing connections.
	// Defaults to 1 minute.
	StopTimeout time.Duration

	// KillTimeout is the duration before which we completely give up and abort
	// even though we still have connected clients. This is useful when a large
	// number of client connections exist and closing them can take a long time.
	// Note, this is in addition to the StopTimeout. Defaults to 1 minute.
	KillTimeout time.Duration

	// Stats is optional. If provided, it will be used to record various metrics.
	Stats stats.Client

	// Clock allows for testing timing related functionality. Do not specify this
	// in production code.
	Clock clock.Clock
}

// Serve provides the low-level API which is useful if you're creating your own
// net.Listener.
func (h HTTP) Serve(s *http.Server, l net.Listener) Server {
	stopTimeout := h.StopTimeout
	if stopTimeout == 0 {
		stopTimeout = defaultStopTimeout
	}
	killTimeout := h.KillTimeout
	if killTimeout == 0 {
		killTimeout = defaultKillTimeout
	}
	klock := h.Clock
	if klock == nil {
		klock = clock.New()
	}

	ss := &server{
		stopTimeout:  stopTimeout,
		killTimeout:  killTimeout,
		stats:        h.Stats,
		clock:        klock,
		oldConnState: s.ConnState,
		listener:     l,
		server:       s,
		serveDone:    make(chan struct{}),
		serveErr:     make(chan error, 1),
		new:          make(chan net.Conn),
		active:       make(chan net.Conn),
		idle:         make(chan net.Conn),
		closed:       make(chan net.Conn),
		stop:         make(chan chan struct{}),
		kill:         make(chan chan struct{}),
	}
	s.ConnState = ss.connState
	go ss.manage()
	go ss.serve()
	return ss
}

// ListenAndServe returns a Server for the given http.Server. It is equivalent
// to ListendAndServe from the standard library, but returns immediately.
// Requests will be accepted in a background goroutine. If the http.Server has
// a non-nil TLSConfig, a TLS enabled listener will be setup.
func (h HTTP) ListenAndServe(s *http.Server) (Server, error) {
	addr := s.Addr
	if addr == "" {
		if s.TLSConfig == nil {
			addr = ":http"
		} else {
			addr = ":https"
		}
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		stats.BumpSum(h.Stats, "listen.error", 1)
		return nil, err
	}
	if s.TLSConfig != nil {
		l = tls.NewListener(l, s.TLSConfig)
	}
	return h.Serve(s, l), nil
}

// server manages the serving process and allows for gracefully stopping it.
type server struct {
	stopTimeout time.Duration
	killTimeout time.Duration
	stats       stats.Client
	clock       clock.Clock

	oldConnState func(net.Conn, http.ConnState)
	server       *http.Server
	serveDone    chan struct{}
	serveErr     chan error
	listener     net.Listener

	new    chan net.Conn
	active chan net.Conn
	idle   chan net.Conn
	closed chan net.Conn
	stop   chan chan struct{}
	kill   chan chan struct{}

	stopOnce sync.Once
	stopErr  error
}

func (s *server) connState(c net.Conn, cs http.ConnState) {
	if s.oldConnState != nil {
		s.oldConnState(c, cs)
	}

	switch cs {
	case http.StateNew:
		s.new <- c
	case http.StateActive:
		s.active <- c
	case http.StateIdle:
		s.idle <- c
	case http.StateHijacked, http.StateClosed:
		s.closed <- c
	}
}

func (s *server) manage() {
	defer func() {
		close(s.new)
		close(s.active)
		close(s.idle)
		close(s.closed)
		close(s.stop)
		close(s.kill)
	}()

	var stopDone chan struct{}

	conns := map[net.Conn]http.ConnState{}
	var countNew, countActive, countIdle float64

	// decConn decrements the count associated with the current state of the
	// given connection.
	decConn := func(c net.Conn) {
		switch conns[c] {
		default:
			panic(fmt.Errorf("unknown existing connection: %s", c))
		case http.StateNew:
			countNew--
		case http.StateActive:
			countActive--
		case http.StateIdle:
			countIdle--
		}
	}

	// setup a ticker to report various values every minute. if we don't have a
	// Stats implementation provided, we Stop it so it never ticks.
	statsTicker := s.clock.Ticker(time.Minute)
	if s.stats == nil {
		statsTicker.Stop()
	}

	for {
		select {
		case <-statsTicker.C:
			// we'll only get here when s.stats is not nil
			s.stats.BumpAvg("http-state.new", countNew)
			s.stats.BumpAvg("http-state.active", countActive)
			s.stats.BumpAvg("http-state.idle", countIdle)
			s.stats.BumpAvg("http-state.total", countNew+countActive+countIdle)
		case c := <-s.new:
			conns[c] = http.StateNew
			countNew++
		case c := <-s.active:
			decConn(c)
			countActive++

			conns[c] = http.StateActive
		case c := <-s.idle:
			decConn(c)
			countIdle++

			conns[c] = http.StateIdle

			// if we're already stopping, close it
			if stopDone != nil {
				c.Close()
			}
		case c := <-s.closed:
			stats.BumpSum(s.stats, "conn.closed", 1)
			decConn(c)
			delete(conns, c)

			// if we're waiting to stop and are all empty, we just closed the last
			// connection and we're done.
			if stopDone != nil && len(conns) == 0 {
				close(stopDone)
				return
			}
		case stopDone = <-s.stop:
			// if we're already all empty, we're already done
			if len(conns) == 0 {
				close(stopDone)
				return
			}

			// close current idle connections right away
			for c, cs := range conns {
				if cs == http.StateIdle {
					c.Close()
				}
			}

			// continue the loop and wait for all the ConnState updates which will
			// eventually close(stopDone) and return from this goroutine.

		case killDone := <-s.kill:
			// force close all connections
			stats.BumpSum(s.stats, "kill.conn.count", float64(len(conns)))
			for c := range conns {
				c.Close()
			}

			// don't block the kill.
			close(killDone)

			// continue the loop and we wait for all the ConnState updates and will
			// return from this goroutine when we're all done. otherwise we'll try to
			// send those ConnState updates on closed channels.

		}
	}
}

func (s *server) serve() {
	stats.BumpSum(s.stats, "serve", 1)
	s.serveErr <- s.server.Serve(s.listener)
	close(s.serveDone)
	close(s.serveErr)
}

func (s *server) Wait() error {
	if err := <-s.serveErr; !isUseOfClosedError(err) {
		return err
	}
	return nil
}

func (s *server) Stop() error {
	s.stopOnce.Do(func() {
		defer stats.BumpTime(s.stats, "stop.time").End()
		stats.BumpSum(s.stats, "stop", 1)

		// first disable keep-alive for new connections
		s.server.SetKeepAlivesEnabled(false)

		// then close the listener so new connections can't connect come thru
		closeErr := s.listener.Close()
		<-s.serveDone

		// then trigger the background goroutine to stop and wait for it
		stopDone := make(chan struct{})
		s.stop <- stopDone

		// wait for stop
		select {
		case <-stopDone:
		case <-s.clock.After(s.stopTimeout):
			defer stats.BumpTime(s.stats, "kill.time").End()
			stats.BumpSum(s.stats, "kill", 1)

			// stop timed out, wait for kill
			killDone := make(chan struct{})
			s.kill <- killDone
			select {
			case <-killDone:
			case <-s.clock.After(s.killTimeout):
				// kill timed out, give up
				stats.BumpSum(s.stats, "kill.timeout", 1)
			}
		}

		if closeErr != nil && !isUseOfClosedError(closeErr) {
			stats.BumpSum(s.stats, "listener.close.error", 1)
			s.stopErr = closeErr
		}
	})
	return s.stopErr
}

func isUseOfClosedError(err error) bool {
	if err == nil {
		return false
	}
	if opErr, ok := err.(*net.OpError); ok {
		err = opErr.Err
	}
	return err.Error() == "use of closed network connection"
}

// ListenAndServe is a convenience function to serve and wait for a SIGTERM
// or SIGINT before shutting down.
func ListenAndServe(s *http.Server, hd *HTTP) error {
	if hd == nil {
		hd = &HTTP{}
	}
	hs, err := hd.ListenAndServe(s)
	if err != nil {
		return err
	}
	log.Printf("serving on http://%s/ with pid %d\n", s.Addr, os.Getpid())

	waiterr := make(chan error, 1)
	go func() {
		defer close(waiterr)
		waiterr <- hs.Wait()
	}()

	signals := make(chan os.Signal, 10)
	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

	select {
	case err := <-waiterr:
		if err != nil {
			return err
		}
	case s := <-signals:
		signal.Stop(signals)
		log.Printf("signal received: %s\n", s)
		if err := hs.Stop(); err != nil {
			return err
		}
		if err := <-waiterr; err != nil {
			return err
		}
	}
	log.Println("exiting")
	return nil
}
<commit_msg>Fix typo in comment<commit_after>// Package httpdown provides http.ConnState enabled graceful termination of
// http.Server.
package httpdown

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/facebookgo/clock"
	"github.com/facebookgo/stats"
)

const (
	defaultStopTimeout = time.Minute
	defaultKillTimeout = time.Minute
)

// A Server allows encapsulates the process of accepting new connections and
// serving them, and gracefully shutting down the listener without dropping
// active connections.
type Server interface {
	// Wait waits for the serving loop to finish. This will happen when Stop is
	// called, at which point it returns no error, or if there is an error in the
	// serving loop. You must call Wait after calling Serve or ListenAndServe.
	Wait() error

	// Stop stops the listener. It will block until all connections have been
	// closed.
	Stop() error
}

// HTTP defines the configuration for serving a http.Server. Multiple calls to
// Serve or ListenAndServe can be made on the same HTTP instance. The default
// timeouts of 1 minute each result in a maximum of 2 minutes before a Stop()
// returns.
type HTTP struct {
	// StopTimeout is the duration before we begin force closing connections.
	// Defaults to 1 minute.
	StopTimeout time.Duration

	// KillTimeout is the duration before which we completely give up and abort
	// even though we still have connected clients. This is useful when a large
	// number of client connections exist and closing them can take a long time.
	// Note, this is in addition to the StopTimeout. Defaults to 1 minute.
	KillTimeout time.Duration

	// Stats is optional. If provided, it will be used to record various metrics.
	Stats stats.Client

	// Clock allows for testing timing related functionality. Do not specify this
	// in production code.
	Clock clock.Clock
}

// Serve provides the low-level API which is useful if you're creating your own
// net.Listener.
func (h HTTP) Serve(s *http.Server, l net.Listener) Server {
	stopTimeout := h.StopTimeout
	if stopTimeout == 0 {
		stopTimeout = defaultStopTimeout
	}
	killTimeout := h.KillTimeout
	if killTimeout == 0 {
		killTimeout = defaultKillTimeout
	}
	klock := h.Clock
	if klock == nil {
		klock = clock.New()
	}

	ss := &server{
		stopTimeout:  stopTimeout,
		killTimeout:  killTimeout,
		stats:        h.Stats,
		clock:        klock,
		oldConnState: s.ConnState,
		listener:     l,
		server:       s,
		serveDone:    make(chan struct{}),
		serveErr:     make(chan error, 1),
		new:          make(chan net.Conn),
		active:       make(chan net.Conn),
		idle:         make(chan net.Conn),
		closed:       make(chan net.Conn),
		stop:         make(chan chan struct{}),
		kill:         make(chan chan struct{}),
	}
	s.ConnState = ss.connState
	go ss.manage()
	go ss.serve()
	return ss
}

// ListenAndServe returns a Server for the given http.Server. It is equivalent
// to ListenAndServe from the standard library, but returns immediately.
// Requests will be accepted in a background goroutine. If the http.Server has
// a non-nil TLSConfig, a TLS enabled listener will be setup.
func (h HTTP) ListenAndServe(s *http.Server) (Server, error) {
	addr := s.Addr
	if addr == "" {
		if s.TLSConfig == nil {
			addr = ":http"
		} else {
			addr = ":https"
		}
	}
	l, err := net.Listen("tcp", addr)
	if err != nil {
		stats.BumpSum(h.Stats, "listen.error", 1)
		return nil, err
	}
	if s.TLSConfig != nil {
		l = tls.NewListener(l, s.TLSConfig)
	}
	return h.Serve(s, l), nil
}

// server manages the serving process and allows for gracefully stopping it.
type server struct {
	stopTimeout time.Duration
	killTimeout time.Duration
	stats       stats.Client
	clock       clock.Clock

	oldConnState func(net.Conn, http.ConnState)
	server       *http.Server
	serveDone    chan struct{}
	serveErr     chan error
	listener     net.Listener

	new    chan net.Conn
	active chan net.Conn
	idle   chan net.Conn
	closed chan net.Conn
	stop   chan chan struct{}
	kill   chan chan struct{}

	stopOnce sync.Once
	stopErr  error
}

func (s *server) connState(c net.Conn, cs http.ConnState) {
	if s.oldConnState != nil {
		s.oldConnState(c, cs)
	}

	switch cs {
	case http.StateNew:
		s.new <- c
	case http.StateActive:
		s.active <- c
	case http.StateIdle:
		s.idle <- c
	case http.StateHijacked, http.StateClosed:
		s.closed <- c
	}
}

func (s *server) manage() {
	defer func() {
		close(s.new)
		close(s.active)
		close(s.idle)
		close(s.closed)
		close(s.stop)
		close(s.kill)
	}()

	var stopDone chan struct{}

	conns := map[net.Conn]http.ConnState{}
	var countNew, countActive, countIdle float64

	// decConn decrements the count associated with the current state of the
	// given connection.
	decConn := func(c net.Conn) {
		switch conns[c] {
		default:
			panic(fmt.Errorf("unknown existing connection: %s", c))
		case http.StateNew:
			countNew--
		case http.StateActive:
			countActive--
		case http.StateIdle:
			countIdle--
		}
	}

	// setup a ticker to report various values every minute. if we don't have a
	// Stats implementation provided, we Stop it so it never ticks.
	statsTicker := s.clock.Ticker(time.Minute)
	if s.stats == nil {
		statsTicker.Stop()
	}

	for {
		select {
		case <-statsTicker.C:
			// we'll only get here when s.stats is not nil
			s.stats.BumpAvg("http-state.new", countNew)
			s.stats.BumpAvg("http-state.active", countActive)
			s.stats.BumpAvg("http-state.idle", countIdle)
			s.stats.BumpAvg("http-state.total", countNew+countActive+countIdle)
		case c := <-s.new:
			conns[c] = http.StateNew
			countNew++
		case c := <-s.active:
			decConn(c)
			countActive++

			conns[c] = http.StateActive
		case c := <-s.idle:
			decConn(c)
			countIdle++

			conns[c] = http.StateIdle

			// if we're already stopping, close it
			if stopDone != nil {
				c.Close()
			}
		case c := <-s.closed:
			stats.BumpSum(s.stats, "conn.closed", 1)
			decConn(c)
			delete(conns, c)

			// if we're waiting to stop and are all empty, we just closed the last
			// connection and we're done.
			if stopDone != nil && len(conns) == 0 {
				close(stopDone)
				return
			}
		case stopDone = <-s.stop:
			// if we're already all empty, we're already done
			if len(conns) == 0 {
				close(stopDone)
				return
			}

			// close current idle connections right away
			for c, cs := range conns {
				if cs == http.StateIdle {
					c.Close()
				}
			}

			// continue the loop and wait for all the ConnState updates which will
			// eventually close(stopDone) and return from this goroutine.

		case killDone := <-s.kill:
			// force close all connections
			stats.BumpSum(s.stats, "kill.conn.count", float64(len(conns)))
			for c := range conns {
				c.Close()
			}

			// don't block the kill.
			close(killDone)

			// continue the loop and we wait for all the ConnState updates and will
			// return from this goroutine when we're all done. otherwise we'll try to
			// send those ConnState updates on closed channels.

		}
	}
}

func (s *server) serve() {
	stats.BumpSum(s.stats, "serve", 1)
	s.serveErr <- s.server.Serve(s.listener)
	close(s.serveDone)
	close(s.serveErr)
}

func (s *server) Wait() error {
	if err := <-s.serveErr; !isUseOfClosedError(err) {
		return err
	}
	return nil
}

func (s *server) Stop() error {
	s.stopOnce.Do(func() {
		defer stats.BumpTime(s.stats, "stop.time").End()
		stats.BumpSum(s.stats, "stop", 1)

		// first disable keep-alive for new connections
		s.server.SetKeepAlivesEnabled(false)

		// then close the listener so new connections can't connect come thru
		closeErr := s.listener.Close()
		<-s.serveDone

		// then trigger the background goroutine to stop and wait for it
		stopDone := make(chan struct{})
		s.stop <- stopDone

		// wait for stop
		select {
		case <-stopDone:
		case <-s.clock.After(s.stopTimeout):
			defer stats.BumpTime(s.stats, "kill.time").End()
			stats.BumpSum(s.stats, "kill", 1)

			// stop timed out, wait for kill
			killDone := make(chan struct{})
			s.kill <- killDone
			select {
			case <-killDone:
			case <-s.clock.After(s.killTimeout):
				// kill timed out, give up
				stats.BumpSum(s.stats, "kill.timeout", 1)
			}
		}

		if closeErr != nil && !isUseOfClosedError(closeErr) {
			stats.BumpSum(s.stats, "listener.close.error", 1)
			s.stopErr = closeErr
		}
	})
	return s.stopErr
}

func isUseOfClosedError(err error) bool {
	if err == nil {
		return false
	}
	if opErr, ok := err.(*net.OpError); ok {
		err = opErr.Err
	}
	return err.Error() == "use of closed network connection"
}

// ListenAndServe is a convenience function to serve and wait for a SIGTERM
// or SIGINT before shutting down.
func ListenAndServe(s *http.Server, hd *HTTP) error {
	if hd == nil {
		hd = &HTTP{}
	}
	hs, err := hd.ListenAndServe(s)
	if err != nil {
		return err
	}
	log.Printf("serving on http://%s/ with pid %d\n", s.Addr, os.Getpid())

	waiterr := make(chan error, 1)
	go func() {
		defer close(waiterr)
		waiterr <- hs.Wait()
	}()

	signals := make(chan os.Signal, 10)
	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

	select {
	case err := <-waiterr:
		if err != nil {
			return err
		}
	case s := <-signals:
		signal.Stop(signals)
		log.Printf("signal received: %s\n", s)
		if err := hs.Stop(); err != nil {
			return err
		}
		if err := <-waiterr; err != nil {
			return err
		}
	}
	log.Println("exiting")
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 The Httpgzip Authors.
// Use of this source code is governed by an Expat-style
// MIT license that can be found in the LICENSE file.

// Package httpgzip implements an http.Handler wrapper adding gzip
// compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not do a simple string search for
// "gzip" (which will fail to do the correct thing for values such as
// "*" or "identity,gzip;q=0"). It will serve either gzip or identity
// content codings (identity meaning no encoding), or return 406 Not
// Acceptable status if it can do neither.
//
// It works correctly with handlers which honour Range request headers
// (such as http.FileServer) by removing the Range header for requests
// which prefer gzip encoding. This is necessary since Range requests
// apply to the gzipped content but the wrapped handler is not aware
// of the compression when it writes byte ranges. The Accept-Ranges
// header is also stripped from corresponding responses.
//
// For requests which prefer gzip encoding a Content-Type header is
// set using http.DetectContentType if it is not set by the wrapped
// handler.
//
// Gzip implementation
//
// By default, httpgzip uses the standard library gzip
// implementation. To use the optimized gzip implementation from
// https://github.com/klauspost/compress instead, download and install
// httpgzip with the "kpgzip" build tag:
//
//     go get -tags kpgzip github.com/xi2/httpgzip
//
// or simply alter the import line in httpgzip.go.
//
// Thanks
//
// Thanks are due to Klaus Post for his blog post which inspired the
// creation of this package and is recommended reading:
//
//     https://blog.klauspost.com/gzip-performance-for-go-webservers/
package httpgzip

import (
	"bytes"
	"fmt"
	"mime"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/xi2/httpgzip/internal/gzip"
)

// These constants are copied from the gzip package, so that code that
// imports this package does not also have to import the gzip package.
const (
	NoCompression      = gzip.NoCompression
	BestSpeed          = gzip.BestSpeed
	BestCompression    = gzip.BestCompression
	DefaultCompression = gzip.DefaultCompression
)

// DefaultContentTypes is the default list of content types for which
// a Handler considers gzip compression. This list originates from the
// file compression.conf within the Apache configuration found at
// https://html5boilerplate.com/.
var DefaultContentTypes = []string{
	"application/atom+xml",
	"application/font-sfnt",
	"application/javascript",
	"application/json",
	"application/ld+json",
	"application/manifest+json",
	"application/rdf+xml",
	"application/rss+xml",
	"application/schema+json",
	"application/vnd.geo+json",
	"application/vnd.ms-fontobject",
	"application/x-font-ttf",
	"application/x-javascript",
	"application/x-web-app-manifest+json",
	"application/xhtml+xml",
	"application/xml",
	"font/eot",
	"font/opentype",
	"image/bmp",
	"image/svg+xml",
	"image/vnd.microsoft.icon",
	"image/x-icon",
	"text/cache-manifest",
	"text/css",
	"text/html",
	"text/javascript",
	"text/plain",
	"text/vcard",
	"text/vnd.rim.location.xloc",
	"text/vtt",
	"text/x-component",
	"text/x-cross-domain-policy",
	"text/xml",
}

var gzipWriterPools = map[int]*sync.Pool{}

func init() {
	levels := map[int]struct{}{
		DefaultCompression: struct{}{},
		NoCompression:      struct{}{},
	}
	for i := BestSpeed; i <= BestCompression; i++ {
		levels[i] = struct{}{}
	}
	for k := range levels {
		level := k // create new variable for closure
		gzipWriterPools[level] = &sync.Pool{
			New: func() interface{} {
				w, _ := gzip.NewWriterLevel(nil, level)
				return w
			},
		}
	}
}

var gzipBufPool = sync.Pool{
	New: func() interface{} { return new(bytes.Buffer) },
}

// A gzipResponseWriter is a modified http.ResponseWriter. It adds
// gzip compression to certain responses, and there are two cases
// where this is done. Case 1 is when encs only allows gzip encoding
// and forbids identity. Case 2 is when encs prefers gzip encoding,
// the response is at least 512 bytes and the response's content type
// is in ctMap.
//
// A gzipResponseWriter sets the Content-Encoding and Content-Type
// headers when appropriate. It is important to call the Close method
// when writing is finished in order to flush and close the
// gzipResponseWriter. The slice encs must contain only encodings from
// {encGzip,encIdentity} and contain at least one encoding.
//
// If a gzip.Writer is used in order to write a response it will use a
// compression level of level.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus int
	ctMap      map[string]struct{}
	encs       []encoding
	level      int
	gw         *gzip.Writer
	buf        *bytes.Buffer
}

func newGzipResponseWriter(w http.ResponseWriter, ctMap map[string]struct{}, encs []encoding, level int) *gzipResponseWriter {
	buf := gzipBufPool.Get().(*bytes.Buffer)
	buf.Reset()
	return &gzipResponseWriter{
		ResponseWriter: w,
		httpStatus:     http.StatusOK,
		ctMap:          ctMap,
		encs:           encs,
		level:          level,
		buf:            buf}
}

// init gets called by Write once at least 512 bytes have been written
// to the temporary buffer buf, or by Close if it has not yet been
// called. Firstly it determines the content type, either from the
// Content-Type header, or by calling http.DetectContentType on
// buf. Then, if needed, a gzip.Writer is initialized. Lastly,
// appropriate headers are set and the ResponseWriter's WriteHeader
// method is called.
func (w *gzipResponseWriter) init() {
	cth := w.Header().Get("Content-Type")
	var ct string
	if cth != "" {
		ct = cth
	} else {
		ct = http.DetectContentType(w.buf.Bytes())
	}
	var gzipContentType bool
	if mt, _, err := mime.ParseMediaType(ct); err == nil {
		if _, ok := w.ctMap[mt]; ok {
			gzipContentType = true
		}
	}
	var useGzip bool
	if w.Header().Get("Content-Encoding") == "" && w.encs[0] == encGzip {
		if gzipContentType && w.buf.Len() >= 512 || len(w.encs) == 1 {
			useGzip = true
		}
	}
	if useGzip {
		w.gw = gzipWriterPools[w.level].Get().(*gzip.Writer)
		w.gw.Reset(w.ResponseWriter)
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Encoding", "gzip")
	}
	w.Header().Del("Accept-Ranges")
	if cth == "" {
		w.Header().Set("Content-Type", ct)
	}
	w.ResponseWriter.WriteHeader(w.httpStatus)
}

func (w *gzipResponseWriter) Write(p []byte) (int, error) {
	var n, written int
	var err error
	if w.buf != nil {
		written = w.buf.Len()
		_, _ = w.buf.Write(p)
		if w.buf.Len() < 512 {
			return len(p), nil
		}
		w.init()
		p = w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
	}
	switch {
	case w.gw != nil:
		n, err = w.gw.Write(p)
	default:
		n, err = w.ResponseWriter.Write(p)
	}
	n -= written
	if n < 0 {
		n = 0
	}
	return n, err
}

func (w *gzipResponseWriter) WriteHeader(httpStatus int) {
	// postpone WriteHeader call until end of init method
	w.httpStatus = httpStatus
}

func (w *gzipResponseWriter) Close() (err error) {
	if w.buf != nil {
		w.init()
		p := w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
		switch {
		case w.gw != nil:
			_, err = w.gw.Write(p)
		default:
			_, err = w.ResponseWriter.Write(p)
		}
	}
	if w.gw != nil {
		e := w.gw.Close()
		if e != nil && err == nil {
			err = e
		}
		gzipWriterPools[w.level].Put(w.gw)
		w.gw = nil
	}
	return err
}

// An encoding is a supported content coding.
type encoding int

const (
	encIdentity encoding = iota
	encGzip
)

// acceptedEncodings returns the supported content codings that are
// accepted by the request r. It returns a slice of encodings in
// client preference order.
//
// If the Sec-WebSocket-Key header is present then compressed content
// encodings are not considered.
//
// ref: http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html
func acceptedEncodings(r *http.Request) []encoding {
	h := r.Header.Get("Accept-Encoding")
	swk := r.Header.Get("Sec-WebSocket-Key")
	if h == "" {
		return []encoding{encIdentity}
	}
	gzip := float64(-1)    // -1 means not accepted, 0 -> 1 means value of q
	identity := float64(0) // -1 means not accepted, 0 -> 1 means value of q
	for _, s := range strings.Split(h, ",") {
		f := strings.Split(s, ";")
		f0 := strings.ToLower(strings.Trim(f[0], " "))
		q := float64(1.0)
		if len(f) > 1 {
			f1 := strings.ToLower(strings.Trim(f[1], " "))
			if strings.HasPrefix(f1, "q=") {
				if flt, err := strconv.ParseFloat(f1[2:], 32); err == nil {
					if flt >= 0 && flt <= 1 {
						q = flt
					}
				}
			}
		}
		if (f0 == "gzip" || f0 == "*") && q > gzip && swk == "" {
			gzip = q
		}
		if (f0 == "gzip" || f0 == "*") && q == 0 {
			gzip = -1
		}
		if (f0 == "identity" || f0 == "*") && q > identity {
			identity = q
		}
		if (f0 == "identity" || f0 == "*") && q == 0 {
			identity = -1
		}
	}
	switch {
	case gzip == -1 && identity == -1:
		return []encoding{}
	case gzip == -1:
		return []encoding{encIdentity}
	case identity == -1:
		return []encoding{encGzip}
	case identity > gzip:
		return []encoding{encIdentity, encGzip}
	default:
		return []encoding{encGzip, encIdentity}
	}
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding gzip compression to certain responses. There are two cases
// where gzip compression is done. Case 1 is responses whose requests
// only allow gzip encoding and forbid identity encoding (identity
// encoding meaning no encoding). Case 2 is responses whose requests
// prefer gzip encoding, whose size is at least 512 bytes and whose
// content types are in contentTypes. If contentTypes is nil then
// DefaultContentTypes is considered instead.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If a request
// expresses a preference for gzip encoding then any Range headers are
// removed from the request before it is passed through to h and
// Accept-Ranges headers are stripped from corresponding
// responses. This happens regardless of whether gzip encoding is
// eventually used in the response or not.
func NewHandler(h http.Handler, contentTypes []string) http.Handler {
	gzh, _ := NewHandlerLevel(h, contentTypes, DefaultCompression)
	return gzh
}

// NewHandlerLevel is like NewHandler but allows one to specify the
// gzip compression level instead of assuming DefaultCompression.
//
// The compression level can be DefaultCompression, NoCompression, or
// any integer value between BestSpeed and BestCompression
// inclusive. The error returned will be nil if the level is valid.
func NewHandlerLevel(h http.Handler, contentTypes []string, level int) (http.Handler, error) {
	switch {
	case level == DefaultCompression || level == NoCompression:
		// no action needed
	case level < BestSpeed || level > BestCompression:
		return nil, fmt.Errorf(
			"httpgzip: invalid compression level: %d", level)
	}
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	ctMap := map[string]struct{}{}
	for _, ct := range contentTypes {
		ctMap[ct] = struct{}{}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		// return if no acceptable encodings
		if len(encs) == 0 {
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}
		if encs[0] == encGzip {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
			// create new ResponseWriter
			w = newGzipResponseWriter(w, ctMap, encs, level)
			defer w.(*gzipResponseWriter).Close()
		}
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	}), nil
}
<commit_msg>Fix faulty parsing of Accept-Encoding header<commit_after>// Copyright (c) 2015 The Httpgzip Authors.
// Use of this source code is governed by an Expat-style
// MIT license that can be found in the LICENSE file.

// Package httpgzip implements an http.Handler wrapper adding gzip
// compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not do a simple string search for
// "gzip" (which will fail to do the correct thing for values such as
// "*" or "identity,gzip;q=0"). It will serve either gzip or identity
// content codings (identity meaning no encoding), or return 406 Not
// Acceptable status if it can do neither.
//
// It works correctly with handlers which honour Range request headers
// (such as http.FileServer) by removing the Range header for requests
// which prefer gzip encoding. This is necessary since Range requests
// apply to the gzipped content but the wrapped handler is not aware
// of the compression when it writes byte ranges. The Accept-Ranges
// header is also stripped from corresponding responses.
//
// For requests which prefer gzip encoding a Content-Type header is
// set using http.DetectContentType if it is not set by the wrapped
// handler.
//
// Gzip implementation
//
// By default, httpgzip uses the standard library gzip
// implementation. To use the optimized gzip implementation from
// https://github.com/klauspost/compress instead, download and install
// httpgzip with the "kpgzip" build tag:
//
//     go get -tags kpgzip github.com/xi2/httpgzip
//
// or simply alter the import line in httpgzip.go.
//
// Thanks
//
// Thanks are due to Klaus Post for his blog post which inspired the
// creation of this package and is recommended reading:
//
//     https://blog.klauspost.com/gzip-performance-for-go-webservers/
package httpgzip

import (
	"bytes"
	"fmt"
	"mime"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/xi2/httpgzip/internal/gzip"
)

// These constants are copied from the gzip package, so that code that
// imports this package does not also have to import the gzip package.
const (
	NoCompression      = gzip.NoCompression
	BestSpeed          = gzip.BestSpeed
	BestCompression    = gzip.BestCompression
	DefaultCompression = gzip.DefaultCompression
)

// DefaultContentTypes is the default list of content types for which
// a Handler considers gzip compression. This list originates from the
// file compression.conf within the Apache configuration found at
// https://html5boilerplate.com/.
var DefaultContentTypes = []string{
	"application/atom+xml",
	"application/font-sfnt",
	"application/javascript",
	"application/json",
	"application/ld+json",
	"application/manifest+json",
	"application/rdf+xml",
	"application/rss+xml",
	"application/schema+json",
	"application/vnd.geo+json",
	"application/vnd.ms-fontobject",
	"application/x-font-ttf",
	"application/x-javascript",
	"application/x-web-app-manifest+json",
	"application/xhtml+xml",
	"application/xml",
	"font/eot",
	"font/opentype",
	"image/bmp",
	"image/svg+xml",
	"image/vnd.microsoft.icon",
	"image/x-icon",
	"text/cache-manifest",
	"text/css",
	"text/html",
	"text/javascript",
	"text/plain",
	"text/vcard",
	"text/vnd.rim.location.xloc",
	"text/vtt",
	"text/x-component",
	"text/x-cross-domain-policy",
	"text/xml",
}

var gzipWriterPools = map[int]*sync.Pool{}

func init() {
	levels := map[int]struct{}{
		DefaultCompression: struct{}{},
		NoCompression:      struct{}{},
	}
	for i := BestSpeed; i <= BestCompression; i++ {
		levels[i] = struct{}{}
	}
	for k := range levels {
		level := k // create new variable for closure
		gzipWriterPools[level] = &sync.Pool{
			New: func() interface{} {
				w, _ := gzip.NewWriterLevel(nil, level)
				return w
			},
		}
	}
}

var gzipBufPool = sync.Pool{
	New: func() interface{} { return new(bytes.Buffer) },
}

// A gzipResponseWriter is a modified http.ResponseWriter. It adds
// gzip compression to certain responses, and there are two cases
// where this is done. Case 1 is when encs only allows gzip encoding
// and forbids identity. Case 2 is when encs prefers gzip encoding,
// the response is at least 512 bytes and the response's content type
// is in ctMap.
//
// A gzipResponseWriter sets the Content-Encoding and Content-Type
// headers when appropriate. It is important to call the Close method
// when writing is finished in order to flush and close the
// gzipResponseWriter. The slice encs must contain only encodings from
// {encGzip,encIdentity} and contain at least one encoding.
//
// If a gzip.Writer is used in order to write a response it will use a
// compression level of level.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus int
	ctMap      map[string]struct{}
	encs       []encoding
	level      int
	gw         *gzip.Writer
	buf        *bytes.Buffer
}

func newGzipResponseWriter(w http.ResponseWriter, ctMap map[string]struct{}, encs []encoding, level int) *gzipResponseWriter {
	buf := gzipBufPool.Get().(*bytes.Buffer)
	buf.Reset()
	return &gzipResponseWriter{
		ResponseWriter: w,
		httpStatus:     http.StatusOK,
		ctMap:          ctMap,
		encs:           encs,
		level:          level,
		buf:            buf}
}

// init gets called by Write once at least 512 bytes have been written
// to the temporary buffer buf, or by Close if it has not yet been
// called. Firstly it determines the content type, either from the
// Content-Type header, or by calling http.DetectContentType on
// buf. Then, if needed, a gzip.Writer is initialized. Lastly,
// appropriate headers are set and the ResponseWriter's WriteHeader
// method is called.
func (w *gzipResponseWriter) init() {
	cth := w.Header().Get("Content-Type")
	var ct string
	if cth != "" {
		ct = cth
	} else {
		ct = http.DetectContentType(w.buf.Bytes())
	}
	var gzipContentType bool
	if mt, _, err := mime.ParseMediaType(ct); err == nil {
		if _, ok := w.ctMap[mt]; ok {
			gzipContentType = true
		}
	}
	var useGzip bool
	if w.Header().Get("Content-Encoding") == "" && w.encs[0] == encGzip {
		if gzipContentType && w.buf.Len() >= 512 || len(w.encs) == 1 {
			useGzip = true
		}
	}
	if useGzip {
		w.gw = gzipWriterPools[w.level].Get().(*gzip.Writer)
		w.gw.Reset(w.ResponseWriter)
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Encoding", "gzip")
	}
	w.Header().Del("Accept-Ranges")
	if cth == "" {
		w.Header().Set("Content-Type", ct)
	}
	w.ResponseWriter.WriteHeader(w.httpStatus)
}

func (w *gzipResponseWriter) Write(p []byte) (int, error) {
	var n, written int
	var err error
	if w.buf != nil {
		written = w.buf.Len()
		_, _ = w.buf.Write(p)
		if w.buf.Len() < 512 {
			return len(p), nil
		}
		w.init()
		p = w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
	}
	switch {
	case w.gw != nil:
		n, err = w.gw.Write(p)
	default:
		n, err = w.ResponseWriter.Write(p)
	}
	n -= written
	if n < 0 {
		n = 0
	}
	return n, err
}

func (w *gzipResponseWriter) WriteHeader(httpStatus int) {
	// postpone WriteHeader call until end of init method
	w.httpStatus = httpStatus
}

func (w *gzipResponseWriter) Close() (err error) {
	if w.buf != nil {
		w.init()
		p := w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
		switch {
		case w.gw != nil:
			_, err = w.gw.Write(p)
		default:
			_, err = w.ResponseWriter.Write(p)
		}
	}
	if w.gw != nil {
		e := w.gw.Close()
		if e != nil && err == nil {
			err = e
		}
		gzipWriterPools[w.level].Put(w.gw)
		w.gw = nil
	}
	return err
}

// An encoding is a supported content coding.
type encoding int

const (
	encIdentity encoding = iota
	encGzip
)

// acceptedEncodings returns the supported content codings that are
// accepted by the request r. It returns a slice of encodings in
// client preference order.
//
// If the Sec-WebSocket-Key header is present then compressed content
// encodings are not considered.
//
// ref: http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html
func acceptedEncodings(r *http.Request) []encoding {
	h := r.Header.Get("Accept-Encoding")
	swk := r.Header.Get("Sec-WebSocket-Key")
	if h == "" {
		return []encoding{encIdentity}
	}
	gzip := float64(-1)     // q-value: -1 means "not present in header"
	identity := float64(-1) // q-value: -1 means "not present in header"
	any := float64(-1)      // q-value: -1 means "not present in header"
	for _, s := range strings.Split(h, ",") {
		f := strings.Split(s, ";")
		f0 := strings.ToLower(strings.Trim(f[0], " "))
		q := float64(1.0)
		if len(f) > 1 {
			f1 := strings.ToLower(strings.Trim(f[1], " "))
			if strings.HasPrefix(f1, "q=") {
				if flt, err := strconv.ParseFloat(f1[2:], 32); err == nil {
					if flt >= 0 && flt <= 1 {
						q = flt
					}
				}
			}
		}
		if f0 == "gzip" && q > gzip && swk == "" {
			gzip = q
		}
		if f0 == "identity" && q > identity {
			identity = q
		}
		if f0 == "*" && q > any {
			any = q
		}
	}
	if identity == -1 {
		if any >= 0 {
			identity = any
		} else {
			identity = 1
		}
	}
	if gzip == -1 && any >= 0 {
		gzip = any
	}
	switch {
	case gzip <= 0 && identity <= 0:
		return []encoding{}
	case gzip <= 0:
		return []encoding{encIdentity}
	case identity <= 0:
		return []encoding{encGzip}
	case identity > gzip:
		return []encoding{encIdentity, encGzip}
	default:
		return []encoding{encGzip, encIdentity}
	}
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding gzip compression to certain responses. There are two cases
// where gzip compression is done. Case 1 is responses whose requests
// only allow gzip encoding and forbid identity encoding (identity
// encoding meaning no encoding). Case 2 is responses whose requests
// prefer gzip encoding, whose size is at least 512 bytes and whose
// content types are in contentTypes. If contentTypes is nil then
// DefaultContentTypes is considered instead.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If a request
// expresses a preference for gzip encoding then any Range headers are
// removed from the request before it is passed through to h and
// Accept-Ranges headers are stripped from corresponding
// responses. This happens regardless of whether gzip encoding is
// eventually used in the response or not.
func NewHandler(h http.Handler, contentTypes []string) http.Handler {
	gzh, _ := NewHandlerLevel(h, contentTypes, DefaultCompression)
	return gzh
}

// NewHandlerLevel is like NewHandler but allows one to specify the
// gzip compression level instead of assuming DefaultCompression.
//
// The compression level can be DefaultCompression, NoCompression, or
// any integer value between BestSpeed and BestCompression
// inclusive. The error returned will be nil if the level is valid.
func NewHandlerLevel(h http.Handler, contentTypes []string, level int) (http.Handler, error) {
	switch {
	case level == DefaultCompression || level == NoCompression:
		// no action needed
	case level < BestSpeed || level > BestCompression:
		return nil, fmt.Errorf(
			"httpgzip: invalid compression level: %d", level)
	}
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	ctMap := map[string]struct{}{}
	for _, ct := range contentTypes {
		ctMap[ct] = struct{}{}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		// return if no acceptable encodings
		if len(encs) == 0 {
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}
		if encs[0] == encGzip {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
			// create new ResponseWriter
			w = newGzipResponseWriter(w, ctMap, encs, level)
			defer w.(*gzipResponseWriter).Close()
		}
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	}), nil
}
<|endoftext|>


<commit_before>package httpunix

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"sync"
	"time"
)

const Scheme = "http+unix"

type HTTPUnixTransport struct {
	DialTimeout           time.Duration
	RequestTimeout        time.Duration
	ResponseHeaderTimeout time.Duration

	mu sync.Mutex
	// map a URL "hostname" to a UNIX domain socket path
	loc map[string]string
}

func (u *HTTPUnixTransport) RegisterLocation(loc string, path string) {
	u.mu.Lock()
	defer u.mu.Unlock()
	if u.loc == nil {
		u.loc = make(map[string]string)
	}
	if _, exists := u.loc[loc]; exists {
		panic("location " + loc + " already registered")
	}
	u.loc[loc] = path
}

func (t *HTTPUnixTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL == nil {
		return nil, errors.New("http+unix: nil Request.URL")
	}
	if req.URL.Scheme != Scheme {
		return nil, errors.New("unsupported protocol scheme: " + req.URL.Scheme)
	}
	if req.URL.Host == "" {
		return nil, errors.New("http+unix: no Host in request URL")
	}
	t.mu.Lock()
	path, ok := t.loc[req.URL.Host]
	t.mu.Unlock()
	if !ok {
		return nil, errors.New("unknown location: " + req.Host)
	}

	c, err := net.DialTimeout("unix", path, t.DialTimeout)
	if err != nil {
		return nil, err
	}
	r := bufio.NewReader(c)
	if t.RequestTimeout > 0 {
		c.SetWriteDeadline(time.Now().Add(t.RequestTimeout))
	}
	if err := req.Write(c); err != nil {
		return nil, err
	}
	if t.ResponseHeaderTimeout > 0 {
		c.SetReadDeadline(time.Now().Add(t.ResponseHeaderTimeout))
	}
	resp, err := http.ReadResponse(r, req)
	return resp, err
}
<commit_msg>Ensure Transport is a http.RoundTripper<commit_after>package httpunix

import (
	"bufio"
	"errors"
	"net"
	"net/http"
	"sync"
	"time"
)

const Scheme = "http+unix"

type HTTPUnixTransport struct {
	DialTimeout           time.Duration
	RequestTimeout        time.Duration
	ResponseHeaderTimeout time.Duration

	mu sync.Mutex
	// map a URL "hostname" to a UNIX domain socket path
	loc map[string]string
}

func (u *HTTPUnixTransport) RegisterLocation(loc string, path string) {
	u.mu.Lock()
	defer u.mu.Unlock()
	if u.loc == nil {
		u.loc = make(map[string]string)
	}
	if _, exists := u.loc[loc]; exists {
		panic("location " + loc + " already registered")
	}
	u.loc[loc] = path
}

var _ http.RoundTripper = (*HTTPUnixTransport)(nil)

func (t *HTTPUnixTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL == nil {
		return nil, errors.New("http+unix: nil Request.URL")
	}
	if req.URL.Scheme != Scheme {
		return nil, errors.New("unsupported protocol scheme: " + req.URL.Scheme)
	}
	if req.URL.Host == "" {
		return nil, errors.New("http+unix: no Host in request URL")
	}
	t.mu.Lock()
	path, ok := t.loc[req.URL.Host]
	t.mu.Unlock()
	if !ok {
		return nil, errors.New("unknown location: " + req.Host)
	}

	c, err := net.DialTimeout("unix", path, t.DialTimeout)
	if err != nil {
		return nil, err
	}
	r := bufio.NewReader(c)
	if t.RequestTimeout > 0 {
		c.SetWriteDeadline(time.Now().Add(t.RequestTimeout))
	}
	if err := req.Write(c); err != nil {
		return nil, err
	}
	if t.ResponseHeaderTimeout > 0 {
		c.SetReadDeadline(time.Now().Add(t.ResponseHeaderTimeout))
	}
	resp, err := http.ReadResponse(r, req)
	return resp, err
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

type Transferable interface {
	Check() (*objectResource, *WrappedError)
	Transfer(CopyCallback) *WrappedError
	Object() *objectResource
	Oid() string
	Size() int64
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	transferc        chan Transferable
	errorc           chan *WrappedError
	watchers         []chan string
	errors           []*WrappedError
	wg               sync.WaitGroup
	workers          int
	files            int
	finished         int64
	size             int64
	authCond         *sync.Cond
	transferables    map[string]Transferable
	bar              *pb.ProgressBar
	clientAuthorized int32
	transferKind     string
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(workers, files int) *TransferQueue {
	return &TransferQueue{
		transferc:     make(chan Transferable, files),
		errorc:        make(chan *WrappedError),
		watchers:      make([]chan string, 0),
		workers:       workers,
		files:         files,
		authCond:      sync.NewCond(&sync.Mutex{}),
		transferables: make(map[string]Transferable),
	}
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.transferables[t.Oid()] = t
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, q.files)
	q.watchers = append(q.watchers, c)
	return c
}

// processIndividual processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) processIndividual() {
	apic := make(chan Transferable, q.files)
	workersReady := make(chan int, q.workers)
	var wg sync.WaitGroup

	for i := 0; i < q.workers; i++ {
		go func() {
			workersReady <- 1
			for t := range apic {
				// If an API authorization has not occured, we wait until we're woken up.
				q.authCond.L.Lock()
				if atomic.LoadInt32(&q.clientAuthorized) == 0 {
					q.authCond.Wait()
				}
				q.authCond.L.Unlock()

				obj, err := t.Check()
				if err != nil {
					q.errorc <- err
					wg.Done()
					continue
				}
				if obj != nil {
					q.wg.Add(1)
					t.SetObject(obj)
					q.transferc <- t
				}
				wg.Done()
			}
		}()
	}

	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, len(q.transferables)))
	q.bar.Start()

	for _, t := range q.transferables {
		wg.Add(1)
		apic <- t
	}

	<-workersReady
	q.authCond.Signal() // Signal the first goroutine to run
	close(apic)
	wg.Wait()

	close(q.transferc)
}

// processBatch processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) processBatch() error {
	transfers := make([]*objectResource, 0, len(q.transferables))
	for _, t := range q.transferables {
		transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
	}

	objects, err := Batch(transfers)
	if err != nil {
		if isNotImplError(err) {
			tracerx.Printf("queue: batch not implemented, disabling")
			configFile := filepath.Join(LocalGitDir, "config")
			git.Config.SetLocal(configFile, "lfs.batch", "false")
		}

		return err
	}

	q.files = 0

	for _, o := range objects {
		if _, ok := o.Links[q.transferKind]; ok {
			// This object needs to be transfered
			if transfer, ok := q.transferables[o.Oid]; ok {
				q.files++
				q.wg.Add(1)
				transfer.SetObject(o)
				q.transferc <- transfer
			}
		}
	}

	close(q.transferc)
	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, q.files))
	q.bar.Start()
	sendApiEvent(apiEventSuccess) // Wake up transfer workers
	return nil
}

// Process starts the transfer queue and displays a progress bar. Process will
// do individual or batch transfers depending on the Config.BatchTransfer() value.
// Process will transfer files sequentially or concurrently depending on the
// Concig.ConcurrentTransfers() value.
func (q *TransferQueue) Process() {
	q.bar = pb.New64(q.size)
	q.bar.SetUnits(pb.U_BYTES)
	q.bar.ShowBar = false

	// This goroutine collects errors returned from transfers
	go func() {
		for err := range q.errorc {
			q.errors = append(q.errors, err)
		}
	}()

	// This goroutine watches for apiEvents. In order to prevent multiple
	// credential requests from happening, the queue is processed sequentially
	// until an API request succeeds (meaning authenication has happened successfully).
	// Once the an API request succeeds, all worker goroutines are woken up and allowed
	// to process transfers. Once a success happens, this goroutine exits.
	go func() {
		for {
			event := <-apiEvent
			switch event {
			case apiEventSuccess:
				atomic.StoreInt32(&q.clientAuthorized, 1)
				q.authCond.Broadcast() // Wake all remaining goroutines
				return
			case apiEventFail:
				q.authCond.Signal() // Wake the next goroutine
			}
		}
	}()

	for i := 0; i < q.workers; i++ {
		// These are the worker goroutines that process transfers
		go func() {
			for transfer := range q.transferc {
				cb := func(total, read int64, current int) error {
					q.bar.Add(current)
					return nil
				}

				if err := transfer.Transfer(cb); err != nil {
					q.errorc <- err
				} else {
					oid := transfer.Oid()
					for _, c := range q.watchers {
						c <- oid
					}
				}

				f := atomic.AddInt64(&q.finished, 1)
				q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.files))
				q.wg.Done()
			}
		}()
	}

	if Config.BatchTransfer() {
		if err := q.processBatch(); err != nil {
			q.processIndividual()
		}
	} else {
		q.processIndividual()
	}

	q.wg.Wait()
	close(q.errorc)
	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.bar.Finish()
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []*WrappedError {
	return q.errors
}
<commit_msg>ララララー ラララー ララ<commit_after>package lfs

import (
	"fmt"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/cheggaaa/pb"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

type Transferable interface {
	Check() (*objectResource, *WrappedError)
	Transfer(CopyCallback) *WrappedError
	Object() *objectResource
	Oid() string
	Size() int64
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	transferc        chan Transferable
	errorc           chan *WrappedError
	watchers         []chan string
	errors           []*WrappedError
	wg               sync.WaitGroup
	workers          int
	files            int
	finished         int64
	size             int64
	authCond         *sync.Cond
	transferables    map[string]Transferable
	bar              *pb.ProgressBar
	clientAuthorized int32
	transferKind     string
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(workers, files int) *TransferQueue {
	return &TransferQueue{
		transferc:     make(chan Transferable, files),
		errorc:        make(chan *WrappedError),
		watchers:      make([]chan string, 0),
		workers:       workers,
		files:         files,
		authCond:      sync.NewCond(&sync.Mutex{}),
		transferables: make(map[string]Transferable),
	}
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.transferables[t.Oid()] = t
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, q.files)
	q.watchers = append(q.watchers, c)
	return c
}

// processIndividual processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) processIndividual() {
	apic := make(chan Transferable, q.files)
	workersReady := make(chan int, q.workers)
	var wg sync.WaitGroup

	for i := 0; i < q.workers; i++ {
		go func() {
			for t := range apic {
				// If an API authorization has not occured, we wait until we're woken up.
				q.authCond.L.Lock()
				if atomic.LoadInt32(&q.clientAuthorized) == 0 {
					workersReady <- 1
					q.authCond.Wait()
				}
				q.authCond.L.Unlock()

				obj, err := t.Check()
				if err != nil {
					q.errorc <- err
					wg.Done()
					continue
				}
				if obj != nil {
					q.wg.Add(1)
					t.SetObject(obj)
					q.transferc <- t
				}
				wg.Done()
			}
		}()
	}

	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, len(q.transferables)))
	q.bar.Start()

	for _, t := range q.transferables {
		wg.Add(1)
		apic <- t
	}

	go func() {
		wg.Wait()
		close(workersReady)
	}()

	<-workersReady
	q.authCond.L.Lock()
	q.authCond.Signal() // Signal the first goroutine to run
	q.authCond.L.Unlock()

	close(apic)
	for _ = range workersReady {
	}

	close(q.transferc)
}

// processBatch processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) processBatch() error {
	transfers := make([]*objectResource, 0, len(q.transferables))
	for _, t := range q.transferables {
		transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
	}

	objects, err := Batch(transfers)
	if err != nil {
		if isNotImplError(err) {
			tracerx.Printf("queue: batch not implemented, disabling")
			configFile := filepath.Join(LocalGitDir, "config")
			git.Config.SetLocal(configFile, "lfs.batch", "false")
		}

		return err
	}

	q.files = 0

	for _, o := range objects {
		if _, ok := o.Links[q.transferKind]; ok {
			// This object needs to be transfered
			if transfer, ok := q.transferables[o.Oid]; ok {
				q.files++
				q.wg.Add(1)
				transfer.SetObject(o)
				q.transferc <- transfer
			}
		}
	}

	close(q.transferc)
	q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, q.files))
	q.bar.Start()
	sendApiEvent(apiEventSuccess) // Wake up transfer workers
	return nil
}

// Process starts the transfer queue and displays a progress bar. Process will
// do individual or batch transfers depending on the Config.BatchTransfer() value.
// Process will transfer files sequentially or concurrently depending on the
// Concig.ConcurrentTransfers() value.
func (q *TransferQueue) Process() {
	q.bar = pb.New64(q.size)
	q.bar.SetUnits(pb.U_BYTES)
	q.bar.ShowBar = false

	// This goroutine collects errors returned from transfers
	go func() {
		for err := range q.errorc {
			q.errors = append(q.errors, err)
		}
	}()

	// This goroutine watches for apiEvents. In order to prevent multiple
	// credential requests from happening, the queue is processed sequentially
	// until an API request succeeds (meaning authenication has happened successfully).
	// Once the an API request succeeds, all worker goroutines are woken up and allowed
	// to process transfers. Once a success happens, this goroutine exits.
	go func() {
		for {
			event := <-apiEvent
			switch event {
			case apiEventSuccess:
				atomic.StoreInt32(&q.clientAuthorized, 1)
				q.authCond.Broadcast() // Wake all remaining goroutines
				return
			case apiEventFail:
				q.authCond.Signal() // Wake the next goroutine
			}
		}
	}()

	for i := 0; i < q.workers; i++ {
		// These are the worker goroutines that process transfers
		go func() {
			for transfer := range q.transferc {
				cb := func(total, read int64, current int) error {
					q.bar.Add(current)
					return nil
				}

				if err := transfer.Transfer(cb); err != nil {
					q.errorc <- err
				} else {
					oid := transfer.Oid()
					for _, c := range q.watchers {
						c <- oid
					}
				}

				f := atomic.AddInt64(&q.finished, 1)
				q.bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.files))
				q.wg.Done()
			}
		}()
	}

	if Config.BatchTransfer() {
		if err := q.processBatch(); err != nil {
			q.processIndividual()
		}
	} else {
		q.processIndividual()
	}

	q.wg.Wait()
	close(q.errorc)
	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.bar.Finish()
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []*WrappedError {
	return q.errors
}
<|endoftext|>