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

// The seeddb command is used to populates a database with an initial set of
// modules.
package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	_ "github.com/jackc/pgx/v4/stdlib" // for pgx driver
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/worker"
	"golang.org/x/sync/errgroup"
)

var (
	seedfile = flag.String("static", "devtools/cmd/seeddb/seed.txt", "filename containing modules for seeding the database")
)

func main() {
	flag.Parse()

	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	log.SetLevel(cfg.LogLevel)

	// Wrap the postgres driver with our own wrapper, which adds OpenCensus instrumentation.
	ddb, err := database.Open("pgx", cfg.DBConnInfo(), "seeddb")
	if err != nil {
		log.Fatalf(ctx, "database.Open for host %s failed with %v", cfg.DBHost, err)
	}
	db := postgres.New(ddb)
	defer db.Close()

	if err := run(ctx, db, cfg.ProxyURL); err != nil {
		log.Fatal(ctx, err)
	}
}

func run(ctx context.Context, db *postgres.DB, proxyURL string) error {
	start := time.Now()

	proxyClient, err := proxy.New(proxyURL)
	if err != nil {
		return err
	}

	sourceClient := source.NewClient(config.SourceTimeout)
	fetchFunc := func(ctx context.Context, modulePath, version string) (int, error) {
		f := &worker.Fetcher{
			ProxyClient:  proxyClient,
			SourceClient: sourceClient,
			DB:           db,
		}
		code, _, err := f.FetchAndUpdateState(ctx, modulePath, version, "")
		return code, err
	}

	seedModules, err := readSeedFile(ctx)
	if err != nil {
		return err
	}

	r := results{}
	g := new(errgroup.Group)
	for _, m := range seedModules {
		m := m

		g.Go(func() error {
			log.Infof(ctx, "Fetch requested: %q %q", m.path, m.version)
			start := time.Now()
			defer func() {
				// Log the duration of this fetch request.
				r.add(m.path, m.version, start)
			}()

			fetchCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
			defer cancel()

			code, err := fetchFunc(fetchCtx, m.path, m.version)
			if err != nil {
				if code == http.StatusNotFound {
					// We expect
					// github.com/jackc/pgx/pgxpool@v3.6.2+incompatible
					// to fail from seed.txt, so that it will redirect to
					// github.com/jackc/pgx/v4/pgxpool in tests.
					return nil
				}
				return err
			}
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	log.Infof(ctx, "Successfully fetched all modules: %v", time.Since(start))

	// Print the time it took to fetch these modules.
	var keys []string
	for k := range r.paths {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		log.Infof(ctx, "%s | %v", k, r.paths[k])
	}
	return nil
}

type results struct {
	paths map[string]time.Duration
	mu    sync.Mutex
}

func (r *results) add(modPath, version string, start time.Time) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.paths == nil {
		r.paths = map[string]time.Duration{}
	}
	r.paths[fmt.Sprintf("%s@%s", modPath, version)] = time.Since(start)
}

type module struct {
	path    string
	version string
}

// readSeedFile reads a file of module versions that we want to fetch for
// seeding the database. Format of the file:
// each line is
//     module@version
func readSeedFile(ctx context.Context) (_ []*module, err error) {
	defer derrors.Wrap(&err, "readSeedFile %q", *seedfile)
	lines, err := readFileLines(*seedfile)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "read %d module versions from %s", len(lines), *seedfile)

	var modules []*module
	for _, l := range lines {
		parts := strings.SplitN(l, "@", 2)
		modules = append(modules, &module{
			path:    parts[0],
			version: parts[1],
		})
	}
	return modules, nil
}

// readFileLines reads filename and returns its lines, trimmed of whitespace.
// Blank lines and lines whose first non-blank character is '#' are omitted.
func readFileLines(filename string) ([]string, error) {
	var lines []string
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		line := strings.TrimSpace(scan.Text())
		if line == "" || line[0] == '#' {
			continue
		}
		lines = append(lines, line)
	}
	if err := scan.Err(); err != nil {
		return nil, err
	}
	return lines, nil
}
<commit_msg>devtools/cmd/seeddb: support experiments<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The seeddb command is used to populates a database with an initial set of
// modules.
package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	_ "github.com/jackc/pgx/v4/stdlib" // for pgx driver
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/config/dynconfig"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/worker"
	"golang.org/x/sync/errgroup"
)

var (
	seedfile = flag.String("static", "devtools/cmd/seeddb/seed.txt", "filename containing modules for seeding the database")
)

func main() {
	flag.Parse()

	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	log.SetLevel(cfg.LogLevel)

	exps, err := fetchExperiments(ctx, cfg)
	if err != nil {
		log.Fatal(ctx, err)
	}
	ctx = experiment.NewContext(ctx, exps...)

	// Wrap the postgres driver with our own wrapper, which adds OpenCensus instrumentation.
	ddb, err := database.Open("pgx", cfg.DBConnInfo(), "seeddb")
	if err != nil {
		log.Fatalf(ctx, "database.Open for host %s failed with %v", cfg.DBHost, err)
	}
	db := postgres.New(ddb)
	defer db.Close()

	if err := run(ctx, db, cfg.ProxyURL); err != nil {
		log.Fatal(ctx, err)
	}
}

func run(ctx context.Context, db *postgres.DB, proxyURL string) error {
	start := time.Now()

	proxyClient, err := proxy.New(proxyURL)
	if err != nil {
		return err
	}

	sourceClient := source.NewClient(config.SourceTimeout)
	fetchFunc := func(ctx context.Context, modulePath, version string) (int, error) {
		f := &worker.Fetcher{
			ProxyClient:  proxyClient,
			SourceClient: sourceClient,
			DB:           db,
		}
		code, _, err := f.FetchAndUpdateState(ctx, modulePath, version, "")
		return code, err
	}

	seedModules, err := readSeedFile(ctx)
	if err != nil {
		return err
	}

	r := results{}
	g := new(errgroup.Group)
	for _, m := range seedModules {
		m := m

		g.Go(func() error {
			log.Infof(ctx, "Fetch requested: %q %q", m.path, m.version)
			start := time.Now()
			defer func() {
				// Log the duration of this fetch request.
				r.add(m.path, m.version, start)
			}()

			fetchCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
			defer cancel()

			code, err := fetchFunc(fetchCtx, m.path, m.version)
			if err != nil {
				if code == http.StatusNotFound {
					// We expect
					// github.com/jackc/pgx/pgxpool@v3.6.2+incompatible
					// to fail from seed.txt, so that it will redirect to
					// github.com/jackc/pgx/v4/pgxpool in tests.
					return nil
				}
				return err
			}
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		return err
	}
	log.Infof(ctx, "Successfully fetched all modules: %v", time.Since(start))

	// Print the time it took to fetch these modules.
	var keys []string
	for k := range r.paths {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		log.Infof(ctx, "%s | %v", k, r.paths[k])
	}
	return nil
}

type results struct {
	paths map[string]time.Duration
	mu    sync.Mutex
}

func (r *results) add(modPath, version string, start time.Time) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.paths == nil {
		r.paths = map[string]time.Duration{}
	}
	r.paths[fmt.Sprintf("%s@%s", modPath, version)] = time.Since(start)
}

type module struct {
	path    string
	version string
}

// readSeedFile reads a file of module versions that we want to fetch for
// seeding the database. Format of the file:
// each line is
//     module@version
func readSeedFile(ctx context.Context) (_ []*module, err error) {
	defer derrors.Wrap(&err, "readSeedFile %q", *seedfile)
	lines, err := readFileLines(*seedfile)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "read %d module versions from %s", len(lines), *seedfile)

	var modules []*module
	for _, l := range lines {
		parts := strings.SplitN(l, "@", 2)
		modules = append(modules, &module{
			path:    parts[0],
			version: parts[1],
		})
	}
	return modules, nil
}

// readFileLines reads filename and returns its lines, trimmed of whitespace.
// Blank lines and lines whose first non-blank character is '#' are omitted.
func readFileLines(filename string) ([]string, error) {
	var lines []string
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		line := strings.TrimSpace(scan.Text())
		if line == "" || line[0] == '#' {
			continue
		}
		lines = append(lines, line)
	}
	if err := scan.Err(); err != nil {
		return nil, err
	}
	return lines, nil
}

