<commit_before>package recipebuilder

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/bbs/models"
	ssh_routes "code.cloudfoundry.org/diego-ssh/routes"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/nsync/helpers"
	"code.cloudfoundry.org/runtimeschema/cc_messages"
)

type BuildpackRecipeBuilder struct {
	logger lager.Logger
	config Config
}

func NewBuildpackRecipeBuilder(logger lager.Logger, config Config) *BuildpackRecipeBuilder {
	return &BuildpackRecipeBuilder{
		logger: logger,
		config: config,
	}
}

func (b *BuildpackRecipeBuilder) BuildTask(task *cc_messages.TaskRequestFromCC) (*models.TaskDefinition, error) {
	logger := b.logger.Session("build-task", lager.Data{"request": task})

	if task.DropletUri == "" {
		logger.Error("missing-droplet-source", ErrDropletSourceMissing)
		return nil, ErrDropletSourceMissing
	}

	if task.DockerPath != "" {
		logger.Error("invalid-docker-path", ErrMultipleAppSources)
		return nil, ErrMultipleAppSources
	}

	downloadAction := &models.DownloadAction{
		From:     task.DropletUri,
		To:       ".",
		CacheKey: "",
		User:     "vcap",
	}

	if task.DropletHash != "" {
		downloadAction.ChecksumAlgorithm = "sha1"
		downloadAction.ChecksumValue = task.DropletHash
	}

	logger.Info("downloadAction", lager.Data{"action": downloadAction})

	runAction := &models.RunAction{
		User:           "vcap",
		Path:           "/tmp/lifecycle/launcher",
		Args:           []string{"app", task.Command, ""},
		Env:            task.EnvironmentVariables,
		LogSource:      task.LogSource,
		ResourceLimits: &models.ResourceLimits{},
	}

	var lifecycle = "buildpack/" + task.RootFs
	lifecyclePath, ok := b.config.Lifecycles[lifecycle]
	if !ok {
		logger.Error("unknown-lifecycle", ErrNoLifecycleDefined, lager.Data{
			"lifecycle": lifecycle,
		})

		return nil, ErrNoLifecycleDefined
	}

	lifecycleURL := lifecycleDownloadURL(lifecyclePath, b.config.FileServerURL)

	cachedDependencies := []*models.CachedDependency{
		&models.CachedDependency{
			From:     lifecycleURL,
			To:       "/tmp/lifecycle",
			CacheKey: fmt.Sprintf("%s-lifecycle", strings.Replace(lifecycle, "/", "-", 1)),
		},
	}

	rootFSPath := models.PreloadedRootFS(task.RootFs)

	taskDefinition := &models.TaskDefinition{
		Privileged:            b.config.PrivilegedContainers,
		LogGuid:               task.LogGuid,
		MemoryMb:              int32(task.MemoryMb),
		DiskMb:                int32(task.DiskMb),
		CpuWeight:             cpuWeight(task.MemoryMb),
		EnvironmentVariables:  task.EnvironmentVariables,
		RootFs:                rootFSPath,
		CompletionCallbackUrl: task.CompletionCallbackUrl,
		Action: models.WrapAction(models.Serial(
			downloadAction,
			runAction,
		)),
		CachedDependencies:            cachedDependencies,
		EgressRules:                   task.EgressRules,
		LegacyDownloadUser:            "vcap",
		TrustedSystemCertificatesPath: TrustedSystemCertificatesPath,
		LogSource:                     task.LogSource,
		VolumeMounts:                  task.VolumeMounts,
	}

	return taskDefinition, nil
}

func (b *BuildpackRecipeBuilder) Build(desiredApp *cc_messages.DesireAppRequestFromCC) (*models.DesiredLRP, error) {
	lrpGuid := desiredApp.ProcessGuid

	buildLogger := b.logger.Session("message-builder")

	if desiredApp.DropletUri == "" {
		buildLogger.Error("desired-app-invalid", ErrDropletSourceMissing, lager.Data{"desired-app": desiredApp})
		return nil, ErrDropletSourceMissing
	}

	if desiredApp.DropletUri != "" && desiredApp.DockerImageUrl != "" {
		buildLogger.Error("desired-app-invalid", ErrMultipleAppSources, lager.Data{"desired-app": desiredApp})
		return nil, ErrMultipleAppSources
	}

	var lifecycle = "buildpack/" + desiredApp.Stack
	lifecyclePath, ok := b.config.Lifecycles[lifecycle]
	if !ok {
		buildLogger.Error("unknown-lifecycle", ErrNoLifecycleDefined, lager.Data{
			"lifecycle": lifecycle,
		})

		return nil, ErrNoLifecycleDefined
	}

	lifecycleURL := lifecycleDownloadURL(lifecyclePath, b.config.FileServerURL)

	rootFSPath := models.PreloadedRootFS(desiredApp.Stack)

	var containerEnvVars []*models.EnvironmentVariable
	containerEnvVars = append(containerEnvVars, &models.EnvironmentVariable{Name: "LANG", Value: DefaultLANG})

	numFiles := DefaultFileDescriptorLimit
	if desiredApp.FileDescriptors != 0 {
		numFiles = desiredApp.FileDescriptors
	}

	var setup []models.ActionInterface
	var actions []models.ActionInterface
	var monitor models.ActionInterface

	cachedDependencies := []*models.CachedDependency{}
	cachedDependencies = append(cachedDependencies, &models.CachedDependency{
		From:     lifecycleURL,
		To:       "/tmp/lifecycle",
		CacheKey: fmt.Sprintf("%s-lifecycle", strings.Replace(lifecycle, "/", "-", 1)),
	})

	desiredAppPorts, err := b.ExtractExposedPorts(desiredApp)
	if err != nil {
		return nil, err
	}

	switch desiredApp.HealthCheckType {
	case cc_messages.PortHealthCheckType, cc_messages.UnspecifiedHealthCheckType:
		monitor = models.Timeout(getParallelAction(desiredAppPorts, "vcap"), 30*time.Second)
	}

	downloadAction := &models.DownloadAction{
		From:     desiredApp.DropletUri,
		To:       ".",
		CacheKey: fmt.Sprintf("droplets-%s", lrpGuid),
		User:     "vcap",
	}

	if desiredApp.DropletHash != "" {
		downloadAction.ChecksumAlgorithm = "sha1"
		downloadAction.ChecksumValue = desiredApp.DropletHash
	}

	setup = append(setup, downloadAction)
	actions = append(actions, &models.RunAction{
		User: "vcap",
		Path: "/tmp/lifecycle/launcher",
		Args: append(
			[]string{"app"},
			desiredApp.StartCommand,
			desiredApp.ExecutionMetadata,
		),
		Env:       createLrpEnv(desiredApp.Environment, desiredAppPorts),
		LogSource: getAppLogSource(desiredApp.LogSource),
		ResourceLimits: &models.ResourceLimits{
			Nofile: &numFiles,
		},
	})

	desiredAppRoutingInfo, err := helpers.CCRouteInfoToRoutes(desiredApp.RoutingInfo, desiredAppPorts)
	if err != nil {
		buildLogger.Error("marshaling-cc-route-info-failed", err)
		return nil, err
	}

	if desiredApp.AllowSSH {
		hostKeyPair, err := b.config.KeyFactory.NewKeyPair(1024)
		if err != nil {
			buildLogger.Error("new-host-key-pair-failed", err)
			return nil, err
		}

		userKeyPair, err := b.config.KeyFactory.NewKeyPair(1024)
		if err != nil {
			buildLogger.Error("new-user-key-pair-failed", err)
			return nil, err
		}

		actions = append(actions, &models.RunAction{
			User: "vcap",
			Path: "/tmp/lifecycle/diego-sshd",
			Args: []string{
				"-address=" + fmt.Sprintf("0.0.0.0:%d", DefaultSSHPort),
				"-hostKey=" + hostKeyPair.PEMEncodedPrivateKey(),
				"-authorizedKey=" + userKeyPair.AuthorizedKey(),
				"-inheritDaemonEnv",
				"-logLevel=fatal",
			},
			Env: createLrpEnv(desiredApp.Environment, desiredAppPorts),
			ResourceLimits: &models.ResourceLimits{
				Nofile: &numFiles,
			},
		})

		sshRoutePayload, err := json.Marshal(ssh_routes.SSHRoute{
			ContainerPort:   2222,
			PrivateKey:      userKeyPair.PEMEncodedPrivateKey(),
			HostFingerprint: hostKeyPair.Fingerprint(),
		})

		if err != nil {
			buildLogger.Error("marshaling-ssh-route-failed", err)
			return nil, err
		}

		sshRouteMessage := json.RawMessage(sshRoutePayload)
		desiredAppRoutingInfo[ssh_routes.DIEGO_SSH] = &sshRouteMessage
		desiredAppPorts = append(desiredAppPorts, DefaultSSHPort)
	}

	setupAction := models.Serial(setup...)
	actionAction := models.Codependent(actions...)

	return &models.DesiredLRP{
		Privileged: b.config.PrivilegedContainers,

		Domain: cc_messages.AppLRPDomain,

		ProcessGuid: lrpGuid,
		Instances:   int32(desiredApp.NumInstances),
		Routes:      &desiredAppRoutingInfo,
		Annotation:  desiredApp.ETag,

		CpuWeight: cpuWeight(desiredApp.MemoryMB),

		MemoryMb: int32(desiredApp.MemoryMB),
		DiskMb:   int32(desiredApp.DiskMB),

		Ports: desiredAppPorts,

		RootFs: rootFSPath,

		LogGuid:   desiredApp.LogGuid,
		LogSource: LRPLogSource,

		MetricsGuid: desiredApp.LogGuid,

		EnvironmentVariables: containerEnvVars,
		CachedDependencies:   cachedDependencies,
		Setup:                models.WrapAction(setupAction),
		Action:               models.WrapAction(actionAction),
		Monitor:              models.WrapAction(monitor),

		StartTimeoutMs: int64(desiredApp.HealthCheckTimeoutInSeconds * 1000),

		EgressRules:        desiredApp.EgressRules,
		Network:            desiredApp.Network,
		LegacyDownloadUser: "vcap",

		TrustedSystemCertificatesPath: TrustedSystemCertificatesPath,
		VolumeMounts:                  desiredApp.VolumeMounts,
	}, nil
}

func (b BuildpackRecipeBuilder) ExtractExposedPorts(desiredApp *cc_messages.DesireAppRequestFromCC) ([]uint32, error) {
	return getDesiredAppPorts(desiredApp.Ports), nil
}
<commit_msg>Remove logging that exposes sensitive droplet info<commit_after>package recipebuilder

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"code.cloudfoundry.org/bbs/models"
	ssh_routes "code.cloudfoundry.org/diego-ssh/routes"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/nsync/helpers"
	"code.cloudfoundry.org/runtimeschema/cc_messages"
)

type BuildpackRecipeBuilder struct {
	logger lager.Logger
	config Config
}

func NewBuildpackRecipeBuilder(logger lager.Logger, config Config) *BuildpackRecipeBuilder {
	return &BuildpackRecipeBuilder{
		logger: logger,
		config: config,
	}
}

