<commit_before>// Copyright 2018 Authors of Cilium
//
// 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 metrics

import (
	"time"

	clientPkg "github.com/cilium/cilium/pkg/client"
	healthClientPkg "github.com/cilium/cilium/pkg/health/client"

	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
)

const (
	updateLatencyMetricsInterval = 30 * time.Second
)

type statusCollector struct {
	ciliumClient *clientPkg.Client
	healthClient *healthClientPkg.Client

	controllersFailingDesc         *prometheus.Desc
	ipAddressesDesc                *prometheus.Desc
	unreachableNodesDesc           *prometheus.Desc
	unreachableHealthEndpointsDesc *prometheus.Desc
}

func newStatusCollector() *statusCollector {
	ciliumClient, err := clientPkg.NewClient("")
	if err != nil {
		log.WithError(err).Error("Error while creating client")
	}

	healthClient, err := healthClientPkg.NewClient("")
	if err != nil {
		log.WithError(err).Error("Error while creating cilium-health client")
	}

	return &statusCollector{
		ciliumClient: ciliumClient,
		healthClient: healthClient,
		controllersFailingDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "controllers_failing"),
			"Number of failing controllers",
			nil, nil,
		),
		ipAddressesDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "ip_addresses"),
			"Number of allocated IP addresses",
			[]string{"family"}, nil,
		),
		unreachableNodesDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "unreachable_nodes"),
			"Number of nodes that cannot be reached",
			nil, nil,
		),
		unreachableHealthEndpointsDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "unreachable_health_endpoints"),
			"Number of health endpoints that cannot be reached",
			nil, nil,
		),
	}
}

func (s *statusCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- s.controllersFailingDesc
	ch <- s.ipAddressesDesc
	ch <- s.unreachableNodesDesc
	ch <- s.unreachableHealthEndpointsDesc
}

func (s *statusCollector) Collect(ch chan<- prometheus.Metric) {
	statusResponse, err := s.ciliumClient.Daemon.GetHealthz(nil)
	if err != nil {
		log.WithError(err).Error("Error while getting Cilium status")
		return
	}

	if statusResponse.Payload == nil {
		return
	}

	// Controllers failing
	controllersFailing := 0

	for _, ctrl := range statusResponse.Payload.Controllers {
		if ctrl.Status == nil {
			continue
		}
		if ctrl.Status.ConsecutiveFailureCount > 0 {
			controllersFailing++
		}
	}

	ch <- prometheus.MustNewConstMetric(
		s.controllersFailingDesc,
		prometheus.GaugeValue,
		float64(controllersFailing),
	)

	// Address count
	ch <- prometheus.MustNewConstMetric(
		s.ipAddressesDesc,
		prometheus.GaugeValue,
		float64(len(statusResponse.Payload.IPAM.IPV4)),
		"ipv4",
	)

	ch <- prometheus.MustNewConstMetric(
		s.ipAddressesDesc,
		prometheus.GaugeValue,
		float64(len(statusResponse.Payload.IPAM.IPV6)),
		"ipv6",
	)

	healthStatusResponse, err := s.healthClient.Connectivity.GetStatus(nil)
	if err != nil {
		log.WithError(err).Error("Error while getting cilium-health status")
		return
	}

	if healthStatusResponse.Payload == nil {
		return
	}

	// Nodes and endpoints healthStatusResponse
	var (
		unreachableNodes     int
		unreachableEndpoints int
	)

	for _, nodeStatus := range healthStatusResponse.Payload.Nodes {
		if !healthClientPkg.PathIsHealthy(nodeStatus.Host.PrimaryAddress) {
			unreachableNodes++
		}
		if nodeStatus.Endpoint != nil && !healthClientPkg.PathIsHealthy(nodeStatus.Endpoint) {
			unreachableEndpoints++
		}
	}

	ch <- prometheus.MustNewConstMetric(
		s.unreachableNodesDesc,
		prometheus.GaugeValue,
		float64(unreachableNodes),
	)

	ch <- prometheus.MustNewConstMetric(
		s.unreachableHealthEndpointsDesc,
		prometheus.GaugeValue,
		float64(unreachableEndpoints),
	)
}
<commit_msg>metrics: Fail with Fatal() when client creation fails<commit_after>// Copyright 2018 Authors of Cilium
//
// 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 metrics

import (
	"time"

	clientPkg "github.com/cilium/cilium/pkg/client"
	healthClientPkg "github.com/cilium/cilium/pkg/health/client"

	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
)

const (
	updateLatencyMetricsInterval = 30 * time.Second
)

type statusCollector struct {
	ciliumClient *clientPkg.Client
	healthClient *healthClientPkg.Client

	controllersFailingDesc         *prometheus.Desc
	ipAddressesDesc                *prometheus.Desc
	unreachableNodesDesc           *prometheus.Desc
	unreachableHealthEndpointsDesc *prometheus.Desc
}

func newStatusCollector() *statusCollector {
	ciliumClient, err := clientPkg.NewClient("")
	if err != nil {
		log.WithError(err).Fatal("Error while creating Cilium API client")
	}

	healthClient, err := healthClientPkg.NewClient("")
	if err != nil {
		log.WithError(err).Fatal("Error while creating cilium-health API client")
	}

	return &statusCollector{
		ciliumClient: ciliumClient,
		healthClient: healthClient,
		controllersFailingDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "controllers_failing"),
			"Number of failing controllers",
			nil, nil,
		),
		ipAddressesDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "ip_addresses"),
			"Number of allocated IP addresses",
			[]string{"family"}, nil,
		),
		unreachableNodesDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "unreachable_nodes"),
			"Number of nodes that cannot be reached",
			nil, nil,
		),
		unreachableHealthEndpointsDesc: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, "", "unreachable_health_endpoints"),
			"Number of health endpoints that cannot be reached",
			nil, nil,
		),
	}
}

func (s *statusCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- s.controllersFailingDesc
	ch <- s.ipAddressesDesc
	ch <- s.unreachableNodesDesc
	ch <- s.unreachableHealthEndpointsDesc
}

func (s *statusCollector) Collect(ch chan<- prometheus.Metric) {
	statusResponse, err := s.ciliumClient.Daemon.GetHealthz(nil)
	if err != nil {
		log.WithError(err).Error("Error while getting Cilium status")
		return
	}

	if statusResponse.Payload == nil {
		return
	}

	// Controllers failing
	controllersFailing := 0

	for _, ctrl := range statusResponse.Payload.Controllers {
		if ctrl.Status == nil {
			continue
		}
		if ctrl.Status.ConsecutiveFailureCount > 0 {
			controllersFailing++
		}
	}

	ch <- prometheus.MustNewConstMetric(
		s.controllersFailingDesc,
		prometheus.GaugeValue,
		float64(controllersFailing),
	)

	// Address count
	ch <- prometheus.MustNewConstMetric(
		s.ipAddressesDesc,
		prometheus.GaugeValue,
		float64(len(statusResponse.Payload.IPAM.IPV4)),
		"ipv4",
	)

	ch <- prometheus.MustNewConstMetric(
		s.ipAddressesDesc,
		prometheus.GaugeValue,
		float64(len(statusResponse.Payload.IPAM.IPV6)),
		"ipv6",
	)

	healthStatusResponse, err := s.healthClient.Connectivity.GetStatus(nil)
	if err != nil {
		log.WithError(err).Error("Error while getting cilium-health status")
		return
	}

	if healthStatusResponse.Payload == nil {
		return
	}

	// Nodes and endpoints healthStatusResponse
	var (
		unreachableNodes     int
		unreachableEndpoints int
	)

	for _, nodeStatus := range healthStatusResponse.Payload.Nodes {
		if !healthClientPkg.PathIsHealthy(nodeStatus.Host.PrimaryAddress) {
			unreachableNodes++
		}
		if nodeStatus.Endpoint != nil && !healthClientPkg.PathIsHealthy(nodeStatus.Endpoint) {
			unreachableEndpoints++
		}
	}

	ch <- prometheus.MustNewConstMetric(
		s.unreachableNodesDesc,
		prometheus.GaugeValue,
		float64(unreachableNodes),
	)

	ch <- prometheus.MustNewConstMetric(
		s.unreachableHealthEndpointsDesc,
		prometheus.GaugeValue,
		float64(unreachableEndpoints),
	)
}
<|endoftext|>


<commit_before>// ClawIO - Scalable Distributed High-Performance Synchronisation and Sharing Service
//
// Copyright (C) 2015  Hugo González Labrador <clawio@hugo.labkode.com>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version. See file COPYNG.

package storage

// AlreadyExistError represents the error ocurred when the resource already exists.
type AlreadyExistError struct {
	Err string
}

func (e *AlreadyExistError) Error() string { return e.Err }

// NotExistError represents the error ocurred when the resource does not exist.
type NotExistError struct {
	Err string
}

func (e *NotExistError) Error() string { return e.Err }

// BadChecksumError represents the error of a missmatch between client and server checksums.
type BadChecksumError struct {
	Err string
}

func (e *BadChecksumError) Error() string {
	return e.Err
}
<commit_msg>Added NotImplementedError to storage<commit_after>// ClawIO - Scalable Distributed High-Performance Synchronisation and Sharing Service
//
// Copyright (C) 2015  Hugo González Labrador <clawio@hugo.labkode.com>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version. See file COPYNG.

package storage

// AlreadyExistError represents the error ocurred when the resource already exists.
type AlreadyExistError struct {
	Err string
}

func (e *AlreadyExistError) Error() string { return e.Err }

// NotExistError represents the error ocurred when the resource does not exist.
type NotExistError struct {
	Err string
}

func (e *NotExistError) Error() string { return e.Err }

// BadChecksumError represents the error of a missmatch between client and server checksums.
type BadChecksumError struct {
	Err string
}

func (e *BadChecksumError) Error() string {
	return e.Err
}

type NotImplementedError struct {
	Err string
}

func (e *NotImplementedError) Error() string {
	return e.Err
}
<|endoftext|>


<commit_before>package sub

import (
	"context"

	"github.com/gopasspw/gopass/pkg/out"
	"github.com/gopasspw/gopass/pkg/store"
	"github.com/gopasspw/gopass/pkg/store/secret"
)

