<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package core

import (
	"crypto/ecdsa"
	"crypto/tls"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/TheThingsNetwork/go-account-lib/cache"
	"github.com/TheThingsNetwork/go-account-lib/claims"
	"github.com/TheThingsNetwork/go-account-lib/tokenkey"
	"github.com/TheThingsNetwork/ttn/api"
	pb_discovery "github.com/TheThingsNetwork/ttn/api/discovery"
	pb_noc "github.com/TheThingsNetwork/ttn/api/noc"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/logging"
	"github.com/TheThingsNetwork/ttn/utils/security"
	"github.com/apex/log"
	"github.com/dgrijalva/jwt-go"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/spf13/viper"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
)

type ComponentInterface interface {
	RegisterRPC(s *grpc.Server)
	Init(c *Component) error
	ValidateNetworkContext(ctx context.Context) (*pb_discovery.Announcement, error)
	ValidateTTNAuthContext(ctx context.Context) (*claims.Claims, error)
}

type ManagementInterface interface {
	RegisterManager(s *grpc.Server)
}

// NewComponent creates a new Component
func NewComponent(ctx log.Interface, serviceName string, announcedAddress string) (*Component, error) {
	go func() {
		memstats := new(runtime.MemStats)
		for range time.Tick(time.Minute) {
			runtime.ReadMemStats(memstats)
			ctx.WithFields(log.Fields{
				"Goroutines": runtime.NumGoroutine(),
				"Memory":     float64(memstats.Alloc) / 1000000,
			}).Debugf("Stats")
		}
	}()

	grpclog.SetLogger(logging.NewGRPCLogger(ctx))

	component := &Component{
		Ctx: ctx,
		Identity: &pb_discovery.Announcement{
			Id:             viper.GetString("id"),
			Description:    viper.GetString("description"),
			ServiceName:    serviceName,
			ServiceVersion: fmt.Sprintf("%s-%s (%s)", viper.GetString("version"), viper.GetString("gitCommit"), viper.GetString("buildDate")),
			NetAddress:     announcedAddress,
		},
		AccessToken: viper.GetString("auth-token"),
		TokenKeyProvider: tokenkey.HTTPProvider(
			viper.GetStringMapString("auth-servers"),
			cache.WriteTroughCacheWithFormat(viper.GetString("key-dir"), "auth-%s.pub"),
		),
	}

	if serviceName != "discovery" {
		var err error
		component.Discovery, err = pb_discovery.NewClient(
			viper.GetString("discovery-server"),
			component.Identity,
			func() string {
				token, _ := component.BuildJWT()
				return token
			},
		)
		if err != nil {
			return nil, err
		}
	}

	if priv, err := security.LoadKeypair(viper.GetString("key-dir")); err == nil {
		component.privateKey = priv

		pubPEM, _ := security.PublicPEM(priv)
		component.Identity.PublicKey = string(pubPEM)

		privPEM, _ := security.PrivatePEM(priv)

		if viper.GetBool("tls") {
			cert, err := security.LoadCert(viper.GetString("key-dir"))
			if err != nil {
				return nil, err
			}
			component.Identity.Certificate = string(cert)

			cer, err := tls.X509KeyPair(cert, privPEM)
			if err != nil {
				return nil, err
			}
			component.tlsConfig = &tls.Config{Certificates: []tls.Certificate{cer}}
		}
	}

	if healthPort := viper.GetInt("health-port"); healthPort > 0 {
		http.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
			switch component.GetStatus() {
			case StatusHealthy:
				w.WriteHeader(200)
				w.Write([]byte("Status is HEALTHY"))
				return
			case StatusUnhealthy:
				w.WriteHeader(503)
				w.Write([]byte("Status is UNHEALTHY"))
				return
			}
		})
		go http.ListenAndServe(fmt.Sprintf(":%d", healthPort), nil)
	}

	if nocAddr := viper.GetString("noc-server"); len(nocAddr) > 0 {
		conn, err := grpc.Dial(nocAddr, append(api.DialOptions, grpc.WithInsecure())...)
		if err != nil {
			return nil, err
		}

		component.Monitor = pb_noc.NewMonitorClient(conn)
	}

	return component, nil
}

// Status indicates the health status of this component
type Status int

const (
	// StatusHealthy indicates a healthy component
	StatusHealthy Status = iota
	// StatusUnhealthy indicates an unhealthy component
	StatusUnhealthy
)

// Component contains the common attributes for all TTN components
type Component struct {
	Identity         *pb_discovery.Announcement
	Discovery        pb_discovery.Client
	Monitor          pb_noc.MonitorClient
	Ctx              log.Interface
	AccessToken      string
	privateKey       *ecdsa.PrivateKey
	tlsConfig        *tls.Config
	TokenKeyProvider tokenkey.Provider
	status           int64
}

// GetStatus gets the health status of the component
func (c *Component) GetStatus() Status {
	return Status(atomic.LoadInt64(&c.status))
}

// SetStatus sets the health status of the component
func (c *Component) SetStatus(status Status) {
	atomic.StoreInt64(&c.status, int64(status))
}

// Discover is used to discover another component
func (c *Component) Discover(serviceName, id string) (*pb_discovery.Announcement, error) {
	res, err := c.Discovery.Get(serviceName, id)
	if err != nil {
		return nil, errors.Wrapf(errors.FromGRPCError(err), "Failed to discover %s/%s", serviceName, id)
	}
	return res, nil
}

// Announce the component to TTN discovery
func (c *Component) Announce() error {
	if c.Identity.Id == "" {
		return errors.NewErrInvalidArgument("Component ID", "can not be empty")
	}
	err := c.Discovery.Announce(c.AccessToken)
	if err != nil {
		return errors.Wrapf(errors.FromGRPCError(err), "Failed to announce this component to TTN discovery: %s", err.Error())
	}
	c.Ctx.Info("ttn: Announced to TTN discovery")

	return nil
}

// UpdateTokenKey updates the OAuth Bearer token key
func (c *Component) UpdateTokenKey() error {
	if c.TokenKeyProvider == nil {
		return errors.NewErrInternal("No public key provider configured for token validation")
	}

	// Set up Auth Server Token Validation
	err := c.TokenKeyProvider.Update()
	if err != nil {
		c.Ctx.Warnf("ttn: Failed to refresh public keys for token validation: %s", err.Error())
	} else {
		c.Ctx.Info("ttn: Got public keys for token validation")
	}

	return nil

}

// ValidateNetworkContext validates the context of a network request (router-broker, broker-handler, etc)
func (c *Component) ValidateNetworkContext(ctx context.Context) (component *pb_discovery.Announcement, err error) {
	defer func() {
		if err != nil {
			time.Sleep(time.Second)
		}
	}()

	md, ok := metadata.FromContext(ctx)
	if !ok {
		err = errors.NewErrInternal("Could not get metadata from context")
		return
	}
	var id, serviceName, token string
	if ids, ok := md["id"]; ok && len(ids) == 1 {
		id = ids[0]
	}
	if id == "" {
		err = errors.NewErrInvalidArgument("Metadata", "id missing")
		return
	}
	if serviceNames, ok := md["service-name"]; ok && len(serviceNames) == 1 {
		serviceName = serviceNames[0]
	}
	if serviceName == "" {
		err = errors.NewErrInvalidArgument("Metadata", "service-name missing")
		return
	}
	if tokens, ok := md["token"]; ok && len(tokens) == 1 {
		token = tokens[0]
	}

	var announcement *pb_discovery.Announcement
	announcement, err = c.Discover(serviceName, id)
	if err != nil {
		return
	}

	if announcement.PublicKey == "" {
		return announcement, nil
	}

	if token == "" {
		err = errors.NewErrInvalidArgument("Metadata", "token missing")
		return
	}

	var claims *jwt.StandardClaims
	claims, err = security.ValidateJWT(token, []byte(announcement.PublicKey))
	if err != nil {
		return
	}
	if claims.Issuer != id {
		err = errors.NewErrInvalidArgument("Metadata", "token was issued by different component id")
		return
	}

	return announcement, nil
}

// ValidateTTNAuthContext gets a token from the context and validates it
func (c *Component) ValidateTTNAuthContext(ctx context.Context) (*claims.Claims, error) {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return nil, errors.NewErrInternal("Could not get metadata from context")
	}
	token, ok := md["token"]
	if !ok || len(token) < 1 {
		return nil, errors.NewErrInvalidArgument("Metadata", "token missing")
	}

	if c.TokenKeyProvider == nil {
		return nil, errors.NewErrInternal("No token provider configured")
	}

	if token[0] == "" {
		return nil, errors.NewErrInvalidArgument("Metadata", "token is empty")
	}

	claims, err := claims.FromToken(c.TokenKeyProvider, token[0])
	if err != nil {
		return nil, err
	}

	return claims, nil
}

func (c *Component) ServerOptions() []grpc.ServerOption {
	unary := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		var peerAddr string
		peer, ok := peer.FromContext(ctx)
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(ctx)
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		logCtx := c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		})
		t := time.Now()
		iface, err := handler(ctx, req)
		logCtx = logCtx.WithField("Duration", time.Now().Sub(t))
		if err != nil {
			err := errors.FromGRPCError(err)
			logCtx.WithField("error", err.Error()).Warn("Could not handle Request")
		} else {
			logCtx.Info("Handled request")
		}
		return iface, err
	}

	stream := func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		var peerAddr string
		peer, ok := peer.FromContext(stream.Context())
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(stream.Context())
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		}).Info("Start stream")
		return handler(srv, stream)
	}

	opts := []grpc.ServerOption{
		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(unary)),
		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(stream)),
	}

	if c.tlsConfig != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(c.tlsConfig)))
	}

	return opts
}

// BuildJWT builds a short-lived JSON Web Token for this component
func (c *Component) BuildJWT() (string, error) {
	if c.privateKey != nil {
		privPEM, err := security.PrivatePEM(c.privateKey)
		if err != nil {
			return "", err
		}
		return security.BuildJWT(c.Identity.Id, 20*time.Second, privPEM)
	}
	return "", nil
}

// GetContext returns a context for outgoing RPC request. If token is "", this function will generate a short lived token from the component
func (c *Component) GetContext(token string) context.Context {
	var serviceName, id, netAddress string
	if c.Identity != nil {
		serviceName = c.Identity.ServiceName
		id = c.Identity.Id
		if token == "" {
			token, _ = c.BuildJWT()
		}
		netAddress = c.Identity.NetAddress
	}
	md := metadata.Pairs(
		"service-name", serviceName,
		"id", id,
		"token", token,
		"net-address", netAddress,
	)
	ctx := metadata.NewContext(context.Background(), md)
	return ctx
}
<commit_msg>core/component.go: refactoring<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package core

import (
	"crypto/ecdsa"
	"crypto/tls"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/TheThingsNetwork/go-account-lib/cache"
	"github.com/TheThingsNetwork/go-account-lib/claims"
	"github.com/TheThingsNetwork/go-account-lib/tokenkey"
	"github.com/TheThingsNetwork/ttn/api"
	pb_discovery "github.com/TheThingsNetwork/ttn/api/discovery"
	pb_noc "github.com/TheThingsNetwork/ttn/api/noc"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/logging"
	"github.com/TheThingsNetwork/ttn/utils/security"
	"github.com/apex/log"
	"github.com/dgrijalva/jwt-go"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/spf13/viper"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
)

type ComponentInterface interface {
	RegisterRPC(s *grpc.Server)
	Init(c *Component) error
	ValidateNetworkContext(ctx context.Context) (*pb_discovery.Announcement, error)
	ValidateTTNAuthContext(ctx context.Context) (*claims.Claims, error)
}

type ManagementInterface interface {
	RegisterManager(s *grpc.Server)
}

