<commit_before>package iso639

// ShortToLong map 2 letter to 3 letter langauge code
// Based on https://www.loc.gov/standards/iso639-2/ISO-639-2_utf-8.txt
// Generated with:
// cat ISO-639-2_utf-8.txt | awk '-F|' '{ printf "\"%s\": \"%s\",\n",$3,$1 }' | grep -v '^"":'
var ShortToLong = map[string]string{
	"aa": "aar",
	"ab": "abk",
	"af": "afr",
	"ak": "aka",
	"sq": "alb",
	"am": "amh",
	"ar": "ara",
	"an": "arg",
	"hy": "arm",
	"as": "asm",
	"av": "ava",
	"ae": "ave",
	"ay": "aym",
	"az": "aze",
	"ba": "bak",
	"bm": "bam",
	"eu": "baq",
	"be": "bel",
	"bn": "ben",
	"bh": "bih",
	"bi": "bis",
	"bs": "bos",
	"br": "bre",
	"bg": "bul",
	"my": "bur",
	"ca": "cat",
	"ch": "cha",
	"ce": "che",
	"zh": "chi",
	"cu": "chu",
	"cv": "chv",
	"kw": "cor",
	"co": "cos",
	"cr": "cre",
	"cs": "cze",
	"da": "dan",
	"dv": "div",
	"nl": "dut",
	"dz": "dzo",
	"en": "eng",
	"eo": "epo",
	"et": "est",
	"ee": "ewe",
	"fo": "fao",
	"fj": "fij",
	"fi": "fin",
	"fr": "fre",
	"fy": "fry",
	"ff": "ful",
	"ka": "geo",
	"de": "ger",
	"gd": "gla",
	"ga": "gle",
	"gl": "glg",
	"gv": "glv",
	"el": "gre",
	"gn": "grn",
	"gu": "guj",
	"ht": "hat",
	"ha": "hau",
	"he": "heb",
	"hz": "her",
	"hi": "hin",
	"ho": "hmo",
	"hr": "hrv",
	"hu": "hun",
	"ig": "ibo",
	"is": "ice",
	"io": "ido",
	"ii": "iii",
	"iu": "iku",
	"ie": "ile",
	"ia": "ina",
	"id": "ind",
	"ik": "ipk",
	"it": "ita",
	"jv": "jav",
	"ja": "jpn",
	"kl": "kal",
	"kn": "kan",
	"ks": "kas",
	"kr": "kau",
	"kk": "kaz",
	"km": "khm",
	"ki": "kik",
	"rw": "kin",
	"ky": "kir",
	"kv": "kom",
	"kg": "kon",
	"ko": "kor",
	"kj": "kua",
	"ku": "kur",
	"lo": "lao",
	"la": "lat",
	"lv": "lav",
	"li": "lim",
	"ln": "lin",
	"lt": "lit",
	"lb": "ltz",
	"lu": "lub",
	"lg": "lug",
	"mk": "mac",
	"mh": "mah",
	"ml": "mal",
	"mi": "mao",
	"mr": "mar",
	"ms": "may",
	"mg": "mlg",
	"mt": "mlt",
	"mn": "mon",
	"na": "nau",
	"nv": "nav",
	"nr": "nbl",
	"nd": "nde",
	"ng": "ndo",
	"ne": "nep",
	"nn": "nno",
	"nb": "nob",
	"no": "nor",
	"ny": "nya",
	"oc": "oci",
	"oj": "oji",
	"or": "ori",
	"om": "orm",
	"os": "oss",
	"pa": "pan",
	"fa": "per",
	"pi": "pli",
	"pl": "pol",
	"pt": "por",
	"ps": "pus",
	"qu": "que",
	"rm": "roh",
	"ro": "rum",
	"rn": "run",
	"ru": "rus",
	"sg": "sag",
	"sa": "san",
	"si": "sin",
	"sk": "slo",
	"sl": "slv",
	"se": "sme",
	"sm": "smo",
	"sn": "sna",
	"sd": "snd",
	"so": "som",
	"st": "sot",
	"es": "spa",
	"sc": "srd",
	"sr": "srp",
	"ss": "ssw",
	"su": "sun",
	"sw": "swa",
	"sv": "swe",
	"ty": "tah",
	"ta": "tam",
	"tt": "tat",
	"te": "tel",
	"tg": "tgk",
	"tl": "tgl",
	"th": "tha",
	"bo": "tib",
	"ti": "tir",
	"to": "ton",
	"tn": "tsn",
	"ts": "tso",
	"tk": "tuk",
	"tr": "tur",
	"tw": "twi",
	"ug": "uig",
	"uk": "ukr",
	"ur": "urd",
	"uz": "uzb",
	"ve": "ven",
	"vi": "vie",
	"vo": "vol",
	"cy": "wel",
	"wa": "wln",
	"wo": "wol",
	"xh": "xho",
	"yi": "yid",
	"yo": "yor",
	"za": "zha",
	"zu": "zul",
}
<commit_msg>Fix typo s/langauge/language<commit_after>package iso639

// ShortToLong map 2 letter to 3 letter language code
// Based on https://www.loc.gov/standards/iso639-2/ISO-639-2_utf-8.txt
// Generated with:
// cat ISO-639-2_utf-8.txt | awk '-F|' '{ printf "\"%s\": \"%s\",\n",$3,$1 }' | grep -v '^"":'
var ShortToLong = map[string]string{
	"aa": "aar",
	"ab": "abk",
	"af": "afr",
	"ak": "aka",
	"sq": "alb",
	"am": "amh",
	"ar": "ara",
	"an": "arg",
	"hy": "arm",
	"as": "asm",
	"av": "ava",
	"ae": "ave",
	"ay": "aym",
	"az": "aze",
	"ba": "bak",
	"bm": "bam",
	"eu": "baq",
	"be": "bel",
	"bn": "ben",
	"bh": "bih",
	"bi": "bis",
	"bs": "bos",
	"br": "bre",
	"bg": "bul",
	"my": "bur",
	"ca": "cat",
	"ch": "cha",
	"ce": "che",
	"zh": "chi",
	"cu": "chu",
	"cv": "chv",
	"kw": "cor",
	"co": "cos",
	"cr": "cre",
	"cs": "cze",
	"da": "dan",
	"dv": "div",
	"nl": "dut",
	"dz": "dzo",
	"en": "eng",
	"eo": "epo",
	"et": "est",
	"ee": "ewe",
	"fo": "fao",
	"fj": "fij",
	"fi": "fin",
	"fr": "fre",
	"fy": "fry",
	"ff": "ful",
	"ka": "geo",
	"de": "ger",
	"gd": "gla",
	"ga": "gle",
	"gl": "glg",
	"gv": "glv",
	"el": "gre",
	"gn": "grn",
	"gu": "guj",
	"ht": "hat",
	"ha": "hau",
	"he": "heb",
	"hz": "her",
	"hi": "hin",
	"ho": "hmo",
	"hr": "hrv",
	"hu": "hun",
	"ig": "ibo",
	"is": "ice",
	"io": "ido",
	"ii": "iii",
	"iu": "iku",
	"ie": "ile",
	"ia": "ina",
	"id": "ind",
	"ik": "ipk",
	"it": "ita",
	"jv": "jav",
	"ja": "jpn",
	"kl": "kal",
	"kn": "kan",
	"ks": "kas",
	"kr": "kau",
	"kk": "kaz",
	"km": "khm",
	"ki": "kik",
	"rw": "kin",
	"ky": "kir",
	"kv": "kom",
	"kg": "kon",
	"ko": "kor",
	"kj": "kua",
	"ku": "kur",
	"lo": "lao",
	"la": "lat",
	"lv": "lav",
	"li": "lim",
	"ln": "lin",
	"lt": "lit",
	"lb": "ltz",
	"lu": "lub",
	"lg": "lug",
	"mk": "mac",
	"mh": "mah",
	"ml": "mal",
	"mi": "mao",
	"mr": "mar",
	"ms": "may",
	"mg": "mlg",
	"mt": "mlt",
	"mn": "mon",
	"na": "nau",
	"nv": "nav",
	"nr": "nbl",
	"nd": "nde",
	"ng": "ndo",
	"ne": "nep",
	"nn": "nno",
	"nb": "nob",
	"no": "nor",
	"ny": "nya",
	"oc": "oci",
	"oj": "oji",
	"or": "ori",
	"om": "orm",
	"os": "oss",
	"pa": "pan",
	"fa": "per",
	"pi": "pli",
	"pl": "pol",
	"pt": "por",
	"ps": "pus",
	"qu": "que",
	"rm": "roh",
	"ro": "rum",
	"rn": "run",
	"ru": "rus",
	"sg": "sag",
	"sa": "san",
	"si": "sin",
	"sk": "slo",
	"sl": "slv",
	"se": "sme",
	"sm": "smo",
	"sn": "sna",
	"sd": "snd",
	"so": "som",
	"st": "sot",
	"es": "spa",
	"sc": "srd",
	"sr": "srp",
	"ss": "ssw",
	"su": "sun",
	"sw": "swa",
	"sv": "swe",
	"ty": "tah",
	"ta": "tam",
	"tt": "tat",
	"te": "tel",
	"tg": "tgk",
	"tl": "tgl",
	"th": "tha",
	"bo": "tib",
	"ti": "tir",
	"to": "ton",
	"tn": "tsn",
	"ts": "tso",
	"tk": "tuk",
	"tr": "tur",
	"tw": "twi",
	"ug": "uig",
	"uk": "ukr",
	"ur": "urd",
	"uz": "uzb",
	"ve": "ven",
	"vi": "vie",
	"vo": "vol",
	"cy": "wel",
	"wa": "wln",
	"wo": "wol",
	"xh": "xho",
	"yi": "yid",
	"yo": "yor",
	"za": "zha",
	"zu": "zul",
}
<|endoftext|>


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

package proxy

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/dzip"
	"golang.org/x/discovery/internal/thirdparty/module"
	"golang.org/x/discovery/internal/thirdparty/semver"
	"golang.org/x/net/context/ctxhttp"
)

const (
	stdlibModulePathProxy = "go.googlesource.com/go.git"
)

// A Client is used by the fetch service to communicate with a module
// proxy. It handles all methods defined by go help goproxy.
type Client struct {
	// URL of the module proxy web server
	url string

	// client used for HTTP requests. It is mutable for testing purposes.
	httpClient *http.Client
}

// A VersionInfo contains metadata about a given version of a module.
type VersionInfo struct {
	Version string
	Time    time.Time
}

// New constructs a *Client using the provided rawurl, which is expected to
// be an absolute URI that can be directly passed to http.Get.
func New(rawurl string) (*Client, error) {
	url, err := url.Parse(rawurl)
	if err != nil {
		return nil, fmt.Errorf("url.Parse(%q): %v", rawurl, err)
	}
	if url.Scheme != "https" {
		return nil, fmt.Errorf("scheme must be https (got %s)", url.Scheme)
	}
	return &Client{url: cleanURL(rawurl), httpClient: http.DefaultClient}, nil
}

// cleanURL trims the rawurl of trailing slashes.
func cleanURL(rawurl string) string {
	return strings.TrimRight(rawurl, "/")
}

// GetInfo makes a request to $GOPROXY/<module>/@v/<version>.info and
// transforms that data into a *VersionInfo.
func (c *Client) GetInfo(ctx context.Context, path, version string) (*VersionInfo, error) {
	v, err := c.getInfo(ctx, path, version)
	if err != nil {
		return nil, err
	}
	if path == "std" {
		v.Version = semver.Canonical(version)
	}
	return v, nil
}

func (c *Client) getInfo(ctx context.Context, path, version string) (*VersionInfo, error) {
	u, err := c.infoURL(path, version)
	if err != nil {
		return nil, fmt.Errorf("infoURL(%q, %q): %v", path, version, err)
	}

	r, err := ctxhttp.Get(ctx, c.httpClient, u)
	if err != nil {
		return nil, fmt.Errorf("ctxhttp.Get(ctx, client, %q): %v", u, err)
	}
	defer r.Body.Close()

	if err := derrors.StatusError(r.StatusCode, "ctxhttp.Get(ctx, client, %q)", u); err != nil {
		return nil, err
	}

	var v VersionInfo
	if err = json.NewDecoder(r.Body).Decode(&v); err != nil {
		return nil, err
	}
	return &v, nil
}