func (b *BuildpackRecipeBuilder) BuildTask(task *cc_messages.TaskRequestFromCC) (*models.TaskDefinition, error) {
	logger := b.logger.Session("build-task", lager.Data{"request": task})

	if task.DropletUri == "" {
		logger.Error("missing-droplet-source", ErrDropletSourceMissing)
		return nil, ErrDropletSourceMissing
	}

	if task.DockerPath != "" {
		logger.Error("invalid-docker-path", ErrMultipleAppSources)
		return nil, ErrMultipleAppSources
	}

	downloadAction := &models.DownloadAction{
		From:     task.DropletUri,
		To:       ".",
		CacheKey: "",
		User:     "vcap",
	}

	if task.DropletHash != "" {
		downloadAction.ChecksumAlgorithm = "sha1"
		downloadAction.ChecksumValue = task.DropletHash
	}

	runAction := &models.RunAction{
		User:           "vcap",
		Path:           "/tmp/lifecycle/launcher",
		Args:           []string{"app", task.Command, ""},
		Env:            task.EnvironmentVariables,
		LogSource:      task.LogSource,
		ResourceLimits: &models.ResourceLimits{},
	}

	var lifecycle = "buildpack/" + task.RootFs
	lifecyclePath, ok := b.config.Lifecycles[lifecycle]
	if !ok {
		logger.Error("unknown-lifecycle", ErrNoLifecycleDefined, lager.Data{
			"lifecycle": lifecycle,
		})

		return nil, ErrNoLifecycleDefined
	}

	lifecycleURL := lifecycleDownloadURL(lifecyclePath, b.config.FileServerURL)

	cachedDependencies := []*models.CachedDependency{
		&models.CachedDependency{
			From:     lifecycleURL,
			To:       "/tmp/lifecycle",
			CacheKey: fmt.Sprintf("%s-lifecycle", strings.Replace(lifecycle, "/", "-", 1)),
		},
	}

	rootFSPath := models.PreloadedRootFS(task.RootFs)

	taskDefinition := &models.TaskDefinition{
		Privileged:            b.config.PrivilegedContainers,
		LogGuid:               task.LogGuid,
		MemoryMb:              int32(task.MemoryMb),
		DiskMb:                int32(task.DiskMb),
		CpuWeight:             cpuWeight(task.MemoryMb),
		EnvironmentVariables:  task.EnvironmentVariables,
		RootFs:                rootFSPath,
		CompletionCallbackUrl: task.CompletionCallbackUrl,
		Action: models.WrapAction(models.Serial(
			downloadAction,
			runAction,
		)),
		CachedDependencies:            cachedDependencies,
		EgressRules:                   task.EgressRules,
		LegacyDownloadUser:            "vcap",
		TrustedSystemCertificatesPath: TrustedSystemCertificatesPath,
		LogSource:                     task.LogSource,
		VolumeMounts:                  task.VolumeMounts,
	}

	return taskDefinition, nil
}

func (b *BuildpackRecipeBuilder) Build(desiredApp *cc_messages.DesireAppRequestFromCC) (*models.DesiredLRP, error) {
	lrpGuid := desiredApp.ProcessGuid

	buildLogger := b.logger.Session("message-builder")

	if desiredApp.DropletUri == "" {
		buildLogger.Error("desired-app-invalid", ErrDropletSourceMissing, lager.Data{"desired-app": desiredApp})
		return nil, ErrDropletSourceMissing
	}

	if desiredApp.DropletUri != "" && desiredApp.DockerImageUrl != "" {
		buildLogger.Error("desired-app-invalid", ErrMultipleAppSources, lager.Data{"desired-app": desiredApp})
		return nil, ErrMultipleAppSources
	}

	var lifecycle = "buildpack/" + desiredApp.Stack
	lifecyclePath, ok := b.config.Lifecycles[lifecycle]
	if !ok {
		buildLogger.Error("unknown-lifecycle", ErrNoLifecycleDefined, lager.Data{
			"lifecycle": lifecycle,
		})

		return nil, ErrNoLifecycleDefined
	}

	lifecycleURL := lifecycleDownloadURL(lifecyclePath, b.config.FileServerURL)

	rootFSPath := models.PreloadedRootFS(desiredApp.Stack)

	var containerEnvVars []*models.EnvironmentVariable
	containerEnvVars = append(containerEnvVars, &models.EnvironmentVariable{Name: "LANG", Value: DefaultLANG})

	numFiles := DefaultFileDescriptorLimit
	if desiredApp.FileDescriptors != 0 {
		numFiles = desiredApp.FileDescriptors
	}

	var setup []models.ActionInterface
	var actions []models.ActionInterface
	var monitor models.ActionInterface

	cachedDependencies := []*models.CachedDependency{}
	cachedDependencies = append(cachedDependencies, &models.CachedDependency{
		From:     lifecycleURL,
		To:       "/tmp/lifecycle",
		CacheKey: fmt.Sprintf("%s-lifecycle", strings.Replace(lifecycle, "/", "-", 1)),
	})

	desiredAppPorts, err := b.ExtractExposedPorts(desiredApp)
	if err != nil {
		return nil, err
	}

	switch desiredApp.HealthCheckType {
	case cc_messages.PortHealthCheckType, cc_messages.UnspecifiedHealthCheckType:
		monitor = models.Timeout(getParallelAction(desiredAppPorts, "vcap"), 30*time.Second)
	}

	downloadAction := &models.DownloadAction{
		From:     desiredApp.DropletUri,
		To:       ".",
		CacheKey: fmt.Sprintf("droplets-%s", lrpGuid),
		User:     "vcap",
	}

	if desiredApp.DropletHash != "" {
		downloadAction.ChecksumAlgorithm = "sha1"
		downloadAction.ChecksumValue = desiredApp.DropletHash
	}

	setup = append(setup, downloadAction)
	actions = append(actions, &models.RunAction{
		User: "vcap",
		Path: "/tmp/lifecycle/launcher",
		Args: append(
			[]string{"app"},
			desiredApp.StartCommand,
			desiredApp.ExecutionMetadata,
		),
		Env:       createLrpEnv(desiredApp.Environment, desiredAppPorts),
		LogSource: getAppLogSource(desiredApp.LogSource),
		ResourceLimits: &models.ResourceLimits{
			Nofile: &numFiles,
		},
	})

	desiredAppRoutingInfo, err := helpers.CCRouteInfoToRoutes(desiredApp.RoutingInfo, desiredAppPorts)
	if err != nil {
		buildLogger.Error("marshaling-cc-route-info-failed", err)
		return nil, err
	}

	if desiredApp.AllowSSH {
		hostKeyPair, err := b.config.KeyFactory.NewKeyPair(1024)
		if err != nil {
			buildLogger.Error("new-host-key-pair-failed", err)
			return nil, err
		}

		userKeyPair, err := b.config.KeyFactory.NewKeyPair(1024)
		if err != nil {
			buildLogger.Error("new-user-key-pair-failed", err)
			return nil, err
		}

		actions = append(actions, &models.RunAction{
			User: "vcap",
			Path: "/tmp/lifecycle/diego-sshd",
			Args: []string{
				"-address=" + fmt.Sprintf("0.0.0.0:%d", DefaultSSHPort),
				"-hostKey=" + hostKeyPair.PEMEncodedPrivateKey(),
				"-authorizedKey=" + userKeyPair.AuthorizedKey(),
				"-inheritDaemonEnv",
				"-logLevel=fatal",
			},
			Env: createLrpEnv(desiredApp.Environment, desiredAppPorts),
			ResourceLimits: &models.ResourceLimits{
				Nofile: &numFiles,
			},
		})

		sshRoutePayload, err := json.Marshal(ssh_routes.SSHRoute{
			ContainerPort:   2222,
			PrivateKey:      userKeyPair.PEMEncodedPrivateKey(),
			HostFingerprint: hostKeyPair.Fingerprint(),
		})

		if err != nil {
			buildLogger.Error("marshaling-ssh-route-failed", err)
			return nil, err
		}

		sshRouteMessage := json.RawMessage(sshRoutePayload)
		desiredAppRoutingInfo[ssh_routes.DIEGO_SSH] = &sshRouteMessage
		desiredAppPorts = append(desiredAppPorts, DefaultSSHPort)
	}

	setupAction := models.Serial(setup...)
	actionAction := models.Codependent(actions...)

	return &models.DesiredLRP{
		Privileged: b.config.PrivilegedContainers,

		Domain: cc_messages.AppLRPDomain,

		ProcessGuid: lrpGuid,
		Instances:   int32(desiredApp.NumInstances),
		Routes:      &desiredAppRoutingInfo,
		Annotation:  desiredApp.ETag,

		CpuWeight: cpuWeight(desiredApp.MemoryMB),

		MemoryMb: int32(desiredApp.MemoryMB),
		DiskMb:   int32(desiredApp.DiskMB),

		Ports: desiredAppPorts,

		RootFs: rootFSPath,

		LogGuid:   desiredApp.LogGuid,
		LogSource: LRPLogSource,

		MetricsGuid: desiredApp.LogGuid,

		EnvironmentVariables: containerEnvVars,
		CachedDependencies:   cachedDependencies,
		Setup:                models.WrapAction(setupAction),
		Action:               models.WrapAction(actionAction),
		Monitor:              models.WrapAction(monitor),

		StartTimeoutMs: int64(desiredApp.HealthCheckTimeoutInSeconds * 1000),

		EgressRules:        desiredApp.EgressRules,
		Network:            desiredApp.Network,
		LegacyDownloadUser: "vcap",

		TrustedSystemCertificatesPath: TrustedSystemCertificatesPath,
		VolumeMounts:                  desiredApp.VolumeMounts,
	}, nil
}

func (b BuildpackRecipeBuilder) ExtractExposedPorts(desiredApp *cc_messages.DesireAppRequestFromCC) ([]uint32, error) {
	return getDesiredAppPorts(desiredApp.Ports), nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package testing

import (
	"time"

	"launchpad.net/juju-core/utils"
)

// ShortWait is a reasonable amount of time to block waiting for something that
// shouldn't actually happen. (as in, the test suite will *actually* wait this
// long before continuing)
const ShortWait = 50 * time.Millisecond

// LongWait is used when something should have already happened, or happens
// quickly, but we want to make sure we just haven't missed it. As in, the test
// suite should proceed without sleeping at all, but just in case. It is long
// so that we don't have spurious failures without actually slowing down the
// test suite
const LongWait = 10 * time.Second

var LongAttempt = &utils.AttemptStrategy{
	Total: LongWait,
	Delay: ShortWait,
}

// SupportedSeries lists the series known to Juju.
var SupportedSeries = []string{"precise", "quantal", "raring", "saucy", "trusty"}
<commit_msg>utopic is now available from distro-info, so needs to be updated<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package testing

import (
	"time"

	"launchpad.net/juju-core/utils"
)

// ShortWait is a reasonable amount of time to block waiting for something that
// shouldn't actually happen. (as in, the test suite will *actually* wait this
// long before continuing)
const ShortWait = 50 * time.Millisecond

// LongWait is used when something should have already happened, or happens
// quickly, but we want to make sure we just haven't missed it. As in, the test
// suite should proceed without sleeping at all, but just in case. It is long
// so that we don't have spurious failures without actually slowing down the
// test suite
const LongWait = 10 * time.Second

var LongAttempt = &utils.AttemptStrategy{
	Total: LongWait,
	Delay: ShortWait,
}

// SupportedSeries lists the series known to Juju.
var SupportedSeries = []string{"precise", "quantal", "raring", "saucy", "trusty", "utopic"}
<|endoftext|>


<commit_before>package testing

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach-go/testserver"

	// Import postgres driver.
	_ "github.com/lib/pq"
)

// application represents a single instance of an application running an ORM and
// exposing an HTTP REST API.
type application struct {
	language string
	orm      string
}

func (app application) name() string {
	return fmt.Sprintf("%s/%s", app.language, app.orm)
}

func (app application) dir() string {
	return fmt.Sprintf("../%s", app.name())
}

func (app application) dbName() string {
	return fmt.Sprintf("company_%s", app.orm)
}

// initTestDatabase launches a test database as a subprocess.
func initTestDatabase(t *testing.T, app application) (*sql.DB, *url.URL, func()) {
	ts, err := testserver.NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

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

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	ts.WaitForInit(db)

	// Create the database if it does not exist.
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + app.dbName()); err != nil {
		t.Fatal(err)
	}

	// Connect to the database again, now with the database in the URL.
	url.Path = app.dbName()
	db, err = sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	return db, url, func() {
		_ = db.Close()
		ts.Stop()
	}
}

type killFunc func()
type restartFunc func() (killFunc, restartFunc)