// NewComponent creates a new Component
func NewComponent(ctx log.Interface, serviceName string, announcedAddress string) (*Component, error) {
	go func() {
		memstats := new(runtime.MemStats)
		for range time.Tick(time.Minute) {
			runtime.ReadMemStats(memstats)
			ctx.WithFields(log.Fields{
				"Goroutines": runtime.NumGoroutine(),
				"Memory":     float64(memstats.Alloc) / 1000000,
			}).Debugf("Stats")
		}
	}()

	grpclog.SetLogger(logging.NewGRPCLogger(ctx))

	component := &Component{
		Ctx: ctx,
		Identity: &pb_discovery.Announcement{
			Id:             viper.GetString("id"),
			Description:    viper.GetString("description"),
			ServiceName:    serviceName,
			ServiceVersion: fmt.Sprintf("%s-%s (%s)", viper.GetString("version"), viper.GetString("gitCommit"), viper.GetString("buildDate")),
			NetAddress:     announcedAddress,
		},
		AccessToken: viper.GetString("auth-token"),
		TokenKeyProvider: tokenkey.HTTPProvider(
			viper.GetStringMapString("auth-servers"),
			cache.WriteTroughCacheWithFormat(viper.GetString("key-dir"), "auth-%s.pub"),
		),
	}

	if serviceName != "discovery" {
		var err error
		component.Discovery, err = pb_discovery.NewClient(
			viper.GetString("discovery-server"),
			component.Identity,
			func() string {
				token, _ := component.BuildJWT()
				return token
			},
		)
		if err != nil {
			return nil, err
		}
	}

	if priv, err := security.LoadKeypair(viper.GetString("key-dir")); err == nil {
		component.privateKey = priv

		pubPEM, _ := security.PublicPEM(priv)
		component.Identity.PublicKey = string(pubPEM)

		privPEM, _ := security.PrivatePEM(priv)

		if viper.GetBool("tls") {
			cert, err := security.LoadCert(viper.GetString("key-dir"))
			if err != nil {
				return nil, err
			}
			component.Identity.Certificate = string(cert)

			cer, err := tls.X509KeyPair(cert, privPEM)
			if err != nil {
				return nil, err
			}
			component.tlsConfig = &tls.Config{Certificates: []tls.Certificate{cer}}
		}
	}

	if healthPort := viper.GetInt("health-port"); healthPort > 0 {
		http.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
			switch component.GetStatus() {
			case StatusHealthy:
				w.WriteHeader(200)
				w.Write([]byte("Status is HEALTHY"))
				return
			case StatusUnhealthy:
				w.WriteHeader(503)
				w.Write([]byte("Status is UNHEALTHY"))
				return
			}
		})
		go http.ListenAndServe(fmt.Sprintf(":%d", healthPort), nil)
	}

	if nocAddr := viper.GetString("noc-server"); nocAddr != "" {
		conn, err := grpc.Dial(nocAddr, append(api.DialOptions, grpc.WithInsecure())...)
		if err != nil {
			return nil, err
		}

		component.Monitor = pb_noc.NewMonitorClient(conn)
	}

	return component, nil
}

// Status indicates the health status of this component
type Status int

const (
	// StatusHealthy indicates a healthy component
	StatusHealthy Status = iota
	// StatusUnhealthy indicates an unhealthy component
	StatusUnhealthy
)

// Component contains the common attributes for all TTN components
type Component struct {
	Identity         *pb_discovery.Announcement
	Discovery        pb_discovery.Client
	Monitor          pb_noc.MonitorClient
	Ctx              log.Interface
	AccessToken      string
	privateKey       *ecdsa.PrivateKey
	tlsConfig        *tls.Config
	TokenKeyProvider tokenkey.Provider
	status           int64
}

// GetStatus gets the health status of the component
func (c *Component) GetStatus() Status {
	return Status(atomic.LoadInt64(&c.status))
}

// SetStatus sets the health status of the component
func (c *Component) SetStatus(status Status) {
	atomic.StoreInt64(&c.status, int64(status))
}

// Discover is used to discover another component
func (c *Component) Discover(serviceName, id string) (*pb_discovery.Announcement, error) {
	res, err := c.Discovery.Get(serviceName, id)
	if err != nil {
		return nil, errors.Wrapf(errors.FromGRPCError(err), "Failed to discover %s/%s", serviceName, id)
	}
	return res, nil
}

// Announce the component to TTN discovery
func (c *Component) Announce() error {
	if c.Identity.Id == "" {
		return errors.NewErrInvalidArgument("Component ID", "can not be empty")
	}
	err := c.Discovery.Announce(c.AccessToken)
	if err != nil {
		return errors.Wrapf(errors.FromGRPCError(err), "Failed to announce this component to TTN discovery: %s", err.Error())
	}
	c.Ctx.Info("ttn: Announced to TTN discovery")

	return nil
}

// UpdateTokenKey updates the OAuth Bearer token key
func (c *Component) UpdateTokenKey() error {
	if c.TokenKeyProvider == nil {
		return errors.NewErrInternal("No public key provider configured for token validation")
	}

	// Set up Auth Server Token Validation
	err := c.TokenKeyProvider.Update()
	if err != nil {
		c.Ctx.Warnf("ttn: Failed to refresh public keys for token validation: %s", err.Error())
	} else {
		c.Ctx.Info("ttn: Got public keys for token validation")
	}

	return nil

}

// ValidateNetworkContext validates the context of a network request (router-broker, broker-handler, etc)
func (c *Component) ValidateNetworkContext(ctx context.Context) (component *pb_discovery.Announcement, err error) {
	defer func() {
		if err != nil {
			time.Sleep(time.Second)
		}
	}()

	md, ok := metadata.FromContext(ctx)
	if !ok {
		err = errors.NewErrInternal("Could not get metadata from context")
		return
	}
	var id, serviceName, token string
	if ids, ok := md["id"]; ok && len(ids) == 1 {
		id = ids[0]
	}
	if id == "" {
		err = errors.NewErrInvalidArgument("Metadata", "id missing")
		return
	}
	if serviceNames, ok := md["service-name"]; ok && len(serviceNames) == 1 {
		serviceName = serviceNames[0]
	}
	if serviceName == "" {
		err = errors.NewErrInvalidArgument("Metadata", "service-name missing")
		return
	}
	if tokens, ok := md["token"]; ok && len(tokens) == 1 {
		token = tokens[0]
	}

	var announcement *pb_discovery.Announcement
	announcement, err = c.Discover(serviceName, id)
	if err != nil {
		return
	}

	if announcement.PublicKey == "" {
		return announcement, nil
	}

	if token == "" {
		err = errors.NewErrInvalidArgument("Metadata", "token missing")
		return
	}

	var claims *jwt.StandardClaims
	claims, err = security.ValidateJWT(token, []byte(announcement.PublicKey))
	if err != nil {
		return
	}
	if claims.Issuer != id {
		err = errors.NewErrInvalidArgument("Metadata", "token was issued by different component id")
		return
	}

	return announcement, nil
}

// ValidateTTNAuthContext gets a token from the context and validates it
func (c *Component) ValidateTTNAuthContext(ctx context.Context) (*claims.Claims, error) {
	md, ok := metadata.FromContext(ctx)
	if !ok {
		return nil, errors.NewErrInternal("Could not get metadata from context")
	}
	token, ok := md["token"]
	if !ok || len(token) < 1 {
		return nil, errors.NewErrInvalidArgument("Metadata", "token missing")
	}

	if c.TokenKeyProvider == nil {
		return nil, errors.NewErrInternal("No token provider configured")
	}

	if token[0] == "" {
		return nil, errors.NewErrInvalidArgument("Metadata", "token is empty")
	}

	claims, err := claims.FromToken(c.TokenKeyProvider, token[0])
	if err != nil {
		return nil, err
	}

	return claims, nil
}

func (c *Component) ServerOptions() []grpc.ServerOption {
	unary := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		var peerAddr string
		peer, ok := peer.FromContext(ctx)
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(ctx)
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		logCtx := c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		})
		t := time.Now()
		iface, err := handler(ctx, req)
		logCtx = logCtx.WithField("Duration", time.Now().Sub(t))
		if err != nil {
			err := errors.FromGRPCError(err)
			logCtx.WithField("error", err.Error()).Warn("Could not handle Request")
		} else {
			logCtx.Info("Handled request")
		}
		return iface, err
	}

	stream := func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		var peerAddr string
		peer, ok := peer.FromContext(stream.Context())
		if ok {
			peerAddr = peer.Addr.String()
		}
		var peerID string
		meta, ok := metadata.FromContext(stream.Context())
		if ok {
			id, ok := meta["id"]
			if ok && len(id) > 0 {
				peerID = id[0]
			}
		}
		c.Ctx.WithFields(log.Fields{
			"CallerID": peerID,
			"CallerIP": peerAddr,
			"Method":   info.FullMethod,
		}).Info("Start stream")
		return handler(srv, stream)
	}

	opts := []grpc.ServerOption{
		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(unary)),
		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(stream)),
	}

	if c.tlsConfig != nil {
		opts = append(opts, grpc.Creds(credentials.NewTLS(c.tlsConfig)))
	}

	return opts
}

// BuildJWT builds a short-lived JSON Web Token for this component
func (c *Component) BuildJWT() (string, error) {
	if c.privateKey != nil {
		privPEM, err := security.PrivatePEM(c.privateKey)
		if err != nil {
			return "", err
		}
		return security.BuildJWT(c.Identity.Id, 20*time.Second, privPEM)
	}
	return "", nil
}

// GetContext returns a context for outgoing RPC request. If token is "", this function will generate a short lived token from the component
func (c *Component) GetContext(token string) context.Context {
	var serviceName, id, netAddress string
	if c.Identity != nil {
		serviceName = c.Identity.ServiceName
		id = c.Identity.Id
		if token == "" {
			token, _ = c.BuildJWT()
		}
		netAddress = c.Identity.NetAddress
	}
	md := metadata.Pairs(
		"service-name", serviceName,
		"id", id,
		"token", token,
		"net-address", netAddress,
	)
	ctx := metadata.NewContext(context.Background(), md)
	return ctx
}
<|endoftext|>


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

// +build acceptance

package acceptance

import (
	"errors"
	"math/rand"
	"strings"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/acceptance/cluster"
	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
)

const longWaitTime = 2 * time.Minute
const shortWaitTime = 20 * time.Second

type checkGossipFunc func(map[string]interface{}) error

// checkGossip fetches the gossip infoStore from each node and invokes the given
// function. The test passes if the function returns 0 for every node,
// retrying for up to the given duration.
func checkGossip(t *testing.T, c cluster.Cluster, d time.Duration,
	f checkGossipFunc) {
	util.SucceedsWithin(t, d, func() error {
		select {
		case <-stopper:
			t.Fatalf("interrupted")
			return nil
		case <-time.After(1 * time.Second):
		}

		for i := 0; i < c.NumNodes(); i++ {
			var m map[string]interface{}
			if err := getJSON(c.URL(i), "/_status/gossip/local", &m); err != nil {
				return err
			}
			infos, ok := m["infos"].(map[string]interface{})
			if !ok {
				return errors.New("no infos yet")
			}
			if err := f(infos); err != nil {
				return util.Errorf("node %d: %s", i, err)
			}
		}

		return nil
	})
}

// hasPeers returns a checkGossipFunc that passes when the given
// number of peers are connected via gossip.
func hasPeers(expected int) checkGossipFunc {
	return func(infos map[string]interface{}) error {
		count := 0
		for k := range infos {
			if strings.HasPrefix(k, "node:") {
				count++
			}
		}
		if count != expected {
			return util.Errorf("expected %d peers, found %d", expected, count)
		}
		return nil
	}
}

// hasSentinel is a checkGossipFunc that passes when the sentinel gossip is present.
func hasSentinel(infos map[string]interface{}) error {
	if _, ok := infos[gossip.KeySentinel]; !ok {
		return util.Errorf("sentinel not found")
	}
	return nil
}

// hasClusterID is a checkGossipFunc that passes when the cluster ID gossip is present.
func hasClusterID(infos map[string]interface{}) error {
	if _, ok := infos[gossip.KeyClusterID]; !ok {
		return util.Errorf("cluster ID not found")
	}
	return nil
}

func TestGossipPeerings(t *testing.T) {
	c := StartCluster(t)
	defer c.AssertAndStop(t)
	num := c.NumNodes()

	deadline := time.Now().Add(*duration)

	waitTime := longWaitTime
	if *duration < waitTime {
		waitTime = shortWaitTime
	}

	for time.Now().Before(deadline) {
		checkGossip(t, c, waitTime, hasPeers(num))

		// Restart the first node.
		log.Infof("restarting node 0")
		if err := c.Restart(0); err != nil {
			t.Fatal(err)
		}
		checkGossip(t, c, waitTime, hasPeers(num))

		// Restart another node (if there is one).
		var pickedNode int
		if num > 1 {
			pickedNode = rand.Intn(num-1) + 1
		}
		log.Infof("restarting node %d", pickedNode)
		if err := c.Restart(pickedNode); err != nil {
			t.Fatal(err)
		}
		checkGossip(t, c, waitTime, hasPeers(num))
	}
}