// Get returns the plaintext of a single key
func (s *Store) Get(ctx context.Context, name string) (store.Secret, error) {
	p := s.passfile(name)

	ciphertext, err := s.storage.Get(ctx, p)
	if err != nil {
		out.Debug(ctx, "File %s not found: %s", p, err)
		return nil, store.ErrNotFound
	}

	content, err := s.crypto.Decrypt(ctx, ciphertext)
	if err != nil {
		out.Debug(ctx, "Decryption failed: %s", err)
		return nil, store.ErrDecrypt
	}

	sec, err := secret.Parse(content)
	if err != nil {
		out.Debug(ctx, "Failed to parse YAML: %s", err)
	}
	return sec, nil
}
<commit_msg>Provide more detailed decryption error messages (#1374)<commit_after>package sub

import (
	"context"

	"github.com/gopasspw/gopass/pkg/out"
	"github.com/gopasspw/gopass/pkg/store"
	"github.com/gopasspw/gopass/pkg/store/secret"
)

// Get returns the plaintext of a single key
func (s *Store) Get(ctx context.Context, name string) (store.Secret, error) {
	p := s.passfile(name)

	ciphertext, err := s.storage.Get(ctx, p)
	if err != nil {
		out.Debug(ctx, "File %s not found: %s", p, err)
		return nil, store.ErrNotFound
	}

	content, err := s.crypto.Decrypt(ctx, ciphertext)
	if err != nil {
		out.Error(ctx, "Decryption failed: %s\n%s", err, string(content))
		return nil, store.ErrDecrypt
	}

	sec, err := secret.Parse(content)
	if err != nil {
		out.Debug(ctx, "Failed to parse YAML: %s", err)
	}
	return sec, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package google

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/logs"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
)

// Option is a functional option for List and Walk.
// TODO: Can we somehow reuse the remote options here?
type Option func(*lister) error

type lister struct {
	auth      authn.Authenticator
	transport http.RoundTripper
	repo      name.Repository
	client    *http.Client
	ctx       context.Context
	userAgent string
}

func newLister(repo name.Repository, options ...Option) (*lister, error) {
	l := &lister{
		auth:      authn.Anonymous,
		transport: http.DefaultTransport,
		repo:      repo,
		ctx:       context.Background(),
	}

	for _, option := range options {
		if err := option(l); err != nil {
			return nil, err
		}
	}

	// Wrap the transport in something that logs requests and responses.
	// It's expensive to generate the dumps, so skip it if we're writing
	// to nothing.
	if logs.Enabled(logs.Debug) {
		l.transport = transport.NewLogger(l.transport)
	}

	// Wrap the transport in something that can retry network flakes.
	l.transport = transport.NewRetry(l.transport)

	// Wrap this last to prevent transport.New from double-wrapping.
	if l.userAgent != "" {
		l.transport = transport.NewUserAgent(l.transport, l.userAgent)
	}

	scopes := []string{repo.Scope(transport.PullScope)}
	tr, err := transport.NewWithContext(l.ctx, repo.Registry, l.auth, l.transport, scopes)
	if err != nil {
		return nil, err
	}

	l.client = &http.Client{Transport: tr}

	return l, nil
}

func (l *lister) list(repo name.Repository) (*Tags, error) {
	uri := &url.URL{
		Scheme: repo.Registry.Scheme(),
		Host:   repo.Registry.RegistryStr(),
		Path:   fmt.Sprintf("/v2/%s/tags/list", repo.RepositoryStr()),
		// ECR returns an error if n > 1000:
		// https://github.com/google/go-containerregistry/issues/681
		RawQuery: "n=1000",
	}

	tags := Tags{}

	// get responses until there is no next page
	for {
		select {
		case <-l.ctx.Done():
			return nil, l.ctx.Err()
		default:
		}

		req, err := http.NewRequest("GET", uri.String(), nil)
		if err != nil {
			return nil, err
		}
		req = req.WithContext(l.ctx)

		resp, err := l.client.Do(req)
		if err != nil {
			return nil, err
		}

		if err := transport.CheckError(resp, http.StatusOK); err != nil {
			return nil, err
		}

		parsed := Tags{}
		if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
			return nil, err
		}

		if err := resp.Body.Close(); err != nil {
			return nil, err
		}

		if len(parsed.Manifests) != 0 {
			// We're dealing with GCR, just return directly.
			return &parsed, nil
		}

		// This isn't GCR, just append the tags and keep paginating.
		logs.Warn.Printf("saw non-google tag listing response, falling back to pagination")
		tags.Tags = append(tags.Tags, parsed.Tags...)

		uri, err = getNextPageURL(resp)
		if err != nil {
			return nil, err
		}
		// no next page
		if uri == nil {
			break
		}
	}

	return &tags, nil
}

// getNextPageURL checks if there is a Link header in a http.Response which
// contains a link to the next page. If yes it returns the url.URL of the next
// page otherwise it returns nil.
func getNextPageURL(resp *http.Response) (*url.URL, error) {
	link := resp.Header.Get("Link")
	if link == "" {
		return nil, nil
	}

	if link[0] != '<' {
		return nil, fmt.Errorf("failed to parse link header: missing '<' in: %s", link)
	}

	end := strings.Index(link, ">")
	if end == -1 {
		return nil, fmt.Errorf("failed to parse link header: missing '>' in: %s", link)
	}
	link = link[1:end]

	linkURL, err := url.Parse(link)
	if err != nil {
		return nil, err
	}
	if resp.Request == nil || resp.Request.URL == nil {
		return nil, nil
	}
	linkURL = resp.Request.URL.ResolveReference(linkURL)
	return linkURL, nil
}

type rawManifestInfo struct {
	Size      string   `json:"imageSizeBytes"`
	MediaType string   `json:"mediaType"`
	Created   string   `json:"timeCreatedMs"`
	Uploaded  string   `json:"timeUploadedMs"`
	Tags      []string `json:"tag"`
}

// ManifestInfo is a Manifests entry is the output of List and Walk.
type ManifestInfo struct {
	Size      uint64    `json:"imageSizeBytes"`
	MediaType string    `json:"mediaType"`
	Created   time.Time `json:"timeCreatedMs"`
	Uploaded  time.Time `json:"timeUploadedMs"`
	Tags      []string  `json:"tag"`
}

func fromUnixMs(ms int64) time.Time {
	sec := ms / 1000
	ns := (ms % 1000) * 1000000
	return time.Unix(sec, ns)
}

func toUnixMs(t time.Time) string {
	return strconv.FormatInt(t.UnixNano()/1000000, 10)
}

// MarshalJSON implements json.Marshaler
func (m ManifestInfo) MarshalJSON() ([]byte, error) {
	return json.Marshal(rawManifestInfo{
		Size:      strconv.FormatUint(m.Size, 10),
		MediaType: m.MediaType,
		Created:   toUnixMs(m.Created),
		Uploaded:  toUnixMs(m.Uploaded),
		Tags:      m.Tags,
	})
}

// UnmarshalJSON implements json.Unmarshaler
func (m *ManifestInfo) UnmarshalJSON(data []byte) error {
	raw := rawManifestInfo{}
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}

	if raw.Size != "" {
		size, err := strconv.ParseUint(string(raw.Size), 10, 64)
		if err != nil {
			return err
		}
		m.Size = size
	}

	if raw.Created != "" {
		created, err := strconv.ParseInt(string(raw.Created), 10, 64)
		if err != nil {
			return err
		}
		m.Created = fromUnixMs(created)
	}

	if raw.Uploaded != "" {
		uploaded, err := strconv.ParseInt(string(raw.Uploaded), 10, 64)
		if err != nil {
			return err
		}
		m.Uploaded = fromUnixMs(uploaded)
	}

	m.MediaType = raw.MediaType
	m.Tags = raw.Tags

	return nil
}

// Tags is the result of List and Walk.
type Tags struct {
	Children  []string                `json:"child"`
	Manifests map[string]ManifestInfo `json:"manifest"`
	Name      string                  `json:"name"`
	Tags      []string                `json:"tags"`
}

// List calls /tags/list for the given repository.
func List(repo name.Repository, options ...Option) (*Tags, error) {
	l, err := newLister(repo, options...)
	if err != nil {
		return nil, err
	}

	return l.list(repo)
}

// WalkFunc is the type of the function called for each repository visited by
// Walk. This implements a similar API to filepath.Walk.
//
// The repo argument contains the argument to Walk as a prefix; that is, if Walk
// is called with "gcr.io/foo", which is a repository containing the repository
// "bar", the walk function will be called with argument "gcr.io/foo/bar".
// The tags and error arguments are the result of calling List on repo.
//
// TODO: Do we want a SkipDir error, as in filepath.WalkFunc?
type WalkFunc func(repo name.Repository, tags *Tags, err error) error

func walk(repo name.Repository, tags *Tags, walkFn WalkFunc, options ...Option) error {
	if tags == nil {
		// This shouldn't happen.
		return fmt.Errorf("tags nil for %q", repo)
	}

	if err := walkFn(repo, tags, nil); err != nil {
		return err
	}

	for _, path := range tags.Children {
		child, err := name.NewRepository(fmt.Sprintf("%s/%s", repo, path), name.StrictValidation)
		if err != nil {
			// We don't expect this ever, so don't pass it through to walkFn.
			return fmt.Errorf("unexpected path failure: %v", err)
		}

		childTags, err := List(child, options...)
		if err != nil {
			if err := walkFn(child, nil, err); err != nil {
				return err
			}
		} else {
			if err := walk(child, childTags, walkFn, options...); err != nil {
				return err
			}
		}
	}

	// We made it!
	return nil
}

// Walk recursively descends repositories, calling walkFn.
func Walk(root name.Repository, walkFn WalkFunc, options ...Option) error {
	tags, err := List(root, options...)
	if err != nil {
		return walkFn(root, nil, err)
	}

	return walk(root, tags, walkFn, options...)
}
<commit_msg>Don't warn about fallback unless we need to (#996)<commit_after>// Copyright 2018 Google LLC All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package google

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/logs"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
)

// Option is a functional option for List and Walk.
// TODO: Can we somehow reuse the remote options here?
type Option func(*lister) error

type lister struct {
	auth      authn.Authenticator
	transport http.RoundTripper
	repo      name.Repository
	client    *http.Client
	ctx       context.Context
	userAgent string
}

func newLister(repo name.Repository, options ...Option) (*lister, error) {
	l := &lister{
		auth:      authn.Anonymous,
		transport: http.DefaultTransport,
		repo:      repo,
		ctx:       context.Background(),
	}

	for _, option := range options {
		if err := option(l); err != nil {
			return nil, err
		}
	}

	// Wrap the transport in something that logs requests and responses.
	// It's expensive to generate the dumps, so skip it if we're writing
	// to nothing.
	if logs.Enabled(logs.Debug) {
		l.transport = transport.NewLogger(l.transport)
	}

	// Wrap the transport in something that can retry network flakes.
	l.transport = transport.NewRetry(l.transport)

	// Wrap this last to prevent transport.New from double-wrapping.
	if l.userAgent != "" {
		l.transport = transport.NewUserAgent(l.transport, l.userAgent)
	}

	scopes := []string{repo.Scope(transport.PullScope)}
	tr, err := transport.NewWithContext(l.ctx, repo.Registry, l.auth, l.transport, scopes)
	if err != nil {
		return nil, err
	}

	l.client = &http.Client{Transport: tr}

	return l, nil
}

func (l *lister) list(repo name.Repository) (*Tags, error) {
	uri := &url.URL{
		Scheme: repo.Registry.Scheme(),
		Host:   repo.Registry.RegistryStr(),
		Path:   fmt.Sprintf("/v2/%s/tags/list", repo.RepositoryStr()),
		// ECR returns an error if n > 1000:
		// https://github.com/google/go-containerregistry/issues/681
		RawQuery: "n=1000",
	}

	tags := Tags{}

	// get responses until there is no next page
	for {
		select {
		case <-l.ctx.Done():
			return nil, l.ctx.Err()
		default:
		}

		req, err := http.NewRequest("GET", uri.String(), nil)
		if err != nil {
			return nil, err
		}
		req = req.WithContext(l.ctx)

		resp, err := l.client.Do(req)
		if err != nil {
			return nil, err
		}

		if err := transport.CheckError(resp, http.StatusOK); err != nil {
			return nil, err
		}

		parsed := Tags{}
		if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
			return nil, err
		}

		if err := resp.Body.Close(); err != nil {
			return nil, err
		}

		if len(parsed.Manifests) != 0 {
			// We're dealing with GCR, just return directly.
			return &parsed, nil
		}

		// This isn't GCR, just append the tags and keep paginating.
		tags.Tags = append(tags.Tags, parsed.Tags...)

		uri, err = getNextPageURL(resp)
		if err != nil {
			return nil, err
		}
		// no next page
		if uri == nil {
			break
		}
		logs.Warn.Printf("saw non-google tag listing response, falling back to pagination")
	}

	return &tags, nil
}