func fetchExperiments(ctx context.Context, cfg *config.Config) ([]string, error) {
	if cfg.DynamicConfigLocation == "" {
		return nil, nil
	}
	dc, err := dynconfig.Read(ctx, cfg.DynamicConfigLocation)
	if err != nil {
		return nil, err
	}
	var exps []string
	for _, e := range dc.Experiments {
		if e.Rollout > 0 {
			exps = append(exps, e.Name)
		}
	}
	return exps, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 integration

import (
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"testing"
	"time"

	"github.com/golang/glog"

	"k8s.io/apimachinery/pkg/util/wait"

	restclient "k8s.io/client-go/rest"

	genericserveroptions "k8s.io/apiserver/pkg/server/options"

	"github.com/kubernetes-incubator/service-catalog/cmd/apiserver/app/server"
	_ "github.com/kubernetes-incubator/service-catalog/pkg/apis/servicecatalog/install"
	servicecatalogclient "github.com/kubernetes-incubator/service-catalog/pkg/client/clientset_generated/clientset"
	serverstorage "github.com/kubernetes-incubator/service-catalog/pkg/registry/servicecatalog/server"
	"k8s.io/apimachinery/pkg/runtime"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

type TestServerConfig struct {
	etcdServerList []string
	storageType    serverstorage.StorageType
	emptyObjFunc   func() runtime.Object
}

// NewTestServerConfig is a default constructor for the standard test-apiserver setup
func NewTestServerConfig() *TestServerConfig {
	return &TestServerConfig{
		etcdServerList: []string{"http://localhost:2379"},
	}
}

func withConfigGetFreshApiserverAndClient(
	t *testing.T,
	serverConfig *TestServerConfig,
) (servicecatalogclient.Interface,
	*restclient.Config,
	func(),
) {
	securePort := rand.Intn(31743) + 1024
	secureAddr := fmt.Sprintf("https://localhost:%d", securePort)
	stopCh := make(chan struct{})
	serverFailed := make(chan struct{})
	shutdownServer := func() {
		t.Logf("Shutting down server on port: %d", securePort)
		close(stopCh)
	}

	t.Logf("Starting server on port: %d", securePort)
	certDir, _ := ioutil.TempDir("", "service-catalog-integration")
	secureServingOptions := genericserveroptions.NewSecureServingOptions()
	// start the server in the background
	go func() {
		var etcdOptions *server.EtcdOptions
		if serverstorage.StorageTypeEtcd == serverConfig.storageType {
			etcdOptions = server.NewEtcdOptions()
			etcdOptions.StorageConfig.ServerList = serverConfig.etcdServerList
			etcdOptions.EtcdOptions.StorageConfig.Prefix = fmt.Sprintf("%s-%08X", server.DefaultEtcdPathPrefix, rand.Int31())
		} else {
			t.Fatal("no storage type specified")
		}

		options := &server.ServiceCatalogServerOptions{
			StorageTypeString:       serverConfig.storageType.String(),
			GenericServerRunOptions: genericserveroptions.NewServerRunOptions(),
			AdmissionOptions:        genericserveroptions.NewAdmissionOptions(),
			SecureServingOptions:    secureServingOptions,
			EtcdOptions:             etcdOptions,
			AuthenticationOptions:   genericserveroptions.NewDelegatingAuthenticationOptions(),
			AuthorizationOptions:    genericserveroptions.NewDelegatingAuthorizationOptions(),
			AuditOptions:            genericserveroptions.NewAuditOptions(),
			DisableAuth:             true,
			StandaloneMode:          true, // this must be true because we have no kube server for integration.
		}
		options.SecureServingOptions.BindPort = securePort
		options.SecureServingOptions.ServerCert.CertDirectory = certDir

		if err := server.RunServer(options, stopCh); err != nil {
			close(serverFailed)
			t.Fatalf("Error in bringing up the server: %v", err)
		}
	}()

	if err := waitForApiserverUp(secureAddr, serverFailed); err != nil {
		t.Fatalf("%v", err)
	}

	config := &restclient.Config{}
	config.Host = secureAddr
	config.Insecure = true
	config.CertFile = secureServingOptions.ServerCert.CertKey.CertFile
	config.KeyFile = secureServingOptions.ServerCert.CertKey.KeyFile
	clientset, err := servicecatalogclient.NewForConfig(config)
	if nil != err {
		t.Fatal("can't make the client from the config", err)
	}
	return clientset, config, shutdownServer
}

func getFreshApiserverAndClient(
	t *testing.T,
	storageTypeStr string,
	newEmptyObj func() runtime.Object,
) (servicecatalogclient.Interface, *restclient.Config, func()) {
	var serverStorageType serverstorage.StorageType
	serverStorageType, err := serverstorage.StorageTypeFromString(storageTypeStr)
	if nil != err {
		t.Fatal("non supported storage type")
	}

	serverConfig := &TestServerConfig{
		etcdServerList: []string{"http://localhost:2379"},
		storageType:    serverStorageType,
		emptyObjFunc:   newEmptyObj,
	}
	client, clientConfig, shutdownFunc := withConfigGetFreshApiserverAndClient(t, serverConfig)
	return client, clientConfig, shutdownFunc
}

func waitForApiserverUp(serverURL string, stopCh <-chan struct{}) error {
	interval := 1 * time.Second
	timeout := 30 * time.Second
	startWaiting := time.Now()
	tries := 0
	return wait.PollImmediate(interval, timeout,
		func() (bool, error) {
			select {
			// we've been told to stop, so no reason to keep going
			case <-stopCh:
				return true, fmt.Errorf("apiserver failed")
			default:
				glog.Infof("Waiting for : %#v", serverURL)
				tr := &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				}
				c := &http.Client{Transport: tr}
				_, err := c.Get(serverURL)
				if err == nil {
					glog.Infof("Found server after %v tries and duration %v",
						tries, time.Since(startWaiting))
					return true, nil
				}
				tries++
				return false, nil
			}
		},
	)
}
<commit_msg>silence go-restful logging output (#1622)<commit_after>/*
Copyright 2017 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 integration

import (
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"testing"
	"time"

	restfullog "github.com/emicklei/go-restful/log"
	"github.com/golang/glog"

	"k8s.io/apimachinery/pkg/util/wait"

	restclient "k8s.io/client-go/rest"

	genericserveroptions "k8s.io/apiserver/pkg/server/options"

	"github.com/kubernetes-incubator/service-catalog/cmd/apiserver/app/server"
	_ "github.com/kubernetes-incubator/service-catalog/pkg/apis/servicecatalog/install"
	servicecatalogclient "github.com/kubernetes-incubator/service-catalog/pkg/client/clientset_generated/clientset"
	serverstorage "github.com/kubernetes-incubator/service-catalog/pkg/registry/servicecatalog/server"
	"k8s.io/apimachinery/pkg/runtime"
)

func init() {
	rand.Seed(time.Now().UnixNano())
	// silence the go-restful webservices swagger logger
	restfullog.SetLogger(log.New(ioutil.Discard, "[restful]", log.LstdFlags|log.Lshortfile))
}

type TestServerConfig struct {
	etcdServerList []string
	storageType    serverstorage.StorageType
	emptyObjFunc   func() runtime.Object
}

// NewTestServerConfig is a default constructor for the standard test-apiserver setup
func NewTestServerConfig() *TestServerConfig {
	return &TestServerConfig{
		etcdServerList: []string{"http://localhost:2379"},
	}
}

func withConfigGetFreshApiserverAndClient(
	t *testing.T,
	serverConfig *TestServerConfig,
) (servicecatalogclient.Interface,
	*restclient.Config,
	func(),
) {
	securePort := rand.Intn(31743) + 1024
	secureAddr := fmt.Sprintf("https://localhost:%d", securePort)
	stopCh := make(chan struct{})
	serverFailed := make(chan struct{})
	shutdownServer := func() {
		t.Logf("Shutting down server on port: %d", securePort)
		close(stopCh)
	}

	t.Logf("Starting server on port: %d", securePort)
	certDir, _ := ioutil.TempDir("", "service-catalog-integration")
	secureServingOptions := genericserveroptions.NewSecureServingOptions()
	// start the server in the background
	go func() {
		var etcdOptions *server.EtcdOptions
		if serverstorage.StorageTypeEtcd == serverConfig.storageType {
			etcdOptions = server.NewEtcdOptions()
			etcdOptions.StorageConfig.ServerList = serverConfig.etcdServerList
			etcdOptions.EtcdOptions.StorageConfig.Prefix = fmt.Sprintf("%s-%08X", server.DefaultEtcdPathPrefix, rand.Int31())
		} else {
			t.Fatal("no storage type specified")
		}

		options := &server.ServiceCatalogServerOptions{
			StorageTypeString:       serverConfig.storageType.String(),
			GenericServerRunOptions: genericserveroptions.NewServerRunOptions(),
			AdmissionOptions:        genericserveroptions.NewAdmissionOptions(),
			SecureServingOptions:    secureServingOptions,
			EtcdOptions:             etcdOptions,
			AuthenticationOptions:   genericserveroptions.NewDelegatingAuthenticationOptions(),
			AuthorizationOptions:    genericserveroptions.NewDelegatingAuthorizationOptions(),
			AuditOptions:            genericserveroptions.NewAuditOptions(),
			DisableAuth:             true,
			StandaloneMode:          true, // this must be true because we have no kube server for integration.
		}
		options.SecureServingOptions.BindPort = securePort
		options.SecureServingOptions.ServerCert.CertDirectory = certDir

		if err := server.RunServer(options, stopCh); err != nil {
			close(serverFailed)
			t.Fatalf("Error in bringing up the server: %v", err)
		}
	}()

	if err := waitForApiserverUp(secureAddr, serverFailed); err != nil {
		t.Fatalf("%v", err)
	}

	config := &restclient.Config{}
	config.Host = secureAddr
	config.Insecure = true
	config.CertFile = secureServingOptions.ServerCert.CertKey.CertFile
	config.KeyFile = secureServingOptions.ServerCert.CertKey.KeyFile
	clientset, err := servicecatalogclient.NewForConfig(config)
	if nil != err {
		t.Fatal("can't make the client from the config", err)
	}
	return clientset, config, shutdownServer
}

func getFreshApiserverAndClient(
	t *testing.T,
	storageTypeStr string,
	newEmptyObj func() runtime.Object,
) (servicecatalogclient.Interface, *restclient.Config, func()) {
	var serverStorageType serverstorage.StorageType
	serverStorageType, err := serverstorage.StorageTypeFromString(storageTypeStr)
	if nil != err {
		t.Fatal("non supported storage type")
	}

	serverConfig := &TestServerConfig{
		etcdServerList: []string{"http://localhost:2379"},
		storageType:    serverStorageType,
		emptyObjFunc:   newEmptyObj,
	}
	client, clientConfig, shutdownFunc := withConfigGetFreshApiserverAndClient(t, serverConfig)
	return client, clientConfig, shutdownFunc
}

func waitForApiserverUp(serverURL string, stopCh <-chan struct{}) error {
	interval := 1 * time.Second
	timeout := 30 * time.Second
	startWaiting := time.Now()
	tries := 0
	return wait.PollImmediate(interval, timeout,
		func() (bool, error) {
			select {
			// we've been told to stop, so no reason to keep going
			case <-stopCh:
				return true, fmt.Errorf("apiserver failed")
			default:
				glog.Infof("Waiting for : %#v", serverURL)
				tr := &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				}
				c := &http.Client{Transport: tr}
				_, err := c.Get(serverURL)
				if err == nil {
					glog.Infof("Found server after %v tries and duration %v",
						tries, time.Since(startWaiting))
					return true, nil
				}
				tries++
				return false, nil
			}
		},
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2021 Gravitational, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
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 bot

import (
	"context"
	"fmt"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/gravitational/teleport/.github/workflows/ci"
	"github.com/gravitational/teleport/.github/workflows/ci/pkg/environment"

	"github.com/google/go-github/v37/github"
	"github.com/gravitational/trace"
)

// Check checks if all the reviewers have approved the pull request in the current context.
func (c *Bot) Check(ctx context.Context) error {
	pr := c.Environment.Metadata
	if c.Environment.IsInternal(pr.Author) {
		return c.checkInternal(ctx)
	}
	return c.checkExternal(ctx)
}

// checkInternal is called to check if a PR reviewed and approved by the
// required reviewers for internal contributors. Unlike approvals for
// external contributors, approvals from internal team members will not be
// invalidated when new changes are pushed to the PR.
func (c *Bot) checkInternal(ctx context.Context) error {
	pr := c.Environment.Metadata
	// Remove any stale workflow runs. As only the current workflow run should
	// be shown because it is the workflow that reflects the correct state of the pull request.
	//
	// Note: This is run for all workflow runs triggered by an event from an internal contributor,
	// but has to be run again in cron workflow because workflows triggered by external contributors do not
	// grant the Github actions token the correct permissions to dismiss workflow runs.
	err := c.dismissStaleWorkflowRuns(ctx, pr.RepoOwner, pr.RepoName, pr.BranchName)
	if err != nil {
		return trace.Wrap(err)
	}
	mostRecentReviews, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	// If an admin approves a pull request, allow the check workflow run to pass.
	// This is used in the event the pull request needs the bypassing
	// of the required reviewers.
	if repoAdminHasApproved(mostRecentReviews) {
		return nil
	}
	log.Printf("Checking if %v has approvals from the required reviewers %+v", pr.Author, c.Environment.GetReviewersForAuthor(pr.Author))
	err = hasRequiredApprovals(mostRecentReviews, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// repoAdminHasApproved checks that at least one listed admin has approved.
func repoAdminHasApproved(reviews map[string]review) bool {
	for _, adminName := range ci.RepoAdmins {
		if hasApproved(adminName, reviews) {
			return true
		}
	}
	return false
}

// checkExternal is called to check if a PR reviewed and approved by the
// required reviewers for external contributors. Approvals for external
// contributors are dismissed when new changes are pushed to the PR. The only
// case in which reviews are not dismissed is if they are from GitHub and
// only update the PR.
func (c *Bot) checkExternal(ctx context.Context) error {
	var obsoleteReviews map[string]review
	var validReviews map[string]review

	pr := c.Environment.Metadata
	mostRecentReviews, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	validReviews, obsoleteReviews = splitReviews(pr.HeadSHA, mostRecentReviews)
	// External contributions require tighter scrutiny than team
	// contributions. As such reviews from previous pushes must
	// not carry over to when new changes are added. Github does
	// not do this automatically, so we must dismiss the reviews
	// manually if there is a file change.
	if err = c.hasFileChangeFromLastApprovedReview(ctx); err != nil {
		err = c.invalidateApprovals(ctx, obsoleteReviews)
		if err != nil {
			return trace.Wrap(err)
		}
	} else {
		// If there are no file changes between current commit and commit where all
		// reviewers have approved, then all most recent reviews are valid.
		validReviews = mostRecentReviews
	}
	log.Printf("Checking if %v has approvals from the required reviewers %+v", pr.Author, c.Environment.GetReviewersForAuthor(pr.Author))
	err = hasRequiredApprovals(validReviews, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// splitReviews splits a list of reviews into two lists: `valid` (those reviews that refer to
// the current PR head revision) and `obsolete` (those that do not)
func splitReviews(headSHA string, reviews map[string]review) (valid, obsolete map[string]review) {
	valid = make(map[string]review)
	obsolete = make(map[string]review)
	for _, r := range reviews {
		if r.commitID == headSHA {
			valid[r.name] = r
		} else {
			obsolete[r.name] = r
		}
	}
	return valid, obsolete
}

// hasRequiredApprovals determines if all required reviewers have approved.
func hasRequiredApprovals(mostRecentReviews map[string]review, required []string) error {
	if len(mostRecentReviews) == 0 {
		return trace.BadParameter("pull request has no approvals")
	}
	var waitingOnApprovalsFrom []string
	for _, requiredReviewer := range required {
		ok := hasApproved(requiredReviewer, mostRecentReviews)
		if !ok {
			waitingOnApprovalsFrom = append(waitingOnApprovalsFrom, requiredReviewer)
		}
	}
	if len(waitingOnApprovalsFrom) > 0 {
		return trace.BadParameter("required reviewers have not yet approved, waiting on approval(s) from %v", waitingOnApprovalsFrom)
	}
	return nil
}

func (c *Bot) getMostRecentReviews(ctx context.Context) (map[string]review, error) {
	env := c.Environment
	pr := c.Environment.Metadata
	reviews, _, err := env.Client.PullRequests.ListReviews(ctx, pr.RepoOwner,
		pr.RepoName,
		pr.Number,
		&github.ListOptions{})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	currentReviewsSlice := []review{}
	for _, rev := range reviews {
		// Because PRs can be submitted by anyone, input here is attacker controlled
		// and do strict validation of input.
		err := validateReviewFields(rev)
		if err != nil {
			return nil, trace.Wrap(err)
		}
		currReview := review{
			name:        *rev.User.Login,
			status:      *rev.State,
			commitID:    *rev.CommitID,
			id:          *rev.ID,
			submittedAt: rev.SubmittedAt,
		}
		currentReviewsSlice = append(currentReviewsSlice, currReview)
	}
	return mostRecent(currentReviewsSlice), nil
}

// review is a pull request review
type review struct {
	name        string
	status      string
	commitID    string
	id          int64
	submittedAt *time.Time
}

// validateReviewFields validates required fields exist and passes them
// through a restrictive allow list (alphanumerics only). This is done to
// mitigate impact of attacker controlled input (the PR).
func validateReviewFields(review *github.PullRequestReview) error {
	switch {
	case review.ID == nil:
		return trace.Errorf("review ID is nil. review: %+v", review)
	case review.State == nil:
		return trace.Errorf("review State is nil. review: %+v", review)
	case review.CommitID == nil:
		return trace.Errorf("review CommitID is nil. review: %+v", review)
	case review.SubmittedAt == nil:
		return trace.Errorf("review SubmittedAt is nil. review: %+v", review)
	case review.User.Login == nil:
		return trace.Errorf("reviewer User.Login is nil. review: %+v", review)
	}
	if err := validateField(*review.State); err != nil {
		return trace.Errorf("review ID err: %v", err)
	}
	if err := validateField(*review.CommitID); err != nil {
		return trace.Errorf("commit ID err: %v", err)
	}
	if err := validateField(*review.User.Login); err != nil {
		return trace.Errorf("user login err: %v", err)
	}
	return nil
}

// mostRecent returns a list of the most recent review from each required reviewer.
func mostRecent(currentReviews []review) map[string]review {
	mostRecentReviews := make(map[string]review)
	for _, rev := range currentReviews {
		val, ok := mostRecentReviews[rev.name]
		if !ok {
			mostRecentReviews[rev.name] = rev
		} else {
			setTime := val.submittedAt
			currTime := rev.submittedAt
			if currTime.After(*setTime) {
				mostRecentReviews[rev.name] = rev
			}
		}
	}
	return mostRecentReviews
}

// hasApproved determines if the reviewer has submitted an approval
// for the pull request.
func hasApproved(reviewer string, reviews map[string]review) bool {
	for _, rev := range reviews {
		if rev.name == reviewer && rev.status == ci.Approved {
			return true
		}
	}
	return false
}

// dimissMessage returns the dimiss message when a review is dismissed
func dismissMessage(pr *environment.Metadata, required []string) string {
	var sb strings.Builder
	sb.WriteString("New commit pushed, please re-review ")
	for _, reviewer := range required {
		sb.WriteString(fmt.Sprintf("@%s ", reviewer))
	}
	return strings.TrimSpace(sb.String())
}

// hasFileChangeFromLastApproved checks if there is a file change from the last commit all
// reviewers approved (if all reviewers approved at a commit) to the current HEAD.
func (c *Bot) hasFileChangeFromLastApprovedReview(ctx context.Context) error {
	pr := c.Environment.Metadata
	lastReviewCommitID, err := c.getLastApprovedReviewCommitID(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	mostRecent, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	// Make sure all approvals are at the same commit.
	err = hasAllRequiredApprovalsAtCommit(lastReviewCommitID, mostRecent, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	// Check for any differences
	err = c.hasFileDiff(ctx, lastReviewCommitID, pr.HeadSHA)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// getLastApprovedReviewCommitID gets the last review's commit ID (last review where a commit was approved).
func (c *Bot) getLastApprovedReviewCommitID(ctx context.Context) (string, error) {
	pr := c.Environment.Metadata
	clt := c.Environment.Client
	reviews, _, err := clt.PullRequests.ListReviews(ctx, pr.RepoOwner, pr.RepoName, pr.Number, &github.ListOptions{})
	if err != nil {
		return "", trace.Wrap(err)
	}
	if len(reviews) == 0 {
		return "", trace.NotFound("pull request has no reviews")
	}

	// Sort reviews from newest to oldest.
	sort.Slice(reviews, func(i, j int) bool {
		time1, time2 := reviews[i].SubmittedAt, reviews[j].SubmittedAt
		return time2.Before(*time1)
	})
	var lastApprovedReview *github.PullRequestReview
	// Find last approved review.
	for _, review := range reviews {
		if review.State == nil {
			continue
		}
		if *review.State == ci.Approved {
			lastApprovedReview = review
			break
		}
	}
	if lastApprovedReview == nil {
		return "", trace.NotFound("no approved reviews found")
	}
	if lastApprovedReview.CommitID == nil {
		return "", trace.NotFound("commit ID not found")
	}
	return *lastApprovedReview.CommitID, nil
}

// hasFileDiff compares two commits and checks if there are changes.
func (c *Bot) hasFileDiff(ctx context.Context, base, head string) error {
	pr := c.Environment.Metadata
	clt := c.Environment.Client
	comparison, _, err := clt.Repositories.CompareCommits(ctx, pr.RepoOwner, pr.RepoName, base, head)
	if err != nil {
		return trace.Wrap(err)
	}
	if len(comparison.Files) != 0 {
		return trace.Errorf("detected file change")
	}
	return nil
}

func hasAllRequiredApprovalsAtCommit(commitSHA string, reviews map[string]review, required []string) error {
	for _, requiredReviewer := range required {
		review, ok := reviews[requiredReviewer]
		if !ok {
			return trace.BadParameter("all reviewers have not approved")
		}
		if review.commitID != commitSHA {
			return trace.Errorf("all reviewers have not approved at %s", commitSHA)
		}
	}
	return nil
}

// invalidateApprovals dismisses all approved reviews on a pull request.
func (c *Bot) invalidateApprovals(ctx context.Context, reviews map[string]review) error {
	pr := c.Environment.Metadata
	msg := dismissMessage(pr, c.Environment.GetReviewersForAuthor(pr.Author))
	for _, v := range reviews {
		if pr.HeadSHA != v.commitID && v.status != ci.Commented {
			_, _, err := c.Environment.Client.PullRequests.DismissReview(ctx,
				pr.RepoOwner,
				pr.RepoName,
				pr.Number,
				v.id,
				&github.PullRequestReviewDismissalRequest{Message: &msg},
			)
			if err != nil {
				return trace.Wrap(err)
			}
		}
	}
	return nil
}
<commit_msg>Re-Request Reviews When Approvals Are Invalidated (#9037)<commit_after>/*
Copyright 2021 Gravitational, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
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 bot

import (
	"context"
	"fmt"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/gravitational/teleport/.github/workflows/ci"
	"github.com/gravitational/teleport/.github/workflows/ci/pkg/environment"

	"github.com/google/go-github/v37/github"
	"github.com/gravitational/trace"
)

// Check checks if all the reviewers have approved the pull request in the current context.
func (c *Bot) Check(ctx context.Context) error {
	pr := c.Environment.Metadata
	if c.Environment.IsInternal(pr.Author) {
		return c.checkInternal(ctx)
	}
	return c.checkExternal(ctx)
}

// checkInternal is called to check if a PR reviewed and approved by the
// required reviewers for internal contributors. Unlike approvals for
// external contributors, approvals from internal team members will not be
// invalidated when new changes are pushed to the PR.
func (c *Bot) checkInternal(ctx context.Context) error {
	pr := c.Environment.Metadata
	// Remove any stale workflow runs. As only the current workflow run should
	// be shown because it is the workflow that reflects the correct state of the pull request.
	//
	// Note: This is run for all workflow runs triggered by an event from an internal contributor,
	// but has to be run again in cron workflow because workflows triggered by external contributors do not
	// grant the Github actions token the correct permissions to dismiss workflow runs.
	err := c.dismissStaleWorkflowRuns(ctx, pr.RepoOwner, pr.RepoName, pr.BranchName)
	if err != nil {
		return trace.Wrap(err)
	}
	mostRecentReviews, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	// If an admin approves a pull request, allow the check workflow run to pass.
	// This is used in the event the pull request needs the bypassing
	// of the required reviewers.
	if repoAdminHasApproved(mostRecentReviews) {
		return nil
	}
	log.Printf("Checking if %v has approvals from the required reviewers %+v", pr.Author, c.Environment.GetReviewersForAuthor(pr.Author))
	err = hasRequiredApprovals(mostRecentReviews, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// repoAdminHasApproved checks that at least one listed admin has approved.
func repoAdminHasApproved(reviews map[string]review) bool {
	for _, adminName := range ci.RepoAdmins {
		if hasApproved(adminName, reviews) {
			return true
		}
	}
	return false
}

// checkExternal is called to check if a PR reviewed and approved by the
// required reviewers for external contributors. Approvals for external
// contributors are dismissed when new changes are pushed to the PR. The only
// case in which reviews are not dismissed is if they are from GitHub and
// only update the PR.
func (c *Bot) checkExternal(ctx context.Context) error {
	var obsoleteReviews map[string]review
	var validReviews map[string]review

	pr := c.Environment.Metadata
	mostRecentReviews, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	validReviews, obsoleteReviews = splitReviews(pr.HeadSHA, mostRecentReviews)
	// External contributions require tighter scrutiny than team
	// contributions. As such reviews from previous pushes must
	// not carry over to when new changes are added. Github does
	// not do this automatically, so we must dismiss the reviews
	// manually if there is a file change.
	if err = c.hasFileChangeFromLastApprovedReview(ctx); err != nil {
		err = c.invalidateApprovals(ctx, obsoleteReviews)
		if err != nil {
			return trace.Wrap(err)
		}
	} else {
		// If there are no file changes between current commit and commit where all
		// reviewers have approved, then all most recent reviews are valid.
		validReviews = mostRecentReviews
	}
	log.Printf("Checking if %v has approvals from the required reviewers %+v", pr.Author, c.Environment.GetReviewersForAuthor(pr.Author))
	err = hasRequiredApprovals(validReviews, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// splitReviews splits a list of reviews into two lists: `valid` (those reviews that refer to
// the current PR head revision) and `obsolete` (those that do not)
func splitReviews(headSHA string, reviews map[string]review) (valid, obsolete map[string]review) {
	valid = make(map[string]review)
	obsolete = make(map[string]review)
	for _, r := range reviews {
		if r.commitID == headSHA {
			valid[r.name] = r
		} else {
			obsolete[r.name] = r
		}
	}
	return valid, obsolete
}

// hasRequiredApprovals determines if all required reviewers have approved.
func hasRequiredApprovals(mostRecentReviews map[string]review, required []string) error {
	if len(mostRecentReviews) == 0 {
		return trace.BadParameter("pull request has no approvals")
	}
	var waitingOnApprovalsFrom []string
	for _, requiredReviewer := range required {
		ok := hasApproved(requiredReviewer, mostRecentReviews)
		if !ok {
			waitingOnApprovalsFrom = append(waitingOnApprovalsFrom, requiredReviewer)
		}
	}
	if len(waitingOnApprovalsFrom) > 0 {
		return trace.BadParameter("required reviewers have not yet approved, waiting on approval(s) from %v", waitingOnApprovalsFrom)
	}
	return nil
}

func (c *Bot) getMostRecentReviews(ctx context.Context) (map[string]review, error) {
	env := c.Environment
	pr := c.Environment.Metadata
	reviews, _, err := env.Client.PullRequests.ListReviews(ctx, pr.RepoOwner,
		pr.RepoName,
		pr.Number,
		&github.ListOptions{})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	currentReviewsSlice := []review{}
	for _, rev := range reviews {
		// Because PRs can be submitted by anyone, input here is attacker controlled
		// and do strict validation of input.
		err := validateReviewFields(rev)
		if err != nil {
			return nil, trace.Wrap(err)
		}
		currReview := review{
			name:        *rev.User.Login,
			status:      *rev.State,
			commitID:    *rev.CommitID,
			id:          *rev.ID,
			submittedAt: rev.SubmittedAt,
		}
		currentReviewsSlice = append(currentReviewsSlice, currReview)
	}
	return mostRecent(currentReviewsSlice), nil
}

// review is a pull request review
type review struct {
	name        string
	status      string
	commitID    string
	id          int64
	submittedAt *time.Time
}

// validateReviewFields validates required fields exist and passes them
// through a restrictive allow list (alphanumerics only). This is done to
// mitigate impact of attacker controlled input (the PR).
func validateReviewFields(review *github.PullRequestReview) error {
	switch {
	case review.ID == nil:
		return trace.Errorf("review ID is nil. review: %+v", review)
	case review.State == nil:
		return trace.Errorf("review State is nil. review: %+v", review)
	case review.CommitID == nil:
		return trace.Errorf("review CommitID is nil. review: %+v", review)
	case review.SubmittedAt == nil:
		return trace.Errorf("review SubmittedAt is nil. review: %+v", review)
	case review.User.Login == nil:
		return trace.Errorf("reviewer User.Login is nil. review: %+v", review)
	}
	if err := validateField(*review.State); err != nil {
		return trace.Errorf("review ID err: %v", err)
	}
	if err := validateField(*review.CommitID); err != nil {
		return trace.Errorf("commit ID err: %v", err)
	}
	if err := validateField(*review.User.Login); err != nil {
		return trace.Errorf("user login err: %v", err)
	}
	return nil
}

// mostRecent returns a list of the most recent review from each required reviewer.
func mostRecent(currentReviews []review) map[string]review {
	mostRecentReviews := make(map[string]review)
	for _, rev := range currentReviews {
		val, ok := mostRecentReviews[rev.name]
		if !ok {
			mostRecentReviews[rev.name] = rev
		} else {
			setTime := val.submittedAt
			currTime := rev.submittedAt
			if currTime.After(*setTime) {
				mostRecentReviews[rev.name] = rev
			}
		}
	}
	return mostRecentReviews
}

// hasApproved determines if the reviewer has submitted an approval
// for the pull request.
func hasApproved(reviewer string, reviews map[string]review) bool {
	for _, rev := range reviews {
		if rev.name == reviewer && rev.status == ci.Approved {
			return true
		}
	}
	return false
}

// dimissMessage returns the dimiss message when a review is dismissed
func dismissMessage(pr *environment.Metadata, required []string) string {
	var sb strings.Builder
	sb.WriteString("New commit pushed, please re-review ")
	for _, reviewer := range required {
		sb.WriteString(fmt.Sprintf("@%s ", reviewer))
	}
	return strings.TrimSpace(sb.String())
}

// hasFileChangeFromLastApproved checks if there is a file change from the last commit all
// reviewers approved (if all reviewers approved at a commit) to the current HEAD.
func (c *Bot) hasFileChangeFromLastApprovedReview(ctx context.Context) error {
	pr := c.Environment.Metadata
	lastReviewCommitID, err := c.getLastApprovedReviewCommitID(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	mostRecent, err := c.getMostRecentReviews(ctx)
	if err != nil {
		return trace.Wrap(err)
	}
	// Make sure all approvals are at the same commit.
	err = hasAllRequiredApprovalsAtCommit(lastReviewCommitID, mostRecent, c.Environment.GetReviewersForAuthor(pr.Author))
	if err != nil {
		return trace.Wrap(err)
	}
	// Check for any differences
	err = c.hasFileDiff(ctx, lastReviewCommitID, pr.HeadSHA)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

// getLastApprovedReviewCommitID gets the last review's commit ID (last review where a commit was approved).
func (c *Bot) getLastApprovedReviewCommitID(ctx context.Context) (string, error) {
	pr := c.Environment.Metadata
	clt := c.Environment.Client
	reviews, _, err := clt.PullRequests.ListReviews(ctx, pr.RepoOwner, pr.RepoName, pr.Number, &github.ListOptions{})
	if err != nil {
		return "", trace.Wrap(err)
	}
	if len(reviews) == 0 {
		return "", trace.NotFound("pull request has no reviews")
	}

	// Sort reviews from newest to oldest.
	sort.Slice(reviews, func(i, j int) bool {
		time1, time2 := reviews[i].SubmittedAt, reviews[j].SubmittedAt
		return time2.Before(*time1)
	})
	var lastApprovedReview *github.PullRequestReview
	// Find last approved review.
	for _, review := range reviews {
		if review.State == nil {
			continue
		}
		if *review.State == ci.Approved {
			lastApprovedReview = review
			break
		}
	}
	if lastApprovedReview == nil {
		return "", trace.NotFound("no approved reviews found")
	}
	if lastApprovedReview.CommitID == nil {
		return "", trace.NotFound("commit ID not found")
	}
	return *lastApprovedReview.CommitID, nil
}

// hasFileDiff compares two commits and checks if there are changes.
func (c *Bot) hasFileDiff(ctx context.Context, base, head string) error {
	pr := c.Environment.Metadata
	clt := c.Environment.Client
	comparison, _, err := clt.Repositories.CompareCommits(ctx, pr.RepoOwner, pr.RepoName, base, head)
	if err != nil {
		return trace.Wrap(err)
	}
	if len(comparison.Files) != 0 {
		return trace.Errorf("detected file change")
	}
	return nil
}

func hasAllRequiredApprovalsAtCommit(commitSHA string, reviews map[string]review, required []string) error {
	for _, requiredReviewer := range required {
		review, ok := reviews[requiredReviewer]
		if !ok {
			return trace.BadParameter("all reviewers have not approved")
		}
		if review.commitID != commitSHA {
			return trace.Errorf("all reviewers have not approved at %s", commitSHA)
		}
	}
	return nil
}

// invalidateApprovals dismisses all approved reviews on a pull request.
func (c *Bot) invalidateApprovals(ctx context.Context, reviews map[string]review) error {
	pr := c.Environment.Metadata
	clt := c.Environment.Client
	msg := dismissMessage(pr, c.Environment.GetReviewersForAuthor(pr.Author))
	for _, v := range reviews {
		if pr.HeadSHA != v.commitID && v.status != ci.Commented {
			_, _, err := clt.PullRequests.DismissReview(ctx,
				pr.RepoOwner,
				pr.RepoName,
				pr.Number,
				v.id,
				&github.PullRequestReviewDismissalRequest{Message: &msg},
			)
			if err != nil {
				return trace.Wrap(err)
			}
		}
	}
	// Re-assign reviewers when dismissing so the
	// pull request shows up in their review requests again.
	return c.Assign(ctx)
}
<|endoftext|>


<commit_before>package tests_test

import (
	"context"
	"fmt"
	"strings"

	"kubevirt.io/kubevirt/tests/framework/cleanup"

	expect "github.com/google/goexpect"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/kubevirt/tests/util"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"

	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/console"
	. "kubevirt.io/kubevirt/tests/framework/matcher"
	"kubevirt.io/kubevirt/tests/libnode"
)

var _ = Describe("[Serial][sig-compute]MediatedDevices", func() {
	var err error
	var virtClient kubecli.KubevirtClient

	BeforeEach(func() {
		virtClient, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)
	})
	checkAllMDEVCreated := func(mdevTypeName string, expectedInstancesCount int) {
		By(fmt.Sprintf("Checking the number of created mdev types, should be %d of %s type ", expectedInstancesCount, mdevTypeName))
		check := fmt.Sprintf(`set -x
		files_num=$(ls -A /sys/bus/mdev/devices/| wc -l)
		if [[ $files_num != %d ]] ; then
		  echo "failed, not enough mdevs of type %[2]s has been created"
		  exit 1
		fi
		for x in $(ls -A /sys/bus/mdev/devices/); do
		  type_name=$(basename "$(readlink -f /sys/bus/mdev/devices/$x/mdev_type)")
		  if [[ "$type_name" != "%[2]s" ]]; then
		     echo "failed, not all mdevs of type %[2]s"
		     exit 1
		  fi
		  exit 0
		done`, expectedInstancesCount, mdevTypeName)
		testPod := tests.RenderPod("test-pod", []string{"/bin/bash", "-c"}, []string{check})
		testPod, err = virtClient.CoreV1().Pods(util.NamespaceTestDefault).Create(context.Background(), testPod, metav1.CreateOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisPod(testPod), 120).Should(BeInPhase(k8sv1.PodSucceeded))
	}
	checkAllMDEVRemoved := func() {
		check := fmt.Sprintf(`set -x
		files_num=$(ls -A /sys/bus/mdev/devices/| wc -l)
		if [[ $files_num != 0 ]] ; then
		  echo "failed, not all mdevs removed"
		  exit 1
		fi
	        exit 0`)
		testPod := tests.RenderPod("test-pod", []string{"/bin/bash", "-c"}, []string{check})
		testPod, err = virtClient.CoreV1().Pods(util.NamespaceTestDefault).Create(context.Background(), testPod, metav1.CreateOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisPod(testPod), 120).Should(BeInPhase(k8sv1.PodSucceeded))
		Eventually(func() bool {
			nodes, err := virtClient.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{})
			Expect(err).ToNot(HaveOccurred())
			for _, node := range nodes.Items {
				for key, amount := range node.Status.Capacity {
					if strings.HasPrefix(string(key), "nvidia.com/") && !amount.IsZero() {
						return true
					}
				}
			}
			return false
		}, 60).Should(BeFalse(), "wait for the kubelet to stop promoting unconfigured devices")
	}
	Context("with mediated devices configuration", func() {
		var vmi *v1.VirtualMachineInstance
		var deviceName string = "nvidia.com/GRID_T4-1B"
		var mdevSelector string = "GRID T4-1B"
		var updatedDeviceName string = "nvidia.com/GRID_T4-2B"
		var updatedMdevSelector string = "GRID T4-2B"
		var parentDeviceID string = "10de:1eb8"
		var desiredMdevTypeName string = "nvidia-222"
		var expectedInstancesNum int = 16
		var config v1.KubeVirtConfiguration
		var mdevTestLabel = "mdevTestLabel1"

		BeforeEach(func() {
			tests.BeforeTestCleanup()
			kv := util.GetCurrentKv(virtClient)

			By("Creating a configuration for mediated devices")
			config = kv.Spec.Configuration
			config.DeveloperConfiguration.FeatureGates = []string{virtconfig.GPUGate}
			config.MediatedDevicesConfiguration = &v1.MediatedDevicesConfiguration{
				MediatedDevicesTypes: []string{desiredMdevTypeName},
			}
			config.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: mdevSelector,
						ResourceName:     deviceName,
					},
					{
						MDEVNameSelector: updatedMdevSelector,
						ResourceName:     updatedDeviceName,
					},
				},
			}
			tests.UpdateKubeVirtConfigValueAndWait(config)

			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVCreated(desiredMdevTypeName, expectedInstancesNum)
		})

		cleanupConfiguredMdevs := func() {
			By("Deleting the VMI")
			Expect(virtClient.VirtualMachineInstance(vmi.Namespace).Delete(vmi.Name, &metav1.DeleteOptions{})).To(Succeed(), "Should delete VMI")
			By("Creating a configuration for mediated devices")
			config.MediatedDevicesConfiguration = &v1.MediatedDevicesConfiguration{}
			tests.UpdateKubeVirtConfigValueAndWait(config)
			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVRemoved()
		}

		AfterEach(func() {
			cleanupConfiguredMdevs()
		})

		It("Should successfully passthrough a mediated device", func() {

			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu1",
					DeviceName: deviceName,
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)
			Expect(console.LoginToFedora(vmi)).To(Succeed())

			By("Making sure the device is present inside the VMI")
			Expect(console.SafeExpectBatch(vmi, []expect.Batcher{
				&expect.BSnd{S: "grep -c " + strings.Replace(parentDeviceID, ":", "", 1) + " /proc/bus/pci/devices\n"},
				&expect.BExp{R: console.RetValue("1")},
			}, 250)).To(Succeed(), "Device not found")

			domXml, err := tests.GetRunningVirtualMachineInstanceDomainXML(virtClient, vmi)
			Expect(err).ToNot(HaveOccurred())
			// make sure that one mdev has display and ramfb on
			By("Maiking sure that a boot display is enabled")
			Expect(domXml).To(MatchRegexp(`<hostdev .*display=.?on.?`), "Display should be on")
			Expect(domXml).To(MatchRegexp(`<hostdev .*ramfb=.?on.?`), "RamFB should be on")
		})
		It("Should successfully passthrough a mediated device with a disabled display", func() {
			_false := false
			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu2",
					DeviceName: deviceName,
					VirtualGPUOptions: &v1.VGPUOptions{
						Display: &v1.VGPUDisplayOptions{
							Enabled: &_false,
						},
					},
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)

			domXml, err := tests.GetRunningVirtualMachineInstanceDomainXML(virtClient, vmi)
			Expect(err).ToNot(HaveOccurred())
			// make sure that another mdev explicitly turned off its display
			By("Maiking sure that a boot display is disabled")
			Expect(domXml).ToNot(MatchRegexp(`<hostdev .*display=.?on.?`), "Display should not be enabled")
			Expect(domXml).ToNot(MatchRegexp(`<hostdev .*ramfb=.?on.?`), "RamFB should not be enabled")
		})
		It("Should override default mdev configuration on a specific node", func() {
			newDesiredMdevTypeName := "nvidia-223"
			newExpectedInstancesNum := 8
			By("Creating a configuration for mediated devices")
			config.MediatedDevicesConfiguration.NodeMediatedDeviceTypes = []v1.NodeMediatedDeviceTypesConfig{
				{
					NodeSelector: map[string]string{
						cleanup.TestLabelForNamespace(util.NamespaceTestDefault): mdevTestLabel,
					},
					MediatedDevicesTypes: []string{
						"nvidia-223",
					},
				},
			}
			tests.UpdateKubeVirtConfigValueAndWait(config)
			By("Verify that the default mdev configuration didn't change")
			checkAllMDEVCreated(desiredMdevTypeName, expectedInstancesNum)

			By("Adding a mdevTestLabel1 that should trigger mdev config change")
			// There should be only one node in this lane
			singleNode := util.GetAllSchedulableNodes(virtClient).Items[0]
			libnode.AddLabelToNode(singleNode.Name, cleanup.TestLabelForNamespace(util.NamespaceTestDefault), mdevTestLabel)

			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu1",
					DeviceName: updatedDeviceName,
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)
			Expect(console.LoginToFedora(vmi)).To(Succeed())

			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVCreated(newDesiredMdevTypeName, newExpectedInstancesNum)
		})
	})
})
<commit_msg>e2e: quarantine a few flaky mdev tests<commit_after>package tests_test

import (
	"context"
	"fmt"
	"strings"

	"kubevirt.io/kubevirt/tests/framework/cleanup"

	expect "github.com/google/goexpect"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/kubevirt/tests/util"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"

	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/console"
	. "kubevirt.io/kubevirt/tests/framework/matcher"
	"kubevirt.io/kubevirt/tests/libnode"
)

var _ = Describe("[Serial][sig-compute]MediatedDevices", func() {
	var err error
	var virtClient kubecli.KubevirtClient

	BeforeEach(func() {
		virtClient, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)
	})
	checkAllMDEVCreated := func(mdevTypeName string, expectedInstancesCount int) {
		By(fmt.Sprintf("Checking the number of created mdev types, should be %d of %s type ", expectedInstancesCount, mdevTypeName))
		check := fmt.Sprintf(`set -x
		files_num=$(ls -A /sys/bus/mdev/devices/| wc -l)
		if [[ $files_num != %d ]] ; then
		  echo "failed, not enough mdevs of type %[2]s has been created"
		  exit 1
		fi
		for x in $(ls -A /sys/bus/mdev/devices/); do
		  type_name=$(basename "$(readlink -f /sys/bus/mdev/devices/$x/mdev_type)")
		  if [[ "$type_name" != "%[2]s" ]]; then
		     echo "failed, not all mdevs of type %[2]s"
		     exit 1
		  fi
		  exit 0
		done`, expectedInstancesCount, mdevTypeName)
		testPod := tests.RenderPod("test-pod", []string{"/bin/bash", "-c"}, []string{check})
		testPod, err = virtClient.CoreV1().Pods(util.NamespaceTestDefault).Create(context.Background(), testPod, metav1.CreateOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisPod(testPod), 120).Should(BeInPhase(k8sv1.PodSucceeded))
	}
	checkAllMDEVRemoved := func() {
		check := fmt.Sprintf(`set -x
		files_num=$(ls -A /sys/bus/mdev/devices/| wc -l)
		if [[ $files_num != 0 ]] ; then
		  echo "failed, not all mdevs removed"
		  exit 1
		fi
	        exit 0`)
		testPod := tests.RenderPod("test-pod", []string{"/bin/bash", "-c"}, []string{check})
		testPod, err = virtClient.CoreV1().Pods(util.NamespaceTestDefault).Create(context.Background(), testPod, metav1.CreateOptions{})
		Expect(err).ToNot(HaveOccurred())
		Eventually(ThisPod(testPod), 120).Should(BeInPhase(k8sv1.PodSucceeded))
		Eventually(func() bool {
			nodes, err := virtClient.CoreV1().Nodes().List(context.Background(), metav1.ListOptions{})
			Expect(err).ToNot(HaveOccurred())
			for _, node := range nodes.Items {
				for key, amount := range node.Status.Capacity {
					if strings.HasPrefix(string(key), "nvidia.com/") && !amount.IsZero() {
						return true
					}
				}
			}
			return false
		}, 60).Should(BeFalse(), "wait for the kubelet to stop promoting unconfigured devices")
	}
	Context("with mediated devices configuration", func() {
		var vmi *v1.VirtualMachineInstance
		var deviceName string = "nvidia.com/GRID_T4-1B"
		var mdevSelector string = "GRID T4-1B"
		var updatedDeviceName string = "nvidia.com/GRID_T4-2B"
		var updatedMdevSelector string = "GRID T4-2B"
		var parentDeviceID string = "10de:1eb8"
		var desiredMdevTypeName string = "nvidia-222"
		var expectedInstancesNum int = 16
		var config v1.KubeVirtConfiguration
		var mdevTestLabel = "mdevTestLabel1"

		BeforeEach(func() {
			tests.BeforeTestCleanup()
			kv := util.GetCurrentKv(virtClient)

			By("Creating a configuration for mediated devices")
			config = kv.Spec.Configuration
			config.DeveloperConfiguration.FeatureGates = []string{virtconfig.GPUGate}
			config.MediatedDevicesConfiguration = &v1.MediatedDevicesConfiguration{
				MediatedDevicesTypes: []string{desiredMdevTypeName},
			}
			config.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: mdevSelector,
						ResourceName:     deviceName,
					},
					{
						MDEVNameSelector: updatedMdevSelector,
						ResourceName:     updatedDeviceName,
					},
				},
			}
			tests.UpdateKubeVirtConfigValueAndWait(config)

			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVCreated(desiredMdevTypeName, expectedInstancesNum)
		})

		cleanupConfiguredMdevs := func() {
			By("Deleting the VMI")
			Expect(virtClient.VirtualMachineInstance(vmi.Namespace).Delete(vmi.Name, &metav1.DeleteOptions{})).To(Succeed(), "Should delete VMI")
			By("Creating a configuration for mediated devices")
			config.MediatedDevicesConfiguration = &v1.MediatedDevicesConfiguration{}
			tests.UpdateKubeVirtConfigValueAndWait(config)
			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVRemoved()
		}

		AfterEach(func() {
			cleanupConfiguredMdevs()
		})

		It("[QUARANTINE]Should successfully passthrough a mediated device", func() {

			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu1",
					DeviceName: deviceName,
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)
			Expect(console.LoginToFedora(vmi)).To(Succeed())

			By("Making sure the device is present inside the VMI")
			Expect(console.SafeExpectBatch(vmi, []expect.Batcher{
				&expect.BSnd{S: "grep -c " + strings.Replace(parentDeviceID, ":", "", 1) + " /proc/bus/pci/devices\n"},
				&expect.BExp{R: console.RetValue("1")},
			}, 250)).To(Succeed(), "Device not found")

			domXml, err := tests.GetRunningVirtualMachineInstanceDomainXML(virtClient, vmi)
			Expect(err).ToNot(HaveOccurred())
			// make sure that one mdev has display and ramfb on
			By("Maiking sure that a boot display is enabled")
			Expect(domXml).To(MatchRegexp(`<hostdev .*display=.?on.?`), "Display should be on")
			Expect(domXml).To(MatchRegexp(`<hostdev .*ramfb=.?on.?`), "RamFB should be on")
		})
		It("Should successfully passthrough a mediated device with a disabled display", func() {
			_false := false
			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu2",
					DeviceName: deviceName,
					VirtualGPUOptions: &v1.VGPUOptions{
						Display: &v1.VGPUDisplayOptions{
							Enabled: &_false,
						},
					},
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)

			domXml, err := tests.GetRunningVirtualMachineInstanceDomainXML(virtClient, vmi)
			Expect(err).ToNot(HaveOccurred())
			// make sure that another mdev explicitly turned off its display
			By("Maiking sure that a boot display is disabled")
			Expect(domXml).ToNot(MatchRegexp(`<hostdev .*display=.?on.?`), "Display should not be enabled")
			Expect(domXml).ToNot(MatchRegexp(`<hostdev .*ramfb=.?on.?`), "RamFB should not be enabled")
		})
		It("[QUARANTINE]Should override default mdev configuration on a specific node", func() {
			newDesiredMdevTypeName := "nvidia-223"
			newExpectedInstancesNum := 8
			By("Creating a configuration for mediated devices")
			config.MediatedDevicesConfiguration.NodeMediatedDeviceTypes = []v1.NodeMediatedDeviceTypesConfig{
				{
					NodeSelector: map[string]string{
						cleanup.TestLabelForNamespace(util.NamespaceTestDefault): mdevTestLabel,
					},
					MediatedDevicesTypes: []string{
						"nvidia-223",
					},
				},
			}
			tests.UpdateKubeVirtConfigValueAndWait(config)
			By("Verify that the default mdev configuration didn't change")
			checkAllMDEVCreated(desiredMdevTypeName, expectedInstancesNum)

			By("Adding a mdevTestLabel1 that should trigger mdev config change")
			// There should be only one node in this lane
			singleNode := util.GetAllSchedulableNodes(virtClient).Items[0]
			libnode.AddLabelToNode(singleNode.Name, cleanup.TestLabelForNamespace(util.NamespaceTestDefault), mdevTestLabel)

			By("Creating a Fedora VMI")
			vmi = tests.NewRandomFedoraVMIWithGuestAgent()
			vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1G")
			vGPUs := []v1.GPU{
				{
					Name:       "gpu1",
					DeviceName: updatedDeviceName,
				},
			}
			vmi.Spec.Domain.Devices.GPUs = vGPUs
			createdVmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
			Expect(err).ToNot(HaveOccurred())
			vmi = createdVmi
			tests.WaitForSuccessfulVMIStart(vmi)
			Expect(console.LoginToFedora(vmi)).To(Succeed())

			By("Verifying that an expected amount of devices has been created")
			checkAllMDEVCreated(newDesiredMdevTypeName, newExpectedInstancesNum)
		})
	})
})
<|endoftext|>


<commit_before>package deploymentconfigs

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/golang/glog"

	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	kapi "k8s.io/kubernetes/pkg/apis/core"

	appsapi "github.com/openshift/origin/pkg/apps/apis/apps"
	appsclient "github.com/openshift/origin/pkg/apps/generated/internalclientset/typed/apps/internalversion"
	triggerapi "github.com/openshift/origin/pkg/image/apis/image/v1/trigger"
	"github.com/openshift/origin/pkg/image/trigger"
)

func indicesForContainerNames(containers []kapi.Container, names []string) []int {
	var index []int
	for _, name := range names {
		for i, container := range containers {
			if name == container.Name {
				index = append(index, i)
			}
		}
	}
	return index
}

func namesInclude(names []string, name string) bool {
	for _, n := range names {
		if name == n {
			return true
		}
	}
	return false
}

// calculateDeploymentConfigTrigger resolves a particular trigger against the deployment config and extracts a list of triggers.
// It will silently ignore triggers that do not point to valid container. It returns empty if no triggers can be found.
func calculateDeploymentConfigTrigger(t appsapi.DeploymentTriggerPolicy, dc *appsapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	if t.ImageChangeParams == nil {
		return nil
	}
	from := t.ImageChangeParams.From
	if from.Kind != "ImageStreamTag" || len(from.Name) == 0 {
		return nil
	}

	var triggers []triggerapi.ObjectFieldTrigger

	// add one trigger per init container and container
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.Containers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.containers[@name==\"%s\"].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.InitContainers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.initContainers[@name==\"%s\"].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	return triggers
}

// calculateDeploymentConfigTriggers transforms a deployment config into a set of image change triggers. If retrieveChanges
// is true an array of the latest state of the triggers will be returned.
func calculateDeploymentConfigTriggers(dc *appsapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	var triggers []triggerapi.ObjectFieldTrigger
	for _, t := range dc.Spec.Triggers {
		addedTriggers := calculateDeploymentConfigTrigger(t, dc)
		triggers = append(triggers, addedTriggers...)
	}
	return triggers
}

// deploymentConfigTriggerIndexer converts deployment config events into entries for the trigger cache, and
// also calculates the latest state of the changes on the object.
type deploymentConfigTriggerIndexer struct {
	prefix string
}

func NewDeploymentConfigTriggerIndexer(prefix string) trigger.Indexer {
	return deploymentConfigTriggerIndexer{prefix: prefix}
}

func (i deploymentConfigTriggerIndexer) Index(obj, old interface{}) (string, *trigger.CacheEntry, cache.DeltaType, error) {
	var (
		triggers []triggerapi.ObjectFieldTrigger
		dc       *appsapi.DeploymentConfig
		change   cache.DeltaType
	)
	switch {
	case obj != nil && old == nil:
		// added
		dc = obj.(*appsapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Added
	case old != nil && obj == nil:
		// deleted
		switch deleted := old.(type) {
		case *appsapi.DeploymentConfig:
			dc = deleted
			triggers = calculateDeploymentConfigTriggers(dc)
		case cache.DeletedFinalStateUnknown:
			// don't process triggers as the state could be stale
			glog.V(4).Infof("skipping trigger calculation for deleted deploymentconfig %q", deleted.Key)
		}
		change = cache.Deleted
	default:
		// updated
		dc = obj.(*appsapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		oldTriggers := calculateDeploymentConfigTriggers(old.(*appsapi.DeploymentConfig))
		switch {
		case len(oldTriggers) == 0:
			change = cache.Added
		case !reflect.DeepEqual(oldTriggers, triggers):
			change = cache.Updated
		}
	}

	if len(triggers) > 0 {
		key := i.prefix + dc.Namespace + "/" + dc.Name
		return key, &trigger.CacheEntry{
			Key:       key,
			Namespace: dc.Namespace,
			Triggers:  triggers,
		}, change, nil
	}
	return "", nil, change, nil
}

// DeploymentConfigReactor converts image trigger changes into updates on deployments.
type DeploymentConfigReactor struct {
	Client appsclient.DeploymentConfigsGetter
}

// UpdateDeploymentConfigImages sets the latest image value from all triggers onto each container, returning false if
// one or more triggers could not be resolved yet or an error. The returned dc will be copied if mutated.
func UpdateDeploymentConfigImages(dc *appsapi.DeploymentConfig, tagRetriever trigger.TagRetriever) (*appsapi.DeploymentConfig, bool, error) {
	var updated *appsapi.DeploymentConfig

	// copy the object and reset dc to the copy
	copyObject := func() {
		if updated != nil {
			return
		}
		dc = dc.DeepCopy()
		updated = dc
	}

	for i, t := range dc.Spec.Triggers {
		p := t.ImageChangeParams
		if p == nil || p.From.Kind != "ImageStreamTag" {
			continue
		}
		if !p.Automatic {
			continue
		}

		namespace := p.From.Namespace
		if len(namespace) == 0 {
			namespace = dc.Namespace
		}

		ref, _, ok := tagRetriever.ImageStreamTag(namespace, p.From.Name)
		if !ok && len(p.LastTriggeredImage) == 0 {
			glog.V(4).Infof("trigger %#v in deployment %s is not resolveable", p, dc.Name)
			return nil, false, nil
		}
		if ref == p.LastTriggeredImage {
			continue
		}

		if len(ref) == 0 {
			ref = p.LastTriggeredImage
		}

		if p.LastTriggeredImage != ref {
			copyObject()
			p = dc.Spec.Triggers[i].ImageChangeParams
			p.LastTriggeredImage = ref
		}

		for i, c := range dc.Spec.Template.Spec.InitContainers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.InitContainers[i]
			container.Image = ref
		}

		for i, c := range dc.Spec.Template.Spec.Containers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.Containers[i]
			container.Image = ref
		}
	}
	return updated, true, nil
}

// ImageChanged is passed a deployment config and a set of changes.
func (r *DeploymentConfigReactor) ImageChanged(obj runtime.Object, tagRetriever trigger.TagRetriever) error {
	dc := obj.(*appsapi.DeploymentConfig)
	newDC := dc.DeepCopy()

	updated, resolvable, err := UpdateDeploymentConfigImages(newDC, tagRetriever)
	if err != nil {
		return err
	}
	if !resolvable {
		if glog.V(4) {
			glog.Infof("Ignoring changes to deployment config %s, has unresolved images: %s", dc.Name, printDeploymentTriggers(newDC.Spec.Triggers))
		}
		return nil
	}
	if updated == nil {
		glog.V(4).Infof("Deployment config %s has not changed", dc.Name)
		return nil
	}
	glog.V(4).Infof("Deployment config %s has changed", dc.Name)
	_, err = r.Client.DeploymentConfigs(dc.Namespace).Update(updated)
	return err
}

func printDeploymentTriggers(triggers []appsapi.DeploymentTriggerPolicy) string {
	var values []string
	for _, t := range triggers {
		if t.ImageChangeParams == nil {
			continue
		}
		values = append(values, fmt.Sprintf("[from=%s last=%s]", t.ImageChangeParams.From.Name, t.ImageChangeParams.LastTriggeredImage))
	}
	return strings.Join(values, ", ")
}
<commit_msg>Fix DC image reactor to reconcile on DC dc.Spec.Template.Spec.Containers changes<commit_after>package deploymentconfigs

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/golang/glog"

	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	kapi "k8s.io/kubernetes/pkg/apis/core"

	appsapi "github.com/openshift/origin/pkg/apps/apis/apps"
	appsclient "github.com/openshift/origin/pkg/apps/generated/internalclientset/typed/apps/internalversion"
	triggerapi "github.com/openshift/origin/pkg/image/apis/image/v1/trigger"
	"github.com/openshift/origin/pkg/image/trigger"
)

func indicesForContainerNames(containers []kapi.Container, names []string) []int {
	var index []int
	for _, name := range names {
		for i, container := range containers {
			if name == container.Name {
				index = append(index, i)
			}
		}
	}
	return index
}

func namesInclude(names []string, name string) bool {
	for _, n := range names {
		if name == n {
			return true
		}
	}
	return false
}

// calculateDeploymentConfigTrigger resolves a particular trigger against the deployment config and extracts a list of triggers.
// It will silently ignore triggers that do not point to valid container. It returns empty if no triggers can be found.
func calculateDeploymentConfigTrigger(t appsapi.DeploymentTriggerPolicy, dc *appsapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	if t.ImageChangeParams == nil {
		return nil
	}
	from := t.ImageChangeParams.From
	if from.Kind != "ImageStreamTag" || len(from.Name) == 0 {
		return nil
	}

	var triggers []triggerapi.ObjectFieldTrigger

	// add one trigger per init container and container
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.Containers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.containers[@name==\"%s\"].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	for _, index := range indicesForContainerNames(dc.Spec.Template.Spec.InitContainers, t.ImageChangeParams.ContainerNames) {
		fieldPath := fmt.Sprintf("spec.template.spec.initContainers[@name==\"%s\"].image", dc.Spec.Template.Spec.Containers[index].Name)
		triggers = append(triggers, triggerapi.ObjectFieldTrigger{
			From: triggerapi.ObjectReference{
				Name:       from.Name,
				Namespace:  from.Namespace,
				Kind:       from.Kind,
				APIVersion: from.APIVersion,
			},
			FieldPath: fieldPath,
			Paused:    !t.ImageChangeParams.Automatic,
		})
	}
	return triggers
}

// calculateDeploymentConfigTriggers transforms a deployment config into a set of image change triggers. If retrieveChanges
// is true an array of the latest state of the triggers will be returned.
func calculateDeploymentConfigTriggers(dc *appsapi.DeploymentConfig) []triggerapi.ObjectFieldTrigger {
	var triggers []triggerapi.ObjectFieldTrigger
	for _, t := range dc.Spec.Triggers {
		addedTriggers := calculateDeploymentConfigTrigger(t, dc)
		triggers = append(triggers, addedTriggers...)
	}
	return triggers
}

// deploymentConfigTriggerIndexer converts deployment config events into entries for the trigger cache, and
// also calculates the latest state of the changes on the object.
type deploymentConfigTriggerIndexer struct {
	prefix string
}

func NewDeploymentConfigTriggerIndexer(prefix string) trigger.Indexer {
	return deploymentConfigTriggerIndexer{prefix: prefix}
}

func (i deploymentConfigTriggerIndexer) Index(obj, old interface{}) (string, *trigger.CacheEntry, cache.DeltaType, error) {
	var (
		triggers []triggerapi.ObjectFieldTrigger
		dc       *appsapi.DeploymentConfig
		change   cache.DeltaType
	)
	switch {
	case obj != nil && old == nil:
		// added
		dc = obj.(*appsapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		change = cache.Added
	case old != nil && obj == nil:
		// deleted
		switch deleted := old.(type) {
		case *appsapi.DeploymentConfig:
			dc = deleted
			triggers = calculateDeploymentConfigTriggers(dc)
		case cache.DeletedFinalStateUnknown:
			// don't process triggers as the state could be stale
			glog.V(4).Infof("skipping trigger calculation for deleted deploymentconfig %q", deleted.Key)
		}
		change = cache.Deleted
	default:
		// updated
		dc = obj.(*appsapi.DeploymentConfig)
		oldDC := old.(*appsapi.DeploymentConfig)
		triggers = calculateDeploymentConfigTriggers(dc)
		oldTriggers := calculateDeploymentConfigTriggers(oldDC)
		switch {
		case len(oldTriggers) == 0:
			change = cache.Added
		case !reflect.DeepEqual(oldTriggers, triggers):
			change = cache.Updated
		// We need to react on image changes as well. Image names could change,
		// images could be set to different value or resetted to "" e.g. by oc apply
		// and we need to make sure those changes get reconciled by re-resolving images
		case !reflect.DeepEqual(dc.Spec.Template.Spec.Containers, oldDC.Spec.Template.Spec.Containers):
			change = cache.Updated
		}
	}

	if len(triggers) > 0 {
		key := i.prefix + dc.Namespace + "/" + dc.Name
		return key, &trigger.CacheEntry{
			Key:       key,
			Namespace: dc.Namespace,
			Triggers:  triggers,
		}, change, nil
	}
	return "", nil, change, nil
}

// DeploymentConfigReactor converts image trigger changes into updates on deployments.
type DeploymentConfigReactor struct {
	Client appsclient.DeploymentConfigsGetter
}

// UpdateDeploymentConfigImages sets the latest image value from all triggers onto each container, returning false if
// one or more triggers could not be resolved yet or an error. The returned dc will be copied if mutated.
func UpdateDeploymentConfigImages(dc *appsapi.DeploymentConfig, tagRetriever trigger.TagRetriever) (*appsapi.DeploymentConfig, bool, error) {
	var updated *appsapi.DeploymentConfig

	// copy the object and reset dc to the copy
	copyObject := func() {
		if updated != nil {
			return
		}
		dc = dc.DeepCopy()
		updated = dc
	}

	for i, t := range dc.Spec.Triggers {
		p := t.ImageChangeParams
		if p == nil || p.From.Kind != "ImageStreamTag" {
			continue
		}
		if !p.Automatic {
			continue
		}

		namespace := p.From.Namespace
		if len(namespace) == 0 {
			namespace = dc.Namespace
		}

		ref, _, ok := tagRetriever.ImageStreamTag(namespace, p.From.Name)
		if !ok && len(p.LastTriggeredImage) == 0 {
			glog.V(4).Infof("trigger %#v in deployment %s is not resolveable", p, dc.Name)
			return nil, false, nil
		}

		if len(ref) == 0 {
			ref = p.LastTriggeredImage
		}

		if p.LastTriggeredImage != ref {
			copyObject()
			p = dc.Spec.Triggers[i].ImageChangeParams
			p.LastTriggeredImage = ref
		}

		for i, c := range dc.Spec.Template.Spec.InitContainers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.InitContainers[i]
			container.Image = ref
		}

		for i, c := range dc.Spec.Template.Spec.Containers {
			if !namesInclude(p.ContainerNames, c.Name) || c.Image == ref {
				continue
			}
			copyObject()
			container := &dc.Spec.Template.Spec.Containers[i]
			container.Image = ref
		}
	}
	return updated, true, nil
}

// ImageChanged is passed a deployment config and a set of changes.
func (r *DeploymentConfigReactor) ImageChanged(obj runtime.Object, tagRetriever trigger.TagRetriever) error {
	dc := obj.(*appsapi.DeploymentConfig)
	newDC := dc.DeepCopy()

	updated, resolvable, err := UpdateDeploymentConfigImages(newDC, tagRetriever)
	if err != nil {
		return err
	}
	if !resolvable {
		if glog.V(4) {
			glog.Infof("Ignoring changes to deployment config %s, has unresolved images: %s", dc.Name, printDeploymentTriggers(newDC.Spec.Triggers))
		}
		return nil
	}
	if updated == nil {
		glog.V(4).Infof("Deployment config %s has not changed", dc.Name)
		return nil
	}
	glog.V(4).Infof("Deployment config %s has changed", dc.Name)
	_, err = r.Client.DeploymentConfigs(dc.Namespace).Update(updated)
	return err
}

func printDeploymentTriggers(triggers []appsapi.DeploymentTriggerPolicy) string {
	var values []string
	for _, t := range triggers {
		if t.ImageChangeParams == nil {
			continue
		}
		values = append(values, fmt.Sprintf("[from=%s last=%s]", t.ImageChangeParams.From.Name, t.ImageChangeParams.LastTriggeredImage))
	}
	return strings.Join(values, ", ")
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel_test

import (
	"strings"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func TestActiveCallReq(t *testing.T) {
	t.Skip("Test skipped due to unreliable way to test for protocol errors")

	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// Note: This test cannot use log verification as the duplicate ID causes a log.
	// It does not use a verified server, as it leaks a message exchange due to the
	// modification of IDs in the relay.
	opts := testutils.NewOpts().DisableLogVerification()
	testutils.WithServer(t, opts, func(ch *Channel, hostPort string) {
		gotCall := make(chan struct{})
		unblock := make(chan struct{})

		testutils.RegisterFunc(ch, "blocked", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			gotCall <- struct{}{}
			<-unblock
			return &raw.Res{}, nil
		})

		relayFunc := func(outgoing bool, frame *Frame) *Frame {
			if outgoing && frame.Header.ID == 3 {
				frame.Header.ID = 2
			}
			return frame
		}

		relayHostPort, closeRelay := testutils.FrameRelay(t, hostPort, relayFunc)
		defer closeRelay()

		firstComplete := make(chan struct{})
		go func() {
			// This call will block until we close unblock.
			raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
			close(firstComplete)
		}()

		// Wait for the first call to be received by the server
		<-gotCall

		// Make a new call, which should fail
		_, _, _, err := raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		assert.Error(t, err, "Expect error")
		assert.True(t, strings.Contains(err.Error(), "already active"),
			"expected already active error, got %v", err)

		close(unblock)
		<-firstComplete
	})
}

func TestInboundConnection(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		testutils.RegisterFunc(ch, "test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			c, _ := InboundConnection(CurrentCall(ctx))
			assert.Equal(t, hostPort, c.RemotePeerInfo().HostPort, "Unexpected host port")
			return &raw.Res{}, nil
		})

		_, _, _, err := raw.Call(ctx, ch, hostPort, ch.PeerInfo().ServiceName, "test", nil, nil)
		require.NoError(t, err, "Call failed")
	})
}