// initORMApp launches an ORM application as a subprocess.
func initORMApp(t *testing.T, app application, dbURL *url.URL) (killFunc, restartFunc) {
	addrFlag := fmt.Sprintf("ADDR=%s", dbURL.String())
	args := []string{"make", "start", "-C", app.dir(), addrFlag}

	cmd := exec.Command(args[0], args[1:]...)

	// Set up stderr to display to console and store in a buffer, so we can later
	// verify that it's clean.
	stderr := new(bytes.Buffer)
	cmd.Stderr = io.MultiWriter(stderr, os.Stderr)

	// make will launch the application in a child process, and this is the most
	// straightforward way to kill all ancestors.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	killed := false
	killCmd := func() {
		if !killed {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			waitForAppExit()
		}
		killed = true
	}

	if err := cmd.Start(); err != nil {
		killCmd()
		t.Fatal(err)
	}
	if cmd.Process != nil {
		log.Printf("process %d started: %s", cmd.Process.Pid, strings.Join(args, " "))
	}

	if err := waitForInit(app); err != nil {
		killCmd()
		t.Fatalf("error waiting for http server initialization: %v stderr=%s", err, stderr.String())
	}
	if s := stderr.String(); len(s) > 0 {
		killCmd()
	}

	restartCmd := func() (killFunc, restartFunc) {
		killCmd()
		return initORMApp(t, app, dbURL)
	}

	return killCmd, restartCmd
}

// waitForInit retries until a connection is successfully established.
func waitForInit(app application) error {
	const maxWait = 3 * time.Minute
	const waitDelay = 250 * time.Millisecond
	const maxWaitLoops = int(maxWait / waitDelay)

	var err error
	var api apiHandler
	for i := 0; i < maxWaitLoops; i++ {
		if err = api.ping(app.name()); err == nil {
			return err
		}
		log.Printf("waitForInit: %v", err)
		time.Sleep(waitDelay)
	}
	return err
}

// waitForExit waits indefinitely for the HTTP port of the ORM to stop
// listening.
func waitForAppExit() {
	const waitDelay = time.Second
	var api apiHandler
	for {
		if !api.canDial() {
			break
		}
		log.Print("waiting for app to exit")
		time.Sleep(waitDelay)
	}
}

func testORM(t *testing.T, language, orm string) {
	app := application{
		language: language,
		orm:      orm,
	}

	db, dbURL, stopDB := initTestDatabase(t, app)
	defer stopDB()

	stopApp, restartApp := initORMApp(t, app, dbURL)
	defer stopApp()

	td := testDriver{
		db:     db,
		dbName: app.dbName(),
	}

	// Test that the correct tables were generated.
	t.Run("GeneratedTables", td.TestGeneratedTables)

	// Test that the correct columns in those tables were generated.
	t.Run("GeneratedColumns", parallelTestGroup{
		"CustomersTable":     td.TestGeneratedCustomersTableColumns,
		"ProductsTable":      td.TestGeneratedProductsTableColumns,
		"OrdersTable":        td.TestGeneratedOrdersTableColumns,
		"OrderProductsTable": td.TestGeneratedOrderProductsTableColumns,
	}.T)

	// Test that the tables begin empty.
	t.Run("EmptyTables", parallelTestGroup{
		"CustomersTable":     td.TestCustomersEmpty,
		"ProductsTable":      td.TestProductsTableEmpty,
		"OrdersTable":        td.TestOrdersTableEmpty,
		"OrderProductsTable": td.TestOrderProductsTableEmpty,
	}.T)

	// Test that the API returns empty sets for each collection.
	t.Run("RetrieveFromAPIBeforeCreation", parallelTestGroup{
		"Customers": td.TestRetrieveCustomersBeforeCreation,
		"Products":  td.TestRetrieveProductsBeforeCreation,
		"Orders":    td.TestRetrieveOrdersBeforeCreation,
	}.T)

	// Test the creation of initial objects.
	t.Run("CreateCustomer", td.TestCreateCustomer)
	t.Run("CreateProduct", td.TestCreateProduct)

	// Test that the API returns what we just created.
	t.Run("RetrieveFromAPIAfterInitialCreation", parallelTestGroup{
		"Customers": td.TestRetrieveCustomerAfterCreation,
		"Products":  td.TestRetrieveProductAfterCreation,
	}.T)

	// Test the creation of dependent objects.
	t.Run("CreateOrder", td.TestCreateOrder)

	// Test that the API returns what we just created.
	t.Run("RetrieveFromAPIAfterDependentCreation", parallelTestGroup{
		"Order": td.TestRetrieveProductAfterCreation,
	}.T)

	// Restart the application.
	stopApp, restartApp = restartApp()
	defer stopApp()

	// Test that the API still returns all created objects.
	t.Run("RetrieveFromAPIAfterRestart", parallelTestGroup{
		"Customers": td.TestRetrieveCustomerAfterCreation,
		"Products":  td.TestRetrieveProductAfterCreation,
		"Order":     td.TestRetrieveProductAfterCreation,
	}.T)
}

func TestGORM(t *testing.T) {
	testORM(t, "go", "gorm")
}

func TestHibernate(t *testing.T) {
	testORM(t, "java", "hibernate")
}
<commit_msg>testing: don't die if there's stuff in stderr<commit_after>package testing

import (
	"bytes"
	"database/sql"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach-go/testserver"

	// Import postgres driver.
	_ "github.com/lib/pq"
)

// application represents a single instance of an application running an ORM and
// exposing an HTTP REST API.
type application struct {
	language string
	orm      string
}

func (app application) name() string {
	return fmt.Sprintf("%s/%s", app.language, app.orm)
}

func (app application) dir() string {
	return fmt.Sprintf("../%s", app.name())
}

func (app application) dbName() string {
	return fmt.Sprintf("company_%s", app.orm)
}

// initTestDatabase launches a test database as a subprocess.
func initTestDatabase(t *testing.T, app application) (*sql.DB, *url.URL, func()) {
	ts, err := testserver.NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

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

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	ts.WaitForInit(db)

	// Create the database if it does not exist.
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + app.dbName()); err != nil {
		t.Fatal(err)
	}

	// Connect to the database again, now with the database in the URL.
	url.Path = app.dbName()
	db, err = sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	return db, url, func() {
		_ = db.Close()
		ts.Stop()
	}
}

type killFunc func()
type restartFunc func() (killFunc, restartFunc)

// initORMApp launches an ORM application as a subprocess.
func initORMApp(t *testing.T, app application, dbURL *url.URL) (killFunc, restartFunc) {
	addrFlag := fmt.Sprintf("ADDR=%s", dbURL.String())
	args := []string{"make", "start", "-C", app.dir(), addrFlag}

	cmd := exec.Command(args[0], args[1:]...)

	// Set up stderr to display to console and store in a buffer, so we can later
	// verify that it's clean.
	stderr := new(bytes.Buffer)
	cmd.Stderr = io.MultiWriter(stderr, os.Stderr)

	// make will launch the application in a child process, and this is the most
	// straightforward way to kill all ancestors.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	killed := false
	killCmd := func() {
		if !killed {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			waitForAppExit()
		}
		killed = true
	}

	if err := cmd.Start(); err != nil {
		killCmd()
		t.Fatal(err)
	}
	if cmd.Process != nil {
		log.Printf("process %d started: %s", cmd.Process.Pid, strings.Join(args, " "))
	}

	if err := waitForInit(app); err != nil {
		killCmd()
		t.Fatalf("error waiting for http server initialization: %v stderr=%s", err, stderr.String())
	}

	restartCmd := func() (killFunc, restartFunc) {
		killCmd()
		return initORMApp(t, app, dbURL)
	}

	return killCmd, restartCmd
}

// waitForInit retries until a connection is successfully established.
func waitForInit(app application) error {
	const maxWait = 3 * time.Minute
	const waitDelay = 250 * time.Millisecond
	const maxWaitLoops = int(maxWait / waitDelay)

	var err error
	var api apiHandler
	for i := 0; i < maxWaitLoops; i++ {
		if err = api.ping(app.name()); err == nil {
			return err
		}
		log.Printf("waitForInit: %v", err)
		time.Sleep(waitDelay)
	}
	return err
}

// waitForExit waits indefinitely for the HTTP port of the ORM to stop
// listening.
func waitForAppExit() {
	const waitDelay = time.Second
	var api apiHandler
	for {
		if !api.canDial() {
			break
		}
		log.Print("waiting for app to exit")
		time.Sleep(waitDelay)
	}
}

func testORM(t *testing.T, language, orm string) {
	app := application{
		language: language,
		orm:      orm,
	}

	db, dbURL, stopDB := initTestDatabase(t, app)
	defer stopDB()

	stopApp, restartApp := initORMApp(t, app, dbURL)
	defer stopApp()

	td := testDriver{
		db:     db,
		dbName: app.dbName(),
	}

	// Test that the correct tables were generated.
	t.Run("GeneratedTables", td.TestGeneratedTables)

	// Test that the correct columns in those tables were generated.
	t.Run("GeneratedColumns", parallelTestGroup{
		"CustomersTable":     td.TestGeneratedCustomersTableColumns,
		"ProductsTable":      td.TestGeneratedProductsTableColumns,
		"OrdersTable":        td.TestGeneratedOrdersTableColumns,
		"OrderProductsTable": td.TestGeneratedOrderProductsTableColumns,
	}.T)

	// Test that the tables begin empty.
	t.Run("EmptyTables", parallelTestGroup{
		"CustomersTable":     td.TestCustomersEmpty,
		"ProductsTable":      td.TestProductsTableEmpty,
		"OrdersTable":        td.TestOrdersTableEmpty,
		"OrderProductsTable": td.TestOrderProductsTableEmpty,
	}.T)

	// Test that the API returns empty sets for each collection.
	t.Run("RetrieveFromAPIBeforeCreation", parallelTestGroup{
		"Customers": td.TestRetrieveCustomersBeforeCreation,
		"Products":  td.TestRetrieveProductsBeforeCreation,
		"Orders":    td.TestRetrieveOrdersBeforeCreation,
	}.T)

	// Test the creation of initial objects.
	t.Run("CreateCustomer", td.TestCreateCustomer)
	t.Run("CreateProduct", td.TestCreateProduct)

	// Test that the API returns what we just created.
	t.Run("RetrieveFromAPIAfterInitialCreation", parallelTestGroup{
		"Customers": td.TestRetrieveCustomerAfterCreation,
		"Products":  td.TestRetrieveProductAfterCreation,
	}.T)

	// Test the creation of dependent objects.
	t.Run("CreateOrder", td.TestCreateOrder)

	// Test that the API returns what we just created.
	t.Run("RetrieveFromAPIAfterDependentCreation", parallelTestGroup{
		"Order": td.TestRetrieveProductAfterCreation,
	}.T)

	// Restart the application.
	stopApp, restartApp = restartApp()
	defer stopApp()

	// Test that the API still returns all created objects.
	t.Run("RetrieveFromAPIAfterRestart", parallelTestGroup{
		"Customers": td.TestRetrieveCustomerAfterCreation,
		"Products":  td.TestRetrieveProductAfterCreation,
		"Order":     td.TestRetrieveProductAfterCreation,
	}.T)
}

func TestGORM(t *testing.T) {
	testORM(t, "go", "gorm")
}

func TestHibernate(t *testing.T) {
	testORM(t, "java", "hibernate")
}
<|endoftext|>


<commit_before>package internal

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/influxdata/influxdb/v2/kit/platform"
	"github.com/influxdata/influxdb/v2/pkg/durablequeue"
	"github.com/influxdata/influxdb/v2/replications/metrics"
	"github.com/influxdata/influxdb/v2/replications/remotewrite"
	"go.uber.org/zap"
)

type remoteWriter interface {
	Write([]byte) error
}

type replicationQueue struct {
	id           platform.ID
	queue        *durablequeue.Queue
	wg           sync.WaitGroup
	done         chan struct{}
	receive      chan struct{}
	logger       *zap.Logger
	metrics      *metrics.ReplicationsMetrics
	remoteWriter remoteWriter
}

type durableQueueManager struct {
	replicationQueues map[platform.ID]*replicationQueue
	logger            *zap.Logger
	queuePath         string
	mutex             sync.RWMutex
	metrics           *metrics.ReplicationsMetrics
	configStore       remotewrite.HttpConfigStore
}

var errStartup = errors.New("startup tasks for replications durable queue management failed, see server logs for details")
var errShutdown = errors.New("shutdown tasks for replications durable queues failed, see server logs for details")

// NewDurableQueueManager creates a new durableQueueManager struct, for managing durable queues associated with
// replication streams.
func NewDurableQueueManager(log *zap.Logger, queuePath string, metrics *metrics.ReplicationsMetrics, configStore remotewrite.HttpConfigStore) *durableQueueManager {
	replicationQueues := make(map[platform.ID]*replicationQueue)

	os.MkdirAll(queuePath, 0777)

	return &durableQueueManager{
		replicationQueues: replicationQueues,
		logger:            log,
		queuePath:         queuePath,
		metrics:           metrics,
		configStore:       configStore,
	}
}