// getNextPageURL checks if there is a Link header in a http.Response which
// contains a link to the next page. If yes it returns the url.URL of the next
// page otherwise it returns nil.
func getNextPageURL(resp *http.Response) (*url.URL, error) {
	link := resp.Header.Get("Link")
	if link == "" {
		return nil, nil
	}

	if link[0] != '<' {
		return nil, fmt.Errorf("failed to parse link header: missing '<' in: %s", link)
	}

	end := strings.Index(link, ">")
	if end == -1 {
		return nil, fmt.Errorf("failed to parse link header: missing '>' in: %s", link)
	}
	link = link[1:end]

	linkURL, err := url.Parse(link)
	if err != nil {
		return nil, err
	}
	if resp.Request == nil || resp.Request.URL == nil {
		return nil, nil
	}
	linkURL = resp.Request.URL.ResolveReference(linkURL)
	return linkURL, nil
}

type rawManifestInfo struct {
	Size      string   `json:"imageSizeBytes"`
	MediaType string   `json:"mediaType"`
	Created   string   `json:"timeCreatedMs"`
	Uploaded  string   `json:"timeUploadedMs"`
	Tags      []string `json:"tag"`
}

// ManifestInfo is a Manifests entry is the output of List and Walk.
type ManifestInfo struct {
	Size      uint64    `json:"imageSizeBytes"`
	MediaType string    `json:"mediaType"`
	Created   time.Time `json:"timeCreatedMs"`
	Uploaded  time.Time `json:"timeUploadedMs"`
	Tags      []string  `json:"tag"`
}

func fromUnixMs(ms int64) time.Time {
	sec := ms / 1000
	ns := (ms % 1000) * 1000000
	return time.Unix(sec, ns)
}

func toUnixMs(t time.Time) string {
	return strconv.FormatInt(t.UnixNano()/1000000, 10)
}

// MarshalJSON implements json.Marshaler
func (m ManifestInfo) MarshalJSON() ([]byte, error) {
	return json.Marshal(rawManifestInfo{
		Size:      strconv.FormatUint(m.Size, 10),
		MediaType: m.MediaType,
		Created:   toUnixMs(m.Created),
		Uploaded:  toUnixMs(m.Uploaded),
		Tags:      m.Tags,
	})
}

// UnmarshalJSON implements json.Unmarshaler
func (m *ManifestInfo) UnmarshalJSON(data []byte) error {
	raw := rawManifestInfo{}
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}

	if raw.Size != "" {
		size, err := strconv.ParseUint(string(raw.Size), 10, 64)
		if err != nil {
			return err
		}
		m.Size = size
	}

	if raw.Created != "" {
		created, err := strconv.ParseInt(string(raw.Created), 10, 64)
		if err != nil {
			return err
		}
		m.Created = fromUnixMs(created)
	}

	if raw.Uploaded != "" {
		uploaded, err := strconv.ParseInt(string(raw.Uploaded), 10, 64)
		if err != nil {
			return err
		}
		m.Uploaded = fromUnixMs(uploaded)
	}

	m.MediaType = raw.MediaType
	m.Tags = raw.Tags

	return nil
}

// Tags is the result of List and Walk.
type Tags struct {
	Children  []string                `json:"child"`
	Manifests map[string]ManifestInfo `json:"manifest"`
	Name      string                  `json:"name"`
	Tags      []string                `json:"tags"`
}

// List calls /tags/list for the given repository.
func List(repo name.Repository, options ...Option) (*Tags, error) {
	l, err := newLister(repo, options...)
	if err != nil {
		return nil, err
	}

	return l.list(repo)
}

// WalkFunc is the type of the function called for each repository visited by
// Walk. This implements a similar API to filepath.Walk.
//
// The repo argument contains the argument to Walk as a prefix; that is, if Walk
// is called with "gcr.io/foo", which is a repository containing the repository
// "bar", the walk function will be called with argument "gcr.io/foo/bar".
// The tags and error arguments are the result of calling List on repo.
//
// TODO: Do we want a SkipDir error, as in filepath.WalkFunc?
type WalkFunc func(repo name.Repository, tags *Tags, err error) error

func walk(repo name.Repository, tags *Tags, walkFn WalkFunc, options ...Option) error {
	if tags == nil {
		// This shouldn't happen.
		return fmt.Errorf("tags nil for %q", repo)
	}

	if err := walkFn(repo, tags, nil); err != nil {
		return err
	}

	for _, path := range tags.Children {
		child, err := name.NewRepository(fmt.Sprintf("%s/%s", repo, path), name.StrictValidation)
		if err != nil {
			// We don't expect this ever, so don't pass it through to walkFn.
			return fmt.Errorf("unexpected path failure: %v", err)
		}

		childTags, err := List(child, options...)
		if err != nil {
			if err := walkFn(child, nil, err); err != nil {
				return err
			}
		} else {
			if err := walk(child, childTags, walkFn, options...); err != nil {
				return err
			}
		}
	}

	// We made it!
	return nil
}

// Walk recursively descends repositories, calling walkFn.
func Walk(root name.Repository, walkFn WalkFunc, options ...Option) error {
	tags, err := List(root, options...)
	if err != nil {
		return walkFn(root, nil, err)
	}

	return walk(root, tags, walkFn, options...)
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/google/trillian"
	"github.com/google/trillian/crypto"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/storage/testonly"
)

// rehashTest encapsulates one test case for the rehasher in isolation. Input data like the storage
// hashes and revisions can be arbitrary but the nodes should have distinct values
type rehashTest struct {
	desc    string
	index   int64
	nodes   []storage.Node
	fetches []merkle.NodeFetch
	output  trillian.Proof
}

// An arbitrary tree revision to be used in tests
const testTreeRevision int64 = 3

// Raw hashes for dummy storage nodes
var h1 = th.HashLeaf([]byte("Hash 1"))
var h2 = th.HashLeaf([]byte("Hash 2"))
var h3 = th.HashLeaf([]byte("Hash 3"))
var h4 = th.HashLeaf([]byte("Hash 4"))
var h5 = th.HashLeaf([]byte("Hash 5"))

// And the dummy nodes themselves
var sn1 = storage.Node{NodeID: storage.NewNodeIDFromHash(h1), Hash: h1, NodeRevision: 11}
var sn2 = storage.Node{NodeID: storage.NewNodeIDFromHash(h2), Hash: h2, NodeRevision: 22}
var sn3 = storage.Node{NodeID: storage.NewNodeIDFromHash(h3), Hash: h3, NodeRevision: 33}
var sn4 = storage.Node{NodeID: storage.NewNodeIDFromHash(h4), Hash: h4, NodeRevision: 44}
var sn5 = storage.Node{NodeID: storage.NewNodeIDFromHash(h5), Hash: h5, NodeRevision: 55}

// And the output proof nodes expected for them if they are passed through without rehashing
var n1 = &trillian.Node{NodeHash: h1, NodeId: mustMarshalNodeID(sn1.NodeID), NodeRevision: sn1.NodeRevision}
var n2 = &trillian.Node{NodeHash: h2, NodeId: mustMarshalNodeID(sn2.NodeID), NodeRevision: sn2.NodeRevision}
var n3 = &trillian.Node{NodeHash: h3, NodeId: mustMarshalNodeID(sn3.NodeID), NodeRevision: sn3.NodeRevision}
var n4 = &trillian.Node{NodeHash: h4, NodeId: mustMarshalNodeID(sn4.NodeID), NodeRevision: sn4.NodeRevision}
var n5 = &trillian.Node{NodeHash: h5, NodeId: mustMarshalNodeID(sn5.NodeID), NodeRevision: sn5.NodeRevision}

// Nodes containing composite hashes. They don't have node ids or revisions as they're recomputed
var n1n2 = &trillian.Node{NodeHash: th.HashChildren(h2, h1)}
var n2n3 = &trillian.Node{NodeHash: th.HashChildren(h3, h2)}
var n2n3n4 = &trillian.Node{NodeHash: th.HashChildren(h4, th.HashChildren(h3, h2))}
var n4n5 = &trillian.Node{NodeHash: th.HashChildren(h5, h4)}

func TestRehasher(t *testing.T) {
	var rehashTests = []rehashTest{
		{
			desc:    "no rehash",
			index:   126,
			nodes:   []storage.Node{sn1, sn2, sn3},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: false}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 126,
				ProofNode: []*trillian.Node{n1, n2, n3},
			},
		},
		{
			desc:    "single rehash",
			index:   999,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}, {Rehash: false}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 999,
				ProofNode: []*trillian.Node{n1, n2n3, n4, n5},
			},
		},
		{
			desc:    "single rehash at end",
			index:   11,
			nodes:   []storage.Node{sn1, sn2, sn3},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}},
			output: trillian.Proof{
				LeafIndex: 11,
				ProofNode: []*trillian.Node{n1, n2n3},
			},
		},
		{
			desc:    "single rehash multiple nodes",
			index:   23,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}, {Rehash: true}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 23,
				ProofNode: []*trillian.Node{n1, n2n3n4, n5},
			},
		},
		{
			desc:    "multiple rehash",
			index:   45,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: true}, {Rehash: true}, {Rehash: false}, {Rehash: true}, {Rehash: true}},
			output: trillian.Proof{
				LeafIndex: 45,
				ProofNode: []*trillian.Node{n1n2, n3, n4n5},
			},
		},
	}

	for _, rehashTest := range rehashTests {
		r := newRehasher()
		for i, node := range rehashTest.nodes {
			r.process(node, rehashTest.fetches[i])
		}

		want := rehashTest.output
		got, err := r.rehashedProof(rehashTest.index)

		if err != nil {
			t.Fatalf("rehash test %s unexpected error: %v", rehashTest.desc, err)
		}

		if !proto.Equal(&got, &want) {
			t.Errorf("rehash test %s:\ngot: %v\nwant: %v", rehashTest.desc, got, want)
		}
	}
}

func TestTree32InclusionProofFetchAll(t *testing.T) {
	for ts := 2; ts <= 32; ts++ {
		mt := treeAtSize(ts)
		r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
			{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, ts-1), ExpectedRoot: expectedRootAtSize(mt)},
		})

		for s := int64(2); s <= int64(ts); s++ {
			for l := int64(0); l < s; l++ {
				fetches, err := merkle.CalcInclusionProofNodeAddresses(s, l, int64(ts), 64)
				if err != nil {
					t.Fatal(err)
				}

				proof, err := fetchNodesAndBuildProof(r, testTreeRevision, int64(l), fetches)
				if err != nil {
					t.Fatal(err)
				}

				// We use +1 here because of the 1 based leaf indexing of this implementation
				refProof := mt.PathToRootAtSnapshot(l+1, s)

				if got, want := len(proof.ProofNode), len(refProof); got != want {
					t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", ts, s, l, got, want, fetches, refProof)
				}

				for i := 0; i < len(proof.ProofNode); i++ {
					if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
						t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", ts, s, l, i, got, want, len(proof.ProofNode), fetches)
					}
				}
			}
		}
	}
}