// infoURL constructs a url for a request to $GOPROXY/<module>/@v/<version>.info.
// If the path is "std", the version will be converted to the corresponding go
// tag. For example, "v1.12.5" will be converted to "go1.12.5" and "v1.12.0"
// will be converted to "go1.12".
func (c *Client) infoURL(path, version string) (string, error) {
	if path == "std" {
		path = stdlibModulePathProxy
		version = fmt.Sprintf("go%s", strings.TrimSuffix(strings.TrimPrefix(version, "v"), ".0"))
	}
	path, version, err := encodeModulePathAndVersion(path, version)
	if err != nil {
		return "", fmt.Errorf("encodeModulePathAndVersion(%q, %q): %v", path, version, err)
	}
	return fmt.Sprintf("%s/%s/@v/%s.info", c.url, path, version), nil
}

// GetZip makes a request to $GOPROXY/<path>/@v/<version>.zip and transforms
// that data into a *zip.Reader. <version> is obtained by first making a
// request to $GOPROXY/<path>/@v/<requestedVersion>.info to obtained the valid
// semantic version.
func (c *Client) GetZip(ctx context.Context, path, requestedVersion string) (*zip.Reader, error) {
	info, err := c.getInfo(ctx, path, requestedVersion)
	if err != nil {
		return nil, err
	}
	zipReader, err := c.getZip(ctx, path, info.Version)
	if err != nil {
		return nil, err
	}
	if path == "std" {
		return createGoZipReader(zipReader, info.Version, requestedVersion)
	}
	return zipReader, nil
}

// getZip makes a request to $GOPROXY/<path>/@v/<version>.zip and transforms
// that data into a *zip.Reader.
func (c *Client) getZip(ctx context.Context, path, version string) (*zip.Reader, error) {
	u, err := c.zipURL(path, version)
	if err != nil {
		return nil, fmt.Errorf("zipURL(%q, %q): %v", path, version, err)
	}

	r, err := ctxhttp.Get(ctx, c.httpClient, u)
	if err != nil {
		return nil, fmt.Errorf("ctxhttp.Get(ctx, nil, %q): %v", u, err)
	}
	defer r.Body.Close()

	if err := derrors.StatusError(r.StatusCode, "HTTP error from proxy for %q: %d", u, r.StatusCode); err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("get(ctx, %q): %v", u, err)
	}
	zipReader, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
	if err != nil {
		return nil, fmt.Errorf("get(ctx, %q): %v", u, err)
	}
	return zipReader, nil
}

// zipURL constructs a url for a request to $GOPROXY/<module>/@v/<version>.zip.
func (c *Client) zipURL(path, version string) (string, error) {
	if path == "std" {
		path = stdlibModulePathProxy
	}
	path, version, err := encodeModulePathAndVersion(path, version)
	if err != nil {
		return "", fmt.Errorf("encodeModulePathAndVersion(%q, %q): %v", path, version, err)
	}
	return fmt.Sprintf("%s/%s/@v/%s.zip", c.url, path, version), nil
}

// createGoZipReader returns a *zip.Reader containing the README, LICENSE and
// contents of the src/ directory for a zip obtained from a request to
// $GOPROXY/go.googlesource.com/go.git/@v/<version>.zip. The filenames returned
// will be trimmed of the prefix go.googlesource.com@<infoVersion> or
// go.googlesource.com@<infoVersion>/src. The prefix std@<goVersion> will be
// added to each of the resulting filenames.
func createGoZipReader(r *zip.Reader, infoVersion string, goVersion string) (*zip.Reader, error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)

	for _, file := range r.File {
		stdlibFilePrefix := fmt.Sprintf("%s@%s/", stdlibModulePathProxy, infoVersion)
		if !strings.HasPrefix(file.Name, stdlibFilePrefix+"src/") &&
			!strings.HasPrefix(file.Name, stdlibFilePrefix+"README") &&
			!strings.HasPrefix(file.Name, stdlibFilePrefix+"LICENSE") {
			continue
		}

		// Trim stdlibFilePrefix from README and LICENSE, and
		// stdlibFilePrefix+"src/" from files in the src/ directory.
		fileName := fmt.Sprintf("std@%s/%s", goVersion, strings.TrimPrefix(strings.TrimPrefix(file.Name, stdlibFilePrefix), "src/"))
		f, err := w.Create(fileName)
		if err != nil {
			return nil, fmt.Errorf("w.Create(%q): %v", file.Name, err)
		}

		contents, err := dzip.ReadZipFile(file)
		if err != nil {
			log.Printf("zip.ReadZipFile(%q): %v", file.Name, err)
			continue
		}

		if _, err = f.Write(contents); err != nil {
			return nil, fmt.Errorf("f.Write: %v", err)
		}
	}

	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("w.Close(): %v", err)
	}

	zipReader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		return nil, fmt.Errorf("zip.NewReader: %v", err)
	}
	return zipReader, nil
}

func encodeModulePathAndVersion(path, version string) (string, string, error) {
	encodedPath, err := module.EncodePath(path)
	if err != nil {
		return "", "", fmt.Errorf("module.EncodePath(%q): %v", path, err)
	}
	encodedVersion, err := module.EncodeVersion(version)
	if err != nil {
		return "", "", fmt.Errorf("module.EncodeVersion(%q): %v", version, err)
	}
	return encodedPath, encodedVersion, nil
}
<commit_msg>internal/proxy: do not fetch go.googlesource.com/go.git<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package proxy

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/dzip"
	"golang.org/x/discovery/internal/thirdparty/module"
	"golang.org/x/discovery/internal/thirdparty/semver"
	"golang.org/x/net/context/ctxhttp"
)

const (
	stdlibModulePathProxy = "go.googlesource.com/go.git"
)

// A Client is used by the fetch service to communicate with a module
// proxy. It handles all methods defined by go help goproxy.
type Client struct {
	// URL of the module proxy web server
	url string

	// client used for HTTP requests. It is mutable for testing purposes.
	httpClient *http.Client
}

// A VersionInfo contains metadata about a given version of a module.
type VersionInfo struct {
	Version string
	Time    time.Time
}

// New constructs a *Client using the provided rawurl, which is expected to
// be an absolute URI that can be directly passed to http.Get.
func New(rawurl string) (*Client, error) {
	url, err := url.Parse(rawurl)
	if err != nil {
		return nil, fmt.Errorf("url.Parse(%q): %v", rawurl, err)
	}
	if url.Scheme != "https" {
		return nil, fmt.Errorf("scheme must be https (got %s)", url.Scheme)
	}
	return &Client{url: cleanURL(rawurl), httpClient: http.DefaultClient}, nil
}

// cleanURL trims the rawurl of trailing slashes.
func cleanURL(rawurl string) string {
	return strings.TrimRight(rawurl, "/")
}

// GetInfo makes a request to $GOPROXY/<module>/@v/<version>.info and
// transforms that data into a *VersionInfo.
func (c *Client) GetInfo(ctx context.Context, path, version string) (*VersionInfo, error) {
	v, err := c.getInfo(ctx, path, version)
	if err != nil {
		return nil, err
	}
	if path == "std" {
		v.Version = semver.Canonical(version)
	}
	return v, nil
}

func (c *Client) getInfo(ctx context.Context, path, version string) (*VersionInfo, error) {
	if path == stdlibModulePathProxy {
		return nil, derrors.NotAcceptable("%q is not an accepted module. Fetch std instead.", stdlibModulePathProxy)
	}

	u, err := c.infoURL(path, version)
	if err != nil {
		return nil, fmt.Errorf("infoURL(%q, %q): %v", path, version, err)
	}

	r, err := ctxhttp.Get(ctx, c.httpClient, u)
	if err != nil {
		return nil, fmt.Errorf("ctxhttp.Get(ctx, client, %q): %v", u, err)
	}
	defer r.Body.Close()

	if err := derrors.StatusError(r.StatusCode, "ctxhttp.Get(ctx, client, %q)", u); err != nil {
		return nil, err
	}

	var v VersionInfo
	if err = json.NewDecoder(r.Body).Decode(&v); err != nil {
		return nil, err
	}
	return &v, nil
}

// infoURL constructs a url for a request to $GOPROXY/<module>/@v/<version>.info.
// If the path is "std", the version will be converted to the corresponding go
// tag. For example, "v1.12.5" will be converted to "go1.12.5" and "v1.12.0"
// will be converted to "go1.12".
func (c *Client) infoURL(path, version string) (string, error) {
	if path == "std" {
		path = stdlibModulePathProxy
		version = fmt.Sprintf("go%s", strings.TrimSuffix(strings.TrimPrefix(version, "v"), ".0"))
	}
	path, version, err := encodeModulePathAndVersion(path, version)
	if err != nil {
		return "", fmt.Errorf("encodeModulePathAndVersion(%q, %q): %v", path, version, err)
	}
	return fmt.Sprintf("%s/%s/@v/%s.info", c.url, path, version), nil
}

// GetZip makes a request to $GOPROXY/<path>/@v/<version>.zip and transforms
// that data into a *zip.Reader. <version> is obtained by first making a
// request to $GOPROXY/<path>/@v/<requestedVersion>.info to obtained the valid
// semantic version.
func (c *Client) GetZip(ctx context.Context, path, requestedVersion string) (*zip.Reader, error) {
	info, err := c.getInfo(ctx, path, requestedVersion)
	if err != nil {
		return nil, err
	}
	zipReader, err := c.getZip(ctx, path, info.Version)
	if err != nil {
		return nil, err
	}
	if path == "std" {
		return createGoZipReader(zipReader, info.Version, requestedVersion)
	}
	return zipReader, nil
}

// getZip makes a request to $GOPROXY/<path>/@v/<version>.zip and transforms
// that data into a *zip.Reader.
func (c *Client) getZip(ctx context.Context, path, version string) (*zip.Reader, error) {
	u, err := c.zipURL(path, version)
	if err != nil {
		return nil, fmt.Errorf("zipURL(%q, %q): %v", path, version, err)
	}

	r, err := ctxhttp.Get(ctx, c.httpClient, u)
	if err != nil {
		return nil, fmt.Errorf("ctxhttp.Get(ctx, nil, %q): %v", u, err)
	}
	defer r.Body.Close()

	if err := derrors.StatusError(r.StatusCode, "HTTP error from proxy for %q: %d", u, r.StatusCode); err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("get(ctx, %q): %v", u, err)
	}
	zipReader, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
	if err != nil {
		return nil, fmt.Errorf("get(ctx, %q): %v", u, err)
	}
	return zipReader, nil
}

// zipURL constructs a url for a request to $GOPROXY/<module>/@v/<version>.zip.
func (c *Client) zipURL(path, version string) (string, error) {
	if path == "std" {
		path = stdlibModulePathProxy
	}
	path, version, err := encodeModulePathAndVersion(path, version)
	if err != nil {
		return "", fmt.Errorf("encodeModulePathAndVersion(%q, %q): %v", path, version, err)
	}
	return fmt.Sprintf("%s/%s/@v/%s.zip", c.url, path, version), nil
}

// createGoZipReader returns a *zip.Reader containing the README, LICENSE and
// contents of the src/ directory for a zip obtained from a request to
// $GOPROXY/go.googlesource.com/go.git/@v/<version>.zip. The filenames returned
// will be trimmed of the prefix go.googlesource.com@<infoVersion> or
// go.googlesource.com@<infoVersion>/src. The prefix std@<goVersion> will be
// added to each of the resulting filenames.
func createGoZipReader(r *zip.Reader, infoVersion string, goVersion string) (*zip.Reader, error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)

	for _, file := range r.File {
		stdlibFilePrefix := fmt.Sprintf("%s@%s/", stdlibModulePathProxy, infoVersion)
		if !strings.HasPrefix(file.Name, stdlibFilePrefix+"src/") &&
			!strings.HasPrefix(file.Name, stdlibFilePrefix+"README") &&
			!strings.HasPrefix(file.Name, stdlibFilePrefix+"LICENSE") {
			continue
		}

		// Trim stdlibFilePrefix from README and LICENSE, and
		// stdlibFilePrefix+"src/" from files in the src/ directory.
		fileName := fmt.Sprintf("std@%s/%s", goVersion, strings.TrimPrefix(strings.TrimPrefix(file.Name, stdlibFilePrefix), "src/"))
		f, err := w.Create(fileName)
		if err != nil {
			return nil, fmt.Errorf("w.Create(%q): %v", file.Name, err)
		}

		contents, err := dzip.ReadZipFile(file)
		if err != nil {
			log.Printf("zip.ReadZipFile(%q): %v", file.Name, err)
			continue
		}

		if _, err = f.Write(contents); err != nil {
			return nil, fmt.Errorf("f.Write: %v", err)
		}
	}

	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("w.Close(): %v", err)
	}

	zipReader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		return nil, fmt.Errorf("zip.NewReader: %v", err)
	}
	return zipReader, nil
}