// TestGossipRestart verifies that the gossip network can be
// re-bootstrapped after a time when all nodes were down
// simultaneously.
func TestGossipRestart(t *testing.T) {
	// This already replicates the first range (in the local setup).
	// The replication of the first range is important: as long as the
	// first range only exists on one node, that node can trivially
	// acquire the leader lease. Once the range is replicated, however,
	// nodes must be able to discover each other over gossip before the
	// lease can be acquired.
	c := StartCluster(t)
	defer c.AssertAndStop(t)
	num := c.NumNodes()

	deadline := time.Now().Add(*duration)

	waitTime := longWaitTime
	if *duration < waitTime {
		waitTime = shortWaitTime
	}

	for time.Now().Before(deadline) {
		log.Infof("waiting for initial gossip connections")
		checkGossip(t, c, waitTime, hasPeers(num))
		checkGossip(t, c, waitTime, hasClusterID)
		checkGossip(t, c, waitTime, hasSentinel)

		log.Infof("killing all nodes")
		for i := 0; i < num; i++ {
			if err := c.Kill(i); err != nil {
				t.Fatal(err)
			}
		}

		log.Infof("restarting all nodes")
		for i := 0; i < num; i++ {
			if err := c.Restart(i); err != nil {
				t.Fatal(err)
			}
		}

		log.Infof("waiting for gossip to be connected")
		checkGossip(t, c, waitTime, hasPeers(num))
		checkGossip(t, c, waitTime, hasClusterID)
		checkGossip(t, c, waitTime, hasSentinel)

		for i := 0; i < num; i++ {
			db, dbStopper := makeClient(t, c.ConnString(i))
			if i == 0 {
				if err := db.Del("count"); err != nil {
					t.Fatal(err)
				}
			}
			var kv client.KeyValue
			if err := db.Txn(func(txn *client.Txn) error {
				var err error
				kv, err = txn.Inc("count", 1)
				return err
			}); err != nil {
				t.Fatal(err)
			} else if v := kv.ValueInt(); v != int64(i+1) {
				t.Fatalf("unexpected value %d for write #%d (expected %d)", v, i, i+1)
			}
			dbStopper.Stop()
		}
	}
}
<commit_msg>disable flaky TestGossip{Restart,Peerings}<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Peter Mattis (peter@cockroachlabs.com)

// +build acceptance

package acceptance

import (
	"errors"
	"math/rand"
	"strings"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/acceptance/cluster"
	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
)

const longWaitTime = 2 * time.Minute
const shortWaitTime = 20 * time.Second

type checkGossipFunc func(map[string]interface{}) error

// checkGossip fetches the gossip infoStore from each node and invokes the given
// function. The test passes if the function returns 0 for every node,
// retrying for up to the given duration.
func checkGossip(t *testing.T, c cluster.Cluster, d time.Duration,
	f checkGossipFunc) {
	util.SucceedsWithin(t, d, func() error {
		select {
		case <-stopper:
			t.Fatalf("interrupted")
			return nil
		case <-time.After(1 * time.Second):
		}

		for i := 0; i < c.NumNodes(); i++ {
			var m map[string]interface{}
			if err := getJSON(c.URL(i), "/_status/gossip/local", &m); err != nil {
				return err
			}
			infos, ok := m["infos"].(map[string]interface{})
			if !ok {
				return errors.New("no infos yet")
			}
			if err := f(infos); err != nil {
				return util.Errorf("node %d: %s", i, err)
			}
		}

		return nil
	})
}

// hasPeers returns a checkGossipFunc that passes when the given
// number of peers are connected via gossip.
func hasPeers(expected int) checkGossipFunc {
	return func(infos map[string]interface{}) error {
		count := 0
		for k := range infos {
			if strings.HasPrefix(k, "node:") {
				count++
			}
		}
		if count != expected {
			return util.Errorf("expected %d peers, found %d", expected, count)
		}
		return nil
	}
}

// hasSentinel is a checkGossipFunc that passes when the sentinel gossip is present.
func hasSentinel(infos map[string]interface{}) error {
	if _, ok := infos[gossip.KeySentinel]; !ok {
		return util.Errorf("sentinel not found")
	}
	return nil
}

// hasClusterID is a checkGossipFunc that passes when the cluster ID gossip is present.
func hasClusterID(infos map[string]interface{}) error {
	if _, ok := infos[gossip.KeyClusterID]; !ok {
		return util.Errorf("cluster ID not found")
	}
	return nil
}

func TestGossipPeerings(t *testing.T) {
	t.Skip("#3611")
	c := StartCluster(t)
	defer c.AssertAndStop(t)
	num := c.NumNodes()

	deadline := time.Now().Add(*duration)

	waitTime := longWaitTime
	if *duration < waitTime {
		waitTime = shortWaitTime
	}

	for time.Now().Before(deadline) {
		checkGossip(t, c, waitTime, hasPeers(num))

		// Restart the first node.
		log.Infof("restarting node 0")
		if err := c.Restart(0); err != nil {
			t.Fatal(err)
		}
		checkGossip(t, c, waitTime, hasPeers(num))

		// Restart another node (if there is one).
		var pickedNode int
		if num > 1 {
			pickedNode = rand.Intn(num-1) + 1
		}
		log.Infof("restarting node %d", pickedNode)
		if err := c.Restart(pickedNode); err != nil {
			t.Fatal(err)
		}
		checkGossip(t, c, waitTime, hasPeers(num))
	}
}

// TestGossipRestart verifies that the gossip network can be
// re-bootstrapped after a time when all nodes were down
// simultaneously.
func TestGossipRestart(t *testing.T) {
	t.Skip("#3611")
	// This already replicates the first range (in the local setup).
	// The replication of the first range is important: as long as the
	// first range only exists on one node, that node can trivially
	// acquire the leader lease. Once the range is replicated, however,
	// nodes must be able to discover each other over gossip before the
	// lease can be acquired.
	c := StartCluster(t)
	defer c.AssertAndStop(t)
	num := c.NumNodes()

	deadline := time.Now().Add(*duration)

	waitTime := longWaitTime
	if *duration < waitTime {
		waitTime = shortWaitTime
	}

	for time.Now().Before(deadline) {
		log.Infof("waiting for initial gossip connections")
		checkGossip(t, c, waitTime, hasPeers(num))
		checkGossip(t, c, waitTime, hasClusterID)
		checkGossip(t, c, waitTime, hasSentinel)

		log.Infof("killing all nodes")
		for i := 0; i < num; i++ {
			if err := c.Kill(i); err != nil {
				t.Fatal(err)
			}
		}

		log.Infof("restarting all nodes")
		for i := 0; i < num; i++ {
			if err := c.Restart(i); err != nil {
				t.Fatal(err)
			}
		}

		log.Infof("waiting for gossip to be connected")
		checkGossip(t, c, waitTime, hasPeers(num))
		checkGossip(t, c, waitTime, hasClusterID)
		checkGossip(t, c, waitTime, hasSentinel)

		for i := 0; i < num; i++ {
			db, dbStopper := makeClient(t, c.ConnString(i))
			if i == 0 {
				if err := db.Del("count"); err != nil {
					t.Fatal(err)
				}
			}
			var kv client.KeyValue
			if err := db.Txn(func(txn *client.Txn) error {
				var err error
				kv, err = txn.Inc("count", 1)
				return err
			}); err != nil {
				t.Fatal(err)
			} else if v := kv.ValueInt(); v != int64(i+1) {
				t.Fatalf("unexpected value %d for write #%d (expected %d)", v, i, i+1)
			}
			dbStopper.Stop()
		}
	}
}
<|endoftext|>


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

package hugolib

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestGroupFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", pageContent, "page2.md", pageContent).
		WithTemplatesAdded("index.html", `
{{ $cool := .Site.RegularPages | group "cool" }}
{{ $cool.Key }}: {{ len $cool.Pages }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html", "cool: 2")
}

func TestSliceFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20)).
		WithTemplatesAdded("index.html", `
{{ $cool := first 1 .Site.RegularPages | group "cool" }}
{{ $blue := after 1 .Site.RegularPages | group "blue" }}
{{ $weightedPages := index (index .Site.Taxonomies "tags") "blue" }}

{{ $p1 := index .Site.RegularPages 0 }}{{ $p2 := index .Site.RegularPages 1 }}
{{ $wp1 := index $weightedPages 0 }}{{ $wp2 := index $weightedPages 1 }}

{{ $pages := slice $p1 $p2 }}
{{ $pageGroups := slice $cool $blue }}
{{ $weighted := slice $wp1 $wp2 }}

{{ printf "pages:%d:%T:%v/%v" (len $pages) $pages (index $pages 0) (index $pages 1) }}
{{ printf "pageGroups:%d:%T:%v/%v" (len $pageGroups) $pageGroups (index (index $pageGroups 0).Pages 0) (index (index $pageGroups 1).Pages 0)}}
{{ printf "weightedPages:%d::%T:%v" (len $weighted) $weighted $weighted | safeHTML }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html",
		"pages:2:page.Pages:Page(/page1.md)/Page(/page2.md)",
		"pageGroups:2:page.PagesGroup:Page(/page1.md)/Page(/page2.md)",
		`weightedPages:2::page.WeightedPages:[WeightedPage(10,"Page") WeightedPage(20,"Page")]`)
}

func TestUnionFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20),
			"page3.md", fmt.Sprintf(pageContent, 30)).
		WithTemplatesAdded("index.html", `
{{ $unionPages := first 2 .Site.RegularPages | union .Site.RegularPages  }}
{{ $unionWeightedPages := .Site.Taxonomies.tags.blue | union .Site.Taxonomies.tags.green }}
{{ printf "unionPages: %T %d" $unionPages (len $unionPages) }}
{{ printf "unionWeightedPages: %T %d" $unionWeightedPages (len $unionWeightedPages) }}
`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 3)

	b.AssertFileContent("public/index.html",
		"unionPages: page.Pages 3",
		"unionWeightedPages: page.WeightedPages 6")
}

func TestAppendFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20)).
		WithTemplatesAdded("index.html", `

{{ $p1 := index .Site.RegularPages 0 }}{{ $p2 := index .Site.RegularPages 1 }}

{{ $pages := slice }}

{{ if true }}
	{{ $pages = $pages | append $p2 $p1 }}
{{ end }}
{{ $appendPages := .Site.Pages | append .Site.RegularPages }}
{{ $appendStrings := slice "a" "b" | append "c" "d" "e" }}
{{ $appendStringsSlice := slice "a" "b" "c" | append (slice "c" "d") }}

{{ printf "pages:%d:%T:%v/%v" (len $pages) $pages (index $pages 0) (index $pages 1)  }}
{{ printf "appendPages:%d:%T:%v/%v" (len $appendPages) $appendPages (index $appendPages 0).Kind (index $appendPages 8).Kind  }}
{{ printf "appendStrings:%T:%v"  $appendStrings $appendStrings  }}
{{ printf "appendStringsSlice:%T:%v"  $appendStringsSlice $appendStringsSlice }}

{{/* add some slightly related funcs to check what types we get */}}
{{ $u :=  $appendStrings | union $appendStringsSlice }}
{{ $i :=  $appendStrings | intersect $appendStringsSlice }}
{{ printf "union:%T:%v" $u $u  }}
{{ printf "intersect:%T:%v" $i $i }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html",
		"pages:2:page.Pages:Page(/page2.md)/Page(/page1.md)",
		"appendPages:9:page.Pages:home/page",
		"appendStrings:[]string:[a b c d e]",
		"appendStringsSlice:[]string:[a b c c d]",
		"union:[]string:[a b c d e]",
		"intersect:[]string:[a b c d]",
	)
}
<commit_msg>hugolib: Add some integration tests for in/uniq using Pages<commit_after>// Copyright 2019 The Hugo Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package hugolib

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestGroupFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", pageContent, "page2.md", pageContent).
		WithTemplatesAdded("index.html", `
{{ $cool := .Site.RegularPages | group "cool" }}
{{ $cool.Key }}: {{ len $cool.Pages }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html", "cool: 2")
}

func TestSliceFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20)).
		WithTemplatesAdded("index.html", `
{{ $cool := first 1 .Site.RegularPages | group "cool" }}
{{ $blue := after 1 .Site.RegularPages | group "blue" }}
{{ $weightedPages := index (index .Site.Taxonomies "tags") "blue" }}

{{ $p1 := index .Site.RegularPages 0 }}{{ $p2 := index .Site.RegularPages 1 }}
{{ $wp1 := index $weightedPages 0 }}{{ $wp2 := index $weightedPages 1 }}

{{ $pages := slice $p1 $p2 }}
{{ $pageGroups := slice $cool $blue }}
{{ $weighted := slice $wp1 $wp2 }}

{{ printf "pages:%d:%T:%v/%v" (len $pages) $pages (index $pages 0) (index $pages 1) }}
{{ printf "pageGroups:%d:%T:%v/%v" (len $pageGroups) $pageGroups (index (index $pageGroups 0).Pages 0) (index (index $pageGroups 1).Pages 0)}}
{{ printf "weightedPages:%d::%T:%v" (len $weighted) $weighted $weighted | safeHTML }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html",
		"pages:2:page.Pages:Page(/page1.md)/Page(/page2.md)",
		"pageGroups:2:page.PagesGroup:Page(/page1.md)/Page(/page2.md)",
		`weightedPages:2::page.WeightedPages:[WeightedPage(10,"Page") WeightedPage(20,"Page")]`)
}

func TestUnionFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20),
			"page3.md", fmt.Sprintf(pageContent, 30)).
		WithTemplatesAdded("index.html", `
{{ $unionPages := first 2 .Site.RegularPages | union .Site.RegularPages  }}
{{ $unionWeightedPages := .Site.Taxonomies.tags.blue | union .Site.Taxonomies.tags.green }}
{{ printf "unionPages: %T %d" $unionPages (len $unionPages) }}
{{ printf "unionWeightedPages: %T %d" $unionWeightedPages (len $unionWeightedPages) }}
`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 3)

	b.AssertFileContent("public/index.html",
		"unionPages: page.Pages 3",
		"unionWeightedPages: page.WeightedPages 6")
}

func TestCollectionsFuncs(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20),
			"page3.md", fmt.Sprintf(pageContent, 30)).
		WithTemplatesAdded("index.html", `
{{ $uniqPages := first 2 .Site.RegularPages | append .Site.RegularPages | uniq  }}
{{ $inTrue := in .Site.RegularPages (index .Site.RegularPages 1)  }}
{{ $inFalse := in .Site.RegularPages (.Site.Home)  }}

{{ printf "uniqPages: %T %d" $uniqPages (len $uniqPages) }}
{{ printf "inTrue: %t" $inTrue }}
{{ printf "inFalse: %t" $inFalse  }}
`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 3)

	b.AssertFileContent("public/index.html",
		"uniqPages: page.Pages 3",
		"inTrue: true",
		"inFalse: false",
	)
}

func TestAppendFunc(t *testing.T) {
	assert := require.New(t)

	pageContent := `
---
title: "Page"
tags: ["blue", "green"]
tags_weight: %d
---

`
	b := newTestSitesBuilder(t)
	b.WithSimpleConfigFile().
		WithContent("page1.md", fmt.Sprintf(pageContent, 10), "page2.md", fmt.Sprintf(pageContent, 20)).
		WithTemplatesAdded("index.html", `

{{ $p1 := index .Site.RegularPages 0 }}{{ $p2 := index .Site.RegularPages 1 }}

{{ $pages := slice }}

{{ if true }}
	{{ $pages = $pages | append $p2 $p1 }}
{{ end }}
{{ $appendPages := .Site.Pages | append .Site.RegularPages }}
{{ $appendStrings := slice "a" "b" | append "c" "d" "e" }}
{{ $appendStringsSlice := slice "a" "b" "c" | append (slice "c" "d") }}

{{ printf "pages:%d:%T:%v/%v" (len $pages) $pages (index $pages 0) (index $pages 1)  }}
{{ printf "appendPages:%d:%T:%v/%v" (len $appendPages) $appendPages (index $appendPages 0).Kind (index $appendPages 8).Kind  }}
{{ printf "appendStrings:%T:%v"  $appendStrings $appendStrings  }}
{{ printf "appendStringsSlice:%T:%v"  $appendStringsSlice $appendStringsSlice }}

{{/* add some slightly related funcs to check what types we get */}}
{{ $u :=  $appendStrings | union $appendStringsSlice }}
{{ $i :=  $appendStrings | intersect $appendStringsSlice }}
{{ printf "union:%T:%v" $u $u  }}
{{ printf "intersect:%T:%v" $i $i }}

`)
	b.CreateSites().Build(BuildCfg{})

	assert.Equal(1, len(b.H.Sites))
	require.Len(t, b.H.Sites[0].RegularPages(), 2)

	b.AssertFileContent("public/index.html",
		"pages:2:page.Pages:Page(/page2.md)/Page(/page1.md)",
		"appendPages:9:page.Pages:home/page",
		"appendStrings:[]string:[a b c d e]",
		"appendStringsSlice:[]string:[a b c c d]",
		"union:[]string:[a b c d e]",
		"intersect:[]string:[a b c d]",
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package nodeports

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/scheduler/framework"
)

// NodePorts is a plugin that checks if a node has free ports for the requested pod ports.
type NodePorts struct{}

var _ framework.PreFilterPlugin = &NodePorts{}
var _ framework.FilterPlugin = &NodePorts{}

const (
	// Name is the name of the plugin used in the plugin registry and configurations.
	Name = "NodePorts"

	// preFilterStateKey is the key in CycleState to NodePorts pre-computed data.
	// Using the name of the plugin will likely help us avoid collisions with other plugins.
	preFilterStateKey = "PreFilter" + Name

	// ErrReason when node ports aren't available.
	ErrReason = "node(s) didn't have free ports for the requested pod ports"
)

type preFilterState []*v1.ContainerPort

// Clone the prefilter state.
func (s preFilterState) Clone() framework.StateData {
	// The state is not impacted by adding/removing existing pods, hence we don't need to make a deep copy.
	return s
}

// Name returns name of the plugin. It is used in logs, etc.
func (pl *NodePorts) Name() string {
	return Name
}

// getContainerPorts returns the used host ports of Pods: if 'port' was used, a 'port:true' pair
// will be in the result; but it does not resolve port conflict.
func getContainerPorts(pods ...*v1.Pod) []*v1.ContainerPort {
	ports := []*v1.ContainerPort{}
	for _, pod := range pods {
		for j := range pod.Spec.Containers {
			container := &pod.Spec.Containers[j]
			for k := range container.Ports {
				ports = append(ports, &container.Ports[k])
			}
		}
	}
	return ports
}

// PreFilter invoked at the prefilter extension point.
func (pl *NodePorts) PreFilter(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod) *framework.Status {
	s := getContainerPorts(pod)
	cycleState.Write(preFilterStateKey, preFilterState(s))
	return nil
}

// PreFilterExtensions do not exist for this plugin.
func (pl *NodePorts) PreFilterExtensions() framework.PreFilterExtensions {
	return nil
}

func getPreFilterState(cycleState *framework.CycleState) (preFilterState, error) {
	c, err := cycleState.Read(preFilterStateKey)
	if err != nil {
		// preFilterState doesn't exist, likely PreFilter wasn't invoked.
		return nil, fmt.Errorf("reading %q from cycleState: %w", preFilterStateKey, err)
	}

	s, ok := c.(preFilterState)
	if !ok {
		return nil, fmt.Errorf("%+v  convert to nodeports.preFilterState error", c)
	}
	return s, nil
}

// Filter invoked at the filter extension point.
func (pl *NodePorts) Filter(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {
	wantPorts, err := getPreFilterState(cycleState)
	if err != nil {
		return framework.AsStatus(err)
	}

	fits := fitsPorts(wantPorts, nodeInfo)
	if !fits {
		return framework.NewStatus(framework.Unschedulable, ErrReason)
	}

	return nil
}

// Fits checks if the pod fits the node.
func Fits(pod *v1.Pod, nodeInfo *framework.NodeInfo) bool {
	return fitsPorts(getContainerPorts(pod), nodeInfo)
}

func fitsPorts(wantPorts []*v1.ContainerPort, nodeInfo *framework.NodeInfo) bool {
	// try to see whether existingPorts and wantPorts will conflict or not
	existingPorts := nodeInfo.UsedPorts
	for _, cp := range wantPorts {
		if existingPorts.CheckConflict(cp.HostIP, string(cp.Protocol), cp.HostPort) {
			return false
		}
	}
	return true
}

// New initializes a new plugin and returns it.
func New(_ runtime.Object, _ framework.Handle) (framework.Plugin, error) {
	return &NodePorts{}, nil
}
<commit_msg>Implemented function EventsToRegister for NodePorts<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package nodeports

import (
	"context"
	"fmt"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/scheduler/framework"
)

// NodePorts is a plugin that checks if a node has free ports for the requested pod ports.
type NodePorts struct{}

var _ framework.PreFilterPlugin = &NodePorts{}
var _ framework.FilterPlugin = &NodePorts{}
var _ framework.EnqueueExtensions = &NodePorts{}

const (
	// Name is the name of the plugin used in the plugin registry and configurations.
	Name = "NodePorts"

	// preFilterStateKey is the key in CycleState to NodePorts pre-computed data.
	// Using the name of the plugin will likely help us avoid collisions with other plugins.
	preFilterStateKey = "PreFilter" + Name

	// ErrReason when node ports aren't available.
	ErrReason = "node(s) didn't have free ports for the requested pod ports"
)

type preFilterState []*v1.ContainerPort

// Clone the prefilter state.
func (s preFilterState) Clone() framework.StateData {
	// The state is not impacted by adding/removing existing pods, hence we don't need to make a deep copy.
	return s
}

// Name returns name of the plugin. It is used in logs, etc.
func (pl *NodePorts) Name() string {
	return Name
}

// getContainerPorts returns the used host ports of Pods: if 'port' was used, a 'port:true' pair
// will be in the result; but it does not resolve port conflict.
func getContainerPorts(pods ...*v1.Pod) []*v1.ContainerPort {
	ports := []*v1.ContainerPort{}
	for _, pod := range pods {
		for j := range pod.Spec.Containers {
			container := &pod.Spec.Containers[j]
			for k := range container.Ports {
				ports = append(ports, &container.Ports[k])
			}
		}
	}
	return ports
}

// PreFilter invoked at the prefilter extension point.
func (pl *NodePorts) PreFilter(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod) *framework.Status {
	s := getContainerPorts(pod)
	cycleState.Write(preFilterStateKey, preFilterState(s))
	return nil
}

// PreFilterExtensions do not exist for this plugin.
func (pl *NodePorts) PreFilterExtensions() framework.PreFilterExtensions {
	return nil
}

func getPreFilterState(cycleState *framework.CycleState) (preFilterState, error) {
	c, err := cycleState.Read(preFilterStateKey)
	if err != nil {
		// preFilterState doesn't exist, likely PreFilter wasn't invoked.
		return nil, fmt.Errorf("reading %q from cycleState: %w", preFilterStateKey, err)
	}

	s, ok := c.(preFilterState)
	if !ok {
		return nil, fmt.Errorf("%+v  convert to nodeports.preFilterState error", c)
	}
	return s, nil
}

// EventsToRegister returns the possible events that may make a Pod
// failed by this plugin schedulable.
func (pl *NodePorts) EventsToRegister() []framework.ClusterEvent {
	return []framework.ClusterEvent{
		// Due to immutable fields `spec.containers[*].ports`, pod update events are ignored.
		{Resource: framework.Pod, ActionType: framework.Delete},
		{Resource: framework.Node, ActionType: framework.Add},
	}
}

// Filter invoked at the filter extension point.
func (pl *NodePorts) Filter(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {
	wantPorts, err := getPreFilterState(cycleState)
	if err != nil {
		return framework.AsStatus(err)
	}

	fits := fitsPorts(wantPorts, nodeInfo)
	if !fits {
		return framework.NewStatus(framework.Unschedulable, ErrReason)
	}

	return nil
}

// Fits checks if the pod fits the node.
func Fits(pod *v1.Pod, nodeInfo *framework.NodeInfo) bool {
	return fitsPorts(getContainerPorts(pod), nodeInfo)
}

func fitsPorts(wantPorts []*v1.ContainerPort, nodeInfo *framework.NodeInfo) bool {
	// try to see whether existingPorts and wantPorts will conflict or not
	existingPorts := nodeInfo.UsedPorts
	for _, cp := range wantPorts {
		if existingPorts.CheckConflict(cp.HostIP, string(cp.Protocol), cp.HostPort) {
			return false
		}
	}
	return true
}

// New initializes a new plugin and returns it.
func New(_ runtime.Object, _ framework.Handle) (framework.Plugin, error) {
	return &NodePorts{}, nil
}
<|endoftext|>


<commit_before>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/bitly/go-simplejson"
	"github.com/pkg/errors"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"

	gocontext "context"
)

// HTTPJobQueue is a JobQueue that uses http
type HTTPJobQueue struct {
	processorPool *ProcessorPool
	jobBoardURL   *url.URL
	site          string
	providerName  string
	queue         string
	workerID      string
	buildJobChan  chan Job

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

type httpFetchJobsRequest struct {
	Jobs []string `json:"jobs"`
}

type httpFetchJobsResponse struct {
	Jobs []string `json:"jobs"`
}

type jobBoardErrorResponse struct {
	Type          string `json:"@type"`
	Error         string `json:"error"`
	UpstreamError string `json:"upstream_error,omitempty"`
}

// NewHTTPJobQueue creates a new job-board job queue
func NewHTTPJobQueue(pool *ProcessorPool, jobBoardURL *url.URL, site, providerName, queue, workerID string) (*HTTPJobQueue, error) {
	return &HTTPJobQueue{
		processorPool: pool,
		jobBoardURL:   jobBoardURL,
		site:          site,
		providerName:  providerName,
		queue:         queue,
		workerID:      workerID,
	}, nil
}

// Jobs consumes new jobs from job-board
func (q *HTTPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	logger := context.LoggerFromContext(ctx)
	if q.buildJobChan != nil {
		return q.buildJobChan, nil
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		for {
			logger.Debug("fetching job ids")
			jobIds, err := q.fetchJobs(ctx)
			if err != nil {
				logger.WithField("err", err).Warn("continuing after failing to get job ids")
				time.Sleep(time.Second)
				continue
			}
			for _, id := range jobIds {
				go func(id uint64) {
					logger.WithField("job_id", id).Debug("fetching complete job")
					buildJob, err := q.fetchJob(ctx, id)
					if err != nil {
						logger.WithField("err", err).Warn("breaking after failing to get complete job")
						return
					}
					logger.WithField("job", buildJob).Debug("sending job to output channel")
					buildJobChan <- buildJob
				}(id)
			}

			select {
			case <-time.After(time.Second):
				logger.Debug("jobs loop again after 1s sleep")
				continue
			case <-ctx.Done():
				logger.WithField("err", ctx.Err()).Warn("returning from jobs loop due to context done")
				q.buildJobChan = nil
				return
			}
		}
	}()

	q.buildJobChan = buildJobChan
	return outChan, nil
}

func (q *HTTPJobQueue) fetchJobs(ctx gocontext.Context) ([]uint64, error) {
	logger := context.LoggerFromContext(ctx)
	fetchRequestPayload := &httpFetchJobsRequest{Jobs: []string{}}
	numWaiting := 0
	q.processorPool.Each(func(i int, p *Processor) {
		switch p.CurrentStatus {
		case "processing":
			fetchRequestPayload.Jobs = append(fetchRequestPayload.Jobs, strconv.FormatUint(p.LastJobID, 10))
		case "waiting", "new":
			numWaiting++
		}
	})

	jobIdsJSON, err := json.Marshal(fetchRequestPayload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal job board jobs request payload")
	}

	u := *q.jobBoardURL

	query := u.Query()
	query.Add("count", strconv.Itoa(numWaiting))
	query.Add("queue", q.queue)

	u.Path = "/jobs"
	u.RawQuery = query.Encode()

	client := &http.Client{}

	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(jobIdsJSON))
	if err != nil {
		return nil, errors.Wrap(err, "failed to create job board jobs request")
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", q.workerID)
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "failed to make job board jobs request")
	}

	fetchResponsePayload := &httpFetchJobsResponse{}
	err = json.NewDecoder(resp.Body).Decode(&fetchResponsePayload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to decode job board jobs response")
	}

	var jobIds []uint64
	for _, strID := range fetchResponsePayload.Jobs {
		alreadyRunning := false
		for _, prevStrID := range fetchRequestPayload.Jobs {
			if strID == prevStrID {
				alreadyRunning = true
			}
		}
		if alreadyRunning {
			logger.WithField("job_id", strID).Debug("skipping running job")
			continue
		}

		id, err := strconv.ParseUint(strID, 10, 64)
		if err != nil {
			return nil, errors.Wrap(err, "failed to parse job ID")
		}
		jobIds = append(jobIds, id)
	}

	return jobIds, nil
}

func (q *HTTPJobQueue) fetchJob(ctx gocontext.Context, id uint64) (Job, error) {
	buildJob := &httpJob{
		payload: &httpJobPayload{
			Data: &JobPayload{},
		},
		startAttributes: &backend.StartAttributes{},

		jobBoardURL: q.jobBoardURL,
		site:        q.site,
		workerID:    q.workerID,
	}
	startAttrs := &httpJobPayloadStartAttrs{
		Data: &jobPayloadStartAttrs{
			Config: &backend.StartAttributes{},
		},
	}

	u := *q.jobBoardURL
	u.Path = fmt.Sprintf("/jobs/%d", id)

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, errors.Wrap(err, "couldn't make job board job request")
	}

	// TODO: ensure infrastructure is not synonymous with providerName since
	// there's the possibility that a provider has multiple infrastructures, which
	// is expected to be the case with the future cloudbrain provider.
	req.Header.Add("Travis-Infrastructure", q.providerName)
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", q.workerID)
	req = req.WithContext(ctx)

	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "error making job board job request")
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(err, "error reading body from job board job request")
	}

	if resp.StatusCode != http.StatusOK {
		var errorResp jobBoardErrorResponse
		err := json.Unmarshal(body, &errorResp)
		if err != nil {
			return nil, errors.Wrapf(err, "job board job fetch request errored with status %d and didn't send an error response", resp.StatusCode)
		}

		return nil, errors.Errorf("job board job fetch request errored with status %d: %s", resp.StatusCode, errorResp.Error)
	}

	err = json.Unmarshal(body, buildJob.payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal job board payload")
	}

	err = json.Unmarshal(body, &startAttrs)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal start attributes from job board")
	}

	rawPayload, err := simplejson.NewJson(body)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse raw payload with simplejson")
	}
	buildJob.rawPayload = rawPayload.Get("data")

	buildJob.startAttributes = startAttrs.Data.Config
	buildJob.startAttributes.VMType = buildJob.payload.Data.VMType
	buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)

	return buildJob, nil
}

// Cleanup does not do anything!
func (q *HTTPJobQueue) Cleanup() error {
	return nil
}
<commit_msg>Add more debug logging around HTTP job ID fetching<commit_after>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/bitly/go-simplejson"
	"github.com/pkg/errors"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"

	gocontext "context"
)

// HTTPJobQueue is a JobQueue that uses http
type HTTPJobQueue struct {
	processorPool *ProcessorPool
	jobBoardURL   *url.URL
	site          string
	providerName  string
	queue         string
	workerID      string
	buildJobChan  chan Job

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

type httpFetchJobsRequest struct {
	Jobs []string `json:"jobs"`
}

type httpFetchJobsResponse struct {
	Jobs []string `json:"jobs"`
}

type jobBoardErrorResponse struct {
	Type          string `json:"@type"`
	Error         string `json:"error"`
	UpstreamError string `json:"upstream_error,omitempty"`
}

// NewHTTPJobQueue creates a new job-board job queue
func NewHTTPJobQueue(pool *ProcessorPool, jobBoardURL *url.URL, site, providerName, queue, workerID string) (*HTTPJobQueue, error) {
	return &HTTPJobQueue{
		processorPool: pool,
		jobBoardURL:   jobBoardURL,
		site:          site,
		providerName:  providerName,
		queue:         queue,
		workerID:      workerID,
	}, nil
}

// Jobs consumes new jobs from job-board
func (q *HTTPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	logger := context.LoggerFromContext(ctx)
	if q.buildJobChan != nil {
		return q.buildJobChan, nil
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		for {
			logger.Debug("fetching job ids")
			jobIds, err := q.fetchJobs(ctx)
			if err != nil {
				logger.WithField("err", err).Warn("continuing after failing to get job ids")
				time.Sleep(time.Second)
				continue
			}
			for _, id := range jobIds {
				go func(id uint64) {
					logger.WithField("job_id", id).Debug("fetching complete job")
					buildJob, err := q.fetchJob(ctx, id)
					if err != nil {
						logger.WithField("err", err).Warn("breaking after failing to get complete job")
						return
					}
					logger.WithField("job", buildJob).Debug("sending job to output channel")
					buildJobChan <- buildJob
				}(id)
			}

			select {
			case <-time.After(time.Second):
				logger.Debug("jobs loop again after 1s sleep")
				continue
			case <-ctx.Done():
				logger.WithField("err", ctx.Err()).Warn("returning from jobs loop due to context done")
				q.buildJobChan = nil
				return
			}
		}
	}()

	q.buildJobChan = buildJobChan
	return outChan, nil
}

func (q *HTTPJobQueue) fetchJobs(ctx gocontext.Context) ([]uint64, error) {
	logger := context.LoggerFromContext(ctx)
	fetchRequestPayload := &httpFetchJobsRequest{Jobs: []string{}}
	numWaiting := 0
	q.processorPool.Each(func(i int, p *Processor) {
		switch p.CurrentStatus {
		case "processing":
			fetchRequestPayload.Jobs = append(fetchRequestPayload.Jobs, strconv.FormatUint(p.LastJobID, 10))
		case "waiting", "new":
			numWaiting++
		}
	})

	jobIdsJSON, err := json.Marshal(fetchRequestPayload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to marshal job board jobs request payload")
	}

	u := *q.jobBoardURL

	query := u.Query()
	query.Add("count", strconv.Itoa(numWaiting))
	query.Add("queue", q.queue)

	u.Path = "/jobs"
	u.RawQuery = query.Encode()

	client := &http.Client{}

	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(jobIdsJSON))
	if err != nil {
		return nil, errors.Wrap(err, "failed to create job board jobs request")
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", q.workerID)
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "failed to make job board jobs request")
	}

	fetchResponsePayload := &httpFetchJobsResponse{}
	err = json.NewDecoder(resp.Body).Decode(&fetchResponsePayload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to decode job board jobs response")
	}

	logger.WithField("jobs", fetchResponsePayload.Jobs).Debug("fetched raw jobs")
	var jobIds []uint64
	for _, strID := range fetchResponsePayload.Jobs {
		alreadyRunning := false
		for _, prevStrID := range fetchRequestPayload.Jobs {
			if strID == prevStrID {
				alreadyRunning = true
			}
		}
		if alreadyRunning {
			logger.WithField("job_id", strID).Debug("skipping running job")
			continue
		}

		id, err := strconv.ParseUint(strID, 10, 64)
		if err != nil {
			return nil, errors.Wrap(err, "failed to parse job ID")
		}
		jobIds = append(jobIds, id)
	}

	logger.WithField("jobs", jobIds).Debug("returning filtered job IDs")
	return jobIds, nil
}

func (q *HTTPJobQueue) fetchJob(ctx gocontext.Context, id uint64) (Job, error) {
	buildJob := &httpJob{
		payload: &httpJobPayload{
			Data: &JobPayload{},
		},
		startAttributes: &backend.StartAttributes{},

		jobBoardURL: q.jobBoardURL,
		site:        q.site,
		workerID:    q.workerID,
	}
	startAttrs := &httpJobPayloadStartAttrs{
		Data: &jobPayloadStartAttrs{
			Config: &backend.StartAttributes{},
		},
	}

	u := *q.jobBoardURL
	u.Path = fmt.Sprintf("/jobs/%d", id)

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, errors.Wrap(err, "couldn't make job board job request")
	}

	// TODO: ensure infrastructure is not synonymous with providerName since
	// there's the possibility that a provider has multiple infrastructures, which
	// is expected to be the case with the future cloudbrain provider.
	req.Header.Add("Travis-Infrastructure", q.providerName)
	req.Header.Add("Travis-Site", q.site)
	req.Header.Add("From", q.workerID)
	req = req.WithContext(ctx)

	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "error making job board job request")
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(err, "error reading body from job board job request")
	}

	if resp.StatusCode != http.StatusOK {
		var errorResp jobBoardErrorResponse
		err := json.Unmarshal(body, &errorResp)
		if err != nil {
			return nil, errors.Wrapf(err, "job board job fetch request errored with status %d and didn't send an error response", resp.StatusCode)
		}

		return nil, errors.Errorf("job board job fetch request errored with status %d: %s", resp.StatusCode, errorResp.Error)
	}

	err = json.Unmarshal(body, buildJob.payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal job board payload")
	}

	err = json.Unmarshal(body, &startAttrs)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal start attributes from job board")
	}

	rawPayload, err := simplejson.NewJson(body)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse raw payload with simplejson")
	}
	buildJob.rawPayload = rawPayload.Get("data")

	buildJob.startAttributes = startAttrs.Data.Config
	buildJob.startAttributes.VMType = buildJob.payload.Data.VMType
	buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)

	return buildJob, nil
}

// Cleanup does not do anything!
func (q *HTTPJobQueue) Cleanup() error {
	return nil
}
<|endoftext|>


<commit_before>package portworx

import (
	"fmt"
	"strings"

	crdv1 "github.com/kubernetes-incubator/external-storage/snapshot/pkg/apis/crd/v1"
	"github.com/kubernetes-incubator/external-storage/snapshot/pkg/controller/snapshotter"
	snapshotVolume "github.com/kubernetes-incubator/external-storage/snapshot/pkg/volume"
	"github.com/libopenstorage/openstorage/api"
	clusterclient "github.com/libopenstorage/openstorage/api/client/cluster"
	volumeclient "github.com/libopenstorage/openstorage/api/client/volume"
	"github.com/libopenstorage/openstorage/cluster"
	"github.com/libopenstorage/openstorage/volume"
	storkvolume "github.com/libopenstorage/stork/drivers/volume"
	"github.com/libopenstorage/stork/pkg/errors"
	"github.com/libopenstorage/stork/pkg/k8sutils"
	"github.com/libopenstorage/stork/pkg/snapshot"
	"github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TODO: Make some of these configurable
const (
	// driverName is the name of the portworx driver implementation
	driverName = "pxd"

	// serviceName is the name of the portworx service
	serviceName = "portworx-service"

	// namespace is the kubernetes namespace in which portworx daemon set runs
	namespace = "kube-system"

	// provisionerName is the name for the driver provisioner
	provisionerName = "kubernetes.io/portworx-volume"

	// pvcProvisionerAnnotation is the annotation on PVC which has the provisioner name
	pvcProvisionerAnnotation = "volume.beta.kubernetes.io/storage-provisioner"

	// pvcNameLabel is the key of the label used to store the PVC name
	pvcNameLabel = "pvc"
)

type portworx struct {
	clusterManager cluster.Cluster
	volDriver      volume.VolumeDriver
}

func (p *portworx) String() string {
	return driverName
}

func (p *portworx) Init(_ interface{}) error {
	var endpoint string
	svc, err := k8sutils.GetService(serviceName, namespace)
	if err == nil {
		endpoint = svc.Spec.ClusterIP
	} else {
		return fmt.Errorf("Failed to get k8s service spec: %v", err)
	}

	if len(endpoint) == 0 {
		return fmt.Errorf("Failed to get endpoint for portworx volume driver")
	}

	logrus.Infof("Using %v as endpoint for portworx volume driver", endpoint)
	clnt, err := clusterclient.NewClusterClient("http://"+endpoint+":9001", "v1")
	if err != nil {
		return err
	}
	p.clusterManager = clusterclient.ClusterManager(clnt)

	clnt, err = volumeclient.NewDriverClient("http://"+endpoint+":9001", "pxd", "", "stork")
	if err != nil {
		return err
	}
	p.volDriver = volumeclient.VolumeDriver(clnt)
	return err
}

func (p *portworx) InspectVolume(volumeID string) (*storkvolume.Info, error) {
	vols, err := p.volDriver.Inspect([]string{volumeID})
	if err != nil {
		return nil, &ErrFailedToInspectVolume{
			ID:    volumeID,
			Cause: fmt.Sprintf("Volume inspect returned err: %v", err),
		}
	}

	if len(vols) == 0 {
		return nil, &errors.ErrNotFound{
			ID:   volumeID,
			Type: "Volume",
		}
	}

	info := &storkvolume.Info{}
	info.VolumeID = vols[0].Id
	info.VolumeName = vols[0].Locator.Name
	for _, rset := range vols[0].ReplicaSets {
		for _, node := range rset.Nodes {
			info.DataNodes = append(info.DataNodes, node)
		}
	}
	if vols[0].Source != nil {
		info.ParentID = vols[0].Source.Parent
	}
	return info, nil
}

func (p *portworx) mapNodeStatus(status api.Status) storkvolume.NodeStatus {
	switch status {
	case api.Status_STATUS_NONE:
		fallthrough
	case api.Status_STATUS_INIT:
		fallthrough
	case api.Status_STATUS_OFFLINE:
		fallthrough
	case api.Status_STATUS_ERROR:
		fallthrough
	case api.Status_STATUS_NOT_IN_QUORUM:
		fallthrough
	case api.Status_STATUS_DECOMMISSION:
		fallthrough
	case api.Status_STATUS_MAINTENANCE:
		fallthrough
	case api.Status_STATUS_NEEDS_REBOOT:
		return storkvolume.NodeOffline

	case api.Status_STATUS_OK:
		fallthrough
	case api.Status_STATUS_STORAGE_DOWN:
		return storkvolume.NodeOnline

	case api.Status_STATUS_STORAGE_DEGRADED:
		fallthrough
	case api.Status_STATUS_STORAGE_REBALANCE:
		fallthrough
	case api.Status_STATUS_STORAGE_DRIVE_REPLACE:
		return storkvolume.NodeDegraded
	default:
		return storkvolume.NodeOffline
	}
}

func (p *portworx) GetNodes() ([]*storkvolume.NodeInfo, error) {
	cluster, err := p.clusterManager.Enumerate()
	if err != nil {
		return nil, &ErrFailedToGetNodes{
			Cause: err.Error(),
		}
	}

	var nodes []*storkvolume.NodeInfo
	for _, n := range cluster.Nodes {
		nodeInfo := &storkvolume.NodeInfo{
			ID:       n.Id,
			Hostname: strings.ToLower(n.Hostname),
			Status:   p.mapNodeStatus(n.Status),
		}
		nodeInfo.IPs = append(nodeInfo.IPs, n.MgmtIp)
		nodeInfo.IPs = append(nodeInfo.IPs, n.DataIp)

		nodes = append(nodes, nodeInfo)
	}
	return nodes, nil
}

func (p *portworx) GetPodVolumes(pod *v1.Pod) ([]*storkvolume.Info, error) {
	var volumes []*storkvolume.Info
	for _, volume := range pod.Spec.Volumes {
		volumeName := ""
		if volume.PersistentVolumeClaim != nil {
			pvc, err := k8sutils.GetPVC(volume.PersistentVolumeClaim.ClaimName, pod.Namespace)
			if err != nil {
				return nil, err
			}

			storageClassName := k8sutils.GetStorageClassName(pvc)
			if storageClassName == "" {
				logrus.Debugf("Empty StorageClass in PVC %v for pod %v, ignoring",
					pvc.Name, pod.Name)
				continue
			}

			provisioner := ""
			// Try getting the provisioner from the Storage class. If that has been
			// deleted, check for the provisioner in the PVC annotation
			storageClass, err := k8sutils.GetStorageClass(storageClassName, pod.Namespace)
			if kerrors.IsNotFound(err) {
				if val, ok := pvc.Annotations[pvcProvisionerAnnotation]; ok {
					provisioner = val
				}
			} else if err != nil {
				return nil, err
			} else {
				provisioner = storageClass.Provisioner
			}

			if provisioner != provisionerName && provisioner != snapshotcontroller.GetProvisionerName() {
				logrus.Debugf("Provisioner in Storageclass not Portworx or from the snapshot Provisioner, ignoring")
				continue
			}

			if pvc.Status.Phase == v1.ClaimPending {
				return nil, &storkvolume.ErrPVCPending{
					Name: volume.PersistentVolumeClaim.ClaimName,
				}
			}
			volumeName = pvc.Spec.VolumeName
		} else if volume.PortworxVolume != nil {
			volumeName = volume.PortworxVolume.VolumeID
		}

		if volumeName != "" {
			volumeInfo, err := p.InspectVolume(volumeName)
			if err != nil {
				return nil, err
			}
			volumes = append(volumes, volumeInfo)
		}
	}
	return volumes, nil
}

func (p *portworx) GetSnapshotPlugin() snapshotVolume.Plugin {
	return p
}

func (p *portworx) SnapshotCreate(pv *v1.PersistentVolume, tags *map[string]string) (*crdv1.VolumeSnapshotDataSource, *[]crdv1.VolumeSnapshotCondition, error) {
	if pv == nil || pv.Spec.PortworxVolume == nil {
		return nil, nil, fmt.Errorf("Invalid PV: %v", pv)
	}
	spec := &pv.Spec
	volumeID := spec.PortworxVolume.VolumeID

	logrus.Debugf("SnapshotCreate for pv: %+v \n tags: %v", pv, tags)
	locator := &api.VolumeLocator{
		Name: (*tags)[snapshotter.CloudSnapshotCreatedForVolumeSnapshotNameTag],
	}
	snapshotID, err := p.volDriver.Snapshot(volumeID, true, locator)
	if err != nil {
		return nil, nil, err
	}

	return &crdv1.VolumeSnapshotDataSource{
		PortworxSnapshot: &crdv1.PortworxVolumeSnapshotSource{
			SnapshotID: snapshotID,
		},
	}, nil, nil
}

func (p *portworx) SnapshotDelete(snapshot *crdv1.VolumeSnapshotDataSource, _ *v1.PersistentVolume) error {
	if snapshot == nil || snapshot.PortworxSnapshot == nil {
		return fmt.Errorf("Invalid Snaphsot source %v", snapshot)
	}
	return p.volDriver.Delete(snapshot.PortworxSnapshot.SnapshotID)
}

func (p *portworx) SnapshotRestore(
	snapshotData *crdv1.VolumeSnapshotData,
	pvc *v1.PersistentVolumeClaim,
	pvName string,
	parameters map[string]string,
) (*v1.PersistentVolumeSource, map[string]string, error) {
	if snapshotData == nil || snapshotData.Spec.PortworxSnapshot == nil {
		return nil, nil, fmt.Errorf("Invalid Snapshot spec")
	}
	if pvc == nil {
		return nil, nil, fmt.Errorf("Invalid PVC spec")
	}

	snapID := snapshotData.Spec.PortworxSnapshot.SnapshotID

	logrus.Debugf("SnapshotRestore for pvc: %+v", pvc)
	locator := &api.VolumeLocator{
		Name: "pvc-" + string(pvc.UID),
		VolumeLabels: map[string]string{
			pvcNameLabel: pvc.Name,
		},
	}
	volumeID, err := p.volDriver.Snapshot(snapID, false, locator)
	if err != nil {
		return nil, nil, err
	}

	vols, err := p.volDriver.Inspect([]string{volumeID})
	if err != nil {
		return nil, nil, &ErrFailedToInspectVolume{
			ID:    volumeID,
			Cause: fmt.Sprintf("Volume inspect returned err: %v", err),
		}
	}

	if len(vols) == 0 {
		return nil, nil, &errors.ErrNotFound{
			ID:   volumeID,
			Type: "Volume",
		}
	}

	pv := &v1.PersistentVolumeSource{
		PortworxVolume: &v1.PortworxVolumeSource{
			VolumeID: volumeID,
			FSType:   vols[0].Format.String(),
			ReadOnly: vols[0].Readonly,
		},
	}

	labels := make(map[string]string)

	return pv, labels, nil
}

func (p *portworx) DescribeSnapshot(snapshotData *crdv1.VolumeSnapshotData) (*[]crdv1.VolumeSnapshotCondition, bool, error) {
	if snapshotData == nil || snapshotData.Spec.PortworxSnapshot == nil {
		return nil, false, fmt.Errorf("Invalid VolumeSnapshotDataSource: %v", snapshotData)
	}
	snapshotID := snapshotData.Spec.PortworxSnapshot.SnapshotID
	_, err := p.InspectVolume(snapshotID)
	if err != nil {
		return nil, false, err
	}
	var snapConditions []crdv1.VolumeSnapshotCondition
	snapConditions = []crdv1.VolumeSnapshotCondition{
		{
			Type:               crdv1.VolumeSnapshotConditionReady,
			Status:             v1.ConditionTrue,
			Message:            "Snapshot created successfully and it is ready",
			LastTransitionTime: metav1.Now(),
		},
	}
	return &snapConditions, true, err
}

// TODO: Implement FindSnapshot
func (p *portworx) FindSnapshot(tags *map[string]string) (*crdv1.VolumeSnapshotDataSource, *[]crdv1.VolumeSnapshotCondition, error) {
	return nil, nil, nil
}

func (p *portworx) VolumeDelete(pv *v1.PersistentVolume) error {
	if pv == nil || pv.Spec.PortworxVolume == nil {
		return fmt.Errorf("Invalid PV: %v", pv)
	}
	return p.volDriver.Delete(pv.Spec.PortworxVolume.VolumeID)
}

func init() {
	if err := storkvolume.Register(driverName, &portworx{}); err != nil {
		logrus.Panicf("Error registering portworx volume driver: %v", err)
	}
}
<commit_msg>Add pvc namespace label to the volume on restoring snapshot<commit_after>package portworx

import (
	"fmt"
	"strings"

	crdv1 "github.com/kubernetes-incubator/external-storage/snapshot/pkg/apis/crd/v1"
	"github.com/kubernetes-incubator/external-storage/snapshot/pkg/controller/snapshotter"
	snapshotVolume "github.com/kubernetes-incubator/external-storage/snapshot/pkg/volume"
	"github.com/libopenstorage/openstorage/api"
	clusterclient "github.com/libopenstorage/openstorage/api/client/cluster"
	volumeclient "github.com/libopenstorage/openstorage/api/client/volume"
	"github.com/libopenstorage/openstorage/cluster"
	"github.com/libopenstorage/openstorage/volume"
	storkvolume "github.com/libopenstorage/stork/drivers/volume"
	"github.com/libopenstorage/stork/pkg/errors"
	"github.com/libopenstorage/stork/pkg/k8sutils"
	"github.com/libopenstorage/stork/pkg/snapshot"
	"github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TODO: Make some of these configurable
const (
	// driverName is the name of the portworx driver implementation
	driverName = "pxd"

	// serviceName is the name of the portworx service
	serviceName = "portworx-service"

	// namespace is the kubernetes namespace in which portworx daemon set runs
	namespace = "kube-system"

	// provisionerName is the name for the driver provisioner
	provisionerName = "kubernetes.io/portworx-volume"

	// pvcProvisionerAnnotation is the annotation on PVC which has the provisioner name
	pvcProvisionerAnnotation = "volume.beta.kubernetes.io/storage-provisioner"

	// pvcNameLabel is the key of the label used to store the PVC name
	pvcNameLabel = "pvc"

	// pvcNamespaceLabel is the key of the label used to store the PVC namespace
	pvcNamespaceLabel = "namespace"
)

type portworx struct {
	clusterManager cluster.Cluster
	volDriver      volume.VolumeDriver
}

func (p *portworx) String() string {
	return driverName
}

func (p *portworx) Init(_ interface{}) error {
	var endpoint string
	svc, err := k8sutils.GetService(serviceName, namespace)
	if err == nil {
		endpoint = svc.Spec.ClusterIP
	} else {
		return fmt.Errorf("Failed to get k8s service spec: %v", err)
	}

	if len(endpoint) == 0 {
		return fmt.Errorf("Failed to get endpoint for portworx volume driver")
	}

	logrus.Infof("Using %v as endpoint for portworx volume driver", endpoint)
	clnt, err := clusterclient.NewClusterClient("http://"+endpoint+":9001", "v1")
	if err != nil {
		return err
	}
	p.clusterManager = clusterclient.ClusterManager(clnt)

	clnt, err = volumeclient.NewDriverClient("http://"+endpoint+":9001", "pxd", "", "stork")
	if err != nil {
		return err
	}
	p.volDriver = volumeclient.VolumeDriver(clnt)
	return err
}

func (p *portworx) InspectVolume(volumeID string) (*storkvolume.Info, error) {
	vols, err := p.volDriver.Inspect([]string{volumeID})
	if err != nil {
		return nil, &ErrFailedToInspectVolume{
			ID:    volumeID,
			Cause: fmt.Sprintf("Volume inspect returned err: %v", err),
		}
	}

	if len(vols) == 0 {
		return nil, &errors.ErrNotFound{
			ID:   volumeID,
			Type: "Volume",
		}
	}

	info := &storkvolume.Info{}
	info.VolumeID = vols[0].Id
	info.VolumeName = vols[0].Locator.Name
	for _, rset := range vols[0].ReplicaSets {
		for _, node := range rset.Nodes {
			info.DataNodes = append(info.DataNodes, node)
		}
	}
	if vols[0].Source != nil {
		info.ParentID = vols[0].Source.Parent
	}
	return info, nil
}

func (p *portworx) mapNodeStatus(status api.Status) storkvolume.NodeStatus {
	switch status {
	case api.Status_STATUS_NONE:
		fallthrough
	case api.Status_STATUS_INIT:
		fallthrough
	case api.Status_STATUS_OFFLINE:
		fallthrough
	case api.Status_STATUS_ERROR:
		fallthrough
	case api.Status_STATUS_NOT_IN_QUORUM:
		fallthrough
	case api.Status_STATUS_DECOMMISSION:
		fallthrough
	case api.Status_STATUS_MAINTENANCE:
		fallthrough
	case api.Status_STATUS_NEEDS_REBOOT:
		return storkvolume.NodeOffline

	case api.Status_STATUS_OK:
		fallthrough
	case api.Status_STATUS_STORAGE_DOWN:
		return storkvolume.NodeOnline

	case api.Status_STATUS_STORAGE_DEGRADED:
		fallthrough
	case api.Status_STATUS_STORAGE_REBALANCE:
		fallthrough
	case api.Status_STATUS_STORAGE_DRIVE_REPLACE:
		return storkvolume.NodeDegraded
	default:
		return storkvolume.NodeOffline
	}
}

func (p *portworx) GetNodes() ([]*storkvolume.NodeInfo, error) {
	cluster, err := p.clusterManager.Enumerate()
	if err != nil {
		return nil, &ErrFailedToGetNodes{
			Cause: err.Error(),
		}
	}

	var nodes []*storkvolume.NodeInfo
	for _, n := range cluster.Nodes {
		nodeInfo := &storkvolume.NodeInfo{
			ID:       n.Id,
			Hostname: strings.ToLower(n.Hostname),
			Status:   p.mapNodeStatus(n.Status),
		}
		nodeInfo.IPs = append(nodeInfo.IPs, n.MgmtIp)
		nodeInfo.IPs = append(nodeInfo.IPs, n.DataIp)

		nodes = append(nodes, nodeInfo)
	}
	return nodes, nil
}

func (p *portworx) GetPodVolumes(pod *v1.Pod) ([]*storkvolume.Info, error) {
	var volumes []*storkvolume.Info
	for _, volume := range pod.Spec.Volumes {
		volumeName := ""
		if volume.PersistentVolumeClaim != nil {
			pvc, err := k8sutils.GetPVC(volume.PersistentVolumeClaim.ClaimName, pod.Namespace)
			if err != nil {
				return nil, err
			}

			storageClassName := k8sutils.GetStorageClassName(pvc)
			if storageClassName == "" {
				logrus.Debugf("Empty StorageClass in PVC %v for pod %v, ignoring",
					pvc.Name, pod.Name)
				continue
			}

			provisioner := ""
			// Try getting the provisioner from the Storage class. If that has been
			// deleted, check for the provisioner in the PVC annotation
			storageClass, err := k8sutils.GetStorageClass(storageClassName, pod.Namespace)
			if kerrors.IsNotFound(err) {
				if val, ok := pvc.Annotations[pvcProvisionerAnnotation]; ok {
					provisioner = val
				}
			} else if err != nil {
				return nil, err
			} else {
				provisioner = storageClass.Provisioner
			}

			if provisioner != provisionerName && provisioner != snapshotcontroller.GetProvisionerName() {
				logrus.Debugf("Provisioner in Storageclass not Portworx or from the snapshot Provisioner, ignoring")
				continue
			}

			if pvc.Status.Phase == v1.ClaimPending {
				return nil, &storkvolume.ErrPVCPending{
					Name: volume.PersistentVolumeClaim.ClaimName,
				}
			}
			volumeName = pvc.Spec.VolumeName
		} else if volume.PortworxVolume != nil {
			volumeName = volume.PortworxVolume.VolumeID
		}

		if volumeName != "" {
			volumeInfo, err := p.InspectVolume(volumeName)
			if err != nil {
				return nil, err
			}
			volumes = append(volumes, volumeInfo)
		}
	}
	return volumes, nil
}

func (p *portworx) GetSnapshotPlugin() snapshotVolume.Plugin {
	return p
}

func (p *portworx) SnapshotCreate(pv *v1.PersistentVolume, tags *map[string]string) (*crdv1.VolumeSnapshotDataSource, *[]crdv1.VolumeSnapshotCondition, error) {
	if pv == nil || pv.Spec.PortworxVolume == nil {
		return nil, nil, fmt.Errorf("Invalid PV: %v", pv)
	}
	spec := &pv.Spec
	volumeID := spec.PortworxVolume.VolumeID

	logrus.Debugf("SnapshotCreate for pv: %+v \n tags: %v", pv, tags)
	locator := &api.VolumeLocator{
		Name: (*tags)[snapshotter.CloudSnapshotCreatedForVolumeSnapshotNameTag],
	}
	snapshotID, err := p.volDriver.Snapshot(volumeID, true, locator)
	if err != nil {
		return nil, nil, err
	}

	return &crdv1.VolumeSnapshotDataSource{
		PortworxSnapshot: &crdv1.PortworxVolumeSnapshotSource{
			SnapshotID: snapshotID,
		},
	}, nil, nil
}

func (p *portworx) SnapshotDelete(snapshot *crdv1.VolumeSnapshotDataSource, _ *v1.PersistentVolume) error {
	if snapshot == nil || snapshot.PortworxSnapshot == nil {
		return fmt.Errorf("Invalid Snaphsot source %v", snapshot)
	}
	return p.volDriver.Delete(snapshot.PortworxSnapshot.SnapshotID)
}

func (p *portworx) SnapshotRestore(
	snapshotData *crdv1.VolumeSnapshotData,
	pvc *v1.PersistentVolumeClaim,
	pvName string,
	parameters map[string]string,
) (*v1.PersistentVolumeSource, map[string]string, error) {
	if snapshotData == nil || snapshotData.Spec.PortworxSnapshot == nil {
		return nil, nil, fmt.Errorf("Invalid Snapshot spec")
	}
	if pvc == nil {
		return nil, nil, fmt.Errorf("Invalid PVC spec")
	}

	snapID := snapshotData.Spec.PortworxSnapshot.SnapshotID

	logrus.Debugf("SnapshotRestore for pvc: %+v", pvc)
	locator := &api.VolumeLocator{
		Name: "pvc-" + string(pvc.UID),
		VolumeLabels: map[string]string{
			pvcNameLabel:      pvc.Name,
			pvcNamespaceLabel: pvc.Namespace,
		},
	}
	volumeID, err := p.volDriver.Snapshot(snapID, false, locator)
	if err != nil {
		return nil, nil, err
	}

	vols, err := p.volDriver.Inspect([]string{volumeID})
	if err != nil {
		return nil, nil, &ErrFailedToInspectVolume{
			ID:    volumeID,
			Cause: fmt.Sprintf("Volume inspect returned err: %v", err),
		}
	}

	if len(vols) == 0 {
		return nil, nil, &errors.ErrNotFound{
			ID:   volumeID,
			Type: "Volume",
		}
	}

	pv := &v1.PersistentVolumeSource{
		PortworxVolume: &v1.PortworxVolumeSource{
			VolumeID: volumeID,
			FSType:   vols[0].Format.String(),
			ReadOnly: vols[0].Readonly,
		},
	}

	labels := make(map[string]string)

	return pv, labels, nil
}

func (p *portworx) DescribeSnapshot(snapshotData *crdv1.VolumeSnapshotData) (*[]crdv1.VolumeSnapshotCondition, bool, error) {
	if snapshotData == nil || snapshotData.Spec.PortworxSnapshot == nil {
		return nil, false, fmt.Errorf("Invalid VolumeSnapshotDataSource: %v", snapshotData)
	}
	snapshotID := snapshotData.Spec.PortworxSnapshot.SnapshotID
	_, err := p.InspectVolume(snapshotID)
	if err != nil {
		return nil, false, err
	}
	var snapConditions []crdv1.VolumeSnapshotCondition
	snapConditions = []crdv1.VolumeSnapshotCondition{
		{
			Type:               crdv1.VolumeSnapshotConditionReady,
			Status:             v1.ConditionTrue,
			Message:            "Snapshot created successfully and it is ready",
			LastTransitionTime: metav1.Now(),
		},
	}
	return &snapConditions, true, err
}

// TODO: Implement FindSnapshot
func (p *portworx) FindSnapshot(tags *map[string]string) (*crdv1.VolumeSnapshotDataSource, *[]crdv1.VolumeSnapshotCondition, error) {
	return nil, nil, nil
}

func (p *portworx) VolumeDelete(pv *v1.PersistentVolume) error {
	if pv == nil || pv.Spec.PortworxVolume == nil {
		return fmt.Errorf("Invalid PV: %v", pv)
	}
	return p.volDriver.Delete(pv.Spec.PortworxVolume.VolumeID)
}

func init() {
	if err := storkvolume.Register(driverName, &portworx{}); err != nil {
		logrus.Panicf("Error registering portworx volume driver: %v", err)
	}
}
<|endoftext|>


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

package logging

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/dustin/go-humanize"
	"gopkg.in/natefinch/lumberjack.v2"
)

func init() {
	caddy.RegisterModule(FileWriter{})
}

// FileWriter can write logs to files. By default, log files
// are rotated ("rolled") when they get large, and old log
// files get deleted, to ensure that the process does not
// exhaust disk space.
type FileWriter struct {
	// Filename is the name of the file to write.
	Filename string `json:"filename,omitempty"`

	// Roll toggles log rolling or rotation, which is
	// enabled by default.
	Roll *bool `json:"roll,omitempty"`

	// When a log file reaches approximately this size,
	// it will be rotated.
	RollSizeMB int `json:"roll_size_mb,omitempty"`

	// Whether to compress rolled files. Default: true
	RollCompress *bool `json:"roll_gzip,omitempty"`

	// Whether to use local timestamps in rolled filenames.
	// Default: false
	RollLocalTime bool `json:"roll_local_time,omitempty"`

	// The maximum number of rolled log files to keep.
	// Default: 10
	RollKeep int `json:"roll_keep,omitempty"`

	// How many days to keep rolled log files. Default: 90
	RollKeepDays int `json:"roll_keep_days,omitempty"`
}

// CaddyModule returns the Caddy module information.
func (FileWriter) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "caddy.logging.writers.file",
		New: func() caddy.Module { return new(FileWriter) },
	}
}

// Provision sets up the module
func (fw *FileWriter) Provision(ctx caddy.Context) error {
	// Replace placeholder in filename
	repl := caddy.NewReplacer()
	filename, err := repl.ReplaceOrErr(fw.Filename, true, true)
	if err != nil {
		return fmt.Errorf("invalid filename for log file: %v", err)
	}

	fw.Filename = filename
	return nil
}

func (fw FileWriter) String() string {
	fpath, err := filepath.Abs(fw.Filename)
	if err == nil {
		return fpath
	}
	return fw.Filename
}

// WriterKey returns a unique key representing this fw.
func (fw FileWriter) WriterKey() string {
	return "file:" + fw.Filename
}

// OpenWriter opens a new file writer.
func (fw FileWriter) OpenWriter() (io.WriteCloser, error) {
	// roll log files by default
	if fw.Roll == nil || *fw.Roll {
		if fw.RollSizeMB == 0 {
			fw.RollSizeMB = 100
		}
		if fw.RollCompress == nil {
			compress := true
			fw.RollCompress = &compress
		}
		if fw.RollKeep == 0 {
			fw.RollKeep = 10
		}
		if fw.RollKeepDays == 0 {
			fw.RollKeepDays = 90
		}

		return &lumberjack.Logger{
			Filename:   fw.Filename,
			MaxSize:    fw.RollSizeMB,
			MaxAge:     fw.RollKeepDays,
			MaxBackups: fw.RollKeep,
			LocalTime:  fw.RollLocalTime,
			Compress:   *fw.RollCompress,
		}, nil
	}

	// otherwise just open a regular file
	return os.OpenFile(fw.Filename, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
}

// UnmarshalCaddyfile sets up the module from Caddyfile tokens. Syntax:
//
//     file <filename> {
//         roll_disabled
//         roll_size     <size>
//         roll_keep     <num>
//         roll_keep_for <days>
//     }
//
// The roll_size value will be rounded down to number of megabytes (MiB).
// The roll_keep_for duration will be rounded down to number of days.
func (fw *FileWriter) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		if !d.NextArg() {
			return d.ArgErr()
		}
		fw.Filename = d.Val()
		if d.NextArg() {
			return d.ArgErr()
		}

		for d.NextBlock(0) {
			switch d.Val() {
			case "roll_disabled":
				var f bool
				fw.Roll = &f
				if d.NextArg() {
					return d.ArgErr()
				}

			case "roll_size":
				var sizeStr string
				if !d.AllArgs(&sizeStr) {
					return d.ArgErr()
				}
				size, err := humanize.ParseBytes(sizeStr)
				if err != nil {
					return d.Errf("parsing size: %v", err)
				}
				fw.RollSizeMB = int(size) / 1024 / 1024

			case "roll_keep":
				var keepStr string
				if !d.AllArgs(&keepStr) {
					return d.ArgErr()
				}
				keep, err := strconv.Atoi(keepStr)
				if err != nil {
					return d.Errf("parsing roll_keep number: %v", err)
				}
				fw.RollKeep = keep

			case "roll_keep_for":
				var keepForStr string
				if !d.AllArgs(&keepForStr) {
					return d.ArgErr()
				}
				keepFor, err := time.ParseDuration(keepForStr)
				if err != nil {
					return d.Errf("parsing roll_keep_for duration: %v", err)
				}
				fw.RollKeepDays = int(keepFor.Hours()) / 24
			}
		}
	}
	return nil
}

// Interface guards
var (
	_ caddy.Provisioner     = (*FileWriter)(nil)
	_ caddy.WriterOpener    = (*FileWriter)(nil)
	_ caddyfile.Unmarshaler = (*FileWriter)(nil)
)
<commit_msg>logging: Fix off-by-one for roll size MB from Caddyfile<commit_after>// Copyright 2015 Matthew Holt and The Caddy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package logging

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/dustin/go-humanize"
	"gopkg.in/natefinch/lumberjack.v2"
)

func init() {
	caddy.RegisterModule(FileWriter{})
}

// FileWriter can write logs to files. By default, log files
// are rotated ("rolled") when they get large, and old log
// files get deleted, to ensure that the process does not
// exhaust disk space.
type FileWriter struct {
	// Filename is the name of the file to write.
	Filename string `json:"filename,omitempty"`

	// Roll toggles log rolling or rotation, which is
	// enabled by default.
	Roll *bool `json:"roll,omitempty"`

	// When a log file reaches approximately this size,
	// it will be rotated.
	RollSizeMB int `json:"roll_size_mb,omitempty"`

	// Whether to compress rolled files. Default: true
	RollCompress *bool `json:"roll_gzip,omitempty"`

	// Whether to use local timestamps in rolled filenames.
	// Default: false
	RollLocalTime bool `json:"roll_local_time,omitempty"`

	// The maximum number of rolled log files to keep.
	// Default: 10
	RollKeep int `json:"roll_keep,omitempty"`

	// How many days to keep rolled log files. Default: 90
	RollKeepDays int `json:"roll_keep_days,omitempty"`
}

// CaddyModule returns the Caddy module information.
func (FileWriter) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "caddy.logging.writers.file",
		New: func() caddy.Module { return new(FileWriter) },
	}
}

// Provision sets up the module
func (fw *FileWriter) Provision(ctx caddy.Context) error {
	// Replace placeholder in filename
	repl := caddy.NewReplacer()
	filename, err := repl.ReplaceOrErr(fw.Filename, true, true)
	if err != nil {
		return fmt.Errorf("invalid filename for log file: %v", err)
	}

	fw.Filename = filename
	return nil
}

func (fw FileWriter) String() string {
	fpath, err := filepath.Abs(fw.Filename)
	if err == nil {
		return fpath
	}
	return fw.Filename
}

// WriterKey returns a unique key representing this fw.
func (fw FileWriter) WriterKey() string {
	return "file:" + fw.Filename
}

// OpenWriter opens a new file writer.
func (fw FileWriter) OpenWriter() (io.WriteCloser, error) {
	// roll log files by default
	if fw.Roll == nil || *fw.Roll {
		if fw.RollSizeMB == 0 {
			fw.RollSizeMB = 100
		}
		if fw.RollCompress == nil {
			compress := true
			fw.RollCompress = &compress
		}
		if fw.RollKeep == 0 {
			fw.RollKeep = 10
		}
		if fw.RollKeepDays == 0 {
			fw.RollKeepDays = 90
		}

		return &lumberjack.Logger{
			Filename:   fw.Filename,
			MaxSize:    fw.RollSizeMB,
			MaxAge:     fw.RollKeepDays,
			MaxBackups: fw.RollKeep,
			LocalTime:  fw.RollLocalTime,
			Compress:   *fw.RollCompress,
		}, nil
	}

	// otherwise just open a regular file
	return os.OpenFile(fw.Filename, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
}

// UnmarshalCaddyfile sets up the module from Caddyfile tokens. Syntax:
//
//     file <filename> {
//         roll_disabled
//         roll_size     <size>
//         roll_keep     <num>
//         roll_keep_for <days>
//     }
//
// The roll_size value will be rounded down to number of megabytes (MiB).
// The roll_keep_for duration will be rounded down to number of days.
func (fw *FileWriter) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		if !d.NextArg() {
			return d.ArgErr()
		}
		fw.Filename = d.Val()
		if d.NextArg() {
			return d.ArgErr()
		}

		for d.NextBlock(0) {
			switch d.Val() {
			case "roll_disabled":
				var f bool
				fw.Roll = &f
				if d.NextArg() {
					return d.ArgErr()
				}

			case "roll_size":
				var sizeStr string
				if !d.AllArgs(&sizeStr) {
					return d.ArgErr()
				}
				size, err := humanize.ParseBytes(sizeStr)
				if err != nil {
					return d.Errf("parsing size: %v", err)
				}
				fw.RollSizeMB = int(size)/1024/1024 + 1

			case "roll_keep":
				var keepStr string
				if !d.AllArgs(&keepStr) {
					return d.ArgErr()
				}
				keep, err := strconv.Atoi(keepStr)
				if err != nil {
					return d.Errf("parsing roll_keep number: %v", err)
				}
				fw.RollKeep = keep

			case "roll_keep_for":
				var keepForStr string
				if !d.AllArgs(&keepForStr) {
					return d.ArgErr()
				}
				keepFor, err := time.ParseDuration(keepForStr)
				if err != nil {
					return d.Errf("parsing roll_keep_for duration: %v", err)
				}
				fw.RollKeepDays = int(keepFor.Hours()) / 24
			}
		}
	}
	return nil
}

// Interface guards
var (
	_ caddy.Provisioner     = (*FileWriter)(nil)
	_ caddy.WriterOpener    = (*FileWriter)(nil)
	_ caddyfile.Unmarshaler = (*FileWriter)(nil)
)
<|endoftext|>