// InitializeQueue creates and opens a new durable queue which is associated with a replication stream.
func (qm *durableQueueManager) InitializeQueue(replicationID platform.ID, maxQueueSizeBytes int64) error {
	qm.mutex.Lock()
	defer qm.mutex.Unlock()

	// Check for duplicate replication ID
	if _, exists := qm.replicationQueues[replicationID]; exists {
		return fmt.Errorf("durable queue already exists for replication ID %q", replicationID)
	}

	// Set up path for new queue on disk
	dir := filepath.Join(
		qm.queuePath,
		replicationID.String(),
	)
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}

	// Create a new durable queue
	newQueue, err := durablequeue.NewQueue(
		dir,
		maxQueueSizeBytes,
		durablequeue.DefaultSegmentSize,
		&durablequeue.SharedCount{},
		durablequeue.MaxWritesPending,
		func(bytes []byte) error {
			return nil
		},
	)

	if err != nil {
		return err
	}

	// Open the new queue
	if err := newQueue.Open(); err != nil {
		return err
	}

	// Map new durable queue and scanner to its corresponding replication stream via replication ID
	rq := qm.newReplicationQueue(replicationID, newQueue)
	qm.replicationQueues[replicationID] = rq
	rq.Open()

	qm.logger.Debug("Created new durable queue for replication stream",
		zap.String("id", replicationID.String()), zap.String("path", dir))

	return nil
}

func (rq *replicationQueue) Open() {
	rq.wg.Add(1)
	go rq.run()
}

func (rq *replicationQueue) Close() error {
	close(rq.receive)
	close(rq.done)
	rq.wg.Wait() // wait for goroutine to finish processing all messages
	return rq.queue.Close()
}

func (rq *replicationQueue) run() {
	defer rq.wg.Done()

	for {
		select {
		case <-rq.done: // end the goroutine when done is messaged
			return
		case <-rq.receive: // run the scanner on data append
			// Receive channel has a buffer to prevent a potential race condition where rq.SendWrite has reached EOF and will
			// return false, but data is queued after evaluating the scanner and before the loop is ready to select on the
			// receive channel again. This would result in data remaining unprocessed in the queue until the next send on the
			// receive channel since the send to the receive channel in qm.EnqueueData is non-blocking. The buffer ensures
			// that rq.SendWrite will be called again in this situation and not leave data in the queue. Outside of this
			// specific scenario, the buffer might result in an extra call to rq.SendWrite that will immediately return on
			// EOF.
			for rq.SendWrite() {
			}
		}
	}
}

// SendWrite processes data enqueued into the durablequeue.Queue.
// SendWrite is responsible for processing all data in the queue at the time of calling.
// Network errors will be handled by the remote writer.
func (rq *replicationQueue) SendWrite() bool {
	// Any error in creating the scanner should exit the loop in run()
	// Either it is io.EOF indicating no data, or some other failure in making
	// the Scanner object that we don't know how to handle.
	scan, err := rq.queue.NewScanner()
	if err != nil {
		if err != io.EOF {
			rq.logger.Error("Error creating replications queue scanner", zap.Error(err))
		}
		return false
	}

	for scan.Next() {
		// An io.EOF error here indicates that there is no more data
		// left to process, and is an expected error.
		if scan.Err() == io.EOF {
			break
		}

		// Any other here indicates a problem, so we log the error and
		// drop the data with a call to scan.Advance() later.
		if scan.Err() != nil {
			rq.logger.Info("Segment read error.", zap.Error(scan.Err()))
			// TODO: Add metrics collection for dropped data here.
			break
		}

		// An error here indicates an unhandlable error. Data is not corrupt, and
		// the remote write is not retryable.
		if err = rq.remoteWriter.Write(scan.Bytes()); err != nil {
			rq.logger.Error("Error in replication stream", zap.Error(err))
			return false
		}
	}

	// Update metrics after the call to scan.Advance()
	defer func() {
		rq.metrics.Dequeue(rq.id, rq.queue.TotalBytes())
	}()

	if _, err = scan.Advance(); err != nil {
		if err != io.EOF {
			rq.logger.Error("Error in replication queue scanner", zap.Error(err))
		}
		return false
	}
	return true
}