func encodeModulePathAndVersion(path, version string) (string, string, error) {
	encodedPath, err := module.EncodePath(path)
	if err != nil {
		return "", "", fmt.Errorf("module.EncodePath(%q): %v", path, err)
	}
	encodedVersion, err := module.EncodeVersion(version)
	if err != nil {
		return "", "", fmt.Errorf("module.EncodeVersion(%q): %v", version, err)
	}
	return encodedPath, encodedVersion, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/blob"
	"github.com/jacobsa/comeback/internal/dag"
	"github.com/jacobsa/comeback/internal/fs"
	"github.com/jacobsa/comeback/internal/repr"
	"github.com/jacobsa/comeback/internal/state"
	"github.com/jacobsa/timeutil"
)

const fileChunkSize = 1 << 24

// Create a dag.Visitor for *fsNode that does the following for each node N:
//
//  *  Ensure that nodes are only regular files, directories, and symlinks.
//
//  *  For files, consult the supplied score map to find a list of scores. If
//     the score map doesn't hit, write the file to the blob store to obtain a
//     list of scores, and update the score map.
//
//  *  For directories, write a listing to blob store to obtain a list of
//     scores.
//
//  *  Write all nodes to the supplied channel.
//
func newVisitor(
	chunkSize int,
	basePath string,
	scoreMap state.ScoreMap,
	blobStore blob.Store,
	clock timeutil.Clock,
	logger *log.Logger,
	visitedNodes chan<- *fsNode) (v dag.Visitor) {
	v = &visitor{
		chunkSize:    chunkSize,
		basePath:     basePath,
		scoreMap:     scoreMap,
		blobStore:    blobStore,
		clock:        clock,
		logger:       logger,
		visitedNodes: visitedNodes,
	}

	return
}

type visitor struct {
	chunkSize    int
	basePath     string
	scoreMap     state.ScoreMap
	blobStore    blob.Store
	clock        timeutil.Clock
	logger       *log.Logger
	visitedNodes chan<- *fsNode

	mu sync.Mutex

	// A freelist of objects that it may be beneficial to reuse from call to
	// call.
	//
	// GUARDED_BY(mu)
	freelist []*reusableObjects
}

type reusableObjects struct {
	// A buffer into which we read file chunks.
	//
	// INVARIANT: len(buf) is always the visitor's chunk size, and cap(buf) is a
	// bit more, to cover magic bytes added by repr.MarshalFile.
	buf []byte

	// A request for storing a blob. Reusing this allows us to reuse any internal
	// state that the blob store may cache.
	storeReq blob.StoreRequest
}

func (v *visitor) Visit(ctx context.Context, untyped dag.Node) (err error) {
	// Check the type of the node.
	n, ok := untyped.(*fsNode)
	if !ok {
		err = fmt.Errorf("Unexpected node type: %T", untyped)
		return
	}

	v.logger.Print(n.RelPath)

	// Check that the node is of a supported type.
	switch n.Info.Type {
	case fs.TypeFile:
	case fs.TypeDirectory:
	case fs.TypeSymlink:

	default:
		err = fmt.Errorf("Unsupported node type: %v", n.Info.Type)
		return
	}

	// Ensure that the node has scores set, if it needs to.
	err = v.setScores(ctx, n)
	if err != nil {
		err = fmt.Errorf("setScores: %v", err)
		return
	}

	// Pass on the node.
	select {
	case v.visitedNodes <- n:
	case <-ctx.Done():
		err = ctx.Err()
		return
	}

	return
}

func (v *visitor) setScores(
	ctx context.Context,
	n *fsNode) (err error) {
	// If the node already has scores, we're done.
	if n.Info.Scores != nil {
		return
	}

	// Files and directories are the only interesting cases.
	switch n.Info.Type {
	case fs.TypeFile:
		n.Info.Scores, err = v.saveFile(ctx, n)
		if err != nil {
			err = fmt.Errorf("saveFile(%q): %v", n.RelPath, err)
			return
		}

	case fs.TypeDirectory:
		n.Info.Scores, err = v.saveDir(ctx, n.Children)
		if err != nil {
			err = fmt.Errorf("saveDir: %v", err)
			return
		}
	}

	return
}

// Guarantees non-nil result when successful, even for empty list of scores.
func (v *visitor) saveFile(
	ctx context.Context,
	n *fsNode) (scores []blob.Score, err error) {
	// Can we short circuit here using the score map?
	scoreMapKey := makeScoreMapKey(n, v.clock)
	if scoreMapKey != nil {
		scores = v.scoreMap.Get(*scoreMapKey)
		if scores != nil {
			return
		}
	}

	// Ensure that our result will be non-nil, even for the empty list.
	scores = make([]blob.Score, 0, 1)

	// Open the file for reading.
	f, err := os.Open(path.Join(v.basePath, n.RelPath))
	if err != nil {
		err = fmt.Errorf("Open: %v", err)
		return
	}

	defer f.Close()

	// Process a chunk at a time.
	ro := v.getReusableObjects()
	defer v.putReusableObjects(ro)

	buf := ro.buf
	storeReq := &ro.storeReq

loop:
	for {
		// Read some data.
		var n int
		n, err = io.ReadFull(f, buf)

		switch {
		case err == io.EOF:
			// EOF means we're done.
			err = nil
			break loop

		case err == io.ErrUnexpectedEOF:
			// A short read is fine.
			err = nil

		case err != nil:
			err = fmt.Errorf("Read: %v", err)
			return
		}

		// Encapsulate the data so it can be identified as a file chunk.
		var chunk []byte
		chunk, err = repr.MarshalFile(buf[:n])
		if err != nil {
			err = fmt.Errorf("MarshalFile: %v", err)
			return
		}

		// Write out the blob.
		storeReq.Blob = chunk

		var s blob.Score
		s, err = v.blobStore.Store(ctx, storeReq)
		if err != nil {
			err = fmt.Errorf("Store: %v", err)
			return
		}

		scores = append(scores, s)
	}

	// Update the score map if the file is eligible.
	if scoreMapKey != nil {
		v.scoreMap.Set(*scoreMapKey, scores)
	}

	return
}

func (v *visitor) saveDir(
	ctx context.Context,
	children []*fsNode) (scores []blob.Score, err error) {
	ro := v.getReusableObjects()
	defer v.putReusableObjects(ro)

	storeReq := &ro.storeReq

	// Set up a list of directory entries.
	var entries []*fs.FileInfo
	for _, child := range children {
		entries = append(entries, &child.Info)
	}

	// Create a blob describing the directory's contents.
	b, err := repr.MarshalDir(entries)
	if err != nil {
		err = fmt.Errorf("MarshalDir: %v", err)
		return
	}

	// Write out the blob.
	storeReq.Blob = b

	s, err := v.blobStore.Store(ctx, storeReq)
	if err != nil {
		err = fmt.Errorf("Store: %v", err)
		return
	}

	scores = []blob.Score{s}
	return
}

// LOCKS_EXCLUDED(v.mu)
func (v *visitor) getReusableObjects() (ro *reusableObjects) {
	// Check the freelist.
	{
		v.mu.Lock()
		l := len(v.freelist)
		if l > 0 {
			ro = v.freelist[l-1]
			v.freelist = v.freelist[:l-1]
		}
		v.mu.Unlock()
	}

	if ro != nil {
		return
	}

	// Otherwise, allocate.
	ro = &reusableObjects{
		buf: make([]byte, v.chunkSize, v.chunkSize+v.chunkSize/2),
	}

	return
}

// LOCKS_EXCLUDED(v.mu)
func (v *visitor) putReusableObjects(ro *reusableObjects) {
	v.mu.Lock()
	v.freelist = append(v.freelist, ro)
	v.mu.Unlock()
}
<commit_msg>Include more context in saveDir errors.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/blob"
	"github.com/jacobsa/comeback/internal/dag"
	"github.com/jacobsa/comeback/internal/fs"
	"github.com/jacobsa/comeback/internal/repr"
	"github.com/jacobsa/comeback/internal/state"
	"github.com/jacobsa/timeutil"
)

const fileChunkSize = 1 << 24

// Create a dag.Visitor for *fsNode that does the following for each node N:
//
//  *  Ensure that nodes are only regular files, directories, and symlinks.
//
//  *  For files, consult the supplied score map to find a list of scores. If
//     the score map doesn't hit, write the file to the blob store to obtain a
//     list of scores, and update the score map.
//
//  *  For directories, write a listing to blob store to obtain a list of
//     scores.
//
//  *  Write all nodes to the supplied channel.
//
func newVisitor(
	chunkSize int,
	basePath string,
	scoreMap state.ScoreMap,
	blobStore blob.Store,
	clock timeutil.Clock,
	logger *log.Logger,
	visitedNodes chan<- *fsNode) (v dag.Visitor) {
	v = &visitor{
		chunkSize:    chunkSize,
		basePath:     basePath,
		scoreMap:     scoreMap,
		blobStore:    blobStore,
		clock:        clock,
		logger:       logger,
		visitedNodes: visitedNodes,
	}

	return
}

type visitor struct {
	chunkSize    int
	basePath     string
	scoreMap     state.ScoreMap
	blobStore    blob.Store
	clock        timeutil.Clock
	logger       *log.Logger
	visitedNodes chan<- *fsNode

	mu sync.Mutex

	// A freelist of objects that it may be beneficial to reuse from call to
	// call.
	//
	// GUARDED_BY(mu)
	freelist []*reusableObjects
}

type reusableObjects struct {
	// A buffer into which we read file chunks.
	//
	// INVARIANT: len(buf) is always the visitor's chunk size, and cap(buf) is a
	// bit more, to cover magic bytes added by repr.MarshalFile.
	buf []byte

	// A request for storing a blob. Reusing this allows us to reuse any internal
	// state that the blob store may cache.
	storeReq blob.StoreRequest
}

func (v *visitor) Visit(ctx context.Context, untyped dag.Node) (err error) {
	// Check the type of the node.
	n, ok := untyped.(*fsNode)
	if !ok {
		err = fmt.Errorf("Unexpected node type: %T", untyped)
		return
	}

	v.logger.Print(n.RelPath)

	// Check that the node is of a supported type.
	switch n.Info.Type {
	case fs.TypeFile:
	case fs.TypeDirectory:
	case fs.TypeSymlink:

	default:
		err = fmt.Errorf("Unsupported node type: %v", n.Info.Type)
		return
	}

	// Ensure that the node has scores set, if it needs to.
	err = v.setScores(ctx, n)
	if err != nil {
		err = fmt.Errorf("setScores: %v", err)
		return
	}

	// Pass on the node.
	select {
	case v.visitedNodes <- n:
	case <-ctx.Done():
		err = ctx.Err()
		return
	}

	return
}

func (v *visitor) setScores(
	ctx context.Context,
	n *fsNode) (err error) {
	// If the node already has scores, we're done.
	if n.Info.Scores != nil {
		return
	}

	// Files and directories are the only interesting cases.
	switch n.Info.Type {
	case fs.TypeFile:
		n.Info.Scores, err = v.saveFile(ctx, n)
		if err != nil {
			err = fmt.Errorf("saveFile(%q): %v", n.RelPath, err)
			return
		}

	case fs.TypeDirectory:
		n.Info.Scores, err = v.saveDir(ctx, n.Children)
		if err != nil {
			err = fmt.Errorf("saveDir(%q): : %v", n.RelPath, err)
			return
		}
	}

	return
}

// Guarantees non-nil result when successful, even for empty list of scores.
func (v *visitor) saveFile(
	ctx context.Context,
	n *fsNode) (scores []blob.Score, err error) {
	// Can we short circuit here using the score map?
	scoreMapKey := makeScoreMapKey(n, v.clock)
	if scoreMapKey != nil {
		scores = v.scoreMap.Get(*scoreMapKey)
		if scores != nil {
			return
		}
	}

	// Ensure that our result will be non-nil, even for the empty list.
	scores = make([]blob.Score, 0, 1)

	// Open the file for reading.
	f, err := os.Open(path.Join(v.basePath, n.RelPath))
	if err != nil {
		err = fmt.Errorf("Open: %v", err)
		return
	}

	defer f.Close()

	// Process a chunk at a time.
	ro := v.getReusableObjects()
	defer v.putReusableObjects(ro)

	buf := ro.buf
	storeReq := &ro.storeReq

loop:
	for {
		// Read some data.
		var n int
		n, err = io.ReadFull(f, buf)

		switch {
		case err == io.EOF:
			// EOF means we're done.
			err = nil
			break loop

		case err == io.ErrUnexpectedEOF:
			// A short read is fine.
			err = nil

		case err != nil:
			err = fmt.Errorf("Read: %v", err)
			return
		}

		// Encapsulate the data so it can be identified as a file chunk.
		var chunk []byte
		chunk, err = repr.MarshalFile(buf[:n])
		if err != nil {
			err = fmt.Errorf("MarshalFile: %v", err)
			return
		}

		// Write out the blob.
		storeReq.Blob = chunk

		var s blob.Score
		s, err = v.blobStore.Store(ctx, storeReq)
		if err != nil {
			err = fmt.Errorf("Store: %v", err)
			return
		}

		scores = append(scores, s)
	}

	// Update the score map if the file is eligible.
	if scoreMapKey != nil {
		v.scoreMap.Set(*scoreMapKey, scores)
	}

	return
}