func TestTree32InclusionProofFetchMultiBatch(t *testing.T) {
	mt := treeAtSize(32)
	// The reader is built up with multiple batches, 4 batches x 8 leaves each
	r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
		{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, 7), ExpectedRoot: expectedRootAtSize(treeAtSize(8))},
		{TreeRevision: testTreeRevision + 1, Leaves: expandLeaves(8, 15), ExpectedRoot: expectedRootAtSize(treeAtSize(16))},
		{TreeRevision: testTreeRevision + 2, Leaves: expandLeaves(16, 23), ExpectedRoot: expectedRootAtSize(treeAtSize(24))},
		{TreeRevision: testTreeRevision + 3, Leaves: expandLeaves(24, 31), ExpectedRoot: expectedRootAtSize(mt)},
	})

	for s := int64(2); s <= 32; s++ {
		for l := int64(0); l < s; l++ {
			fetches, err := merkle.CalcInclusionProofNodeAddresses(s, l, 32, 64)
			if err != nil {
				t.Fatal(err)
			}

			// Use the highest tree revision that should be available from the node reader
			proof, err := fetchNodesAndBuildProof(r, testTreeRevision+3, l, fetches)
			if err != nil {
				t.Fatal(err)
			}

			// We use +1 here because of the 1 based leaf indexing of this implementation
			refProof := mt.PathToRootAtSnapshot(l+1, s)

			if got, want := len(proof.ProofNode), len(refProof); got != want {
				t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", 32, s, l, got, want, fetches, refProof)
			}

			for i := 0; i < len(proof.ProofNode); i++ {
				if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
					t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", 32, s, l, i, got, want, len(proof.ProofNode), fetches)
				}
			}
		}
	}
}

func TestTree32ConsistencyProofFetchAll(t *testing.T) {
	for ts := 2; ts <= 32; ts++ {
		mt := treeAtSize(ts)
		r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
			{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, ts-1), ExpectedRoot: expectedRootAtSize(mt)},
		})

		for s1 := int64(2); s1 < int64(ts); s1++ {
			for s2 := int64(s1 + 1); s2 < int64(ts); s2++ {
				fetches, err := merkle.CalcConsistencyProofNodeAddresses(s1, s2, int64(ts), 64)
				if err != nil {
					t.Fatal(err)
				}

				proof, err := fetchNodesAndBuildProof(r, testTreeRevision, int64(s1), fetches)
				if err != nil {
					t.Fatal(err)
				}

				refProof := mt.SnapshotConsistency(s1, s2)

				if got, want := len(proof.ProofNode), len(refProof); got != want {
					t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", ts, s1, s2, got, want, fetches, refProof)
				}

				for i := 0; i < len(proof.ProofNode); i++ {
					if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
						t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", ts, s1, s2, i, got, want, len(proof.ProofNode), fetches)
					}
				}
			}
		}
	}
}

func mustMarshalNodeID(nodeID storage.NodeID) []byte {
	idBytes, err := proto.Marshal(nodeID.AsProto())
	if err != nil {
		panic(err)
	}
	return idBytes
}

func expandLeaves(n, m int) []string {
	leaves := make([]string, 0, m-n+1)
	for l := n; l <= m; l++ {
		leaves = append(leaves, fmt.Sprintf("Leaf %d", l))
	}
	return leaves
}

// expectedRootAtSize uses the in memory tree, the tree built with Compact Merkle Tree should
// have the same root.
func expectedRootAtSize(mt *merkle.InMemoryMerkleTree) string {
	return hex.EncodeToString(mt.CurrentRoot().Hash())
}

func treeAtSize(n int) *merkle.InMemoryMerkleTree {
	leaves := expandLeaves(0, n-1)
	mt := merkle.NewInMemoryMerkleTree(merkle.NewRFC6962TreeHasher(crypto.NewSHA256()))
	for _, leaf := range leaves {
		mt.AddLeaf([]byte(leaf))
	}
	return mt
}
<commit_msg>Add a larger test of building proof paths to the fetcher test (#355)<commit_after>package server

import (
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/golang/protobuf/proto"
	"github.com/google/trillian"
	"github.com/google/trillian/crypto"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/storage/testonly"
)

// rehashTest encapsulates one test case for the rehasher in isolation. Input data like the storage
// hashes and revisions can be arbitrary but the nodes should have distinct values
type rehashTest struct {
	desc    string
	index   int64
	nodes   []storage.Node
	fetches []merkle.NodeFetch
	output  trillian.Proof
}

// An arbitrary tree revision to be used in tests
const testTreeRevision int64 = 3

// Raw hashes for dummy storage nodes
var h1 = th.HashLeaf([]byte("Hash 1"))
var h2 = th.HashLeaf([]byte("Hash 2"))
var h3 = th.HashLeaf([]byte("Hash 3"))
var h4 = th.HashLeaf([]byte("Hash 4"))
var h5 = th.HashLeaf([]byte("Hash 5"))

// And the dummy nodes themselves
var sn1 = storage.Node{NodeID: storage.NewNodeIDFromHash(h1), Hash: h1, NodeRevision: 11}
var sn2 = storage.Node{NodeID: storage.NewNodeIDFromHash(h2), Hash: h2, NodeRevision: 22}
var sn3 = storage.Node{NodeID: storage.NewNodeIDFromHash(h3), Hash: h3, NodeRevision: 33}
var sn4 = storage.Node{NodeID: storage.NewNodeIDFromHash(h4), Hash: h4, NodeRevision: 44}
var sn5 = storage.Node{NodeID: storage.NewNodeIDFromHash(h5), Hash: h5, NodeRevision: 55}

// And the output proof nodes expected for them if they are passed through without rehashing
var n1 = &trillian.Node{NodeHash: h1, NodeId: mustMarshalNodeID(sn1.NodeID), NodeRevision: sn1.NodeRevision}
var n2 = &trillian.Node{NodeHash: h2, NodeId: mustMarshalNodeID(sn2.NodeID), NodeRevision: sn2.NodeRevision}
var n3 = &trillian.Node{NodeHash: h3, NodeId: mustMarshalNodeID(sn3.NodeID), NodeRevision: sn3.NodeRevision}
var n4 = &trillian.Node{NodeHash: h4, NodeId: mustMarshalNodeID(sn4.NodeID), NodeRevision: sn4.NodeRevision}
var n5 = &trillian.Node{NodeHash: h5, NodeId: mustMarshalNodeID(sn5.NodeID), NodeRevision: sn5.NodeRevision}

// Nodes containing composite hashes. They don't have node ids or revisions as they're recomputed
var n1n2 = &trillian.Node{NodeHash: th.HashChildren(h2, h1)}
var n2n3 = &trillian.Node{NodeHash: th.HashChildren(h3, h2)}
var n2n3n4 = &trillian.Node{NodeHash: th.HashChildren(h4, th.HashChildren(h3, h2))}
var n4n5 = &trillian.Node{NodeHash: th.HashChildren(h5, h4)}

func TestRehasher(t *testing.T) {
	var rehashTests = []rehashTest{
		{
			desc:    "no rehash",
			index:   126,
			nodes:   []storage.Node{sn1, sn2, sn3},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: false}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 126,
				ProofNode: []*trillian.Node{n1, n2, n3},
			},
		},
		{
			desc:    "single rehash",
			index:   999,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}, {Rehash: false}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 999,
				ProofNode: []*trillian.Node{n1, n2n3, n4, n5},
			},
		},
		{
			desc:    "single rehash at end",
			index:   11,
			nodes:   []storage.Node{sn1, sn2, sn3},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}},
			output: trillian.Proof{
				LeafIndex: 11,
				ProofNode: []*trillian.Node{n1, n2n3},
			},
		},
		{
			desc:    "single rehash multiple nodes",
			index:   23,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: false}, {Rehash: true}, {Rehash: true}, {Rehash: true}, {Rehash: false}},
			output: trillian.Proof{
				LeafIndex: 23,
				ProofNode: []*trillian.Node{n1, n2n3n4, n5},
			},
		},
		{
			desc:    "multiple rehash",
			index:   45,
			nodes:   []storage.Node{sn1, sn2, sn3, sn4, sn5},
			fetches: []merkle.NodeFetch{{Rehash: true}, {Rehash: true}, {Rehash: false}, {Rehash: true}, {Rehash: true}},
			output: trillian.Proof{
				LeafIndex: 45,
				ProofNode: []*trillian.Node{n1n2, n3, n4n5},
			},
		},
	}

	for _, rehashTest := range rehashTests {
		r := newRehasher()
		for i, node := range rehashTest.nodes {
			r.process(node, rehashTest.fetches[i])
		}

		want := rehashTest.output
		got, err := r.rehashedProof(rehashTest.index)

		if err != nil {
			t.Fatalf("rehash test %s unexpected error: %v", rehashTest.desc, err)
		}

		if !proto.Equal(&got, &want) {
			t.Errorf("rehash test %s:\ngot: %v\nwant: %v", rehashTest.desc, got, want)
		}
	}
}

func TestTree813FetchAll(t *testing.T) {
	const ts int64 = 813

	mt := treeAtSize(int(ts))
	r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
		{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, int(ts - 1)), ExpectedRoot: expectedRootAtSize(mt)},
	})

	for l := int64(271); l < ts; l++ {
		fetches, err := merkle.CalcInclusionProofNodeAddresses(ts, l, ts, 64)

		if err != nil {
			t.Fatal(err)
		}

		proof, err := fetchNodesAndBuildProof(r, testTreeRevision, int64(l), fetches)
		if err != nil {
			t.Fatal(err)
		}

		// We use +1 here because of the 1 based leaf indexing of this implementation
		refProof := mt.PathToRootAtSnapshot(l + 1, ts)

		if got, want := len(proof.ProofNode), len(refProof); got != want {
			for i, f := range fetches {
				t.Errorf("Fetch: %d => %s", i, f.NodeID.CoordString())
			}
			t.Fatalf("(%d, %d): got proof len: %d, want: %d: %v\n%v", ts, l, got, want, fetches, refProof)
		}

		for i := 0; i < len(proof.ProofNode); i++ {
			if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
				t.Fatalf("(%d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", ts, l, i, got, want, len(proof.ProofNode), fetches)
			}
		}
	}
}

func TestTree32InclusionProofFetchAll(t *testing.T) {
	for ts := 2; ts <= 32; ts++ {
		mt := treeAtSize(ts)
		r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
			{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, ts-1), ExpectedRoot: expectedRootAtSize(mt)},
		})

		for s := int64(2); s <= int64(ts); s++ {
			for l := int64(0); l < s; l++ {
				fetches, err := merkle.CalcInclusionProofNodeAddresses(s, l, int64(ts), 64)
				if err != nil {
					t.Fatal(err)
				}

				proof, err := fetchNodesAndBuildProof(r, testTreeRevision, int64(l), fetches)
				if err != nil {
					t.Fatal(err)
				}

				// We use +1 here because of the 1 based leaf indexing of this implementation
				refProof := mt.PathToRootAtSnapshot(l+1, s)

				if got, want := len(proof.ProofNode), len(refProof); got != want {
					t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", ts, s, l, got, want, fetches, refProof)
				}

				for i := 0; i < len(proof.ProofNode); i++ {
					if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
						t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", ts, s, l, i, got, want, len(proof.ProofNode), fetches)
					}
				}
			}
		}
	}
}