// DeleteQueue deletes a durable queue and its associated data on disk.
func (qm *durableQueueManager) DeleteQueue(replicationID platform.ID) error {
	qm.mutex.Lock()
	defer qm.mutex.Unlock()

	if _, exist := qm.replicationQueues[replicationID]; !exist {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	rq := qm.replicationQueues[replicationID]

	// Close the queue
	if err := rq.Close(); err != nil {
		return err
	}

	qm.logger.Debug("Closed replication stream durable queue",
		zap.String("id", replicationID.String()), zap.String("path", rq.queue.Dir()))

	// Delete any enqueued, un-flushed data on disk for this queue
	if err := rq.queue.Remove(); err != nil {
		return err
	}

	qm.logger.Debug("Deleted data associated with replication stream durable queue",
		zap.String("id", replicationID.String()), zap.String("path", rq.queue.Dir()))

	// Remove entry from replicationQueues map
	delete(qm.replicationQueues, replicationID)

	return nil
}

// UpdateMaxQueueSize updates the maximum size of a durable queue.
func (qm *durableQueueManager) UpdateMaxQueueSize(replicationID platform.ID, maxQueueSizeBytes int64) error {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	if _, exist := qm.replicationQueues[replicationID]; !exist {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	if err := qm.replicationQueues[replicationID].queue.SetMaxSize(maxQueueSizeBytes); err != nil {
		return err
	}

	return nil
}

// CurrentQueueSizes returns the current size-on-disk for the requested set of durable queues.
func (qm *durableQueueManager) CurrentQueueSizes(ids []platform.ID) (map[platform.ID]int64, error) {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	sizes := make(map[platform.ID]int64, len(ids))

	for _, id := range ids {
		if _, exist := qm.replicationQueues[id]; !exist {
			return nil, fmt.Errorf("durable queue not found for replication ID %q", id)
		}
		sizes[id] = qm.replicationQueues[id].queue.DiskUsage()
	}

	return sizes, nil
}

// StartReplicationQueues updates the durableQueueManager.replicationQueues map, fully removing any partially deleted
// queues (present on disk, but not tracked in sqlite), opening all current queues, and logging info for each.
func (qm *durableQueueManager) StartReplicationQueues(trackedReplications map[platform.ID]int64) error {
	errOccurred := false

	for id, size := range trackedReplications {
		// Re-initialize a queue struct for each replication stream from sqlite
		queue, err := durablequeue.NewQueue(
			filepath.Join(qm.queuePath, id.String()),
			size,
			durablequeue.DefaultSegmentSize,
			&durablequeue.SharedCount{},
			durablequeue.MaxWritesPending,
			func(bytes []byte) error {
				return nil
			},
		)

		if err != nil {
			qm.logger.Error("failed to initialize replication stream durable queue", zap.Error(err))
			errOccurred = true
			continue
		}

		// Open and map the queue struct to its replication ID
		if err := queue.Open(); err != nil {
			qm.logger.Error("failed to open replication stream durable queue", zap.Error(err), zap.String("id", id.String()))
			errOccurred = true
			continue
		} else {
			qm.replicationQueues[id] = qm.newReplicationQueue(id, queue)
			qm.replicationQueues[id].Open()
			qm.logger.Info("Opened replication stream", zap.String("id", id.String()), zap.String("path", queue.Dir()))
		}
	}

	if errOccurred {
		return errStartup
	}

	// Get contents of replicationq directory
	entries, err := os.ReadDir(qm.queuePath)
	if err != nil {
		return err
	}

	for _, entry := range entries {
		// Skip over non-relevant entries (must be a dir named with a replication ID)
		if !entry.IsDir() {
			continue
		}

		id, err := platform.IDFromString(entry.Name())
		if err != nil {
			continue
		}

		// Partial delete found, needs to be fully removed
		if _, exist := qm.replicationQueues[*id]; !exist {
			if err := os.RemoveAll(filepath.Join(qm.queuePath, id.String())); err != nil {
				qm.logger.Error("failed to remove durable queue during partial delete cleanup", zap.Error(err), zap.String("id", id.String()))
				errOccurred = true
			}
		}
	}

	if errOccurred {
		return errStartup
	} else {
		return nil
	}
}

// CloseAll loops through all current replication stream queues and closes them without deleting on-disk resources
func (qm *durableQueueManager) CloseAll() error {
	errOccurred := false

	for id, replicationQueue := range qm.replicationQueues {
		if err := replicationQueue.Close(); err != nil {
			qm.logger.Error("failed to close durable queue", zap.Error(err), zap.String("id", id.String()))
			errOccurred = true
		}
	}

	if errOccurred {
		return errShutdown
	} else {
		return nil
	}
}

// EnqueueData persists a set of bytes to a replication's durable queue.
func (qm *durableQueueManager) EnqueueData(replicationID platform.ID, data []byte, numPoints int) (err error) {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	// Update metrics if data fails to be added to queue
	defer func() {
		if err != nil {
			qm.metrics.EnqueueError(replicationID, len(data), numPoints)
		}
	}()

	rq, ok := qm.replicationQueues[replicationID]
	if !ok {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	if err := rq.queue.Append(data); err != nil {
		return err
	}
	// Update metrics for this replication queue when adding data to the queue.
	qm.metrics.EnqueueData(replicationID, len(data), numPoints, rq.queue.TotalBytes())

	// Send to the replication receive channel if it is not full to activate the queue processing. If the receive channel
	// is full, don't block further writes and return.
	select {
	case qm.replicationQueues[replicationID].receive <- struct{}{}:
	default:
	}

	return nil
}

func (qm *durableQueueManager) newReplicationQueue(id platform.ID, queue *durablequeue.Queue) *replicationQueue {
	logger := qm.logger.With(zap.String("replication_id", id.String()))
	done := make(chan struct{})

	return &replicationQueue{
		id:           id,
		queue:        queue,
		done:         done,
		receive:      make(chan struct{}, 1),
		logger:       logger,
		metrics:      qm.metrics,
		remoteWriter: remotewrite.NewWriter(id, qm.configStore, qm.metrics, logger, done),
	}
}
<commit_msg>fix: advance replications queue after successful remote writes (#22967)<commit_after>package internal

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/influxdata/influxdb/v2/kit/platform"
	"github.com/influxdata/influxdb/v2/pkg/durablequeue"
	"github.com/influxdata/influxdb/v2/replications/metrics"
	"github.com/influxdata/influxdb/v2/replications/remotewrite"
	"go.uber.org/zap"
)

type remoteWriter interface {
	Write([]byte) error
}

type replicationQueue struct {
	id           platform.ID
	queue        *durablequeue.Queue
	wg           sync.WaitGroup
	done         chan struct{}
	receive      chan struct{}
	logger       *zap.Logger
	metrics      *metrics.ReplicationsMetrics
	remoteWriter remoteWriter
}

type durableQueueManager struct {
	replicationQueues map[platform.ID]*replicationQueue
	logger            *zap.Logger
	queuePath         string
	mutex             sync.RWMutex
	metrics           *metrics.ReplicationsMetrics
	configStore       remotewrite.HttpConfigStore
}

var errStartup = errors.New("startup tasks for replications durable queue management failed, see server logs for details")
var errShutdown = errors.New("shutdown tasks for replications durable queues failed, see server logs for details")

// NewDurableQueueManager creates a new durableQueueManager struct, for managing durable queues associated with
// replication streams.
func NewDurableQueueManager(log *zap.Logger, queuePath string, metrics *metrics.ReplicationsMetrics, configStore remotewrite.HttpConfigStore) *durableQueueManager {
	replicationQueues := make(map[platform.ID]*replicationQueue)

	os.MkdirAll(queuePath, 0777)

	return &durableQueueManager{
		replicationQueues: replicationQueues,
		logger:            log,
		queuePath:         queuePath,
		metrics:           metrics,
		configStore:       configStore,
	}
}

// InitializeQueue creates and opens a new durable queue which is associated with a replication stream.
func (qm *durableQueueManager) InitializeQueue(replicationID platform.ID, maxQueueSizeBytes int64) error {
	qm.mutex.Lock()
	defer qm.mutex.Unlock()

	// Check for duplicate replication ID
	if _, exists := qm.replicationQueues[replicationID]; exists {
		return fmt.Errorf("durable queue already exists for replication ID %q", replicationID)
	}

	// Set up path for new queue on disk
	dir := filepath.Join(
		qm.queuePath,
		replicationID.String(),
	)
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}

	// Create a new durable queue
	newQueue, err := durablequeue.NewQueue(
		dir,
		maxQueueSizeBytes,
		durablequeue.DefaultSegmentSize,
		&durablequeue.SharedCount{},
		durablequeue.MaxWritesPending,
		func(bytes []byte) error {
			return nil
		},
	)

	if err != nil {
		return err
	}

	// Open the new queue
	if err := newQueue.Open(); err != nil {
		return err
	}

	// Map new durable queue and scanner to its corresponding replication stream via replication ID
	rq := qm.newReplicationQueue(replicationID, newQueue)
	qm.replicationQueues[replicationID] = rq
	rq.Open()

	qm.logger.Debug("Created new durable queue for replication stream",
		zap.String("id", replicationID.String()), zap.String("path", dir))

	return nil
}

func (rq *replicationQueue) Open() {
	rq.wg.Add(1)
	go rq.run()
}

func (rq *replicationQueue) Close() error {
	close(rq.receive)
	close(rq.done)
	rq.wg.Wait() // wait for goroutine to finish processing all messages
	return rq.queue.Close()
}

func (rq *replicationQueue) run() {
	defer rq.wg.Done()

	for {
		select {
		case <-rq.done: // end the goroutine when done is messaged
			return
		case <-rq.receive: // run the scanner on data append
			// Receive channel has a buffer to prevent a potential race condition where rq.SendWrite has reached EOF and will
			// return false, but data is queued after evaluating the scanner and before the loop is ready to select on the
			// receive channel again. This would result in data remaining unprocessed in the queue until the next send on the
			// receive channel since the send to the receive channel in qm.EnqueueData is non-blocking. The buffer ensures
			// that rq.SendWrite will be called again in this situation and not leave data in the queue. Outside of this
			// specific scenario, the buffer might result in an extra call to rq.SendWrite that will immediately return on
			// EOF.
			for rq.SendWrite() {
			}
		}
	}
}

// SendWrite processes data enqueued into the durablequeue.Queue.
// SendWrite is responsible for processing all data in the queue at the time of calling.
// Network errors will be handled by the remote writer.
func (rq *replicationQueue) SendWrite() bool {
	// Any error in creating the scanner should exit the loop in run()
	// Either it is io.EOF indicating no data, or some other failure in making
	// the Scanner object that we don't know how to handle.
	scan, err := rq.queue.NewScanner()
	if err != nil {
		if !errors.Is(err, io.EOF) {
			rq.logger.Error("Error creating replications queue scanner", zap.Error(err))
		}
		return false
	}

	for scan.Next() {
		if err := scan.Err(); err != nil {
			if errors.Is(err, io.EOF) {
				// An io.EOF error here indicates that there is no more data left to process, and is an expected error.
				return false
			}
			// Any other error here indicates a problem reading the data from the queue, so we log the error and drop the data
			// with a call to scan.Advance() later.
			rq.logger.Info("Segment read error.", zap.Error(scan.Err()))
		}

		if err = rq.remoteWriter.Write(scan.Bytes()); err != nil {
			// An error here indicates an unhandleable remote write error. The scanner will not be advanced.
			rq.logger.Error("Error in replication stream", zap.Error(err))
			return false
		}

		// TODO: As a potential future optimization, Advance() could be called only if a certain amount of time has passed
		// since an Advance(), a certain amount of data has been transferred, a certain number of writes, etc. to avoid an
		// fsync after every remote write.
		if _, err = scan.Advance(); err != nil {
			if err != io.EOF {
				rq.logger.Error("Error in replication queue scanner", zap.Error(err))
			}
			return false
		}
		rq.metrics.Dequeue(rq.id, rq.queue.TotalBytes())
	}

	return true
}

// DeleteQueue deletes a durable queue and its associated data on disk.
func (qm *durableQueueManager) DeleteQueue(replicationID platform.ID) error {
	qm.mutex.Lock()
	defer qm.mutex.Unlock()

	if _, exist := qm.replicationQueues[replicationID]; !exist {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	rq := qm.replicationQueues[replicationID]

	// Close the queue
	if err := rq.Close(); err != nil {
		return err
	}

	qm.logger.Debug("Closed replication stream durable queue",
		zap.String("id", replicationID.String()), zap.String("path", rq.queue.Dir()))

	// Delete any enqueued, un-flushed data on disk for this queue
	if err := rq.queue.Remove(); err != nil {
		return err
	}

	qm.logger.Debug("Deleted data associated with replication stream durable queue",
		zap.String("id", replicationID.String()), zap.String("path", rq.queue.Dir()))

	// Remove entry from replicationQueues map
	delete(qm.replicationQueues, replicationID)

	return nil
}

// UpdateMaxQueueSize updates the maximum size of a durable queue.
func (qm *durableQueueManager) UpdateMaxQueueSize(replicationID platform.ID, maxQueueSizeBytes int64) error {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	if _, exist := qm.replicationQueues[replicationID]; !exist {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	if err := qm.replicationQueues[replicationID].queue.SetMaxSize(maxQueueSizeBytes); err != nil {
		return err
	}

	return nil
}

// CurrentQueueSizes returns the current size-on-disk for the requested set of durable queues.
func (qm *durableQueueManager) CurrentQueueSizes(ids []platform.ID) (map[platform.ID]int64, error) {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	sizes := make(map[platform.ID]int64, len(ids))

	for _, id := range ids {
		if _, exist := qm.replicationQueues[id]; !exist {
			return nil, fmt.Errorf("durable queue not found for replication ID %q", id)
		}
		sizes[id] = qm.replicationQueues[id].queue.DiskUsage()
	}

	return sizes, nil
}

// StartReplicationQueues updates the durableQueueManager.replicationQueues map, fully removing any partially deleted
// queues (present on disk, but not tracked in sqlite), opening all current queues, and logging info for each.
func (qm *durableQueueManager) StartReplicationQueues(trackedReplications map[platform.ID]int64) error {
	errOccurred := false

	for id, size := range trackedReplications {
		// Re-initialize a queue struct for each replication stream from sqlite
		queue, err := durablequeue.NewQueue(
			filepath.Join(qm.queuePath, id.String()),
			size,
			durablequeue.DefaultSegmentSize,
			&durablequeue.SharedCount{},
			durablequeue.MaxWritesPending,
			func(bytes []byte) error {
				return nil
			},
		)

		if err != nil {
			qm.logger.Error("failed to initialize replication stream durable queue", zap.Error(err))
			errOccurred = true
			continue
		}

		// Open and map the queue struct to its replication ID
		if err := queue.Open(); err != nil {
			qm.logger.Error("failed to open replication stream durable queue", zap.Error(err), zap.String("id", id.String()))
			errOccurred = true
			continue
		} else {
			qm.replicationQueues[id] = qm.newReplicationQueue(id, queue)
			qm.replicationQueues[id].Open()
			qm.logger.Info("Opened replication stream", zap.String("id", id.String()), zap.String("path", queue.Dir()))
		}
	}

	if errOccurred {
		return errStartup
	}

	// Get contents of replicationq directory
	entries, err := os.ReadDir(qm.queuePath)
	if err != nil {
		return err
	}

	for _, entry := range entries {
		// Skip over non-relevant entries (must be a dir named with a replication ID)
		if !entry.IsDir() {
			continue
		}

		id, err := platform.IDFromString(entry.Name())
		if err != nil {
			continue
		}

		// Partial delete found, needs to be fully removed
		if _, exist := qm.replicationQueues[*id]; !exist {
			if err := os.RemoveAll(filepath.Join(qm.queuePath, id.String())); err != nil {
				qm.logger.Error("failed to remove durable queue during partial delete cleanup", zap.Error(err), zap.String("id", id.String()))
				errOccurred = true
			}
		}
	}

	if errOccurred {
		return errStartup
	} else {
		return nil
	}
}

// CloseAll loops through all current replication stream queues and closes them without deleting on-disk resources
func (qm *durableQueueManager) CloseAll() error {
	errOccurred := false

	for id, replicationQueue := range qm.replicationQueues {
		if err := replicationQueue.Close(); err != nil {
			qm.logger.Error("failed to close durable queue", zap.Error(err), zap.String("id", id.String()))
			errOccurred = true
		}
	}

	if errOccurred {
		return errShutdown
	} else {
		return nil
	}
}

// EnqueueData persists a set of bytes to a replication's durable queue.
func (qm *durableQueueManager) EnqueueData(replicationID platform.ID, data []byte, numPoints int) (err error) {
	qm.mutex.RLock()
	defer qm.mutex.RUnlock()

	// Update metrics if data fails to be added to queue
	defer func() {
		if err != nil {
			qm.metrics.EnqueueError(replicationID, len(data), numPoints)
		}
	}()

	rq, ok := qm.replicationQueues[replicationID]
	if !ok {
		return fmt.Errorf("durable queue not found for replication ID %q", replicationID)
	}

	if err := rq.queue.Append(data); err != nil {
		return err
	}
	// Update metrics for this replication queue when adding data to the queue.
	qm.metrics.EnqueueData(replicationID, len(data), numPoints, rq.queue.TotalBytes())

	// Send to the replication receive channel if it is not full to activate the queue processing. If the receive channel
	// is full, don't block further writes and return.
	select {
	case qm.replicationQueues[replicationID].receive <- struct{}{}:
	default:
	}

	return nil
}

func (qm *durableQueueManager) newReplicationQueue(id platform.ID, queue *durablequeue.Queue) *replicationQueue {
	logger := qm.logger.With(zap.String("replication_id", id.String()))
	done := make(chan struct{})

	return &replicationQueue{
		id:           id,
		queue:        queue,
		done:         done,
		receive:      make(chan struct{}, 1),
		logger:       logger,
		metrics:      qm.metrics,
		remoteWriter: remotewrite.NewWriter(id, qm.configStore, qm.metrics, logger, done),
	}
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 chartutil

import (
	"log"

	"github.com/mitchellh/copystructure"
	"github.com/pkg/errors"

	"helm.sh/helm/v3/pkg/chart"
)

// CoalesceValues coalesces all of the values in a chart (and its subcharts).
//
// Values are coalesced together using the following rules:
//
//	- Values in a higher level chart always override values in a lower-level
//		dependency chart
//	- Scalar values and arrays are replaced, maps are merged
//	- A chart has access to all of the variables for it, as well as all of
//		the values destined for its dependencies.
func CoalesceValues(chrt *chart.Chart, vals map[string]interface{}) (Values, error) {
	// create a copy of vals and then pass it to coalesce
	// and coalesceDeps, as both will mutate the passed values
	v, err := copystructure.Copy(vals)
	if err != nil {
		return vals, err
	}

	valsCopy := v.(map[string]interface{})
	// if we have an empty map, make sure it is initialized
	if valsCopy == nil {
		valsCopy = make(map[string]interface{})
	}
	return coalesce(chrt, valsCopy)
}

// coalesce coalesces the dest values and the chart values, giving priority to the dest values.
//
// This is a helper function for CoalesceValues.
func coalesce(ch *chart.Chart, dest map[string]interface{}) (map[string]interface{}, error) {
	coalesceValues(ch, dest)
	return coalesceDeps(ch, dest)
}

// coalesceDeps coalesces the dependencies of the given chart.
func coalesceDeps(chrt *chart.Chart, dest map[string]interface{}) (map[string]interface{}, error) {
	for _, subchart := range chrt.Dependencies() {
		if c, ok := dest[subchart.Name()]; !ok {
			// If dest doesn't already have the key, create it.
			dest[subchart.Name()] = make(map[string]interface{})
		} else if !istable(c) {
			return dest, errors.Errorf("type mismatch on %s: %t", subchart.Name(), c)
		}
		if dv, ok := dest[subchart.Name()]; ok {
			dvmap := dv.(map[string]interface{})

			// Get globals out of dest and merge them into dvmap.
			coalesceGlobals(dvmap, dest)

			// Now coalesce the rest of the values.
			var err error
			dest[subchart.Name()], err = coalesce(subchart, dvmap)
			if err != nil {
				return dest, err
			}
		}
	}
	return dest, nil
}

// coalesceGlobals copies the globals out of src and merges them into dest.
//
// For convenience, returns dest.
func coalesceGlobals(dest, src map[string]interface{}) {
	var dg, sg map[string]interface{}

	if destglob, ok := dest[GlobalKey]; !ok {
		dg = make(map[string]interface{})
	} else if dg, ok = destglob.(map[string]interface{}); !ok {
		log.Printf("warning: skipping globals because destination %s is not a table.", GlobalKey)
		return
	}

	if srcglob, ok := src[GlobalKey]; !ok {
		sg = make(map[string]interface{})
	} else if sg, ok = srcglob.(map[string]interface{}); !ok {
		log.Printf("warning: skipping globals because source %s is not a table.", GlobalKey)
		return
	}

	// EXPERIMENTAL: In the past, we have disallowed globals to test tables. This
	// reverses that decision. It may somehow be possible to introduce a loop
	// here, but I haven't found a way. So for the time being, let's allow
	// tables in globals.
	for key, val := range sg {
		if istable(val) {
			vv := copyMap(val.(map[string]interface{}))
			if destv, ok := dg[key]; !ok {
				// Here there is no merge. We're just adding.
				dg[key] = vv
			} else {
				if destvmap, ok := destv.(map[string]interface{}); !ok {
					log.Printf("Conflict: cannot merge map onto non-map for %q. Skipping.", key)
				} else {
					// Basically, we reverse order of coalesce here to merge
					// top-down.
					CoalesceTables(vv, destvmap)
					dg[key] = vv
					continue
				}
			}
		} else if dv, ok := dg[key]; ok && istable(dv) {
			// It's not clear if this condition can actually ever trigger.
			log.Printf("key %s is table. Skipping", key)
			continue
		}
		// TODO: Do we need to do any additional checking on the value?
		dg[key] = val
	}
	dest[GlobalKey] = dg
}

func copyMap(src map[string]interface{}) map[string]interface{} {
	m := make(map[string]interface{}, len(src))
	for k, v := range src {
		m[k] = v
	}
	return m
}

// coalesceValues builds up a values map for a particular chart.
//
// Values in v will override the values in the chart.
func coalesceValues(c *chart.Chart, v map[string]interface{}) {
	for key, val := range c.Values {
		if value, ok := v[key]; ok {
			if value == nil {
				// When the YAML value is null, we remove the value's key.
				// This allows Helm's various sources of values (value files or --set) to
				// remove incompatible keys from any previous chart, file, or set values.
				delete(v, key)
			} else if dest, ok := value.(map[string]interface{}); ok {
				// if v[key] is a table, merge nv's val table into v[key].
				src, ok := val.(map[string]interface{})
				if !ok {
					log.Printf("warning: skipped value for %s: Not a table.", key)
					continue
				}
				// Because v has higher precedence than nv, dest values override src
				// values.
				CoalesceTables(dest, src)
			}
		} else {
			// If the key is not in v, copy it from nv.
			v[key] = val
		}
	}
}

// CoalesceTables merges a source map into a destination map.
//
// dest is considered authoritative.
func CoalesceTables(dst, src map[string]interface{}) map[string]interface{} {
	if dst == nil || src == nil {
		return src
	}
	// Because dest has higher precedence than src, dest values override src
	// values.
	for key, val := range src {
		if dv, ok := dst[key]; ok && dv == nil {
			delete(dst, key)
		} else if !ok {
			dst[key] = val
		} else if istable(val) {
			if istable(dv) {
				CoalesceTables(dv.(map[string]interface{}), val.(map[string]interface{}))
			} else {
				log.Printf("warning: cannot overwrite table with non table for %s (%v)", key, val)
			}
		} else if istable(dv) {
			log.Printf("warning: destination for %s is a table. Ignoring non-table value %v", key, val)
		}
	}
	return dst
}
<commit_msg>fix: removed inaccurate comment (#7937)<commit_after>/*
Copyright The Helm 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 chartutil

import (
	"log"

	"github.com/mitchellh/copystructure"
	"github.com/pkg/errors"

	"helm.sh/helm/v3/pkg/chart"
)

// CoalesceValues coalesces all of the values in a chart (and its subcharts).
//
// Values are coalesced together using the following rules:
//
//	- Values in a higher level chart always override values in a lower-level
//		dependency chart
//	- Scalar values and arrays are replaced, maps are merged
//	- A chart has access to all of the variables for it, as well as all of
//		the values destined for its dependencies.
func CoalesceValues(chrt *chart.Chart, vals map[string]interface{}) (Values, error) {
	v, err := copystructure.Copy(vals)
	if err != nil {
		return vals, err
	}

	valsCopy := v.(map[string]interface{})
	// if we have an empty map, make sure it is initialized
	if valsCopy == nil {
		valsCopy = make(map[string]interface{})
	}
	return coalesce(chrt, valsCopy)
}

// coalesce coalesces the dest values and the chart values, giving priority to the dest values.
//
// This is a helper function for CoalesceValues.
func coalesce(ch *chart.Chart, dest map[string]interface{}) (map[string]interface{}, error) {
	coalesceValues(ch, dest)
	return coalesceDeps(ch, dest)
}

// coalesceDeps coalesces the dependencies of the given chart.
func coalesceDeps(chrt *chart.Chart, dest map[string]interface{}) (map[string]interface{}, error) {
	for _, subchart := range chrt.Dependencies() {
		if c, ok := dest[subchart.Name()]; !ok {
			// If dest doesn't already have the key, create it.
			dest[subchart.Name()] = make(map[string]interface{})
		} else if !istable(c) {
			return dest, errors.Errorf("type mismatch on %s: %t", subchart.Name(), c)
		}
		if dv, ok := dest[subchart.Name()]; ok {
			dvmap := dv.(map[string]interface{})

			// Get globals out of dest and merge them into dvmap.
			coalesceGlobals(dvmap, dest)

			// Now coalesce the rest of the values.
			var err error
			dest[subchart.Name()], err = coalesce(subchart, dvmap)
			if err != nil {
				return dest, err
			}
		}
	}
	return dest, nil
}

// coalesceGlobals copies the globals out of src and merges them into dest.
//
// For convenience, returns dest.
func coalesceGlobals(dest, src map[string]interface{}) {
	var dg, sg map[string]interface{}

	if destglob, ok := dest[GlobalKey]; !ok {
		dg = make(map[string]interface{})
	} else if dg, ok = destglob.(map[string]interface{}); !ok {
		log.Printf("warning: skipping globals because destination %s is not a table.", GlobalKey)
		return
	}

	if srcglob, ok := src[GlobalKey]; !ok {
		sg = make(map[string]interface{})
	} else if sg, ok = srcglob.(map[string]interface{}); !ok {
		log.Printf("warning: skipping globals because source %s is not a table.", GlobalKey)
		return
	}

	// EXPERIMENTAL: In the past, we have disallowed globals to test tables. This
	// reverses that decision. It may somehow be possible to introduce a loop
	// here, but I haven't found a way. So for the time being, let's allow
	// tables in globals.
	for key, val := range sg {
		if istable(val) {
			vv := copyMap(val.(map[string]interface{}))
			if destv, ok := dg[key]; !ok {
				// Here there is no merge. We're just adding.
				dg[key] = vv
			} else {
				if destvmap, ok := destv.(map[string]interface{}); !ok {
					log.Printf("Conflict: cannot merge map onto non-map for %q. Skipping.", key)
				} else {
					// Basically, we reverse order of coalesce here to merge
					// top-down.
					CoalesceTables(vv, destvmap)
					dg[key] = vv
					continue
				}
			}
		} else if dv, ok := dg[key]; ok && istable(dv) {
			// It's not clear if this condition can actually ever trigger.
			log.Printf("key %s is table. Skipping", key)
			continue
		}
		// TODO: Do we need to do any additional checking on the value?
		dg[key] = val
	}
	dest[GlobalKey] = dg
}

func copyMap(src map[string]interface{}) map[string]interface{} {
	m := make(map[string]interface{}, len(src))
	for k, v := range src {
		m[k] = v
	}
	return m
}

// coalesceValues builds up a values map for a particular chart.
//
// Values in v will override the values in the chart.
func coalesceValues(c *chart.Chart, v map[string]interface{}) {
	for key, val := range c.Values {
		if value, ok := v[key]; ok {
			if value == nil {
				// When the YAML value is null, we remove the value's key.
				// This allows Helm's various sources of values (value files or --set) to
				// remove incompatible keys from any previous chart, file, or set values.
				delete(v, key)
			} else if dest, ok := value.(map[string]interface{}); ok {
				// if v[key] is a table, merge nv's val table into v[key].
				src, ok := val.(map[string]interface{})
				if !ok {
					log.Printf("warning: skipped value for %s: Not a table.", key)
					continue
				}
				// Because v has higher precedence than nv, dest values override src
				// values.
				CoalesceTables(dest, src)
			}
		} else {
			// If the key is not in v, copy it from nv.
			v[key] = val
		}
	}
}

// CoalesceTables merges a source map into a destination map.
//
// dest is considered authoritative.
func CoalesceTables(dst, src map[string]interface{}) map[string]interface{} {
	if dst == nil || src == nil {
		return src
	}
	// Because dest has higher precedence than src, dest values override src
	// values.
	for key, val := range src {
		if dv, ok := dst[key]; ok && dv == nil {
			delete(dst, key)
		} else if !ok {
			dst[key] = val
		} else if istable(val) {
			if istable(dv) {
				CoalesceTables(dv.(map[string]interface{}), val.(map[string]interface{}))
			} else {
				log.Printf("warning: cannot overwrite table with non table for %s (%v)", key, val)
			}
		} else if istable(dv) {
			log.Printf("warning: destination for %s is a table. Ignoring non-table value %v", key, val)
		}
	}
	return dst
}
<|endoftext|>


<commit_before>package command

import (
	"strconv"

	"github.com/naoty/todo/todo"
	"github.com/urfave/cli"
)

// Delete is a command to delete todos.
var Delete = cli.Command{
	Name:   "delete",
	Usage:  "Delete todos",
	Action: delete,
}

func delete(c *cli.Context) error {
	if c.NArg() == 0 {
		cli.ShowCommandHelp(c, "delete")
		return nil
	}

	orders := make([]int, len(c.Args()))
	for _, arg := range c.Args() {
		order, err := strconv.Atoi(arg)
		if err == nil {
			orders = append(orders, order)
		}
	}

	path := todoFilePath()
	todos, err := readTodos(path)
	if err != nil {
		return err
	}

	var newTodos []todo.Todo
	for i, todo := range todos {
		order := i + 1
		if !contains(orders, order) {
			newTodos = append(newTodos, todo)
		}
	}

	err = writeTodos(newTodos, path)
	if err != nil {
		return err
	}

	return nil
}

func contains(ns []int, n int) bool {
	for _, e := range ns {
		if e == n {
			return true
		}
	}
	return false
}
<commit_msg>Delete subtodos<commit_after>package command

import (
	"strconv"
	"strings"

	"github.com/naoty/todo/todo"
	"github.com/urfave/cli"
)

// Delete is a command to delete todos.
var Delete = cli.Command{
	Name:   "delete",
	Usage:  "Delete todos",
	Action: delete,
}

func delete(c *cli.Context) error {
	if c.NArg() == 0 {
		cli.ShowCommandHelp(c, "delete")
		return nil
	}

	path := todoFilePath()
	todos, err := readTodos(path)
	if err != nil {
		return err
	}

	for _, arg := range c.Args() {
		var orders []int
		for _, id := range strings.Split(arg, "-") {
			order, err2 := strconv.Atoi(id)
			if err2 == nil {
				orders = append(orders, order)
			}
		}

		todos = deleteTodo(todos, orders)
	}

	err = writeTodos(todos, path)
	if err != nil {
		return err
	}

	return nil
}

func deleteTodo(todos []todo.Todo, orders []int) []todo.Todo {
	if len(orders) == 0 || len(todos) == 0 {
		return todos
	}

	i := orders[0] - 1
	if i >= len(todos) {
		return todos
	}

	if len(orders) == 1 {
		if len(todos) == 1 {
			return []todo.Todo{}
		}

		if i+1 >= len(todos) {
			return todos[:i]
		}

		return append(todos[:i], todos[i+1:]...)
	}

	todo := todos[i]
	todo.Todos = deleteTodo(todo.Todos, orders[1:])
	todos[i] = todo

	return todos
}
<|endoftext|>


<commit_before>// Package cluster provides helix-alike leader/standby model.
//
// It rebalances resources to participants and watch for cluster changes.
//
// The cluster changes might be:
// participants come and go, resource added and deleted, leader change.
package cluster

// RebalanceCallback connects cluster with its caller when leader decides to rebalance.
type RebalanceCallback func(Decision)

// Controller is the contral brain of the cluster, which assigns resource tickets
// to participants.
//
// Controller works by assigning resources to participants if neccessary.
//
// For dbus cluster to work, all dbus instances must share the same configuration,
// using zookeeper as configuration repo is highly suggetsted.
type Controller interface {

	// Start startup the controller and start election across the cluster.
	Start() error

	// Stop closes the controller underlying connection and do all the related cleanup.
	Stop() error
}
<commit_msg>fix typo<commit_after>// Package cluster provides helix-alike leader/standby model.
//
// It rebalances resources to participants and watch for cluster changes.
//
// The cluster changes might be:
// participants come and go, resource added and deleted, leader change.
package cluster

// RebalanceCallback connects cluster with its caller when leader decides to rebalance.
type RebalanceCallback func(Decision)

// Controller is the contral brain of the cluster, which assigns resource tickets
// to participants.
//
// Controller works by assigning resources to participants if necessary.
//
// For dbus cluster to work, all dbus instances must share the same configuration,
// using zookeeper as configuration repo is highly suggetsted.
type Controller interface {

	// Start startup the controller and start election across the cluster.
	Start() error

	// Stop closes the controller underlying connection and do all the related cleanup.
	Stop() error
}
<|endoftext|>


<commit_before>// Copyright 2016 Michael Schenk. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package command

import (
	"flag"
	"strings"
	"time"

	"github.com/git-time-metric/gtm/metric"
	"github.com/git-time-metric/gtm/note"
	"github.com/git-time-metric/gtm/project"
	"github.com/git-time-metric/gtm/report"
	"github.com/git-time-metric/gtm/util"
	"github.com/mitchellh/cli"
)

// StatusCmd containt methods for status command
type StatusCmd struct {
	Ui cli.Ui
}

// NewStatus returns new StatusCmd struct
func NewStatus() (cli.Command, error) {
	return StatusCmd{}, nil
}

// Help returns help for status command
func (c StatusCmd) Help() string {
	helpText := `
Usage: gtm status [options]

  Show pending time for git repositories.

Options:

  -terminal-off=false        Exclude time spent in terminal (Terminal plug-in is required)

  -color=false               Always output color even if no terminal is detected, i.e 'gtm status -color | less -R'

  -total-only=false          Only display total pending time

  -long-duration             If total-only, display total pending time in long duration format

  -tags=""                   Project tags to report status for, i.e --tags tag1,tag2

  -all=false                 Show status for all projects
`
	return strings.TrimSpace(helpText)
}

// Run executes status command with args
func (c StatusCmd) Run(args []string) int {
	var color, terminalOff, totalOnly, all, profile, longDuration bool
	var tags string
	cmdFlags := flag.NewFlagSet("status", flag.ContinueOnError)
	cmdFlags.BoolVar(&color, "color", false, "Always output color even if no terminal is detected. Use this with pagers i.e 'less -R' or 'more -R'")
	cmdFlags.BoolVar(&terminalOff, "terminal-off", false, "Exclude time spent in terminal (Terminal plugin is required)")
	cmdFlags.BoolVar(&totalOnly, "total-only", false, "Only display total time")
	cmdFlags.BoolVar(&longDuration, "long-duration", false, "Display total time in long duration format")
	cmdFlags.StringVar(&tags, "tags", "", "Project tags to show status on")
	cmdFlags.BoolVar(&all, "all", false, "Show status for all projects")
	cmdFlags.BoolVar(&profile, "profile", false, "Enable profiling")
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}
	util.TimeTrackEnable = profile
	defer util.TimeTrack(time.Now(), "status.Run")

	if totalOnly && (all || tags != "") {
		c.Ui.Error("\n-tags and -all options not allowed with -total-only\n")
		return 1
	}

	var (
		err        error
		commitNote note.CommitNote
		out        string
	)

	index, err := project.NewIndex()
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	tagList := []string{}
	if tags != "" {
		tagList = util.Map(strings.Split(tags, ","), strings.TrimSpace)
	}

	projects, err := index.Get(tagList, all)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	options := report.OutputOptions{
		TotalOnly:    totalOnly,
		LongDuration: longDuration,
		TerminalOff:  terminalOff,
		Color:        color}

	for _, projPath := range projects {
		if commitNote, err = metric.Process(true, projPath); err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		o, err := report.Status(commitNote, options, projPath)
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		out += o
	}

	c.Ui.Output(out)
	return 0
}

// Synopsis returns help for status command
func (c StatusCmd) Synopsis() string {
	return "Show pending time"
}
<commit_msg>Output status -time-only without escape sequences<commit_after>// Copyright 2016 Michael Schenk. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package command

import (
	"flag"
	"fmt"
	"strings"
	"time"

	"github.com/git-time-metric/gtm/metric"
	"github.com/git-time-metric/gtm/note"
	"github.com/git-time-metric/gtm/project"
	"github.com/git-time-metric/gtm/report"
	"github.com/git-time-metric/gtm/util"
	"github.com/mitchellh/cli"
)

// StatusCmd containt methods for status command
type StatusCmd struct {
	Ui cli.Ui
}

// NewStatus returns new StatusCmd struct
func NewStatus() (cli.Command, error) {
	return StatusCmd{}, nil
}

// Help returns help for status command
func (c StatusCmd) Help() string {
	helpText := `
Usage: gtm status [options]

  Show pending time for git repositories.

Options:

  -terminal-off=false        Exclude time spent in terminal (Terminal plug-in is required)

  -color=false               Always output color even if no terminal is detected, i.e 'gtm status -color | less -R'

  -total-only=false          Only display total pending time

  -long-duration             If total-only, display total pending time in long duration format

  -tags=""                   Project tags to report status for, i.e --tags tag1,tag2

  -all=false                 Show status for all projects
`
	return strings.TrimSpace(helpText)
}

// Run executes status command with args
func (c StatusCmd) Run(args []string) int {
	var color, terminalOff, totalOnly, all, profile, longDuration bool
	var tags string
	cmdFlags := flag.NewFlagSet("status", flag.ContinueOnError)
	cmdFlags.BoolVar(&color, "color", false, "Always output color even if no terminal is detected. Use this with pagers i.e 'less -R' or 'more -R'")
	cmdFlags.BoolVar(&terminalOff, "terminal-off", false, "Exclude time spent in terminal (Terminal plugin is required)")
	cmdFlags.BoolVar(&totalOnly, "total-only", false, "Only display total time")
	cmdFlags.BoolVar(&longDuration, "long-duration", false, "Display total time in long duration format")
	cmdFlags.StringVar(&tags, "tags", "", "Project tags to show status on")
	cmdFlags.BoolVar(&all, "all", false, "Show status for all projects")
	cmdFlags.BoolVar(&profile, "profile", false, "Enable profiling")
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}
	util.TimeTrackEnable = profile
	defer util.TimeTrack(time.Now(), "status.Run")

	if totalOnly && (all || tags != "") {
		c.Ui.Error("\n-tags and -all options not allowed with -total-only\n")
		return 1
	}

	var (
		err        error
		commitNote note.CommitNote
		out        string
	)

	index, err := project.NewIndex()
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	tagList := []string{}
	if tags != "" {
		tagList = util.Map(strings.Split(tags, ","), strings.TrimSpace)
	}

	projects, err := index.Get(tagList, all)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	options := report.OutputOptions{
		TotalOnly:    totalOnly,
		LongDuration: longDuration,
		TerminalOff:  terminalOff,
		Color:        color}

	for _, projPath := range projects {
		if commitNote, err = metric.Process(true, projPath); err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		o, err := report.Status(commitNote, options, projPath)
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		out += o
	}

	if totalOnly {
		// plain output, no ansi escape sequences
		fmt.Print(out)
	} else {
		c.Ui.Output(out)
	}
	return 0
}

// Synopsis returns help for status command
func (c StatusCmd) Synopsis() string {
	return "Show pending time"
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The HAProxy Ingress Controller 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 controller

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type metrics struct {
	responseTime       *prometheus.HistogramVec
	ctlProcTimeSum     *prometheus.CounterVec
	ctlProcTimeCount   *prometheus.CounterVec
	procSecondsCounter *prometheus.CounterVec
	updatesCounter     *prometheus.CounterVec
	updateSuccessGauge *prometheus.GaugeVec
	certExpireGauge    *prometheus.GaugeVec
	lastTrack          time.Time
}

func createMetrics() *metrics {
	namespace := "haproxyingress"
	metrics := &metrics{
		responseTime: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: namespace,
				Name:      "haproxy_response_time_seconds",
				Help:      "Response time to commands sent via admin socket",
				Buckets:   []float64{.0005, .001, .002, .005, .01},
			},
			[]string{"command"},
		),
		ctlProcTimeSum: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "controller_processing_time_seconds_sum",
				Help:      "Cumulative time in seconds spent on haproxy-ingress tasks",
			},
			[]string{"task"},
		),
		ctlProcTimeCount: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "controller_processing_time_seconds_count",
				Help:      "Cumulative number of haproxy-ingress tasks executed",
			},
			[]string{"task"},
		),
		procSecondsCounter: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "haproxy_processing_seconds_total",
				Help:      "Cumulative time in seconds a single thread spent processing requests, based on Idle_pct.",
			},
			[]string{},
		),
		updatesCounter: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "updates_total",
				Help:      "Cumulative number of Ingress controller updates. Status can be noop, dynamic, full.",
			},
			[]string{"status"},
		),
		updateSuccessGauge: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "update_success",
				Help:      "Whether the last haproxy update was successful.",
			},
			[]string{},
		),
		certExpireGauge: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "cert_expire_date_epoch",
				Help:      "The SSL certificate expiration date in unix epoch time.",
			},
			[]string{"domain", "cn"},
		),
	}
	prometheus.MustRegister(metrics.responseTime)
	prometheus.MustRegister(metrics.ctlProcTimeSum)
	prometheus.MustRegister(metrics.ctlProcTimeCount)
	prometheus.MustRegister(metrics.procSecondsCounter)
	prometheus.MustRegister(metrics.updatesCounter)
	prometheus.MustRegister(metrics.updateSuccessGauge)
	prometheus.MustRegister(metrics.certExpireGauge)
	return metrics
}