func (v *visitor) saveDir(
	ctx context.Context,
	children []*fsNode) (scores []blob.Score, err error) {
	ro := v.getReusableObjects()
	defer v.putReusableObjects(ro)

	storeReq := &ro.storeReq

	// Set up a list of directory entries.
	var entries []*fs.FileInfo
	for _, child := range children {
		entries = append(entries, &child.Info)
	}

	// Create a blob describing the directory's contents.
	b, err := repr.MarshalDir(entries)
	if err != nil {
		err = fmt.Errorf("MarshalDir: %v", err)
		return
	}

	// Write out the blob.
	storeReq.Blob = b

	s, err := v.blobStore.Store(ctx, storeReq)
	if err != nil {
		err = fmt.Errorf("Store: %v", err)
		return
	}

	scores = []blob.Score{s}
	return
}

// LOCKS_EXCLUDED(v.mu)
func (v *visitor) getReusableObjects() (ro *reusableObjects) {
	// Check the freelist.
	{
		v.mu.Lock()
		l := len(v.freelist)
		if l > 0 {
			ro = v.freelist[l-1]
			v.freelist = v.freelist[:l-1]
		}
		v.mu.Unlock()
	}

	if ro != nil {
		return
	}

	// Otherwise, allocate.
	ro = &reusableObjects{
		buf: make([]byte, v.chunkSize, v.chunkSize+v.chunkSize/2),
	}

	return
}