func TestTree32InclusionProofFetchMultiBatch(t *testing.T) {
	mt := treeAtSize(32)
	// The reader is built up with multiple batches, 4 batches x 8 leaves each
	r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
		{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, 7), ExpectedRoot: expectedRootAtSize(treeAtSize(8))},
		{TreeRevision: testTreeRevision + 1, Leaves: expandLeaves(8, 15), ExpectedRoot: expectedRootAtSize(treeAtSize(16))},
		{TreeRevision: testTreeRevision + 2, Leaves: expandLeaves(16, 23), ExpectedRoot: expectedRootAtSize(treeAtSize(24))},
		{TreeRevision: testTreeRevision + 3, Leaves: expandLeaves(24, 31), ExpectedRoot: expectedRootAtSize(mt)},
	})

	for s := int64(2); s <= 32; s++ {
		for l := int64(0); l < s; l++ {
			fetches, err := merkle.CalcInclusionProofNodeAddresses(s, l, 32, 64)
			if err != nil {
				t.Fatal(err)
			}

			// Use the highest tree revision that should be available from the node reader
			proof, err := fetchNodesAndBuildProof(r, testTreeRevision+3, l, fetches)
			if err != nil {
				t.Fatal(err)
			}

			// We use +1 here because of the 1 based leaf indexing of this implementation
			refProof := mt.PathToRootAtSnapshot(l+1, s)

			if got, want := len(proof.ProofNode), len(refProof); got != want {
				t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", 32, s, l, got, want, fetches, refProof)
			}

			for i := 0; i < len(proof.ProofNode); i++ {
				if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
					t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", 32, s, l, i, got, want, len(proof.ProofNode), fetches)
				}
			}
		}
	}
}

func TestTree32ConsistencyProofFetchAll(t *testing.T) {
	for ts := 2; ts <= 32; ts++ {
		mt := treeAtSize(ts)
		r := testonly.NewMultiFakeNodeReaderFromLeaves([]testonly.LeafBatch{
			{TreeRevision: testTreeRevision, Leaves: expandLeaves(0, ts-1), ExpectedRoot: expectedRootAtSize(mt)},
		})

		for s1 := int64(2); s1 < int64(ts); s1++ {
			for s2 := int64(s1 + 1); s2 < int64(ts); s2++ {
				fetches, err := merkle.CalcConsistencyProofNodeAddresses(s1, s2, int64(ts), 64)
				if err != nil {
					t.Fatal(err)
				}

				proof, err := fetchNodesAndBuildProof(r, testTreeRevision, int64(s1), fetches)
				if err != nil {
					t.Fatal(err)
				}

				refProof := mt.SnapshotConsistency(s1, s2)

				if got, want := len(proof.ProofNode), len(refProof); got != want {
					t.Fatalf("(%d, %d, %d): got proof len: %d, want: %d: %v\n%v", ts, s1, s2, got, want, fetches, refProof)
				}

				for i := 0; i < len(proof.ProofNode); i++ {
					if got, want := hex.EncodeToString(proof.ProofNode[i].NodeHash), hex.EncodeToString(refProof[i].Value.Hash()); got != want {
						t.Fatalf("(%d, %d, %d): %d got proof node: %s, want: %s l:%d fetches: %v", ts, s1, s2, i, got, want, len(proof.ProofNode), fetches)
					}
				}
			}
		}
	}
}

func mustMarshalNodeID(nodeID storage.NodeID) []byte {
	idBytes, err := proto.Marshal(nodeID.AsProto())
	if err != nil {
		panic(err)
	}
	return idBytes
}

func expandLeaves(n, m int) []string {
	leaves := make([]string, 0, m-n+1)
	for l := n; l <= m; l++ {
		leaves = append(leaves, fmt.Sprintf("Leaf %d", l))
	}
	return leaves
}

// expectedRootAtSize uses the in memory tree, the tree built with Compact Merkle Tree should
// have the same root.
func expectedRootAtSize(mt *merkle.InMemoryMerkleTree) string {
	return hex.EncodeToString(mt.CurrentRoot().Hash())
}

func treeAtSize(n int) *merkle.InMemoryMerkleTree {
	leaves := expandLeaves(0, n-1)
	mt := merkle.NewInMemoryMerkleTree(merkle.NewRFC6962TreeHasher(crypto.NewSHA256()))
	for _, leaf := range leaves {
		mt.AddLeaf([]byte(leaf))
	}
	return mt
}
<|endoftext|>


<commit_before>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package schedule

import (
	"math/rand"
	"time"

	"github.com/pingcap/pd/server/cache"
	"github.com/pingcap/pd/server/core"
)

// FlowKind is a identify Flow types.
type FlowKind uint32

// Flags for flow.
const (
	WriteFlow FlowKind = iota
	ReadFlow
)

// HotSpotCache is a cache hold hot regions.
type HotSpotCache struct {
	writeFlow cache.Cache
	readFlow  cache.Cache
}

func newHotSpotCache() *HotSpotCache {
	return &HotSpotCache{
		writeFlow: cache.NewCache(statCacheMaxLen, cache.TwoQueueCache),
		readFlow:  cache.NewCache(statCacheMaxLen, cache.TwoQueueCache),
	}
}

// CheckWrite checks the write status, returns whether need update statistics and item.
func (w *HotSpotCache) CheckWrite(region *core.RegionInfo, stores *core.StoresInfo) (bool, *core.RegionStat) {
	var (
		WrittenBytesPerSec uint64
		value              *core.RegionStat
	)
	v, isExist := w.writeFlow.Peek(region.GetId())
	if isExist && !Simulating {
		value = v.(*core.RegionStat)
		interval := time.Since(value.LastUpdateTime).Seconds()
		if interval < minHotRegionReportInterval {
			return false, nil
		}
		WrittenBytesPerSec = uint64(float64(region.WrittenBytes) / interval)
	} else {
		WrittenBytesPerSec = uint64(float64(region.WrittenBytes) / float64(RegionHeartBeatReportInterval))
	}

	hotRegionThreshold := calculateWriteHotThreshold(stores)
	return w.isNeedUpdateStatCache(region, WrittenBytesPerSec, hotRegionThreshold, value, WriteFlow)
}

// CheckRead checks the read status, returns whether need update statistics and item.
func (w *HotSpotCache) CheckRead(region *core.RegionInfo, stores *core.StoresInfo) (bool, *core.RegionStat) {
	var (
		ReadBytesPerSec uint64
		value           *core.RegionStat
	)
	v, isExist := w.readFlow.Peek(region.GetId())
	if isExist && !Simulating {
		value = v.(*core.RegionStat)
		interval := time.Since(value.LastUpdateTime).Seconds()
		if interval < minHotRegionReportInterval {
			return false, nil
		}
		ReadBytesPerSec = uint64(float64(region.ReadBytes) / interval)
	} else {
		ReadBytesPerSec = uint64(float64(region.ReadBytes) / float64(RegionHeartBeatReportInterval))
	}
	hotRegionThreshold := calculateReadHotThreshold(stores)
	return w.isNeedUpdateStatCache(region, ReadBytesPerSec, hotRegionThreshold, value, ReadFlow)
}

func (w *HotSpotCache) incMetrics(name string, kind FlowKind) {
	switch kind {
	case WriteFlow:
		hotCacheStatusGauge.WithLabelValues(name, "write").Inc()
	case ReadFlow:
		hotCacheStatusGauge.WithLabelValues(name, "read").Inc()
	}
}

func calculateWriteHotThreshold(stores *core.StoresInfo) uint64 {
	// hotRegionThreshold is use to pick hot region
	// suppose the number of the hot Regions is statCacheMaxLen
	// and we use total written Bytes past storeHeartBeatReportInterval seconds to divide the number of hot Regions
	// divide 2 because the store reports data about two times than the region record write to rocksdb
	divisor := float64(statCacheMaxLen) * 2
	hotRegionThreshold := uint64(float64(stores.TotalBytesWriteRate()) / divisor)

	if hotRegionThreshold < hotWriteRegionMinFlowRate {
		hotRegionThreshold = hotWriteRegionMinFlowRate
	}
	return hotRegionThreshold
}

func calculateReadHotThreshold(stores *core.StoresInfo) uint64 {
	// hotRegionThreshold is use to pick hot region
	// suppose the number of the hot Regions is statLRUMaxLen
	// and we use total Read Bytes past storeHeartBeatReportInterval seconds to divide the number of hot Regions
	divisor := float64(statCacheMaxLen)
	hotRegionThreshold := uint64(float64(stores.TotalBytesReadRate()) / divisor)

	if hotRegionThreshold < hotReadRegionMinFlowRate {
		hotRegionThreshold = hotReadRegionMinFlowRate
	}
	return hotRegionThreshold
}

const rollingWindowsSize = 5

func (w *HotSpotCache) isNeedUpdateStatCache(region *core.RegionInfo, flowBytes uint64, hotRegionThreshold uint64, oldItem *core.RegionStat, kind FlowKind) (bool, *core.RegionStat) {
	newItem := &core.RegionStat{
		RegionID:       region.GetId(),
		FlowBytes:      flowBytes,
		LastUpdateTime: time.Now(),
		StoreID:        region.Leader.GetStoreId(),
		Version:        region.GetRegionEpoch().GetVersion(),
		AntiCount:      hotRegionAntiCount,
	}

	if oldItem != nil {
		newItem.HotDegree = oldItem.HotDegree + 1
		newItem.Stats = oldItem.Stats
	}
	if flowBytes >= hotRegionThreshold {
		if oldItem == nil {
			w.incMetrics("add_item", kind)
			newItem.Stats = core.NewRollingStats(rollingWindowsSize)
		}
		newItem.Stats.Add(float64(flowBytes))
		return true, newItem
	}
	// smaller than hotReionThreshold
	if oldItem == nil {
		return false, newItem
	}
	if oldItem.AntiCount <= 0 {
		w.incMetrics("remove_item", kind)
		return true, nil
	}
	// eliminate some noise
	newItem.HotDegree = oldItem.HotDegree - 1
	newItem.AntiCount = oldItem.AntiCount - 1
	newItem.Stats.Add(float64(flowBytes))
	return true, newItem
}

// Update updates the cache.
func (w *HotSpotCache) Update(key uint64, item *core.RegionStat, kind FlowKind) {
	switch kind {
	case WriteFlow:
		if item == nil {
			w.writeFlow.Remove(key)
		} else {
			w.writeFlow.Put(key, item)
			w.incMetrics("update_item", kind)
		}
	case ReadFlow:
		if item == nil {
			w.readFlow.Remove(key)
		} else {
			w.readFlow.Put(key, item)
			w.incMetrics("update_item", kind)
		}
	}
}

// RegionStats returns hot items according to kind
func (w *HotSpotCache) RegionStats(kind FlowKind) []*core.RegionStat {
	var elements []*cache.Item
	switch kind {
	case WriteFlow:
		elements = w.writeFlow.Elems()
	case ReadFlow:
		elements = w.readFlow.Elems()
	}
	stats := make([]*core.RegionStat, len(elements))
	for i := range elements {
		stats[i] = elements[i].Value.(*core.RegionStat)
	}
	return stats
}

// RandHotRegionFromStore random picks a hot region in specify store.
func (w *HotSpotCache) RandHotRegionFromStore(storeID uint64, kind FlowKind, hotThreshold int) *core.RegionStat {
	stats := w.RegionStats(kind)
	for _, i := range rand.Perm(len(stats)) {
		if stats[i].HotDegree >= hotThreshold && stats[i].StoreID == storeID {
			return stats[i]
		}
	}
	return nil
}