func (m *metrics) HAProxyShowInfoResponseTime(duration time.Duration) {
	m.responseTime.WithLabelValues("show_info").Observe(duration.Seconds())
}

func (m *metrics) HAProxySetServerResponseTime(duration time.Duration) {
	m.responseTime.WithLabelValues("set_server").Observe(duration.Seconds())
}

func (m *metrics) ControllerProcTime(task string, duration time.Duration) {
	m.ctlProcTimeSum.WithLabelValues(task).Add(duration.Seconds())
	m.ctlProcTimeCount.WithLabelValues(task).Inc()
}

func (m *metrics) AddIdleFactor(idle int) {
	now := time.Now()
	if m.lastTrack.IsZero() {
		m.lastTrack = now
		return
	}
	totalTime := now.Sub(m.lastTrack).Seconds()
	m.lastTrack = now
	m.procSecondsCounter.WithLabelValues().Add(float64(100-idle) * totalTime / 100)
}

func (m *metrics) IncUpdateNoop() {
	m.updatesCounter.WithLabelValues("noop").Inc()
}

func (m *metrics) IncUpdateDynamic() {
	m.updatesCounter.WithLabelValues("dynamic").Inc()
}

func (m *metrics) IncUpdateFull() {
	m.updatesCounter.WithLabelValues("full").Inc()
}

