<commit_before>package image

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	gocontext "context"

	"github.com/cenk/backoff"
	"github.com/pkg/errors"
	workererrors "github.com/travis-ci/worker/errors"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/trace"
)

const (
	imageAPIRequestContentType = "application/x-www-form-urlencoded; boundary=NL"
)

type APISelector struct {
	baseURL *url.URL

	maxInterval    time.Duration
	maxElapsedTime time.Duration
}

func NewAPISelector(u *url.URL) *APISelector {
	return &APISelector{
		baseURL:        u,
		maxInterval:    10 * time.Second,
		maxElapsedTime: time.Minute,
	}
}

func (as *APISelector) SetMaxInterval(maxInterval time.Duration) {
	as.maxInterval = maxInterval
}

func (as *APISelector) SetMaxElapsedTime(maxElapsedTime time.Duration) {
	as.maxElapsedTime = maxElapsedTime
}

func (as *APISelector) Select(ctx gocontext.Context, params *Params) (string, error) {
	tagSets, err := as.buildCandidateTags(params)
	if err != nil {
		return "default", err
	}

	imageName, err := as.queryWithTags(ctx, params.Infra, tagSets)
	if err != nil {
		return "default", err
	}

	if imageName != "" {
		return imageName, nil
	}

	return "default", nil
}

func (as *APISelector) queryWithTags(ctx gocontext.Context, infra string, tags []*tagSet) (string, error) {
	ctx, span := trace.StartSpan(ctx, "APISelector.querywithTags")
	defer span.End()

	bodyLines := []string{}
	lastJobID := uint64(0)
	lastRepo := ""

	for _, ts := range tags {
		qs := url.Values{}
		qs.Set("infra", infra)
		qs.Set("fields[images]", "name")
		qs.Set("limit", "1")
		qs.Set("job_id", fmt.Sprintf("%v", ts.JobID))
		qs.Set("repo", ts.Repo)
		qs.Set("is_default", fmt.Sprintf("%v", ts.IsDefault))
		if len(ts.Tags) > 0 {
			qs.Set("tags", strings.Join(ts.Tags, ","))
		}

		bodyLines = append(bodyLines, qs.Encode())
		lastJobID = ts.JobID
		lastRepo = ts.Repo
	}

	qs := url.Values{}
	qs.Set("infra", infra)
	qs.Set("is_default", "true")
	qs.Set("fields[images]", "name")
	qs.Set("limit", "1")
	qs.Set("job_id", fmt.Sprintf("%v", lastJobID))
	qs.Set("repo", lastRepo)

	bodyLines = append(bodyLines, qs.Encode())

	u, err := url.Parse(as.baseURL.String())
	if err != nil {
		return "", err
	}

	imageResp, err := as.makeImageRequest(ctx, u.String(), bodyLines)
	if err != nil {
		return "", err
	}

	if len(imageResp.Data) == 0 {
		return "", nil
	}

	return imageResp.Data[0].Name, nil
}

func (as *APISelector) makeImageRequest(ctx gocontext.Context, urlString string, bodyLines []string) (*apiSelectorImageResponse, error) {
	var responseBody []byte

	b := backoff.NewExponentialBackOff()
	b.MaxInterval = as.maxInterval
	b.MaxElapsedTime = as.maxElapsedTime

	client := &http.Client{
		Transport: &ochttp.Transport{},
	}

	err := backoff.Retry(func() error {
		req, err := http.NewRequest("POST", urlString, strings.NewReader(strings.Join(bodyLines, "\n")+"\n"))
		if err != nil {
			return err
		}

		req.Header.Add("Content-Type", imageAPIRequestContentType)
		req.WithContext(ctx)
		resp, err := client.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()<commit_msg>declare req again for http client<commit_after>package image

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	gocontext "context"

	"github.com/cenk/backoff"
	"github.com/pkg/errors"
	workererrors "github.com/travis-ci/worker/errors"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/trace"
)

const (
	imageAPIRequestContentType = "application/x-www-form-urlencoded; boundary=NL"
)

type APISelector struct {
	baseURL *url.URL

	maxInterval    time.Duration
	maxElapsedTime time.Duration
}

func NewAPISelector(u *url.URL) *APISelector {
	return &APISelector{
		baseURL:        u,
		maxInterval:    10 * time.Second,
		maxElapsedTime: time.Minute,
	}
}

func (as *APISelector) SetMaxInterval(maxInterval time.Duration) {
	as.maxInterval = maxInterval
}

func (as *APISelector) SetMaxElapsedTime(maxElapsedTime time.Duration) {
	as.maxElapsedTime = maxElapsedTime
}

func (as *APISelector) Select(ctx gocontext.Context, params *Params) (string, error) {
	tagSets, err := as.buildCandidateTags(params)
	if err != nil {
		return "default", err
	}

	imageName, err := as.queryWithTags(ctx, params.Infra, tagSets)
	if err != nil {
		return "default", err
	}

	if imageName != "" {
		return imageName, nil
	}

	return "default", nil
}

func (as *APISelector) queryWithTags(ctx gocontext.Context, infra string, tags []*tagSet) (string, error) {
	ctx, span := trace.StartSpan(ctx, "APISelector.querywithTags")
	defer span.End()

	bodyLines := []string{}
	lastJobID := uint64(0)
	lastRepo := ""

	for _, ts := range tags {
		qs := url.Values{}
		qs.Set("infra", infra)
		qs.Set("fields[images]", "name")
		qs.Set("limit", "1")
		qs.Set("job_id", fmt.Sprintf("%v", ts.JobID))
		qs.Set("repo", ts.Repo)
		qs.Set("is_default", fmt.Sprintf("%v", ts.IsDefault))
		if len(ts.Tags) > 0 {
			qs.Set("tags", strings.Join(ts.Tags, ","))
		}

		bodyLines = append(bodyLines, qs.Encode())
		lastJobID = ts.JobID
		lastRepo = ts.Repo
	}

	qs := url.Values{}
	qs.Set("infra", infra)
	qs.Set("is_default", "true")
	qs.Set("fields[images]", "name")
	qs.Set("limit", "1")
	qs.Set("job_id", fmt.Sprintf("%v", lastJobID))
	qs.Set("repo", lastRepo)

	bodyLines = append(bodyLines, qs.Encode())

	u, err := url.Parse(as.baseURL.String())
	if err != nil {
		return "", err
	}

	imageResp, err := as.makeImageRequest(ctx, u.String(), bodyLines)
	if err != nil {
		return "", err
	}

	if len(imageResp.Data) == 0 {
		return "", nil
	}

	return imageResp.Data[0].Name, nil
}

func (as *APISelector) makeImageRequest(ctx gocontext.Context, urlString string, bodyLines []string) (*apiSelectorImageResponse, error) {
	var responseBody []byte

	b := backoff.NewExponentialBackOff()
	b.MaxInterval = as.maxInterval
	b.MaxElapsedTime = as.maxElapsedTime

	client := &http.Client{
		Transport: &ochttp.Transport{},
	}

	err := backoff.Retry(func() error {
		req, err := http.NewRequest("POST", urlString, strings.NewReader(strings.Join(bodyLines, "\n")+"\n"))
		if err != nil {
			return err
		}

		req.Header.Add("Content-Type", imageAPIRequestContentType)
		req = req.WithContext(ctx)
		resp, err := client.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
{'fraction_non_alphanumeric': 0.16346758967904343, 'fraction_numerical': 0.0031466331025802393, 'mean_word_length': 4.253719008264463, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '620364', 'n_tokens_mistral': 2517, 'n_tokens_neox': 2399, 'n_words': 495}