// CollectMetrics collect the hot cache metrics
func (w *HotSpotCache) CollectMetrics(stores *core.StoresInfo) {
	hotCacheStatusGauge.WithLabelValues("total_length", "write").Set(float64(w.writeFlow.Len()))
	hotCacheStatusGauge.WithLabelValues("total_length", "read").Set(float64(w.readFlow.Len()))
	threshold := calculateWriteHotThreshold(stores)
	hotCacheStatusGauge.WithLabelValues("hotThreshold", "write").Set(float64(threshold))
	threshold = calculateReadHotThreshold(stores)
	hotCacheStatusGauge.WithLabelValues("hotThreshold", "read").Set(float64(threshold))
}

func (w *HotSpotCache) isRegionHot(id uint64, hotThreshold int) bool {
	if stat, ok := w.writeFlow.Peek(id); ok {
		if stat.(*core.RegionStat).HotDegree >= hotThreshold {
			return true
		}
	}
	if stat, ok := w.readFlow.Peek(id); ok {
		return stat.(*core.RegionStat).HotDegree >= hotThreshold
	}
	return false
}
<commit_msg>fix hot write and read (#1150)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package schedule

import (
	"math/rand"
	"time"

	"github.com/pingcap/pd/server/cache"
	"github.com/pingcap/pd/server/core"
)

// FlowKind is a identify Flow types.
type FlowKind uint32

// Flags for flow.
const (
	WriteFlow FlowKind = iota
	ReadFlow
)

// HotSpotCache is a cache hold hot regions.
type HotSpotCache struct {
	writeFlow cache.Cache
	readFlow  cache.Cache
}

func newHotSpotCache() *HotSpotCache {
	return &HotSpotCache{
		writeFlow: cache.NewCache(statCacheMaxLen, cache.TwoQueueCache),
		readFlow:  cache.NewCache(statCacheMaxLen, cache.TwoQueueCache),
	}
}

// CheckWrite checks the write status, returns whether need update statistics and item.
func (w *HotSpotCache) CheckWrite(region *core.RegionInfo, stores *core.StoresInfo) (bool, *core.RegionStat) {
	var (
		WrittenBytesPerSec uint64
		value              *core.RegionStat
	)

	WrittenBytesPerSec = uint64(float64(region.WrittenBytes) / float64(RegionHeartBeatReportInterval))

	v, isExist := w.writeFlow.Peek(region.GetId())
	if isExist {
		value = v.(*core.RegionStat)
		if !Simulating {
			interval := time.Since(value.LastUpdateTime).Seconds()
			if interval < minHotRegionReportInterval {
				return false, nil
			}
			WrittenBytesPerSec = uint64(float64(region.WrittenBytes) / interval)
		}
	}

	hotRegionThreshold := calculateWriteHotThreshold(stores)
	return w.isNeedUpdateStatCache(region, WrittenBytesPerSec, hotRegionThreshold, value, WriteFlow)
}

// CheckRead checks the read status, returns whether need update statistics and item.
func (w *HotSpotCache) CheckRead(region *core.RegionInfo, stores *core.StoresInfo) (bool, *core.RegionStat) {
	var (
		ReadBytesPerSec uint64
		value           *core.RegionStat
	)

	ReadBytesPerSec = uint64(float64(region.ReadBytes) / float64(RegionHeartBeatReportInterval))

	v, isExist := w.readFlow.Peek(region.GetId())
	if isExist {
		value = v.(*core.RegionStat)
		if !Simulating {
			interval := time.Since(value.LastUpdateTime).Seconds()
			if interval < minHotRegionReportInterval {
				return false, nil
			}
			ReadBytesPerSec = uint64(float64(region.ReadBytes) / interval)
		}
	}

	hotRegionThreshold := calculateReadHotThreshold(stores)
	return w.isNeedUpdateStatCache(region, ReadBytesPerSec, hotRegionThreshold, value, ReadFlow)
}

func (w *HotSpotCache) incMetrics(name string, kind FlowKind) {
	switch kind {
	case WriteFlow:
		hotCacheStatusGauge.WithLabelValues(name, "write").Inc()
	case ReadFlow:
		hotCacheStatusGauge.WithLabelValues(name, "read").Inc()
	}
}

func calculateWriteHotThreshold(stores *core.StoresInfo) uint64 {
	// hotRegionThreshold is use to pick hot region
	// suppose the number of the hot Regions is statCacheMaxLen
	// and we use total written Bytes past storeHeartBeatReportInterval seconds to divide the number of hot Regions
	// divide 2 because the store reports data about two times than the region record write to rocksdb
	divisor := float64(statCacheMaxLen) * 2
	hotRegionThreshold := uint64(float64(stores.TotalBytesWriteRate()) / divisor)

	if hotRegionThreshold < hotWriteRegionMinFlowRate {
		hotRegionThreshold = hotWriteRegionMinFlowRate
	}
	return hotRegionThreshold
}

func calculateReadHotThreshold(stores *core.StoresInfo) uint64 {
	// hotRegionThreshold is use to pick hot region
	// suppose the number of the hot Regions is statLRUMaxLen
	// and we use total Read Bytes past storeHeartBeatReportInterval seconds to divide the number of hot Regions
	divisor := float64(statCacheMaxLen)
	hotRegionThreshold := uint64(float64(stores.TotalBytesReadRate()) / divisor)

	if hotRegionThreshold < hotReadRegionMinFlowRate {
		hotRegionThreshold = hotReadRegionMinFlowRate
	}
	return hotRegionThreshold
}

const rollingWindowsSize = 5

func (w *HotSpotCache) isNeedUpdateStatCache(region *core.RegionInfo, flowBytes uint64, hotRegionThreshold uint64, oldItem *core.RegionStat, kind FlowKind) (bool, *core.RegionStat) {
	newItem := &core.RegionStat{
		RegionID:       region.GetId(),
		FlowBytes:      flowBytes,
		LastUpdateTime: time.Now(),
		StoreID:        region.Leader.GetStoreId(),
		Version:        region.GetRegionEpoch().GetVersion(),
		AntiCount:      hotRegionAntiCount,
	}

	if oldItem != nil {
		newItem.HotDegree = oldItem.HotDegree + 1
		newItem.Stats = oldItem.Stats
	}
	if flowBytes >= hotRegionThreshold {
		if oldItem == nil {
			w.incMetrics("add_item", kind)
			newItem.Stats = core.NewRollingStats(rollingWindowsSize)
		}
		newItem.Stats.Add(float64(flowBytes))
		return true, newItem
	}
	// smaller than hotReionThreshold
	if oldItem == nil {
		return false, newItem
	}
	if oldItem.AntiCount <= 0 {
		w.incMetrics("remove_item", kind)
		return true, nil
	}
	// eliminate some noise
	newItem.HotDegree = oldItem.HotDegree - 1
	newItem.AntiCount = oldItem.AntiCount - 1
	newItem.Stats.Add(float64(flowBytes))
	return true, newItem
}

// Update updates the cache.
func (w *HotSpotCache) Update(key uint64, item *core.RegionStat, kind FlowKind) {
	switch kind {
	case WriteFlow:
		if item == nil {
			w.writeFlow.Remove(key)
		} else {
			w.writeFlow.Put(key, item)
			w.incMetrics("update_item", kind)
		}
	case ReadFlow:
		if item == nil {
			w.readFlow.Remove(key)
		} else {
			w.readFlow.Put(key, item)
			w.incMetrics("update_item", kind)
		}
	}
}

// RegionStats returns hot items according to kind
func (w *HotSpotCache) RegionStats(kind FlowKind) []*core.RegionStat {
	var elements []*cache.Item
	switch kind {
	case WriteFlow:
		elements = w.writeFlow.Elems()
	case ReadFlow:
		elements = w.readFlow.Elems()
	}
	stats := make([]*core.RegionStat, len(elements))
	for i := range elements {
		stats[i] = elements[i].Value.(*core.RegionStat)
	}
	return stats
}

// RandHotRegionFromStore random picks a hot region in specify store.
func (w *HotSpotCache) RandHotRegionFromStore(storeID uint64, kind FlowKind, hotThreshold int) *core.RegionStat {
	stats := w.RegionStats(kind)
	for _, i := range rand.Perm(len(stats)) {
		if stats[i].HotDegree >= hotThreshold && stats[i].StoreID == storeID {
			return stats[i]
		}
	}
	return nil
}

// CollectMetrics collect the hot cache metrics
func (w *HotSpotCache) CollectMetrics(stores *core.StoresInfo) {
	hotCacheStatusGauge.WithLabelValues("total_length", "write").Set(float64(w.writeFlow.Len()))
	hotCacheStatusGauge.WithLabelValues("total_length", "read").Set(float64(w.readFlow.Len()))
	threshold := calculateWriteHotThreshold(stores)
	hotCacheStatusGauge.WithLabelValues("hotThreshold", "write").Set(float64(threshold))
	threshold = calculateReadHotThreshold(stores)
	hotCacheStatusGauge.WithLabelValues("hotThreshold", "read").Set(float64(threshold))
}

func (w *HotSpotCache) isRegionHot(id uint64, hotThreshold int) bool {
	if stat, ok := w.writeFlow.Peek(id); ok {
		if stat.(*core.RegionStat).HotDegree >= hotThreshold {
			return true
		}
	}
	if stat, ok := w.readFlow.Peek(id); ok {
		return stat.(*core.RegionStat).HotDegree >= hotThreshold
	}
	return false
}
<|endoftext|>


<commit_before>package calcium

import (
	"archive/tar"
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/docker/distribution/reference"
	enginetypes "github.com/docker/docker/api/types"
	enginecontainer "github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	engineapi "github.com/docker/docker/client"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/registry"
	"github.com/projecteru2/core/cluster"
	"github.com/projecteru2/core/lock"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

// Lock is lock for calcium
func (c *Calcium) Lock(ctx context.Context, name string, timeout int) (lock.DistributedLock, error) {
	lock, err := c.store.CreateLock(name, timeout)
	if err != nil {
		return nil, err
	}
	if err = lock.Lock(ctx); err != nil {
		return nil, err
	}
	return lock, nil
}

func makeResourceSetting(cpu float64, memory int64, cpuMap types.CPUMap, softlimit bool) enginecontainer.Resources {
	resource := enginecontainer.Resources{}
	if cpu > 0 {
		resource.CPUPeriod = cluster.CPUPeriodBase
		resource.CPUQuota = int64(cpu * float64(cluster.CPUPeriodBase))
	}
	if cpuMap != nil && len(cpuMap) > 0 {
		cpuIDs := []string{}
		for cpuID := range cpuMap {
			cpuIDs = append(cpuIDs, cpuID)
		}
		resource.CpusetCpus = strings.Join(cpuIDs, ",")
	}
	if softlimit {
		resource.MemoryReservation = memory
	} else {
		resource.Memory = memory
		resource.MemorySwap = memory
	}
	return resource
}

// image begin
// MakeAuthConfigFromRemote Calculate encoded AuthConfig from registry and eru-core config
// See https://github.com/docker/cli/blob/16cccc30f95c8163f0749eba5a2e80b807041342/cli/command/registry.go#L67
func makeEncodedAuthConfigFromRemote(authConfigs map[string]types.AuthConfig, remote string) (string, error) {
	ref, err := reference.ParseNormalizedNamed(remote)
	if err != nil {
		return "", err
	}

	// Resolve the Repository name from fqn to RepositoryInfo
	repoInfo, err := registry.ParseRepositoryInfo(ref)
	if err != nil {
		return "", err
	}

	serverAddress := repoInfo.Index.Name
	if authConfig, exists := authConfigs[serverAddress]; exists {
		if encodedAuth, err := encodeAuthToBase64(authConfig); err == nil {
			return encodedAuth, nil
		}
		return "", err
	}
	return "dummy", nil
}

// EncodeAuthToBase64 serializes the auth configuration as JSON base64 payload
// See https://github.com/docker/cli/blob/master/cli/command/registry.go#L41
func encodeAuthToBase64(authConfig types.AuthConfig) (string, error) {
	buf, err := json.Marshal(authConfig)
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(buf), nil
}

func makeCPUAndMem(nodes []*types.Node) map[string]types.CPUAndMem {
	r := make(map[string]types.CPUAndMem)
	for _, node := range nodes {
		r[node.Name] = types.CPUAndMem{
			CPUMap: node.CPU,
			MemCap: node.MemCap,
		}
	}
	return r
}

// As the name says,
// blocks until the stream is empty, until we meet EOF
func ensureReaderClosed(stream io.ReadCloser) {
	if stream == nil {
		return
	}
	io.Copy(ioutil.Discard, stream)
	stream.Close()
}

// make mount paths
// 使用volumes, 参数格式跟docker一样
// volumes:
//     - "/foo-data:$SOMEENV/foodata:rw"
func makeMountPaths(opts *types.DeployOptions) ([]string, map[string]struct{}) {
	binds := []string{}
	volumes := make(map[string]struct{})

	var expandENV = func(env string) string {
		envMap := map[string]string{}
		for _, env := range opts.Env {
			parts := strings.Split(env, "=")
			envMap[parts[0]] = parts[1]
		}
		return envMap[env]
	}

	for _, path := range opts.Volumes {
		expanded := os.Expand(path, expandENV)
		parts := strings.Split(expanded, ":")
		if len(parts) == 2 {
			binds = append(binds, fmt.Sprintf("%s:%s:rw", parts[0], parts[1]))
			volumes[parts[1]] = struct{}{}
		} else if len(parts) == 3 {
			binds = append(binds, fmt.Sprintf("%s:%s:%s", parts[0], parts[1], parts[2]))
			volumes[parts[1]] = struct{}{}
		}
	}

	return binds, volumes
}

// 跑存在labels里的exec
// 为什么要存labels呢, 因为下线容器的时候根本不知道entrypoint是啥
func execuateInside(ctx context.Context, client *engineapi.Client, ID, cmd, user string, env []string, privileged bool) ([]byte, error) {
	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := enginetypes.ExecConfig{
		User:         user,
		Cmd:          cmds,
		Privileged:   privileged,
		Env:          env,
		AttachStderr: true,
		AttachStdout: true,
	}
	//TODO should timeout
	//Fuck docker, ctx will not use inside funcs!!
	idResp, err := client.ContainerExecCreate(ctx, ID, execConfig)
	if err != nil {
		return []byte{}, err
	}
	resp, err := client.ContainerExecAttach(ctx, idResp.ID, enginetypes.ExecStartCheck{})
	if err != nil {
		return []byte{}, err
	}
	defer resp.Close()
	stream := utils.FuckDockerStream(ioutil.NopCloser(resp.Reader))
	b, err := ioutil.ReadAll(stream)
	if err != nil {
		return []byte{}, err
	}
	info, err := client.ContainerExecInspect(ctx, idResp.ID)
	if err != nil {
		return []byte{}, err
	}
	if info.ExitCode != 0 {
		return []byte{}, fmt.Errorf("%s", b)
	}
	return b, nil
}

// Pull an image
func pullImage(ctx context.Context, node *types.Node, image, auth string) error {
	log.Debugf("[pullImage] Pulling image %s", image)
	if image == "" {
		return fmt.Errorf("Goddamn empty image, WTF?")
	}

	// check local
	exists := false
	image = utils.NormalizeImageName(image)
	f := filters.NewArgs()
	f.Add("reference", image)
	listOptions := enginetypes.ImageListOptions{
		All:     false,
		Filters: f,
	}
	listResp, err := node.Engine.ImageList(ctx, listOptions)
	if err != nil {
		log.Errorf("[pullImage] Error during check image %s: %v", image, err)
	}

	if len(listResp) > 0 {
		log.Debugf("[pullImage] Image %s exists", image)
		exists = true
	}

	if exists {
		log.Infof("[pullImage] Image cached, skip pulling")
		return nil
	}

	pullOptions := enginetypes.ImagePullOptions{All: false, RegistryAuth: auth}
	outStream, err := node.Engine.ImagePull(ctx, image, pullOptions)
	if err != nil {
		log.Errorf("[pullImage] Error during pulling image %s: %v", image, err)
		return err
	}
	ensureReaderClosed(outStream)
	log.Debugf("[pullImage] Done pulling image %s", image)
	return nil
}

func makeErrorBuildImageMessage(err error) *types.BuildImageMessage {
	return &types.BuildImageMessage{Error: err.Error()}
}

func createTarStream(path string) (io.ReadCloser, error) {
	tarOpts := &archive.TarOptions{
		ExcludePatterns: []string{},
		IncludeFiles:    []string{"."},
		Compression:     archive.Uncompressed,
		NoLchown:        true,
	}
	return archive.TarWithOptions(path, tarOpts)
}

func createTarFileBuffer(filename string, data []byte) (io.Reader, error) {
	buf := &bytes.Buffer{}
	tw := tar.NewWriter(buf)
	defer tw.Close()
	hdr := &tar.Header{
		Name: filename,
		Mode: 0755,
		Size: int64(len(data)),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		return buf, err
	}
	if _, err := tw.Write(data); err != nil {
		return buf, err
	}
	return buf, nil
}

func makeCommonPart(build *types.Build) (string, error) {
	tmpl := template.Must(template.New("common").Parse(commonTmpl))
	out := bytes.Buffer{}
	if err := tmpl.Execute(&out, build); err != nil {
		return "", err
	}
	return out.String(), nil
}

func makeUserPart(opts *types.BuildOptions) (string, error) {
	tmpl := template.Must(template.New("user").Parse(userTmpl))
	out := bytes.Buffer{}
	if err := tmpl.Execute(&out, opts); err != nil {
		return "", err
	}
	return out.String(), nil
}

func makeMainPart(opts *types.BuildOptions, build *types.Build, from string, commands, copys []string) (string, error) {
	var buildTmpl []string
	common, err := makeCommonPart(build)
	if err != nil {
		return "", err
	}
	buildTmpl = append(buildTmpl, from, common)
	if len(copys) > 0 {
		buildTmpl = append(buildTmpl, copys...)
	}
	if len(commands) > 0 {
		buildTmpl = append(buildTmpl, commands...)
	}
	return strings.Join(buildTmpl, "\n"), nil
}

// Dockerfile
func createDockerfile(dockerfile, buildDir string) error {
	f, err := os.Create(filepath.Join(buildDir, "Dockerfile"))
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(dockerfile)
	return err
}

// Image tag
// 格式严格按照 Hub/HubPrefix/appname:version 来
func createImageTag(config types.DockerConfig, appname, version string) string {
	prefix := strings.Trim(config.Namespace, "/")
	if prefix == "" {
		return fmt.Sprintf("%s/%s:%s", config.Hub, appname, version)
	}
	return fmt.Sprintf("%s/%s/%s:%s", config.Hub, prefix, appname, version)
}

// 只要一个image的前面, tag不要
func normalizeImage(image string) string {
	if strings.Contains(image, ":") {
		t := strings.Split(image, ":")
		return t[0]
	}
	return image
}

// 清理一个node上的这个image
// 只清理同名字不同tag的
// 并且保留最新的 count 个
func cleanImageOnNode(ctx context.Context, node *types.Node, image string, count int) error {
	log.Debugf("[cleanImageOnNode] node: %s, image: %s", node.Name, strings.Split(image, ":")[0])
	imgListFilter := filters.NewArgs()
	image = normalizeImage(image)
	imgListFilter.Add("reference", image) // 相同repo的image
	images, err := node.Engine.ImageList(ctx, enginetypes.ImageListOptions{Filters: imgListFilter})
	if err != nil {
		return err
	}

	if len(images) < count {
		return nil
	}

	images = images[count:]
	log.Debugf("[cleanImageOnNode] Delete Images: %v", images)

	for _, image := range images {
		_, err := node.Engine.ImageRemove(ctx, image.ID, enginetypes.ImageRemoveOptions{
			Force:         false,
			PruneChildren: true,
		})
		if err != nil {
			log.Errorf("[cleanImageOnNode] Node %s ImageRemove error: %s, imageID: %s", node.Name, err, image.ID)
		}
	}
	return nil
}

func makeCopyMessage(id, status, name, path string, err error, data io.ReadCloser) *types.CopyMessage {
	return &types.CopyMessage{
		ID:     id,
		Status: status,
		Name:   name,
		Path:   path,
		Error:  err,
		Data:   data,
	}
}

func updateContainer(ctx context.Context, ID string, node *types.Node, config enginecontainer.UpdateConfig) error {
	_, err := node.Engine.ContainerUpdate(ctx, ID, config)
	return err
}

func filterNode(node *types.Node, labels map[string]string) bool {
	if node.Labels == nil && labels == nil {
		return true
	} else if node.Labels == nil && labels != nil {
		return false
	} else if node.Labels != nil && labels == nil {
		return true
	}

	for k, v := range labels {
		if d, ok := node.Labels[k]; !ok {
			return false
		} else if d != v {
			return false
		}
	}
	return true
}

func getNodesInfo(cpuAndMemData map[string]types.CPUAndMem) []types.NodeInfo {
	result := []types.NodeInfo{}
	for nodename, cpuAndMem := range cpuAndMemData {
		n := types.NodeInfo{
			CPUAndMem: cpuAndMem,
			Name:      nodename,
			CPUs:      len(cpuAndMem.CPUMap),
			Capacity:  0,
			Count:     0,
			Deploy:    0,
		}
		result = append(result, n)
	}
	return result
}

func parseStatusKey(key string) (string, string, string, string) {
	parts := strings.Split(key, "/")
	l := len(parts)
	return parts[l-4], parts[l-3], parts[l-2], parts[l-1]
}
<commit_msg>check remote image digest before pull image<commit_after>package calcium

import (
	"archive/tar"
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/docker/distribution/reference"
	enginetypes "github.com/docker/docker/api/types"
	enginecontainer "github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	engineapi "github.com/docker/docker/client"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/registry"
	"github.com/projecteru2/core/cluster"
	"github.com/projecteru2/core/lock"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

// Lock is lock for calcium
func (c *Calcium) Lock(ctx context.Context, name string, timeout int) (lock.DistributedLock, error) {
	lock, err := c.store.CreateLock(name, timeout)
	if err != nil {
		return nil, err
	}
	if err = lock.Lock(ctx); err != nil {
		return nil, err
	}
	return lock, nil
}

func makeResourceSetting(cpu float64, memory int64, cpuMap types.CPUMap, softlimit bool) enginecontainer.Resources {
	resource := enginecontainer.Resources{}
	if cpu > 0 {
		resource.CPUPeriod = cluster.CPUPeriodBase
		resource.CPUQuota = int64(cpu * float64(cluster.CPUPeriodBase))
	}
	if cpuMap != nil && len(cpuMap) > 0 {
		cpuIDs := []string{}
		for cpuID := range cpuMap {
			cpuIDs = append(cpuIDs, cpuID)
		}
		resource.CpusetCpus = strings.Join(cpuIDs, ",")
	}
	if softlimit {
		resource.MemoryReservation = memory
	} else {
		resource.Memory = memory
		resource.MemorySwap = memory
	}
	return resource
}

// image begin
// MakeAuthConfigFromRemote Calculate encoded AuthConfig from registry and eru-core config
// See https://github.com/docker/cli/blob/16cccc30f95c8163f0749eba5a2e80b807041342/cli/command/registry.go#L67
func makeEncodedAuthConfigFromRemote(authConfigs map[string]types.AuthConfig, remote string) (string, error) {
	ref, err := reference.ParseNormalizedNamed(remote)
	if err != nil {
		return "", err
	}

	// Resolve the Repository name from fqn to RepositoryInfo
	repoInfo, err := registry.ParseRepositoryInfo(ref)
	if err != nil {
		return "", err
	}

	serverAddress := repoInfo.Index.Name
	if authConfig, exists := authConfigs[serverAddress]; exists {
		if encodedAuth, err := encodeAuthToBase64(authConfig); err == nil {
			return encodedAuth, nil
		}
		return "", err
	}
	return "dummy", nil
}

// EncodeAuthToBase64 serializes the auth configuration as JSON base64 payload
// See https://github.com/docker/cli/blob/master/cli/command/registry.go#L41
func encodeAuthToBase64(authConfig types.AuthConfig) (string, error) {
	buf, err := json.Marshal(authConfig)
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(buf), nil
}

func makeCPUAndMem(nodes []*types.Node) map[string]types.CPUAndMem {
	r := make(map[string]types.CPUAndMem)
	for _, node := range nodes {
		r[node.Name] = types.CPUAndMem{
			CPUMap: node.CPU,
			MemCap: node.MemCap,
		}
	}
	return r
}

// As the name says,
// blocks until the stream is empty, until we meet EOF
func ensureReaderClosed(stream io.ReadCloser) {
	if stream == nil {
		return
	}
	io.Copy(ioutil.Discard, stream)
	stream.Close()
}

// make mount paths
// 使用volumes, 参数格式跟docker一样
// volumes:
//     - "/foo-data:$SOMEENV/foodata:rw"
func makeMountPaths(opts *types.DeployOptions) ([]string, map[string]struct{}) {
	binds := []string{}
	volumes := make(map[string]struct{})

	var expandENV = func(env string) string {
		envMap := map[string]string{}
		for _, env := range opts.Env {
			parts := strings.Split(env, "=")
			envMap[parts[0]] = parts[1]
		}
		return envMap[env]
	}

	for _, path := range opts.Volumes {
		expanded := os.Expand(path, expandENV)
		parts := strings.Split(expanded, ":")
		if len(parts) == 2 {
			binds = append(binds, fmt.Sprintf("%s:%s:rw", parts[0], parts[1]))
			volumes[parts[1]] = struct{}{}
		} else if len(parts) == 3 {
			binds = append(binds, fmt.Sprintf("%s:%s:%s", parts[0], parts[1], parts[2]))
			volumes[parts[1]] = struct{}{}
		}
	}

	return binds, volumes
}

// 跑存在labels里的exec
// 为什么要存labels呢, 因为下线容器的时候根本不知道entrypoint是啥
func execuateInside(ctx context.Context, client *engineapi.Client, ID, cmd, user string, env []string, privileged bool) ([]byte, error) {
	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := enginetypes.ExecConfig{
		User:         user,
		Cmd:          cmds,
		Privileged:   privileged,
		Env:          env,
		AttachStderr: true,
		AttachStdout: true,
	}
	//TODO should timeout
	//Fuck docker, ctx will not use inside funcs!!
	idResp, err := client.ContainerExecCreate(ctx, ID, execConfig)
	if err != nil {
		return []byte{}, err
	}
	resp, err := client.ContainerExecAttach(ctx, idResp.ID, enginetypes.ExecStartCheck{})
	if err != nil {
		return []byte{}, err
	}
	defer resp.Close()
	stream := utils.FuckDockerStream(ioutil.NopCloser(resp.Reader))
	b, err := ioutil.ReadAll(stream)
	if err != nil {
		return []byte{}, err
	}
	info, err := client.ContainerExecInspect(ctx, idResp.ID)
	if err != nil {
		return []byte{}, err
	}
	if info.ExitCode != 0 {
		return []byte{}, fmt.Errorf("%s", b)
	}
	return b, nil
}

func distributionInspect(ctx context.Context, node *types.Node, image, auth string, digests []string) bool {
	inspect, err := node.Engine.DistributionInspect(ctx, image, auth)
	if err != nil {
		log.Errorf("[distributionInspect] get manifest failed %v", err)
		return false
	}
	remoteDigest := fmt.Sprintf("%s@%s", normalizeImage(image), inspect.Descriptor.Digest.String())

	for _, digest := range digests {
		if digest == remoteDigest {
			log.Debugf("[distributionInspect]\ndigest %s \nremote %s", digest, remoteDigest)
			return true
		}
	}
	return false
}

// Pull an image
func pullImage(ctx context.Context, node *types.Node, image, auth string) error {
	log.Debugf("[pullImage] Pulling image %s", image)
	if image == "" {
		return fmt.Errorf("Goddamn empty image, WTF?")
	}

	// check local
	exists := false
	inspect, _, err := node.Engine.ImageInspectWithRaw(ctx, image)
	if err != nil {
		log.Errorf("[pullImage] Check image failed %v", err)
	} else {
		exists = true
	}

	if exists && distributionInspect(ctx, node, image, auth, inspect.RepoDigests) {
		log.Infof("[pullImage] Image cached, skip pulling")
		return nil
	}

	pullOptions := enginetypes.ImagePullOptions{All: false, RegistryAuth: auth}
	outStream, err := node.Engine.ImagePull(ctx, image, pullOptions)
	if err != nil {
		log.Errorf("[pullImage] Error during pulling image %s: %v", image, err)
		return err
	}
	ensureReaderClosed(outStream)
	log.Debugf("[pullImage] Done pulling image %s", image)
	return nil
}

func makeErrorBuildImageMessage(err error) *types.BuildImageMessage {
	return &types.BuildImageMessage{Error: err.Error()}
}

func createTarStream(path string) (io.ReadCloser, error) {
	tarOpts := &archive.TarOptions{
		ExcludePatterns: []string{},
		IncludeFiles:    []string{"."},
		Compression:     archive.Uncompressed,
		NoLchown:        true,
	}
	return archive.TarWithOptions(path, tarOpts)
}

func createTarFileBuffer(filename string, data []byte) (io.Reader, error) {
	buf := &bytes.Buffer{}
	tw := tar.NewWriter(buf)
	defer tw.Close()
	hdr := &tar.Header{
		Name: filename,
		Mode: 0755,
		Size: int64(len(data)),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		return buf, err
	}
	if _, err := tw.Write(data); err != nil {
		return buf, err
	}
	return buf, nil
}

func makeCommonPart(build *types.Build) (string, error) {
	tmpl := template.Must(template.New("common").Parse(commonTmpl))
	out := bytes.Buffer{}
	if err := tmpl.Execute(&out, build); err != nil {
		return "", err
	}
	return out.String(), nil
}

func makeUserPart(opts *types.BuildOptions) (string, error) {
	tmpl := template.Must(template.New("user").Parse(userTmpl))
	out := bytes.Buffer{}
	if err := tmpl.Execute(&out, opts); err != nil {
		return "", err
	}
	return out.String(), nil
}

func makeMainPart(opts *types.BuildOptions, build *types.Build, from string, commands, copys []string) (string, error) {
	var buildTmpl []string
	common, err := makeCommonPart(build)
	if err != nil {
		return "", err
	}
	buildTmpl = append(buildTmpl, from, common)
	if len(copys) > 0 {
		buildTmpl = append(buildTmpl, copys...)
	}
	if len(commands) > 0 {
		buildTmpl = append(buildTmpl, commands...)
	}
	return strings.Join(buildTmpl, "\n"), nil
}

// Dockerfile
func createDockerfile(dockerfile, buildDir string) error {
	f, err := os.Create(filepath.Join(buildDir, "Dockerfile"))
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(dockerfile)
	return err
}

// Image tag
// 格式严格按照 Hub/HubPrefix/appname:version 来
func createImageTag(config types.DockerConfig, appname, version string) string {
	prefix := strings.Trim(config.Namespace, "/")
	if prefix == "" {
		return fmt.Sprintf("%s/%s:%s", config.Hub, appname, version)
	}
	return fmt.Sprintf("%s/%s/%s:%s", config.Hub, prefix, appname, version)
}

// 只要一个image的前面, tag不要
func normalizeImage(image string) string {
	if strings.Contains(image, ":") {
		t := strings.Split(image, ":")
		return t[0]
	}
	return image
}

// 清理一个node上的这个image
// 只清理同名字不同tag的
// 并且保留最新的 count 个
func cleanImageOnNode(ctx context.Context, node *types.Node, image string, count int) error {
	log.Debugf("[cleanImageOnNode] node: %s, image: %s", node.Name, strings.Split(image, ":")[0])
	imgListFilter := filters.NewArgs()
	image = normalizeImage(image)
	imgListFilter.Add("reference", image) // 相同repo的image
	images, err := node.Engine.ImageList(ctx, enginetypes.ImageListOptions{Filters: imgListFilter})
	if err != nil {
		return err
	}

	if len(images) < count {
		return nil
	}

	images = images[count:]
	log.Debugf("[cleanImageOnNode] Delete Images: %v", images)

	for _, image := range images {
		_, err := node.Engine.ImageRemove(ctx, image.ID, enginetypes.ImageRemoveOptions{
			Force:         false,
			PruneChildren: true,
		})
		if err != nil {
			log.Errorf("[cleanImageOnNode] Node %s ImageRemove error: %s, imageID: %s", node.Name, err, image.ID)
		}
	}
	return nil
}

func makeCopyMessage(id, status, name, path string, err error, data io.ReadCloser) *types.CopyMessage {
	return &types.CopyMessage{
		ID:     id,
		Status: status,
		Name:   name,
		Path:   path,
		Error:  err,
		Data:   data,
	}
}

func updateContainer(ctx context.Context, ID string, node *types.Node, config enginecontainer.UpdateConfig) error {
	_, err := node.Engine.ContainerUpdate(ctx, ID, config)
	return err
}

func filterNode(node *types.Node, labels map[string]string) bool {
	if node.Labels == nil && labels == nil {
		return true
	} else if node.Labels == nil && labels != nil {
		return false
	} else if node.Labels != nil && labels == nil {
		return true
	}

	for k, v := range labels {
		if d, ok := node.Labels[k]; !ok {
			return false
		} else if d != v {
			return false
		}
	}
	return true
}

func getNodesInfo(cpuAndMemData map[string]types.CPUAndMem) []types.NodeInfo {
	result := []types.NodeInfo{}
	for nodename, cpuAndMem := range cpuAndMemData {
		n := types.NodeInfo{
			CPUAndMem: cpuAndMem,
			Name:      nodename,
			CPUs:      len(cpuAndMem.CPUMap),
			Capacity:  0,
			Count:     0,
			Deploy:    0,
		}
		result = append(result, n)
	}
	return result
}

func parseStatusKey(key string) (string, string, string, string) {
	parts := strings.Split(key, "/")
	l := len(parts)
	return parts[l-4], parts[l-3], parts[l-2], parts[l-1]
}
<|endoftext|>


<commit_before><commit_msg>fix: ensure cache settings read/writes are protected by mutex (#3753)<commit_after><|endoftext|>