func (m *metrics) UpdateSuccessful(success bool) {
	value := map[bool]float64{false: 0, true: 1}
	m.updateSuccessGauge.WithLabelValues().Set(value[success])
}

func (m *metrics) SetCertExpireDate(domain, cn string, notAfter *time.Time) {
	if notAfter == nil {
		m.certExpireGauge.DeleteLabelValues(domain, cn)
		return
	}
	m.certExpireGauge.WithLabelValues(domain, cn).Set(float64(notAfter.Unix()))
}
<commit_msg>fix processing count metric name<commit_after>/*
Copyright 2019 The HAProxy Ingress Controller 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 controller

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type metrics struct {
	responseTime       *prometheus.HistogramVec
	ctlProcTimeSum     *prometheus.CounterVec
	ctlProcCount       *prometheus.CounterVec
	procSecondsCounter *prometheus.CounterVec
	updatesCounter     *prometheus.CounterVec
	updateSuccessGauge *prometheus.GaugeVec
	certExpireGauge    *prometheus.GaugeVec
	lastTrack          time.Time
}

func createMetrics() *metrics {
	namespace := "haproxyingress"
	metrics := &metrics{
		responseTime: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: namespace,
				Name:      "haproxy_response_time_seconds",
				Help:      "Response time to commands sent via admin socket",
				Buckets:   []float64{.0005, .001, .002, .005, .01},
			},
			[]string{"command"},
		),
		ctlProcTimeSum: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "controller_processing_time_seconds_sum",
				Help:      "Cumulative time in seconds spent on haproxy-ingress tasks",
			},
			[]string{"task"},
		),
		ctlProcCount: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "controller_processing_count",
				Help:      "Cumulative number of haproxy-ingress tasks executed",
			},
			[]string{"task"},
		),
		procSecondsCounter: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "haproxy_processing_seconds_total",
				Help:      "Cumulative time in seconds a single thread spent processing requests, based on Idle_pct.",
			},
			[]string{},
		),
		updatesCounter: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: namespace,
				Name:      "updates_total",
				Help:      "Cumulative number of Ingress controller updates. Status can be noop, dynamic, full.",
			},
			[]string{"status"},
		),
		updateSuccessGauge: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "update_success",
				Help:      "Whether the last haproxy update was successful.",
			},
			[]string{},
		),
		certExpireGauge: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "cert_expire_date_epoch",
				Help:      "The SSL certificate expiration date in unix epoch time.",
			},
			[]string{"domain", "cn"},
		),
	}
	prometheus.MustRegister(metrics.responseTime)
	prometheus.MustRegister(metrics.ctlProcTimeSum)
	prometheus.MustRegister(metrics.ctlProcCount)
	prometheus.MustRegister(metrics.procSecondsCounter)
	prometheus.MustRegister(metrics.updatesCounter)
	prometheus.MustRegister(metrics.updateSuccessGauge)
	prometheus.MustRegister(metrics.certExpireGauge)
	return metrics
}

func (m *metrics) HAProxyShowInfoResponseTime(duration time.Duration) {
	m.responseTime.WithLabelValues("show_info").Observe(duration.Seconds())
}

func (m *metrics) HAProxySetServerResponseTime(duration time.Duration) {
	m.responseTime.WithLabelValues("set_server").Observe(duration.Seconds())
}

func (m *metrics) ControllerProcTime(task string, duration time.Duration) {
	m.ctlProcTimeSum.WithLabelValues(task).Add(duration.Seconds())
	m.ctlProcCount.WithLabelValues(task).Inc()
}

func (m *metrics) AddIdleFactor(idle int) {
	now := time.Now()
	if m.lastTrack.IsZero() {
		m.lastTrack = now
		return
	}
	totalTime := now.Sub(m.lastTrack).Seconds()
	m.lastTrack = now
	m.procSecondsCounter.WithLabelValues().Add(float64(100-idle) * totalTime / 100)
}

func (m *metrics) IncUpdateNoop() {
	m.updatesCounter.WithLabelValues("noop").Inc()
}

func (m *metrics) IncUpdateDynamic() {
	m.updatesCounter.WithLabelValues("dynamic").Inc()
}

func (m *metrics) IncUpdateFull() {
	m.updatesCounter.WithLabelValues("full").Inc()
}

func (m *metrics) UpdateSuccessful(success bool) {
	value := map[bool]float64{false: 0, true: 1}
	m.updateSuccessGauge.WithLabelValues().Set(value[success])
}

func (m *metrics) SetCertExpireDate(domain, cn string, notAfter *time.Time) {
	if notAfter == nil {
		m.certExpireGauge.DeleteLabelValues(domain, cn)
		return
	}
	m.certExpireGauge.WithLabelValues(domain, cn).Set(float64(notAfter.Unix()))
}
<|endoftext|>


<commit_before>package dhcp6client

import (
	"math/rand"
	"net"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/header"
	"github.com/mdlayher/dhcp6"
)

func newSolicitOptions(mac net.HardwareAddr) (dhcp6.Options, error) {
	// make options: iana
	options := make(dhcp6.Options)

	var id = [4]byte{'r', 'o', 'o', 't'}
	if err := options.Add(dhcp6.OptionIANA, dhcp6.NewIANA(id, 0, 0, nil)); err != nil {
		return nil, err
	}
	if err := options.Add(dhcp6.OptionRapidCommit, nil); err != nil {
		return nil, err
	}
	if err := options.Add(dhcp6.OptionElapsedTime, dhcp6.ElapsedTime(0)); err != nil {
		return nil, err
	}

	// Requesting Default IP TTL and MTU Timeout.
	oro := dhcp6.OptionRequestOption{dhcp6.OptionCode(0x17), dhcp6.OptionCode(0x18)}
	if err := options.Add(dhcp6.OptionORO, oro); err != nil {
		return nil, err
	}

	// HardwareType *should* be 1?
	if err := options.Add(dhcp6.OptionClientID, dhcp6.NewDUIDLL(6, mac)); err != nil {
		return nil, err
	}

	return options, nil
}

func newSolicitPacket(mac net.HardwareAddr) (*dhcp6.Packet, error) {
	options, err := newSolicitOptions(mac)
	if err != nil {
		return nil, err
	}

	return &dhcp6.Packet{
		MessageType:   dhcp6.MessageTypeSolicit,
		TransactionID: [3]byte{0x00, 0x01, 0x02},
		Options:       options,
	}, nil
}

// mac2ipv6 gives the EUI-64 IPv6 address corresponding to the mac.
func mac2ipv6(mac net.HardwareAddr) []byte {
	var v6addr []byte
	v6addr = append(mac[:3], append([]byte{0xff, 0xfe}, mac[3:]...)...)

	// Invert 7th bit from left.
	if v6addr[0]&0x02 == 0x02 {
		v6addr[0] -= 0x02
	} else {
		v6addr[0] += 0x02
	}

	return append([]byte{0xfe, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, v6addr...)
}

// ipv6UDPPacket wraps a dhcp6 packet in a IPv6 and UDP packet for sending on a
// packet socket.
func ipv6UDPPacket(p *dhcp6.Packet, srcMAC net.HardwareAddr) ([]byte, error) {
	pb, err := p.MarshalBinary()
	if err != nil {
		return nil, err
	}

	length := uint16(header.UDPMinimumSize + len(pb))
	ipv6fields := &header.IPv6Fields{
		FlowLabel:     rand.Uint32(),
		PayloadLength: length,
		NextHeader:    uint8(header.UDPProtocolNumber),
		HopLimit:      1,
		SrcAddr:       tcpip.Address(mac2ipv6(srcMAC)),
		DstAddr:       tcpip.Address(net.ParseIP("FF02::1:2").To16()),
	}
	ipv6header := header.IPv6(make([]byte, header.IPv6MinimumSize))
	ipv6header.Encode(ipv6fields)

	udphdr := header.UDP(make([]byte, header.UDPMinimumSize))
	udphdr.Encode(&header.UDPFields{
		SrcPort: srcPort,
		DstPort: dstPort,
		Length:  length,
	})

	xsum := header.Checksum(pb, header.PseudoHeaderChecksum(ipv6header.TransportProtocol(), ipv6fields.SrcAddr, ipv6fields.DstAddr))
	udphdr.SetChecksum(^udphdr.CalculateChecksum(xsum, length))

	pkt := append([]byte(ipv6header), []byte(udphdr)...)
	return append(pkt, pb...), nil
}
<commit_msg>dhcp6client: name option codes used<commit_after>package dhcp6client

import (
	"math/rand"
	"net"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/header"
	"github.com/mdlayher/dhcp6"
)

// Should send a PR to dhcp6 to add these and others.
const (
	OptionDefaultIPTTL dhcp6.OptionCode = 23
	OptionMTUTimeout   dhcp6.OptionCode = 24
)

func newSolicitOptions(mac net.HardwareAddr) (dhcp6.Options, error) {
	// make options: iana
	options := make(dhcp6.Options)

	var id = [4]byte{'r', 'o', 'o', 't'}
	if err := options.Add(dhcp6.OptionIANA, dhcp6.NewIANA(id, 0, 0, nil)); err != nil {
		return nil, err
	}
	if err := options.Add(dhcp6.OptionRapidCommit, nil); err != nil {
		return nil, err
	}
	if err := options.Add(dhcp6.OptionElapsedTime, dhcp6.ElapsedTime(0)); err != nil {
		return nil, err
	}

	oro := dhcp6.OptionRequestOption{OptionDefaultIPTTL, OptionMTUTimeout}
	if err := options.Add(dhcp6.OptionORO, oro); err != nil {
		return nil, err
	}

	// HardwareType *should* be 1?
	if err := options.Add(dhcp6.OptionClientID, dhcp6.NewDUIDLL(6, mac)); err != nil {
		return nil, err
	}

	return options, nil
}

func newSolicitPacket(mac net.HardwareAddr) (*dhcp6.Packet, error) {
	options, err := newSolicitOptions(mac)
	if err != nil {
		return nil, err
	}

	return &dhcp6.Packet{
		MessageType:   dhcp6.MessageTypeSolicit,
		TransactionID: [3]byte{0x00, 0x01, 0x02},
		Options:       options,
	}, nil
}

// mac2ipv6 gives the EUI-64 IPv6 address corresponding to the mac.
func mac2ipv6(mac net.HardwareAddr) []byte {
	var v6addr []byte
	v6addr = append(mac[:3], append([]byte{0xff, 0xfe}, mac[3:]...)...)

	// Invert 7th bit from left.
	if v6addr[0]&0x02 == 0x02 {
		v6addr[0] -= 0x02
	} else {
		v6addr[0] += 0x02
	}

	return append([]byte{0xfe, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, v6addr...)
}

// ipv6UDPPacket wraps a dhcp6 packet in a IPv6 and UDP packet for sending on a
// packet socket.
func ipv6UDPPacket(p *dhcp6.Packet, srcMAC net.HardwareAddr) ([]byte, error) {
	pb, err := p.MarshalBinary()
	if err != nil {
		return nil, err
	}

	length := uint16(header.UDPMinimumSize + len(pb))
	ipv6fields := &header.IPv6Fields{
		FlowLabel:     rand.Uint32(),
		PayloadLength: length,
		NextHeader:    uint8(header.UDPProtocolNumber),
		HopLimit:      1,
		SrcAddr:       tcpip.Address(mac2ipv6(srcMAC)),
		DstAddr:       tcpip.Address(net.ParseIP("FF02::1:2").To16()),
	}
	ipv6header := header.IPv6(make([]byte, header.IPv6MinimumSize))
	ipv6header.Encode(ipv6fields)

	udphdr := header.UDP(make([]byte, header.UDPMinimumSize))
	udphdr.Encode(&header.UDPFields{
		SrcPort: srcPort,
		DstPort: dstPort,
		Length:  length,
	})

	xsum := header.Checksum(pb, header.PseudoHeaderChecksum(ipv6header.TransportProtocol(), ipv6fields.SrcAddr, ipv6fields.DstAddr))
	udphdr.SetChecksum(^udphdr.CalculateChecksum(xsum, length))

	pkt := append([]byte(ipv6header), []byte(udphdr)...)
	return append(pkt, pb...), nil
}
<|endoftext|>