// LOCKS_EXCLUDED(v.mu)
func (v *visitor) putReusableObjects(ro *reusableObjects) {
	v.mu.Lock()
	v.freelist = append(v.freelist, ro)
	v.mu.Unlock()
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package token

import (
	"errors"
	"fmt"
	"io"
	"reflect"

	"github.com/issue9/is"

	"github.com/caixw/apidoc/v7/core"
	"github.com/caixw/apidoc/v7/internal/locale"
	"github.com/caixw/apidoc/v7/internal/node"
)

// Decoder 实现从 p 中解码内容到当前对象的值
type Decoder interface {
	// 从 p 中读取内容并实例化到当前对象中
	//
	// 必须要同时从 p 中读取相应的 EndElement 才能返回。
	// end 表示 EndElement.End 的值。
	//
	// NOTE: 如果是自闭合标签，则不会调用该接口。
	//
	// 接口应该只返回 *core.SyntaxError 作为错误对象。
	DecodeXML(p *Parser, start *StartElement) (end *EndElement, err error)
}

// AttrDecoder 实现从 attr 中解码内容到当前对象的值
type AttrDecoder interface {
	// 解析属性值
	//
	// 接口应该只返回 *core.SyntaxError 作为错误对象。
	DecodeXMLAttr(p *Parser, attr *Attribute) error
}

// Sanitizer 用于验证和修改对象中的数据
type Sanitizer interface {
	// 验证数据是否正确
	//
	// 可以通过 p.NewError 和 p.WithError 返回 *core.SyntaxError 类型的错误
	Sanitize(p *Parser) error
}

var (
	attrDecoderType = reflect.TypeOf((*AttrDecoder)(nil)).Elem()
	decoderType     = reflect.TypeOf((*Decoder)(nil)).Elem()
	sanitizerType   = reflect.TypeOf((*Sanitizer)(nil)).Elem()
)

// Decode 将 p 中的 XML 内容解码至 v 对象中
func Decode(h *core.MessageHandler, p *Parser, v interface{}, namespace string) {
	var hasRoot bool
	for {
		t, r, err := p.Token()
		if errors.Is(err, io.EOF) {
			return
		} else if err != nil {
			h.Error(err)
			return
		}

		switch elem := t.(type) {
		case *StartElement:
			if hasRoot { // 多个根元素
				h.Error(p.NewError(elem.Start, elem.End, "", locale.ErrInvalidXML))
				return
			}
			hasRoot = true

			prefix := findPrefix(elem, namespace)
			vv := node.ParseValue(reflect.ValueOf(v))
			if err := decodeElement(p, elem, vv, prefix); err != nil {
				h.Error(err)
				return
			}
		case *Comment, *String, *Instruction: // 忽略注释和普通的文本内容
		default:
			h.Error(p.NewError(r.Start, r.End, "", locale.ErrInvalidXML))
			return
		}
	}
}

func findPrefix(start *StartElement, namespace string) string {
	for _, attr := range start.Attributes {
		if attr.Name.Prefix.Value == "xmlns" && attr.Value.Value == namespace {
			return attr.Name.Local.Value
		}
	}
	return ""
}

func decode(n *node.Node, p *Parser, start *StartElement, prefix string) (*EndElement, error) {
	if err := decodeAttributes(n, p, start, prefix); err != nil {
		return nil, err
	}

	if start.Close {
		return nil, decodeCheckOmitempty(n, p, start.Start, start.End)
	}

	end, err := decodeElements(n, p, prefix)
	if err != nil {
		return nil, err
	}
	return end, decodeCheckOmitempty(n, p, start.Start, end.End)
}

// 判断 omitempty 属性
func decodeCheckOmitempty(n *node.Node, p *Parser, start, end core.Position) error {
	for _, attr := range n.Attributes {
		if canNotEmpty(attr) {
			return p.NewError(start, end, attr.Name, locale.ErrRequired)
		}
	}
	for _, elem := range n.Elements {
		if canNotEmpty(elem) {
			return p.NewError(start, end, elem.Name, locale.ErrRequired)
		}
	}
	if n.CData != nil && canNotEmpty(n.CData) {
		return p.NewError(start, end, "cdata", locale.ErrRequired)
	}
	if n.Content != nil && canNotEmpty(n.Content) {
		return p.NewError(start, end, "content", locale.ErrRequired)
	}

	return nil
}

// 当前表示的值必须是一个非空值
func canNotEmpty(v *node.Value) bool {
	return v.Name != "" && // cdata 和 content 在未初始化时 name 字段为空值
		!v.Omitempty &&
		(!v.CanInterface() || is.Empty(v.Interface(), true))
}

// 将 start 的属性内容解码到 obj.Attributes 之中
func decodeAttributes(n *node.Node, p *Parser, start *StartElement, prefix string) error {
	if start == nil {
		return nil
	}

	for _, attr := range start.Attributes {
		item, found := n.Attribute(attr.Name.Local.Value)
		if !found || prefix != attr.Name.Prefix.Value { // 未找到或是命名空间不匹配
			continue
		}
		v := node.GetRealValue(item.Value)
		v.Set(reflect.New(v.Type()).Elem())

		var impl bool
		if item.CanInterface() && item.Type().Implements(attrDecoderType) {
			if err := item.Interface().(AttrDecoder).DecodeXMLAttr(p, attr); err != nil {
				return err
			}
			impl = true
		} else if item.CanAddr() {
			pv := item.Addr()
			if pv.CanInterface() && pv.Type().Implements(attrDecoderType) {
				if err := pv.Interface().(AttrDecoder).DecodeXMLAttr(p, attr); err != nil {
					return err
				}
				impl = true
			}
		}

		if !impl {
			panic(fmt.Sprintf("当前属性 %s 未实现 AttrDecoder 接口", attr.Name))
		}

		if err := setAttributeValue(item.Value, item.Usage, p, attr); err != nil {
			return err
		}
	}

	return nil
}

func decodeElements(n *node.Node, p *Parser, prefix string) (*EndElement, error) {
	for {
		t, r, err := p.Token()
		if errors.Is(err, io.EOF) {
			// 应该只有 EndElement 才能返回，否则就不完整的 XML
			return nil, p.NewError(p.Position().Position, p.Position().Position, "", locale.ErrInvalidXML)
		} else if err != nil {
			return nil, err
		}

		switch elem := t.(type) {
		case *EndElement: // 找到当前对象的结束标签
			if (elem.Name.Local.Value == n.Value.Name) && (elem.Name.Prefix.Value == prefix) {
				return elem, nil
			}
			return nil, p.NewError(elem.Start, elem.End, n.Value.Name, locale.ErrNotFoundEndTag)
		case *CData:
			if n.CData != nil {
				setContentValue(n.CData.Value, reflect.ValueOf(elem))
			}
		case *String:
			if n.Content != nil {
				setContentValue(n.Content.Value, reflect.ValueOf(elem))
			}
		case *StartElement:
			item, found := n.Element(elem.Name.Local.Value)
			if !found || prefix != elem.Name.Prefix.Value {
				if err := findEndElement(p, elem); err != nil {
					return nil, err
				}
				break // 忽略不存在的子元素
			}
			if err = decodeElement(p, elem, item, prefix); err != nil {
				return nil, err
			}
		case *Comment: // 忽略注释内容
		default:
			return nil, p.NewError(r.Start, r.End, "", locale.ErrInvalidXML)
		}
	}
}

func setContentValue(target, source reflect.Value) {
	target = node.GetRealValue(target)
	source = node.GetRealValue(source)

	st := source.Type()
	num := st.NumField()
	for i := 0; i < num; i++ {
		name := st.Field(i).Name
		target.FieldByName(name).
			Set(source.Field(i))
	}
}

func decodeElement(p *Parser, start *StartElement, v *node.Value, prefix string) error {
	v.Value = node.GetRealValue(v.Value)
	k := v.Kind()
	switch {
	case k == reflect.Ptr, k == reflect.Func, k == reflect.Chan, k == reflect.Array, node.IsPrimitive(v.Value):
		panic(fmt.Sprintf("%s 是无效的类型", v.Value.Type()))
	case k == reflect.Slice:
		return decodeSlice(p, start, v, prefix)
	}

	end, impl, err := callDecodeXML(v.Value, p, start)
	if !impl {
		end, err = decode(node.New(start.Name.Local.Value, v.Value), p, start, prefix)
	}
	if err != nil {
		return err
	}
	return setTagValue(v.Value, v.Usage, p, start, end)
}

func decodeSlice(p *Parser, start *StartElement, slice *node.Value, prefix string) (err error) {
	// 不相配，表示当前元素找不到与之相配的元素，需要忽略这个元素，
	// 所以要过滤与 start 想匹配的结束符号才算结束。
	if !start.Close && (start.Name.Local.Value != slice.Name) {
		return findEndElement(p, start)
	}

	elem := reflect.New(slice.Type().Elem()).Elem()
	if elem.Kind() == reflect.Ptr && elem.IsNil() {
		elem.Set(reflect.New(elem.Type().Elem()))
	}

	end, impl, err := callDecodeXML(elem, p, start)
	if !impl {
		if node.IsPrimitive(elem) {
			panic(fmt.Sprintf("%s:%s 必须实现 Decoder 接口", slice.Name, elem.Type()))
		}
		end, err = decode(node.New(start.Name.Local.Value, elem), p, start, prefix)
	}
	if err != nil {
		return err
	}

	if err = setTagValue(elem, slice.Usage, p, start, end); err != nil {
		return err
	}
	slice.Value.Set(reflect.Append(slice.Value, elem))
	return nil
}

// 调用 v 的 DecodeXML 接口方法
func callDecodeXML(v reflect.Value, p *Parser, start *StartElement) (end *EndElement, impl bool, err error) {
	if v.CanInterface() && v.Type().Implements(decoderType) {
		if !start.Close {
			end, err = v.Interface().(Decoder).DecodeXML(p, start)
		}
		return end, true, err
	} else if v.CanAddr() {
		pv := v.Addr()
		if pv.CanInterface() && pv.Type().Implements(decoderType) {
			if !start.Close {
				end, err = pv.Interface().(Decoder).DecodeXML(p, start)
			}
			return end, true, err
		}
	}
	return nil, false, nil
}

// 找到与 start 相对应的结束符号位置
func findEndElement(p *Parser, start *StartElement) error {
	level := 0
	for {
		t, _, err := p.Token()
		if errors.Is(err, io.EOF) {
			return p.NewError(start.Start, start.End, start.Name.String(), locale.ErrNotFoundEndTag) // 找不到相配的结束符号
		} else if err != nil {
			return err
		}

		switch elem := t.(type) {
		case *StartElement:
			if elem.Name.Equal(start.Name) {
				level++
			}
		case *EndElement:
			if level == 0 && start.Match(elem) {
				return nil
			}
			level--
		}
	}
}

func setTagValue(v reflect.Value, usage string, p *Parser, start *StartElement, end *EndElement) error {
	v = node.GetRealValue(v)
	if v.Kind() != reflect.Struct {
		panic(fmt.Sprintf("无效的 kind 类型: %s:%s", v.Type(), v.Kind()))
	}

	v.FieldByName(usageKeyName).Set(reflect.ValueOf(usage))
	v.FieldByName(elementTagName).Set(reflect.ValueOf(start.Name))
	if end == nil {
		v.FieldByName(rangeName).Set(reflect.ValueOf(start.Range))
	} else {
		v.FieldByName(rangeName).Set(reflect.ValueOf(core.Range{Start: start.Start, End: end.End}))
		v.FieldByName(elementTagEndName).Set(reflect.ValueOf(end.Name))
	}

	// Sanitize 在最后调用，可以保证 Sanitize 调用中可以取 v.Range 的值
	return callSanitizer(v, p)
}

func setAttributeValue(v reflect.Value, usage string, p *Parser, attr *Attribute) error {
	v = node.GetRealValue(v)
	if v.Kind() != reflect.Struct {
		panic(fmt.Sprintf("无效的 kind 类型: %s:%s", v.Type(), v.Kind()))
	}

	v.FieldByName(rangeName).Set(reflect.ValueOf(attr.Range))
	v.FieldByName(usageKeyName).Set(reflect.ValueOf(usage))
	v.FieldByName(attributeNameName).Set(reflect.ValueOf(attr.Name))

	return callSanitizer(v, p)
}

func callSanitizer(v reflect.Value, p *Parser) error {
	if v.CanInterface() && v.Type().Implements(sanitizerType) {
		return v.Interface().(Sanitizer).Sanitize(p)
	} else if v.CanAddr() {
		pv := v.Addr()
		if pv.CanInterface() && pv.Type().Implements(sanitizerType) {
			return pv.Interface().(Sanitizer).Sanitize(p)
		}
	}
	return nil
}
<commit_msg>refactor(internal/token): 添加 decoder 对象<commit_after>// SPDX-License-Identifier: MIT

package token

import (
	"errors"
	"fmt"
	"io"
	"reflect"

	"github.com/issue9/is"
	"golang.org/x/text/message"

	"github.com/caixw/apidoc/v7/core"
	"github.com/caixw/apidoc/v7/internal/locale"
	"github.com/caixw/apidoc/v7/internal/node"
)

// Decoder 实现从 p 中解码内容到当前对象的值
type Decoder interface {
	// 从 p 中读取内容并实例化到当前对象中
	//
	// 必须要同时从 p 中读取相应的 EndElement 才能返回。
	// end 表示 EndElement.End 的值。
	//
	// NOTE: 如果是自闭合标签，则不会调用该接口。
	//
	// 接口应该只返回 *core.SyntaxError 作为错误对象。
	DecodeXML(p *Parser, start *StartElement) (end *EndElement, err error)
}

// AttrDecoder 实现从 attr 中解码内容到当前对象的值
type AttrDecoder interface {
	// 解析属性值
	//
	// 接口应该只返回 *core.SyntaxError 作为错误对象。
	DecodeXMLAttr(p *Parser, attr *Attribute) error
}

// Sanitizer 用于验证和修改对象中的数据
type Sanitizer interface {
	// 验证数据是否正确
	//
	// 可以通过 p.NewError 和 p.WithError 返回 *core.SyntaxError 类型的错误
	Sanitize(p *Parser) error
}

var (
	attrDecoderType = reflect.TypeOf((*AttrDecoder)(nil)).Elem()
	decoderType     = reflect.TypeOf((*Decoder)(nil)).Elem()
	sanitizerType   = reflect.TypeOf((*Sanitizer)(nil)).Elem()
)

type decoder struct {
	h      *core.MessageHandler
	p      *Parser
	prefix string // 命名空间所表示的前缀
}

func (d *decoder) message(t core.MessageType, start, end core.Position, field string, key message.Reference, v ...interface{}) bool {
	d.h.Message(t, d.p.NewError(start, end, field, key, v...))
	return false
}

func (d *decoder) error(err error) bool {
	d.h.Error(err)
	return false
}

// Decode 将 p 中的 XML 内容解码至 v 对象中
//
// namespace 如果不为空表示当前的 xml 所在的命名空间，只有该命名空间的元素才会被正确识别。
func Decode(h *core.MessageHandler, p *Parser, v interface{}, namespace string) {
	d := &decoder{
		h: h,
		p: p,
	}

	var hasRoot bool
	for {
		t, r, err := p.Token()
		if errors.Is(err, io.EOF) {
			return
		} else if err != nil {
			h.Error(err)
			return
		}

		switch elem := t.(type) {
		case *StartElement:
			if hasRoot { // 多个根元素
				d.message(core.Erro, elem.Start, elem.End, "", locale.ErrInvalidXML)
				return
			}
			hasRoot = true

			d.prefix = findPrefix(elem, namespace)
			if !d.decodeElement(elem, node.ParseValue(reflect.ValueOf(v))) {
				return
			}
		case *Comment, *String, *Instruction: // 忽略注释和普通的文本内容
		default:
			d.message(core.Erro, r.Start, r.End, "", locale.ErrInvalidXML)
			return
		}
	}
}

func findPrefix(start *StartElement, namespace string) string {
	for _, attr := range start.Attributes {
		if attr.Name.Prefix.Value == "xmlns" && attr.Value.Value == namespace {
			return attr.Name.Local.Value
		}
	}
	return ""
}

func (d *decoder) decode(n *node.Node, start *StartElement) (end *EndElement, ok bool) {
	if err := decodeAttributes(n, d.p, start, d.prefix); err != nil {
		return nil, d.error(err)
	}

	if start.Close {
		return nil, d.decodeCheckOmitempty(n, start.Start, start.End)
	}

	if end, ok = d.decodeElements(n); ok {
		return end, d.decodeCheckOmitempty(n, start.Start, end.End)
	}
	return nil, false
}

// 判断 omitempty 属性
func (d *decoder) decodeCheckOmitempty(n *node.Node, start, end core.Position) bool {
	for _, attr := range n.Attributes {
		if canNotEmpty(attr) {
			return d.message(core.Erro, start, end, attr.Name, locale.ErrRequired)
		}
	}
	for _, elem := range n.Elements {
		if canNotEmpty(elem) {
			return d.message(core.Erro, start, end, elem.Name, locale.ErrRequired)
		}
	}
	if n.CData != nil && canNotEmpty(n.CData) {
		return d.message(core.Erro, start, end, "cdata", locale.ErrRequired)
	}
	if n.Content != nil && canNotEmpty(n.Content) {
		return d.message(core.Erro, start, end, "content", locale.ErrRequired)
	}

	return true
}

// 当前表示的值必须是一个非空值
func canNotEmpty(v *node.Value) bool {
	return v.Name != "" && // cdata 和 content 在未初始化时 name 字段为空值
		!v.Omitempty &&
		(!v.CanInterface() || is.Empty(v.Interface(), true))
}

// 将 start 的属性内容解码到 obj.Attributes 之中
func decodeAttributes(n *node.Node, p *Parser, start *StartElement, prefix string) error {
	if start == nil {
		return nil
	}

	for _, attr := range start.Attributes {
		item, found := n.Attribute(attr.Name.Local.Value)
		if !found || prefix != attr.Name.Prefix.Value { // 未找到或是命名空间不匹配
			continue
		}
		v := node.GetRealValue(item.Value)
		v.Set(reflect.New(v.Type()).Elem())

		var impl bool
		if item.CanInterface() && item.Type().Implements(attrDecoderType) {
			if err := item.Interface().(AttrDecoder).DecodeXMLAttr(p, attr); err != nil {
				return err
			}
			impl = true
		} else if item.CanAddr() {
			pv := item.Addr()
			if pv.CanInterface() && pv.Type().Implements(attrDecoderType) {
				if err := pv.Interface().(AttrDecoder).DecodeXMLAttr(p, attr); err != nil {
					return err
				}
				impl = true
			}
		}

		if !impl {
			panic(fmt.Sprintf("当前属性 %s 未实现 AttrDecoder 接口", attr.Name))
		}

		if err := setAttributeValue(item.Value, item.Usage, p, attr); err != nil {
			return err
		}
	}

	return nil
}

func (d *decoder) decodeElements(n *node.Node) (end *EndElement, ok bool) {
	for {
		t, r, err := d.p.Token()
		if errors.Is(err, io.EOF) {
			// 应该只有 EndElement 才能返回，否则就不完整的 XML
			return nil, d.message(core.Erro, d.p.Position().Position, d.p.Position().Position, "", locale.ErrInvalidXML)
		} else if err != nil {
			return nil, d.error(err)
		}

		switch elem := t.(type) {
		case *EndElement: // 找到当前对象的结束标签
			if (elem.Name.Local.Value == n.Value.Name) && (elem.Name.Prefix.Value == d.prefix) {
				return elem, true
			}
			return nil, d.message(core.Erro, elem.Start, elem.End, n.Value.Name, locale.ErrNotFoundEndTag)
		case *CData:
			if n.CData != nil {
				setContentValue(n.CData.Value, reflect.ValueOf(elem))
			}
		case *String:
			if n.Content != nil {
				setContentValue(n.Content.Value, reflect.ValueOf(elem))
			}
		case *StartElement:
			item, found := n.Element(elem.Name.Local.Value)
			if !found || d.prefix != elem.Name.Prefix.Value {
				if err := findEndElement(d.p, elem); err != nil {
					return nil, d.error(err)
				}

				d.message(core.Warn, elem.Name.Start, elem.Name.End, elem.Name.String(), locale.ErrInvalidValue)
				break // 忽略不存在的子元素
			}
			if !d.decodeElement(elem, item) {
				return nil, false
			}
		case *Comment: // 忽略注释内容
		default:
			return nil, d.message(core.Erro, r.Start, r.End, "", locale.ErrInvalidXML)
		}
	}
}

func setContentValue(target, source reflect.Value) {
	target = node.GetRealValue(target)
	source = node.GetRealValue(source)

	st := source.Type()
	num := st.NumField()
	for i := 0; i < num; i++ {
		name := st.Field(i).Name
		target.FieldByName(name).
			Set(source.Field(i))
	}
}

func (d *decoder) decodeElement(start *StartElement, v *node.Value) (ok bool) {
	v.Value = node.GetRealValue(v.Value)
	k := v.Kind()
	switch {
	case k == reflect.Ptr, k == reflect.Func, k == reflect.Chan, k == reflect.Array, node.IsPrimitive(v.Value):
		panic(fmt.Sprintf("%s 是无效的类型", v.Value.Type()))
	case k == reflect.Slice:
		return d.decodeSlice(start, v)
	}

	end, impl, err := callDecodeXML(v.Value, d.p, start)
	if !impl {
		if end, ok = d.decode(node.New(start.Name.Local.Value, v.Value), start); !ok {
			return false
		}
	}
	if err != nil {
		return d.error(err)
	}

	if err := setTagValue(v.Value, v.Usage, d.p, start, end); err != nil {
		return d.error(err)
	}
	return true
}

func (d *decoder) decodeSlice(start *StartElement, slice *node.Value) (ok bool) {
	// 不相配，表示当前元素找不到与之相配的元素，需要忽略这个元素，
	// 所以要过滤与 start 想匹配的结束符号才算结束。
	if !start.Close && (start.Name.Local.Value != slice.Name) {
		if err := findEndElement(d.p, start); err != nil {
			return d.error(err)
		}
	}

	elem := reflect.New(slice.Type().Elem()).Elem()
	if elem.Kind() == reflect.Ptr && elem.IsNil() {
		elem.Set(reflect.New(elem.Type().Elem()))
	}

	end, impl, err := callDecodeXML(elem, d.p, start)
	if !impl {
		if node.IsPrimitive(elem) {
			panic(fmt.Sprintf("%s:%s 必须实现 Decoder 接口", slice.Name, elem.Type()))
		}
		if end, ok = d.decode(node.New(start.Name.Local.Value, elem), start); !ok {
			return false
		}
	}
	if err != nil {
		return d.error(err)
	}

	if err = setTagValue(elem, slice.Usage, d.p, start, end); err != nil {
		return d.error(err)
	}
	slice.Value.Set(reflect.Append(slice.Value, elem))
	return true
}

// 调用 v 的 DecodeXML 接口方法
//
// 当 impl 为 true 时，err 表示的是 DecodeXML 接口返回的错误，否则 err 永远为 nil
func callDecodeXML(v reflect.Value, p *Parser, start *StartElement) (end *EndElement, impl bool, err error) {
	if v.CanInterface() && v.Type().Implements(decoderType) {
		if !start.Close {
			end, err = v.Interface().(Decoder).DecodeXML(p, start)
		}
		return end, true, err
	} else if v.CanAddr() {
		pv := v.Addr()
		if pv.CanInterface() && pv.Type().Implements(decoderType) {
			if !start.Close {
				end, err = pv.Interface().(Decoder).DecodeXML(p, start)
			}
			return end, true, err
		}
	}
	return nil, false, nil
}

// 找到与 start 相对应的结束符号位置
func findEndElement(p *Parser, start *StartElement) error {
	level := 0
	for {
		t, _, err := p.Token()
		if errors.Is(err, io.EOF) {
			return p.NewError(start.Start, start.End, start.Name.String(), locale.ErrNotFoundEndTag) // 找不到相配的结束符号
		} else if err != nil {
			return err
		}

		switch elem := t.(type) {
		case *StartElement:
			if elem.Name.Equal(start.Name) {
				level++
			}
		case *EndElement:
			if level == 0 && start.Match(elem) {
				return nil
			}
			level--
		}
	}
}

func setTagValue(v reflect.Value, usage string, p *Parser, start *StartElement, end *EndElement) error {
	v = node.GetRealValue(v)
	if v.Kind() != reflect.Struct {
		panic(fmt.Sprintf("无效的 kind 类型: %s:%s", v.Type(), v.Kind()))
	}

	v.FieldByName(usageKeyName).Set(reflect.ValueOf(usage))
	v.FieldByName(elementTagName).Set(reflect.ValueOf(start.Name))
	if end == nil {
		v.FieldByName(rangeName).Set(reflect.ValueOf(start.Range))
	} else {
		v.FieldByName(rangeName).Set(reflect.ValueOf(core.Range{Start: start.Start, End: end.End}))
		v.FieldByName(elementTagEndName).Set(reflect.ValueOf(end.Name))
	}

	// Sanitize 在最后调用，可以保证 Sanitize 调用中可以取 v.Range 的值
	return callSanitizer(v, p)
}

func setAttributeValue(v reflect.Value, usage string, p *Parser, attr *Attribute) error {
	v = node.GetRealValue(v)
	if v.Kind() != reflect.Struct {
		panic(fmt.Sprintf("无效的 kind 类型: %s:%s", v.Type(), v.Kind()))
	}

	v.FieldByName(rangeName).Set(reflect.ValueOf(attr.Range))
	v.FieldByName(usageKeyName).Set(reflect.ValueOf(usage))
	v.FieldByName(attributeNameName).Set(reflect.ValueOf(attr.Name))

	return callSanitizer(v, p)
}

func callSanitizer(v reflect.Value, p *Parser) error {
	if v.CanInterface() && v.Type().Implements(sanitizerType) {
		return v.Interface().(Sanitizer).Sanitize(p)
	} else if v.CanAddr() {
		pv := v.Addr()
		if pv.CanInterface() && pv.Type().Implements(sanitizerType) {
			return pv.Interface().(Sanitizer).Sanitize(p)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// dataloader is an implimentation of facebook's dataloader in go.
// See https://github.com/facebook/dataloader for more information
package dataloader

import (
	"fmt"
	"sync"
	"time"
)

// A `DataLoader` Interface defines a public API for loading data from a particular
// data back-end with unique keys such as the `id` column of a SQL table or
// document name in a MongoDB database, given a batch loading function.
//
// Each `DataLoader` instance should contain a unique memoized cache. Use caution when
// used in long-lived applications or those which serve many users with
// different access permissions and consider creating a new instance per
// web request.
type Interface interface {
	Load(string) Getter
	LoadMany([]string) GetterMany
	Clear(string) Interface
	ClearAll() Interface
	Prime(key string, value interface{}) Interface
}

// BatchFunc is a Function, which when given an Array of keys (string), returns an array of `results`.
// It's important the the length of the input keys must match the length of
// the ouput results.
type BatchFunc func([]string) []*Result

// Result is the data structure that is primarily used by the BatchFunc. It contains the resolved data,
// and any errors that may have occured while fetching the data.
type Result struct {
	Data  interface{}
	Error error
}

// ResultMany is used by the loadMany method. It contains a list of resolved data and a list of erros
// if any occured.
type ResultMany struct {
	Data  []interface{}
	Error []error
}

// Implements the dataloader.Interface.
type Loader struct {
	// the batch function to be used by this loader
	batchFn BatchFunc
	// the maximum batch size. Set to 0 if you want it to be unbounded.
	cap int
	// the internal cache. This packages contains a basic cache implementation but any custom cache
	// implementation could be used as long as it implements the `Cache` interface.
	cache Cache

	// internal channel that is used to batch items
	input chan *batchRequest
	// input mutex
	mu sync.Mutex
}

// type used to on input channel
type batchRequest struct {
	key     string
	channel chan *Result
}

type Getter func() *Result
type GetterMany func() *ResultMany

// NewBatchedLoader constructs a new Loader with given options
func NewBatchedLoader(batchFn BatchFunc, cache Cache, cap int) *Loader {
	return &Loader{
		batchFn: batchFn,
		cache:   cache,
		cap:     cap,
	}
}

// Load load/resolves the given key, returning a channel that will contain the value and error
func (l *Loader) Load(key string) Getter {
	c := make(chan *Result, 1)
	if v, ok := l.cache.Get(key); ok {
		return v.(func() *Result)
	}

	var value *Result
	getter := func() *Result {
		if v, ok := <-c; ok {
			value = v
		}
		return value
	}

	l.cache.Set(key, getter)
	req := &batchRequest{key, c}
	l.addToQueue(req)

	return getter
}

// this help ensure that data return is in same order as data recieved
type result struct {
	*Result
	index int
}

// this help match the error to the key of a specific index
type resultError struct {
	error
	index int
}

// LoadMany loads mulitiple keys, returning a channel that will resolve to an array of values and an error
func (l *Loader) LoadMany(keys []string) GetterMany {
	out := make(chan *ResultMany, 1)
	c := make(chan *result, len(keys))

	for i := range keys {
		go func(i int) {
			getter := l.Load(keys[i])
			c <- &result{getter(), i}
		}(i)
	}

	go func() {
		defer close(out)
		var i int = 1
		outputData := make([]interface{}, len(keys), len(keys))
		outputErrors := make([]error, 0, len(keys))
		for result := range c {
			outputData[result.index] = result.Data
			if result.Error != nil {
				outputErrors = append(outputErrors, resultError{result.Error, result.index})
			}
			if i == len(keys) {
				close(c)
			}
			i++
		}
		out <- &ResultMany{outputData, outputErrors}
	}()

	var value *ResultMany
	getterMany := func() *ResultMany {
		if v, ok := <-out; ok {
			value = v
		}
		return value
	}

	return getterMany
}

// Clear clears the value at `key` from the cache, it it exsits. Returs self for method chaining
func (l *Loader) Clear(key string) Interface {
	l.cache.Delete(key)
	return l
}

// ClearAll clears the entire cache. To be used when some event results in unknown invalidations.
// Returs self for method chaining.
func (l *Loader) ClearAll() Interface {
	l.cache.Clear()
	return l
}

// Adds the provided key and value to the cache. If the key already exists, no change is made.
// Returns self for method chaining
func (l *Loader) Prime(key string, value interface{}) Interface {
	if _, ok := l.cache.Get(key); !ok {
		l.cache.Set(key, value)
	}
	return l
}

// queues item to be batched
func (l *Loader) addToQueue(req *batchRequest) {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.input == nil {
		l.input = make(chan *batchRequest, l.cap)
		go l.batch()
	}

	l.input <- req
}

// execuite the batch of all items in queue
func (l *Loader) batch() {
	var keys []string
	var reqs []*batchRequest

	go l.sleeper()

	for item := range l.input {
		keys = append(keys, item.key)
		reqs = append(reqs, item)
	}

	items := l.batchFn(keys)

	if len(items) != len(reqs) {
		panic(fmt.Errorf("length of keys must match length of responses"))
	}

	for i, req := range reqs {
		req.channel <- items[i]
		close(req.channel)
	}
}

// wait the appropriate amount of time for next batch
func (l *Loader) sleeper() {
	// this will move this goroutine to the back of the callstack
	time.Sleep(1)
	l.mu.Lock()
	defer l.mu.Unlock()

	close(l.input)
	l.input = nil
}
<commit_msg>add lock to value in getter function to prevent nil pointer reference<commit_after>// dataloader is an implimentation of facebook's dataloader in go.
// See https://github.com/facebook/dataloader for more information
package dataloader

import (
	"fmt"
	"sync"
	"time"
)

// A `DataLoader` Interface defines a public API for loading data from a particular
// data back-end with unique keys such as the `id` column of a SQL table or
// document name in a MongoDB database, given a batch loading function.
//
// Each `DataLoader` instance should contain a unique memoized cache. Use caution when
// used in long-lived applications or those which serve many users with
// different access permissions and consider creating a new instance per
// web request.
type Interface interface {
	Load(string) Getter
	LoadMany([]string) GetterMany
	Clear(string) Interface
	ClearAll() Interface
	Prime(key string, value interface{}) Interface
}

// BatchFunc is a Function, which when given an Array of keys (string), returns an array of `results`.
// It's important the the length of the input keys must match the length of
// the ouput results.
type BatchFunc func([]string) []*Result

// Result is the data structure that is primarily used by the BatchFunc. It contains the resolved data,
// and any errors that may have occured while fetching the data.
type Result struct {
	Data  interface{}
	Error error
}

// ResultMany is used by the loadMany method. It contains a list of resolved data and a list of erros
// if any occured.
type ResultMany struct {
	Data  []interface{}
	Error []error
}

// Implements the dataloader.Interface.
type Loader struct {
	// the batch function to be used by this loader
	batchFn BatchFunc
	// the maximum batch size. Set to 0 if you want it to be unbounded.
	cap int
	// the internal cache. This packages contains a basic cache implementation but any custom cache
	// implementation could be used as long as it implements the `Cache` interface.
	cache Cache

	// internal channel that is used to batch items
	input chan *batchRequest
	// input mutex
	mu sync.Mutex
}

// type used to on input channel
type batchRequest struct {
	key     string
	channel chan *Result
}

type Getter func() *Result
type GetterMany func() *ResultMany

// NewBatchedLoader constructs a new Loader with given options
func NewBatchedLoader(batchFn BatchFunc, cache Cache, cap int) *Loader {
	return &Loader{
		batchFn: batchFn,
		cache:   cache,
		cap:     cap,
	}
}

// Load load/resolves the given key, returning a channel that will contain the value and error
func (l *Loader) Load(key string) Getter {
	c := make(chan *Result, 1)
	if v, ok := l.cache.Get(key); ok {
		return v.(func() *Result)
	}

	var value struct {
		value *Result
		lock  sync.Mutex
	}
	getter := func() *Result {
		value.lock.Lock()
		defer value.lock.Unlock()
		if v, ok := <-c; ok {
			value.value = v
		}
		return value.value
	}

	l.cache.Set(key, getter)
	req := &batchRequest{key, c}
	l.addToQueue(req)

	return getter
}

// this help ensure that data return is in same order as data recieved
type result struct {
	*Result
	index int
}

// this help match the error to the key of a specific index
type resultError struct {
	error
	index int
}

// LoadMany loads mulitiple keys, returning a channel that will resolve to an array of values and an error
func (l *Loader) LoadMany(keys []string) GetterMany {
	out := make(chan *ResultMany, 1)
	c := make(chan *result, len(keys))

	for i := range keys {
		go func(i int) {
			getter := l.Load(keys[i])
			c <- &result{getter(), i}
		}(i)
	}

	go func() {
		defer close(out)
		var i int = 1
		outputData := make([]interface{}, len(keys), len(keys))
		outputErrors := make([]error, 0, len(keys))
		for result := range c {
			outputData[result.index] = result.Data
			if result.Error != nil {
				outputErrors = append(outputErrors, resultError{result.Error, result.index})
			}
			if i == len(keys) {
				close(c)
			}
			i++
		}
		out <- &ResultMany{outputData, outputErrors}
	}()

	var value *ResultMany
	getterMany := func() *ResultMany {
		if v, ok := <-out; ok {
			value = v
		}
		return value
	}

	return getterMany
}

// Clear clears the value at `key` from the cache, it it exsits. Returs self for method chaining
func (l *Loader) Clear(key string) Interface {
	l.cache.Delete(key)
	return l
}

// ClearAll clears the entire cache. To be used when some event results in unknown invalidations.
// Returs self for method chaining.
func (l *Loader) ClearAll() Interface {
	l.cache.Clear()
	return l
}

// Adds the provided key and value to the cache. If the key already exists, no change is made.
// Returns self for method chaining
func (l *Loader) Prime(key string, value interface{}) Interface {
	if _, ok := l.cache.Get(key); !ok {
		l.cache.Set(key, value)
	}
	return l
}

// queues item to be batched
func (l *Loader) addToQueue(req *batchRequest) {
	l.mu.Lock()
	defer l.mu.Unlock()

	if l.input == nil {
		l.input = make(chan *batchRequest, l.cap)
		go l.batch()
	}

	l.input <- req
}

// execuite the batch of all items in queue
func (l *Loader) batch() {
	var keys []string
	var reqs []*batchRequest

	go l.sleeper()

	for item := range l.input {
		keys = append(keys, item.key)
		reqs = append(reqs, item)
	}

	items := l.batchFn(keys)

	if len(items) != len(reqs) {
		panic(fmt.Errorf("length of keys must match length of responses"))
	}

	for i, req := range reqs {
		req.channel <- items[i]
		close(req.channel)
	}
}

// wait the appropriate amount of time for next batch
func (l *Loader) sleeper() {
	// this will move this goroutine to the back of the callstack
	time.Sleep(1)
	l.mu.Lock()
	defer l.mu.Unlock()

	close(l.input)
	l.input = nil
}
<|endoftext|>


<commit_before>// Package dataloader is an implimentation of facebook's dataloader in go.
// See https://github.com/facebook/dataloader for more information
package dataloader

import (
	"context"
	"fmt"
	"log"
	"runtime"
	"sync"
	"time"
)

// Interface is a `DataLoader` Interface which defines a public API for loading data from a particular
// data back-end with unique keys such as the `id` column of a SQL table or
// document name in a MongoDB database, given a batch loading function.
//
// Each `DataLoader` instance should contain a unique memoized cache. Use caution when
// used in long-lived applications or those which serve many users with
// different access permissions and consider creating a new instance per
// web request.
type Interface interface {
	Load(context.Context, string) Thunk
	LoadMany(context.Context, []string) ThunkMany
	Clear(string) Interface
	ClearAll() Interface
	Prime(key string, value interface{}) Interface
}

// BatchFunc is a function, which when given a slice of keys (string), returns an slice of `results`.
// It's important that the length of the input keys matches the length of the output results.
//
// The keys passed to this function are guaranteed to be unique
type BatchFunc func(context.Context, []string) []*Result

// Result is the data structure that a BatchFunc returns.
// It contains the resolved data, and any errors that may have occurred while fetching the data.
type Result struct {
	Data  interface{}
	Error error
}

// ResultMany is used by the LoadMany method.
// It contains a list of resolved data and a list of errors.
// The lengths of the data list and error list will match, and elements at each index correspond to each other.
type ResultMany struct {
	Data  []interface{}
	Error []error
}

// Loader implements the dataloader.Interface.
type Loader struct {
	// the batch function to be used by this loader
	batchFn BatchFunc

	// the maximum batch size. Set to 0 if you want it to be unbounded.
	batchCap int

	// the internal cache. This packages contains a basic cache implementation but any custom cache
	// implementation could be used as long as it implements the `Cache` interface.
	cacheLock sync.Mutex
	cache     Cache
	// should we clear the cache on each batch?
	// this would allow batching but no long term caching
	clearCacheOnBatch bool

	// count of queued up items
	count int

	// the maximum input queue size. Set to 0 if you want it to be unbounded.
	inputCap int

	// the amount of time to wait before triggering a batch
	wait time.Duration

	// lock to protect the batching operations
	batchLock sync.Mutex

	// current batcher
	curBatcher *batcher

	// used to close the sleeper of the current batcher
	endSleeper chan bool

	// used by tests to prevent logs
	silent bool

	// can be set to trace calls to dataloader
	tracer Tracer
}

// Thunk is a function that will block until the value (*Result) it contins is resolved.
// After the value it contains is resolved, this function will return the result.
// This function can be called many times, much like a Promise is other languages.
// The value will only need to be resolved once so subsequent calls will return immediately.
type Thunk func() (interface{}, error)

// ThunkMany is much like the Thunk func type but it contains a list of results.
type ThunkMany func() ([]interface{}, []error)

// type used to on input channel
type batchRequest struct {
	key     string
	channel chan *Result
}

// Option allows for configuration of Loader fields.
type Option func(*Loader)

// WithCache sets the BatchedLoader cache. Defaults to InMemoryCache if a Cache is not set.
func WithCache(c Cache) Option {
	return func(l *Loader) {
		l.cache = c
	}
}

// WithBatchCapacity sets the batch capacity. Default is 0 (unbounded).
func WithBatchCapacity(c int) Option {
	return func(l *Loader) {
		l.batchCap = c
	}
}

// WithInputCapacity sets the input capacity. Default is 1000.
func WithInputCapacity(c int) Option {
	return func(l *Loader) {
		l.inputCap = c
	}
}

// WithWait sets the amount of time to wait before triggering a batch.
// Default duration is 16 milliseconds.
func WithWait(d time.Duration) Option {
	return func(l *Loader) {
		l.wait = d
	}
}

// WithClearCacheOnBatch allows batching of items but no long term caching.
// It accomplishes this by clearing the cache after each batch operation.
func WithClearCacheOnBatch() Option {
	return func(l *Loader) {
		l.cacheLock.Lock()
		l.clearCacheOnBatch = true
		l.cacheLock.Unlock()
	}
}

// withSilentLogger turns of log messages. It's used by the tests
func withSilentLogger() Option {
	return func(l *Loader) {
		l.silent = true
	}
}

// WithTracer allows tracing of calls to Load and LoadMany
func WithTracer(tracer Tracer) Option {
	return func(l *Loader) {
		l.tracer = tracer
	}
}

// WithOpenTracingTracer allows tracing of calls to Load and LoadMany
func WithOpenTracingTracer(tracer Tracer) Option {
	return WithTracer(&OpenTracingTracer{})
}

// NewBatchedLoader constructs a new Loader with given options.
func NewBatchedLoader(batchFn BatchFunc, opts ...Option) *Loader {
	loader := &Loader{
		batchFn:  batchFn,
		inputCap: 1000,
		wait:     16 * time.Millisecond,
	}

	// Apply options
	for _, apply := range opts {
		apply(loader)
	}

	// Set defaults
	if loader.cache == nil {
		loader.cache = NewCache()
	}

	if loader.tracer == nil {
		loader.tracer = &NoopTracer{}
	}

	return loader
}

// Load load/resolves the given key, returning a channel that will contain the value and error
func (l *Loader) Load(originalContext context.Context, key string) Thunk {
	ctx, finish := l.tracer.TraceLoad(originalContext, key)
	c := make(chan *Result, 1)
	var result struct {
		mu    sync.RWMutex
		value *Result
	}

	// lock to prevent duplicate keys coming in before item has been added to cache.
	l.cacheLock.Lock()
	if v, ok := l.cache.Get(key); ok {
		defer finish(v)
		defer l.cacheLock.Unlock()
		return v
	}

	thunk := func() (interface{}, error) {
		result.mu.RLock()
		resultNotSet := result.value == nil
		result.mu.RUnlock()

		if resultNotSet {
			result.mu.Lock()
			if v, ok := <-c; ok {
				result.value = v
			}
			result.mu.Unlock()
		}
		result.mu.RLock()
		defer result.mu.RUnlock()
		return result.value.Data, result.value.Error
	}

	l.cache.Set(key, thunk)
	l.cacheLock.Unlock()

	// this is sent to batch fn. It contains the key and the channel to return the
	// the result on
	req := &batchRequest{key, c}

	l.batchLock.Lock()
	// start the batch window if it hasn't already started.
	if l.curBatcher == nil {
		l.curBatcher = l.newBatcher(l.silent, l.tracer)
		// start the current batcher batch function
		go l.curBatcher.batch(ctx)
		// start a sleeper for the current batcher
		l.endSleeper = make(chan bool)
		go l.sleeper(l.curBatcher, l.endSleeper)
	}

	l.curBatcher.input <- req

	// if we need to keep track of the count (max batch), then do so.
	if l.batchCap > 0 {
		l.count++
		// if we hit our limit, force the batch to start
		if l.count == l.batchCap {
			// end the batcher synchronously here because another call to Load
			// may concurrently happen and needs to go to a new batcher.
			l.curBatcher.end()
			// end the sleeper for the current batcher.
			// this is to stop the goroutine without waiting for the
			// sleeper timeout.
			close(l.endSleeper)
			l.reset()
		}
	}
	l.batchLock.Unlock()

	defer finish(thunk)
	return thunk
}

// LoadMany loads mulitiple keys, returning a thunk (type: ThunkMany) that will resolve the keys passed in.
func (l *Loader) LoadMany(originalContext context.Context, keys []string) ThunkMany {
	ctx, finish := l.tracer.TraceLoadMany(originalContext, keys)
	length := len(keys)
	data := make([]interface{}, length)
	errors := make([]error, length)
	c := make(chan *ResultMany, 1)
	wg := sync.WaitGroup{}

	wg.Add(length)
	for i := range keys {
		go func(i int) {
			defer wg.Done()
			thunk := l.Load(ctx, keys[i])
			result, err := thunk()
			data[i] = result
			errors[i] = err
		}(i)
	}

	go func() {
		wg.Wait()
		c <- &ResultMany{data, errors}
		close(c)
	}()

	var result struct {
		mu    sync.RWMutex
		value *ResultMany
	}

	thunkMany := func() ([]interface{}, []error) {
		result.mu.RLock()
		resultNotSet := result.value == nil
		result.mu.RUnlock()

		if resultNotSet {
			result.mu.Lock()
			if v, ok := <-c; ok {
				result.value = v
			}
			result.mu.Unlock()
		}
		result.mu.RLock()
		defer result.mu.RUnlock()
		return result.value.Data, result.value.Error
	}

	defer finish(thunkMany)
	return thunkMany
}

// Clear clears the value at `key` from the cache, it it exsits. Returs self for method chaining
func (l *Loader) Clear(key string) Interface {
	l.cacheLock.Lock()
	l.cache.Delete(key)
	l.cacheLock.Unlock()
	return l
}

// ClearAll clears the entire cache. To be used when some event results in unknown invalidations.
// Returns self for method chaining.
func (l *Loader) ClearAll() Interface {
	l.cacheLock.Lock()
	l.cache.Clear()
	l.cacheLock.Unlock()
	return l
}

// Prime adds the provided key and value to the cache. If the key already exists, no change is made.
// Returns self for method chaining
func (l *Loader) Prime(key string, value interface{}) Interface {
	if _, ok := l.cache.Get(key); !ok {
		thunk := func() (interface{}, error) {
			return value, nil
		}
		l.cache.Set(key, thunk)
	}
	return l
}

func (l *Loader) reset() {
	l.count = 0
	l.curBatcher = nil

	if l.clearCacheOnBatch {
		l.cache.Clear()
	}
}

type batcher struct {
	input    chan *batchRequest
	batchFn  BatchFunc
	finished bool
	silent   bool
	tracer   Tracer
}

// newBatcher returns a batcher for the current requests
// all the batcher methods must be protected by a global batchLock
func (l *Loader) newBatcher(silent bool, tracer Tracer) *batcher {
	return &batcher{
		input:   make(chan *batchRequest, l.inputCap),
		batchFn: l.batchFn,
		silent:  silent,
		tracer:  tracer,
	}
}

// stop receiving input and process batch function
func (b *batcher) end() {
	if !b.finished {
		close(b.input)
		b.finished = true
	}
}

// execute the batch of all items in queue
func (b *batcher) batch(originalContext context.Context) {
	var keys []string
	var reqs []*batchRequest
	var items []*Result
	var panicErr interface{}

	for item := range b.input {
		keys = append(keys, item.key)
		reqs = append(reqs, item)
	}

	ctx, finish := b.tracer.TraceBatch(originalContext, keys)
	defer finish(items)

	func() {
		defer func() {
			if r := recover(); r != nil {
				panicErr = r
				if b.silent {
					return
				}
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				log.Printf("Dataloader: Panic received in batch function:: %v\n%s", panicErr, buf)
			}
		}()
		items = b.batchFn(ctx, keys)
	}()

	if panicErr != nil {
		for _, req := range reqs {
			req.channel <- &Result{Error: fmt.Errorf("Panic received in batch function: %v", panicErr)}
			close(req.channel)
		}
		return
	}

	if len(items) != len(keys) {
		err := &Result{Error: fmt.Errorf(`
			The batch function supplied did not return an array of responses
			the same length as the array of keys.

			Keys:
			%v

			Values:
			%v
		`, keys, items)}

		for _, req := range reqs {
			req.channel <- err
			close(req.channel)
		}

		return
	}

	for i, req := range reqs {
		req.channel <- items[i]
		close(req.channel)
	}
}

// wait the appropriate amount of time for the provided batcher
func (l *Loader) sleeper(b *batcher, close chan bool) {
	select {
	// used by batch to close early. usually triggered by max batch size
	case <-close:
		return
	// this will move this goroutine to the back of the callstack?
	case <-time.After(l.wait):
	}

	// reset
	// this is protected by the batchLock to avoid closing the batcher input
	// channel while Load is inserting a request
	l.batchLock.Lock()
	b.end()

	// We can end here also if the batcher has already been closed and a
	// new one has been created. So reset the loader state only if the batcher
	// is the current one
	if l.curBatcher == b {
		l.reset()
	}
	l.batchLock.Unlock()
}
<commit_msg>fix opentracing tracer<commit_after>// Package dataloader is an implimentation of facebook's dataloader in go.
// See https://github.com/facebook/dataloader for more information
package dataloader

import (
	"context"
	"fmt"
	"log"
	"runtime"
	"sync"
	"time"
)

// Interface is a `DataLoader` Interface which defines a public API for loading data from a particular
// data back-end with unique keys such as the `id` column of a SQL table or
// document name in a MongoDB database, given a batch loading function.
//
// Each `DataLoader` instance should contain a unique memoized cache. Use caution when
// used in long-lived applications or those which serve many users with
// different access permissions and consider creating a new instance per
// web request.
type Interface interface {
	Load(context.Context, string) Thunk
	LoadMany(context.Context, []string) ThunkMany
	Clear(string) Interface
	ClearAll() Interface
	Prime(key string, value interface{}) Interface
}

// BatchFunc is a function, which when given a slice of keys (string), returns an slice of `results`.
// It's important that the length of the input keys matches the length of the output results.
//
// The keys passed to this function are guaranteed to be unique
type BatchFunc func(context.Context, []string) []*Result

// Result is the data structure that a BatchFunc returns.
// It contains the resolved data, and any errors that may have occurred while fetching the data.
type Result struct {
	Data  interface{}
	Error error
}

// ResultMany is used by the LoadMany method.
// It contains a list of resolved data and a list of errors.
// The lengths of the data list and error list will match, and elements at each index correspond to each other.
type ResultMany struct {
	Data  []interface{}
	Error []error
}

// Loader implements the dataloader.Interface.
type Loader struct {
	// the batch function to be used by this loader
	batchFn BatchFunc

	// the maximum batch size. Set to 0 if you want it to be unbounded.
	batchCap int

	// the internal cache. This packages contains a basic cache implementation but any custom cache
	// implementation could be used as long as it implements the `Cache` interface.
	cacheLock sync.Mutex
	cache     Cache
	// should we clear the cache on each batch?
	// this would allow batching but no long term caching
	clearCacheOnBatch bool

	// count of queued up items
	count int

	// the maximum input queue size. Set to 0 if you want it to be unbounded.
	inputCap int

	// the amount of time to wait before triggering a batch
	wait time.Duration

	// lock to protect the batching operations
	batchLock sync.Mutex

	// current batcher
	curBatcher *batcher

	// used to close the sleeper of the current batcher
	endSleeper chan bool

	// used by tests to prevent logs
	silent bool

	// can be set to trace calls to dataloader
	tracer Tracer
}

// Thunk is a function that will block until the value (*Result) it contins is resolved.
// After the value it contains is resolved, this function will return the result.
// This function can be called many times, much like a Promise is other languages.
// The value will only need to be resolved once so subsequent calls will return immediately.
type Thunk func() (interface{}, error)

// ThunkMany is much like the Thunk func type but it contains a list of results.
type ThunkMany func() ([]interface{}, []error)

// type used to on input channel
type batchRequest struct {
	key     string
	channel chan *Result
}

// Option allows for configuration of Loader fields.
type Option func(*Loader)

// WithCache sets the BatchedLoader cache. Defaults to InMemoryCache if a Cache is not set.
func WithCache(c Cache) Option {
	return func(l *Loader) {
		l.cache = c
	}
}

// WithBatchCapacity sets the batch capacity. Default is 0 (unbounded).
func WithBatchCapacity(c int) Option {
	return func(l *Loader) {
		l.batchCap = c
	}
}

// WithInputCapacity sets the input capacity. Default is 1000.
func WithInputCapacity(c int) Option {
	return func(l *Loader) {
		l.inputCap = c
	}
}

// WithWait sets the amount of time to wait before triggering a batch.
// Default duration is 16 milliseconds.
func WithWait(d time.Duration) Option {
	return func(l *Loader) {
		l.wait = d
	}
}

// WithClearCacheOnBatch allows batching of items but no long term caching.
// It accomplishes this by clearing the cache after each batch operation.
func WithClearCacheOnBatch() Option {
	return func(l *Loader) {
		l.cacheLock.Lock()
		l.clearCacheOnBatch = true
		l.cacheLock.Unlock()
	}
}

// withSilentLogger turns of log messages. It's used by the tests
func withSilentLogger() Option {
	return func(l *Loader) {
		l.silent = true
	}
}

// WithTracer allows tracing of calls to Load and LoadMany
func WithTracer(tracer Tracer) Option {
	return func(l *Loader) {
		l.tracer = tracer
	}
}

// WithOpenTracingTracer allows tracing of calls to Load and LoadMany
func WithOpenTracingTracer() Option {
	return WithTracer(&OpenTracingTracer{})
}

// NewBatchedLoader constructs a new Loader with given options.
func NewBatchedLoader(batchFn BatchFunc, opts ...Option) *Loader {
	loader := &Loader{
		batchFn:  batchFn,
		inputCap: 1000,
		wait:     16 * time.Millisecond,
	}

	// Apply options
	for _, apply := range opts {
		apply(loader)
	}

	// Set defaults
	if loader.cache == nil {
		loader.cache = NewCache()
	}

	if loader.tracer == nil {
		loader.tracer = &NoopTracer{}
	}

	return loader
}

// Load load/resolves the given key, returning a channel that will contain the value and error
func (l *Loader) Load(originalContext context.Context, key string) Thunk {
	ctx, finish := l.tracer.TraceLoad(originalContext, key)
	c := make(chan *Result, 1)
	var result struct {
		mu    sync.RWMutex
		value *Result
	}

	// lock to prevent duplicate keys coming in before item has been added to cache.
	l.cacheLock.Lock()
	if v, ok := l.cache.Get(key); ok {
		defer finish(v)
		defer l.cacheLock.Unlock()
		return v
	}

	thunk := func() (interface{}, error) {
		result.mu.RLock()
		resultNotSet := result.value == nil
		result.mu.RUnlock()

		if resultNotSet {
			result.mu.Lock()
			if v, ok := <-c; ok {
				result.value = v
			}
			result.mu.Unlock()
		}
		result.mu.RLock()
		defer result.mu.RUnlock()
		return result.value.Data, result.value.Error
	}

	l.cache.Set(key, thunk)
	l.cacheLock.Unlock()

	// this is sent to batch fn. It contains the key and the channel to return the
	// the result on
	req := &batchRequest{key, c}

	l.batchLock.Lock()
	// start the batch window if it hasn't already started.
	if l.curBatcher == nil {
		l.curBatcher = l.newBatcher(l.silent, l.tracer)
		// start the current batcher batch function
		go l.curBatcher.batch(ctx)
		// start a sleeper for the current batcher
		l.endSleeper = make(chan bool)
		go l.sleeper(l.curBatcher, l.endSleeper)
	}

	l.curBatcher.input <- req

	// if we need to keep track of the count (max batch), then do so.
	if l.batchCap > 0 {
		l.count++
		// if we hit our limit, force the batch to start
		if l.count == l.batchCap {
			// end the batcher synchronously here because another call to Load
			// may concurrently happen and needs to go to a new batcher.
			l.curBatcher.end()
			// end the sleeper for the current batcher.
			// this is to stop the goroutine without waiting for the
			// sleeper timeout.
			close(l.endSleeper)
			l.reset()
		}
	}
	l.batchLock.Unlock()

	defer finish(thunk)
	return thunk
}

// LoadMany loads mulitiple keys, returning a thunk (type: ThunkMany) that will resolve the keys passed in.
func (l *Loader) LoadMany(originalContext context.Context, keys []string) ThunkMany {
	ctx, finish := l.tracer.TraceLoadMany(originalContext, keys)
	length := len(keys)
	data := make([]interface{}, length)
	errors := make([]error, length)
	c := make(chan *ResultMany, 1)
	wg := sync.WaitGroup{}

	wg.Add(length)
	for i := range keys {
		go func(i int) {
			defer wg.Done()
			thunk := l.Load(ctx, keys[i])
			result, err := thunk()
			data[i] = result
			errors[i] = err
		}(i)
	}

	go func() {
		wg.Wait()
		c <- &ResultMany{data, errors}
		close(c)
	}()

	var result struct {
		mu    sync.RWMutex
		value *ResultMany
	}

	thunkMany := func() ([]interface{}, []error) {
		result.mu.RLock()
		resultNotSet := result.value == nil
		result.mu.RUnlock()

		if resultNotSet {
			result.mu.Lock()
			if v, ok := <-c; ok {
				result.value = v
			}
			result.mu.Unlock()
		}
		result.mu.RLock()
		defer result.mu.RUnlock()
		return result.value.Data, result.value.Error
	}

	defer finish(thunkMany)
	return thunkMany
}

// Clear clears the value at `key` from the cache, it it exsits. Returs self for method chaining
func (l *Loader) Clear(key string) Interface {
	l.cacheLock.Lock()
	l.cache.Delete(key)
	l.cacheLock.Unlock()
	return l
}

// ClearAll clears the entire cache. To be used when some event results in unknown invalidations.
// Returns self for method chaining.
func (l *Loader) ClearAll() Interface {
	l.cacheLock.Lock()
	l.cache.Clear()
	l.cacheLock.Unlock()
	return l
}

// Prime adds the provided key and value to the cache. If the key already exists, no change is made.
// Returns self for method chaining
func (l *Loader) Prime(key string, value interface{}) Interface {
	if _, ok := l.cache.Get(key); !ok {
		thunk := func() (interface{}, error) {
			return value, nil
		}
		l.cache.Set(key, thunk)
	}
	return l
}

func (l *Loader) reset() {
	l.count = 0
	l.curBatcher = nil

	if l.clearCacheOnBatch {
		l.cache.Clear()
	}
}

type batcher struct {
	input    chan *batchRequest
	batchFn  BatchFunc
	finished bool
	silent   bool
	tracer   Tracer
}

// newBatcher returns a batcher for the current requests
// all the batcher methods must be protected by a global batchLock
func (l *Loader) newBatcher(silent bool, tracer Tracer) *batcher {
	return &batcher{
		input:   make(chan *batchRequest, l.inputCap),
		batchFn: l.batchFn,
		silent:  silent,
		tracer:  tracer,
	}
}

// stop receiving input and process batch function
func (b *batcher) end() {
	if !b.finished {
		close(b.input)
		b.finished = true
	}
}

// execute the batch of all items in queue
func (b *batcher) batch(originalContext context.Context) {
	var keys []string
	var reqs []*batchRequest
	var items []*Result
	var panicErr interface{}

	for item := range b.input {
		keys = append(keys, item.key)
		reqs = append(reqs, item)
	}

	ctx, finish := b.tracer.TraceBatch(originalContext, keys)
	defer finish(items)

	func() {
		defer func() {
			if r := recover(); r != nil {
				panicErr = r
				if b.silent {
					return
				}
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				log.Printf("Dataloader: Panic received in batch function:: %v\n%s", panicErr, buf)
			}
		}()
		items = b.batchFn(ctx, keys)
	}()

	if panicErr != nil {
		for _, req := range reqs {
			req.channel <- &Result{Error: fmt.Errorf("Panic received in batch function: %v", panicErr)}
			close(req.channel)
		}
		return
	}

	if len(items) != len(keys) {
		err := &Result{Error: fmt.Errorf(`
			The batch function supplied did not return an array of responses
			the same length as the array of keys.

			Keys:
			%v

			Values:
			%v
		`, keys, items)}

		for _, req := range reqs {
			req.channel <- err
			close(req.channel)
		}

		return
	}

	for i, req := range reqs {
		req.channel <- items[i]
		close(req.channel)
	}
}

// wait the appropriate amount of time for the provided batcher
func (l *Loader) sleeper(b *batcher, close chan bool) {
	select {
	// used by batch to close early. usually triggered by max batch size
	case <-close:
		return
	// this will move this goroutine to the back of the callstack?
	case <-time.After(l.wait):
	}

	// reset
	// this is protected by the batchLock to avoid closing the batcher input
	// channel while Load is inserting a request
	l.batchLock.Lock()
	b.end()

	// We can end here also if the batcher has already been closed and a
	// new one has been created. So reset the loader state only if the batcher
	// is the current one
	if l.curBatcher == b {
		l.reset()
	}
	l.batchLock.Unlock()
}
<|endoftext|>