func TestInboundConnection_CallOptions(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	testutils.WithTestServer(t, nil, func(server *testutils.TestServer) {
		server.RegisterFunc("test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			assert.Equal(t, "client", CurrentCall(ctx).CallerName(), "Expected caller name to be passed through")
			return &raw.Res{}, nil
		})

		backendName := server.ServiceName()

		proxyCh := server.NewServer(&testutils.ChannelOpts{ServiceName: "proxy"})
		defer proxyCh.Close()

		subCh := proxyCh.GetSubChannel(backendName)
		subCh.SetHandler(HandlerFunc(func(ctx context.Context, inbound *InboundCall) {
			outbound, err := proxyCh.BeginCall(ctx, server.HostPort(), backendName, inbound.MethodString(), inbound.CallOptions())
			require.NoError(t, err, "Create outbound call failed")
			arg2, arg3, _, err := raw.WriteArgs(outbound, []byte("hello"), []byte("world"))
			require.NoError(t, err, "Write outbound call failed")
			require.NoError(t, raw.WriteResponse(inbound.Response(), &raw.Res{
				Arg2: arg2,
				Arg3: arg3,
			}), "Write response failed")
		}))

		clientCh := server.NewClient(&testutils.ChannelOpts{
			ServiceName: "client",
		})
		defer clientCh.Close()

		_, _, _, err := raw.Call(ctx, clientCh, proxyCh.PeerInfo().HostPort, backendName, "test", nil, nil)
		require.NoError(t, err, "Call through proxy failed")
	})
}
<commit_msg>Fix flaky test TestInboundConnection (#588)<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel_test

import (
	"strings"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func TestActiveCallReq(t *testing.T) {
	t.Skip("Test skipped due to unreliable way to test for protocol errors")

	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// Note: This test cannot use log verification as the duplicate ID causes a log.
	// It does not use a verified server, as it leaks a message exchange due to the
	// modification of IDs in the relay.
	opts := testutils.NewOpts().DisableLogVerification()
	testutils.WithServer(t, opts, func(ch *Channel, hostPort string) {
		gotCall := make(chan struct{})
		unblock := make(chan struct{})

		testutils.RegisterFunc(ch, "blocked", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			gotCall <- struct{}{}
			<-unblock
			return &raw.Res{}, nil
		})

		relayFunc := func(outgoing bool, frame *Frame) *Frame {
			if outgoing && frame.Header.ID == 3 {
				frame.Header.ID = 2
			}
			return frame
		}

		relayHostPort, closeRelay := testutils.FrameRelay(t, hostPort, relayFunc)
		defer closeRelay()

		firstComplete := make(chan struct{})
		go func() {
			// This call will block until we close unblock.
			raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
			close(firstComplete)
		}()

		// Wait for the first call to be received by the server
		<-gotCall

		// Make a new call, which should fail
		_, _, _, err := raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		assert.Error(t, err, "Expect error")
		assert.True(t, strings.Contains(err.Error(), "already active"),
			"expected already active error, got %v", err)

		close(unblock)
		<-firstComplete
	})
}

func TestInboundConnection(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// Disable relay since relays hide host:port on outbound calls.
	opts := testutils.NewOpts().NoRelay()
	testutils.WithTestServer(t, opts, func(ts *testutils.TestServer) {
		s2 := ts.NewServer(nil)

		ts.RegisterFunc("test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			c, _ := InboundConnection(CurrentCall(ctx))
			assert.Equal(t, s2.PeerInfo().HostPort, c.RemotePeerInfo().HostPort, "Unexpected host port")
			return &raw.Res{}, nil
		})

		_, _, _, err := raw.Call(ctx, s2, ts.HostPort(), ts.ServiceName(), "test", nil, nil)
		require.NoError(t, err, "Call failed")
	})
}

func TestInboundConnection_CallOptions(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	testutils.WithTestServer(t, nil, func(server *testutils.TestServer) {
		server.RegisterFunc("test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			assert.Equal(t, "client", CurrentCall(ctx).CallerName(), "Expected caller name to be passed through")
			return &raw.Res{}, nil
		})

		backendName := server.ServiceName()

		proxyCh := server.NewServer(&testutils.ChannelOpts{ServiceName: "proxy"})
		defer proxyCh.Close()

		subCh := proxyCh.GetSubChannel(backendName)
		subCh.SetHandler(HandlerFunc(func(ctx context.Context, inbound *InboundCall) {
			outbound, err := proxyCh.BeginCall(ctx, server.HostPort(), backendName, inbound.MethodString(), inbound.CallOptions())
			require.NoError(t, err, "Create outbound call failed")
			arg2, arg3, _, err := raw.WriteArgs(outbound, []byte("hello"), []byte("world"))
			require.NoError(t, err, "Write outbound call failed")
			require.NoError(t, raw.WriteResponse(inbound.Response(), &raw.Res{
				Arg2: arg2,
				Arg3: arg3,
			}), "Write response failed")
		}))

		clientCh := server.NewClient(&testutils.ChannelOpts{
			ServiceName: "client",
		})
		defer clientCh.Close()

		_, _, _, err := raw.Call(ctx, clientCh, proxyCh.PeerInfo().HostPort, backendName, "test", nil, nil)
		require.NoError(t, err, "Call through proxy failed")
	})
}
<|endoftext|>